from dataclasses import dataclass from datetime import datetime from typing import IO, TYPE_CHECKING, List, Optional, Union from lightning_sdk.api import OrgApi from lightning_sdk.api.billing_api import ActivityFileFormat, BillingActivity, BillingActivityFilterValues, BillingApi from lightning_sdk.api.org_api import MonthlySummary from lightning_sdk.owner import Owner from lightning_sdk.utils.resolve import _resolve_org_name, _resolve_teamspace if TYPE_CHECKING: from lightning_sdk.teamspace import Teamspace @dataclass class BillingActivityFilters: """Filters for narrowing down a billing activity query. Attributes: resource_types: Only include activity for these resource types (e.g. "Studio", "Job"). resource_ids: Only include activity for these specific resource IDs. user_ids: Only include activity generated by these users. start: Only include activity on or after this time. end: Only include activity on or before this time. limit: Maximum number of entries to return. """ resource_types: Optional[List[str]] = None resource_ids: Optional[List[str]] = None user_ids: Optional[List[str]] = None start: Optional[datetime] = None end: Optional[datetime] = None limit: Optional[int] = None @dataclass class BillingActivityCursor: """Pagination cursor for continuing a billing activity query. Attributes: search_after: Only include entries strictly after this time. search_after_resource_id: Resource ID of the last entry from the previous page. Required alongside ``search_after`` to break ties between resources sharing the same time. search_after_resource_type: Resource type of the last entry from the previous page. Required alongside ``search_after`` to break ties between resources sharing the same time. """ search_after: Optional[datetime] = None search_after_resource_id: Optional[str] = None search_after_resource_type: Optional[str] = None class Organization(Owner): """Represents an organization owner of teamspaces and studios. Args: name: the name of the organization Note: Arguments will be automatically inferred from environment variables if possible, unless explicitly specified """ def __init__(self, name: Optional[str] = None) -> None: super().__init__() self._org_api = OrgApi() self._billing_api = BillingApi() if name is None: name = _resolve_org_name(name) if name is None: raise ValueError( "Neither name is provided nor can the organization be inferred from the environment variable!" ) self._org = self._org_api.get_org(name=name) @property def name(self) -> str: """The organization's name.""" return self._org.name @property def id(self) -> str: """The organization's ID.""" return self._org.id @property def default_cloud_account(self) -> Optional[str]: """The organization's preferred cloud account ID, or None if not set. Returns: str | None: The preferred cloud account ID. """ return self._org.preferred_cluster or None def create_teamspace(self, name: str) -> "Teamspace": """Create a new teamspace owned by this organization. Args: name: The name for the new teamspace. Returns: Teamspace: The newly created teamspace. """ from lightning_sdk.teamspace import Teamspace self._org_api.create_teamspace(name, self.id) return Teamspace(name=name, org=self) def get_monthly_summary( self, start: Optional[datetime] = None, end: Optional[datetime] = None, ) -> list[MonthlySummary]: """Returns a monthly summary of credits purchased, used, and remaining. Exactly one of ``start`` and ``end`` must be supplied, or both together: - only ``start``: grabs monthly summaries for months after ``pivot=start``. - only ``end``: grabs monthly summaries for months before ``pivot=end``. - both ``start`` and ``end``: acts as a normal range. Args: start: Start of the time range. If given without ``end``, acts as an "AFTER" pivot. end: End of the time range. If given without ``start``, acts as a "BEFORE" pivot. Returns: A list of :class:`~lightning_sdk.api.org_api.MonthlySummary` dicts, one per month in the requested range:: [ { "period_start": datetime, "period_end": datetime, "total_credits_remaining": float, "total_credits_purchased": float, }, ... ] Raises: ValueError: If neither ``start`` nor ``end`` is provided, if ``start`` is after ``end``, or if an "AFTER" pivot (derived from a lone ``start``) is in the future. """ return self._org_api.get_monthly_summary( self.id, start=start, end=end, ) def _resolve_billing_teamspaces( self, teamspace: Optional[Union[str, "Teamspace"]] = None, teamspaces: Optional[List[Union[str, "Teamspace"]]] = None, ) -> List["Teamspace"]: """Resolve and validate the teamspace(s) to scope a billing query to.""" if teamspace is not None and teamspaces is not None: raise ValueError("Provide at most one of 'teamspace' or 'teamspaces'.") raw_teamspaces: List[Union[str, "Teamspace"]] = [] if teamspace is not None: raw_teamspaces = [teamspace] elif teamspaces is not None: raw_teamspaces = list(teamspaces) resolved_teamspaces: List["Teamspace"] = [] for raw_teamspace in raw_teamspaces: resolved_teamspace = _resolve_teamspace(raw_teamspace, org=self, user=None) if resolved_teamspace is None: raise ValueError(f"Could not resolve teamspace {raw_teamspace!r}.") if isinstance(resolved_teamspace.owner, Organization) and resolved_teamspace.owner.id != self.id: raise ValueError( f"Teamspace '{resolved_teamspace.name}' belongs to organization " f"'{resolved_teamspace.owner.name}', not '{self.name}'." ) resolved_teamspaces.append(resolved_teamspace) return resolved_teamspaces def get_activity( self, teamspace: Optional[Union[str, "Teamspace"]] = None, teamspaces: Optional[List[Union[str, "Teamspace"]]] = None, filters: Optional[BillingActivityFilters] = None, cursor: Optional[BillingActivityCursor] = None, ) -> BillingActivity: """Get billing activity for this organization, or one or more of its teamspaces. Returns paginated, per resource rollup rows keyed by raw IDs (no resolved names). For a resolved-name, non-paginated report meant for reading or exporting, use :meth:`get_session_activity` (one row per session) or :meth:`get_resource_activity` (one row per resource) instead. Args: teamspace: A single teamspace to scope the query to. teamspaces: Multiple teamspaces to scope the query to. Provide at most one of ``teamspace``/``teamspaces``. If neither is given, the query covers the whole organization. filters: Optional filters to narrow down the query. cursor: Optional pagination cursor to continue a previous query. Returns: BillingActivity: The usage report response. """ filters = filters or BillingActivityFilters() cursor = cursor or BillingActivityCursor() resolved_teamspaces = self._resolve_billing_teamspaces(teamspace, teamspaces) return self._billing_api.get_activity( org_id=self.id, project_ids=[t.id for t in resolved_teamspaces] or None, resource_types=filters.resource_types, resource_ids=filters.resource_ids, user_ids=filters.user_ids, start=filters.start, end=filters.end, limit=filters.limit, search_after=cursor.search_after, search_after_resource_id=cursor.search_after_resource_id, search_after_resource_type=cursor.search_after_resource_type, ) def get_activity_filter_values( self, teamspace: Optional[Union[str, "Teamspace"]] = None, ) -> BillingActivityFilterValues: """Get the set of values a billing activity query can be filtered by. Args: teamspace: Optionally scope to a single teamspace. If omitted, covers the whole organization. Returns: BillingActivityFilterValues: The available filter values. """ resolved_teamspaces = self._resolve_billing_teamspaces(teamspace) project_id = resolved_teamspaces[0].id if resolved_teamspaces else None return self._billing_api.get_activity_filter_values(org_id=self.id, project_id=project_id) def get_session_activity( self, writer: IO[str], format: ActivityFileFormat = ActivityFileFormat.JSON, # noqa: A002 teamspace: Optional[Union[str, "Teamspace"]] = None, teamspaces: Optional[List[Union[str, "Teamspace"]]] = None, filters: Optional[BillingActivityFilters] = None, cursor: Optional[BillingActivityCursor] = None, ) -> None: """Get the session-level billing activity report for this organization, as CSV or JSON. One row per session. A resource (e.g. a Studio or Job) can have many sessions within the queried time range, so this report is the finer-grained of the two; use :meth:`get_resource_activity` for one row per resource instead. Args: writer: Writable text stream to write the report to, e.g. an open file, an ``io.StringIO`` buffer, or ``sys.stdout``. To get the JSON as a string rather than writing it to a file, pass an ``io.StringIO()`` and read it back with ``.getvalue()``. format: File format to return the report in, either CSV or JSON. Defaults to JSON. teamspace: A single teamspace to scope the report to. teamspaces: Multiple teamspaces to scope the report to. Provide at most one of ``teamspace``/``teamspaces``. filters: Optional filters to narrow down the report. cursor: Optional pagination cursor to continue a previous query. """ filters = filters or BillingActivityFilters() cursor = cursor or BillingActivityCursor() resolved_teamspaces = self._resolve_billing_teamspaces(teamspace, teamspaces) self._billing_api.get_session_activity( org_id=self.id, format=format, writer=writer, project_ids=[t.id for t in resolved_teamspaces] or None, resource_types=filters.resource_types, resource_ids=filters.resource_ids, user_ids=filters.user_ids, start=filters.start, end=filters.end, limit=filters.limit, search_after=cursor.search_after, search_after_resource_id=cursor.search_after_resource_id, search_after_resource_type=cursor.search_after_resource_type, ) def get_resource_activity( self, writer: IO[str], format: ActivityFileFormat = ActivityFileFormat.JSON, # noqa: A002 teamspace: Optional[Union[str, "Teamspace"]] = None, teamspaces: Optional[List[Union[str, "Teamspace"]]] = None, filters: Optional[BillingActivityFilters] = None, cursor: Optional[BillingActivityCursor] = None, ) -> None: """Get the resource-level billing activity report for this organization, as CSV or JSON. One row per resource (e.g. a Studio or Job) that was active in the queried time range, rather than one row per session; use :meth:`get_session_activity` for the finer-grained, per-session breakdown of a resource's activity. Args: writer: Writable text stream to write the report to, e.g. an open file, an ``io.StringIO`` buffer, or ``sys.stdout``. To get the JSON as a string rather than writing it to a file, pass an ``io.StringIO()`` and read it back with ``.getvalue()``. format: File format to return the report in, either CSV or JSON. Defaults to JSON. teamspace: A single teamspace to scope the report to. teamspaces: Multiple teamspaces to scope the report to. Provide at most one of ``teamspace``/``teamspaces``. filters: Optional filters to narrow down the report. cursor: Optional pagination cursor to continue a previous query. """ filters = filters or BillingActivityFilters() cursor = cursor or BillingActivityCursor() resolved_teamspaces = self._resolve_billing_teamspaces(teamspace, teamspaces) self._billing_api.get_resource_activity( org_id=self.id, format=format, writer=writer, project_ids=[t.id for t in resolved_teamspaces] or None, resource_types=filters.resource_types, resource_ids=filters.resource_ids, user_ids=filters.user_ids, start=filters.start, end=filters.end, limit=filters.limit, search_after=cursor.search_after, search_after_resource_id=cursor.search_after_resource_id, search_after_resource_type=cursor.search_after_resource_type, ) def __repr__(self) -> str: """Returns reader friendly representation.""" return f"Organization(name={self.name})"