lightning-sdk 2025.11.13.post0__py3-none-any.whl

lightning_sdk/mmt/v2.py

@@ -0,0 +1,252 @@
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
+
+from lightning_sdk.api.mmt_api import MMTApiV2
+from lightning_sdk.api.utils import _get_cloud_url
+from lightning_sdk.status import Status
+from lightning_sdk.utils.resolve import _get_org_id
+
+if TYPE_CHECKING:
+    from lightning_sdk.job.job import Job
+    from lightning_sdk.machine import CloudProvider, Machine
+    from lightning_sdk.organization import Organization
+    from lightning_sdk.studio import Studio
+    from lightning_sdk.teamspace import Teamspace
+    from lightning_sdk.user import User
+
+from lightning_sdk.mmt.base import _BaseMMT
+
+
+class _MMTV2(_BaseMMT):
+    """New implementation of Multi-Machine Training."""
+
+    def __init__(
+        self,
+        name: str,
+        teamspace: Union[str, "Teamspace", None] = None,
+        org: Union[str, "Organization", None] = None,
+        user: Union[str, "User", None] = None,
+        *,
+        _fetch_job: bool = True,
+    ) -> None:
+        """Fetch already existing jobs.
+
+        Args:
+            name: the name of the job
+            teamspace: the teamspace the job is part of
+            org: the name of the organization owning the :param`teamspace` in case it is owned by an org
+            user: the name of the user owning the :param`teamspace`
+                in case it is owned directly by a user instead of an org.
+        """
+        self._job_api = MMTApiV2()
+        super().__init__(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=_fetch_job)
+
+    def _submit(
+        self,
+        num_machines: int,
+        machine: Union["Machine", str],
+        command: Optional[str] = None,
+        studio: Optional["Studio"] = None,
+        image: Optional[str] = None,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+        cloud_account: Optional[str] = None,
+        cloud_provider: Optional[Union["CloudProvider", str]] = None,
+        image_credentials: Optional[str] = None,
+        cloud_account_auth: bool = False,
+        entrypoint: str = "sh -c",
+        path_mappings: Optional[Dict[str, str]] = None,
+        max_runtime: Optional[int] = None,
+        artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
+        artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
+        reuse_snapshot: bool = True,
+    ) -> "_MMTV2":
+        """Submit a new multi-machine job to the Lightning AI platform.
+
+        Args:
+            num_machines: The number of machines to run on.
+            machine: The machine type to run the job on. One of {", ".join(_MACHINE_VALUES)}.
+            command: The command to run inside your job. Required if using a studio. Optional if using an image.
+                If not provided for images, will run the container entrypoint and default command.
+            studio: The studio env to run the job with. Mutually exclusive with image.
+            image: The docker image to run the job with. Mutually exclusive with studio.
+            env: Environment variables to set inside the job.
+            interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
+            cloud_account: The cloud account to run the job on.
+                Defaults to the studio cloud account if running with studio compute env.
+                If not provided and `cloud_account_provider` is set, will resolve cluster from this, else
+                will fall back to the teamspaces default cloud account.
+            cloud_account_provider: The provider to select the cloud-account from.
+                If set, must be in agreement with the provider from the cloud_account (if specified).
+                If not specified, falls backto the teamspace default cloud account.
+            image_credentials: The credentials used to pull the image. Required if the image is private.
+                This should be the name of the respective credentials secret created on the Lightning AI platform.
+            cloud_account_auth: Whether to authenticate with the cloud account to pull the image.
+                Required if the registry is part of a cloud provider (e.g. ECR).
+            entrypoint: The entrypoint of your docker container. Defaults to `sh -c` which
+                just runs the provided command in a standard shell.
+                To use the pre-defined entrypoint of the provided image, set this to an empty string.
+                Only applicable when submitting docker jobs.
+            path_mappings: Dictionary of path mappings. The keys are the path inside the container whereas the value
+                represents the data-connection name and the path inside that connection.
+                Should be of form
+                    {
+                        "<CONTAINER_PATH_1>": "<CONNECTION_NAME_1>:<PATH_WITHIN_CONNECTION_1>",
+                        "<CONTAINER_PATH_2>": "<CONNECTION_NAME_2>"
+                    }
+                If the path inside the connection is omitted it's assumed to be the root path of that connection.
+                Only applicable when submitting docker jobs.
+            max_runtime: the duration (in seconds) for which to allocate the machine.
+                Irrelevant for most machines, required for some of the top-end machines on GCP.
+                If in doubt, set it. Won't have an effect on machines not requiring it.
+                Defaults to 3h
+            reuse_snapshot: Whether the job should reuse a Studio snapshot when multiple jobs for the same Studio are
+                submitted. Turning this off may result in longer job startup times. Defaults to True.
+        """
+        # Command is required if Studio is provided to know what to run
+        # Image is mutually exclusive with Studio
+        # Command is optional for Image
+        # Either image or studio must be provided
+        if studio is not None:
+            studio_id = studio._studio.id
+            if image is not None:
+                raise ValueError(
+                    "image and studio are mutually exclusive as both define the environment to run the job in"
+                )
+            if command is None:
+                raise ValueError("command is required when using a studio")
+        else:
+            studio_id = None
+            if image is None:
+                raise ValueError("either image or studio must be provided")
+
+        cloud_account = self._cloud_account_api.resolve_cloud_account(
+            self._teamspace.id,
+            cloud_account=cloud_account,
+            cloud_provider=cloud_provider,
+            default_cloud_account=self._teamspace.default_cloud_account,
+        )
+
+        submitted = self._job_api.submit_job(
+            name=self.name,
+            num_machines=num_machines,
+            command=command,
+            cloud_account=cloud_account,
+            teamspace_id=self._teamspace.id,
+            studio_id=studio_id,
+            image=image,
+            machine=machine,
+            interruptible=interruptible,
+            env=env,
+            image_credentials=image_credentials,
+            cloud_account_auth=cloud_account_auth,
+            entrypoint=entrypoint,
+            path_mappings=path_mappings,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+            max_runtime=max_runtime,
+            reuse_snapshot=reuse_snapshot,
+        )
+        self._job = submitted
+        self._name = submitted.name
+        return self
+
+    @property
+    def machines(self) -> Tuple["Job", ...]:
+        """Returns the sub-jobs for each individual instance."""
+        from lightning_sdk.job import Job
+
+        return tuple(
+            Job(name=j.name, teamspace=self.teamspace)
+            for j in self._job_api.list_mmt_subjobs(self._guaranteed_job.id, self.teamspace.id)
+        )
+
+    def stop(self) -> None:
+        """Stops the job."""
+        if self.status in (Status.Stopped, Status.Completed, Status.Failed):
+            return
+        self._job_api.stop_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id)
+
+    def delete(self) -> None:
+        """Deletes the job.
+
+        Caution: This also deletes all artifacts and snapshots associated with the job.
+        """
+        self._job_api.delete_job(
+            job_id=self._guaranteed_job.id,
+            teamspace_id=self._teamspace.id,
+        )
+
+    @property
+    def status(self) -> "Status":
+        """The current status of the job."""
+        return self._job_api._job_state_to_external(self._latest_job.state)
+
+    @property
+    def artifact_path(self) -> Optional[str]:
+        """Path to the artifacts created by the job within the distributed teamspace filesystem."""
+        # TODO: Since grouping for those is not done yet on the BE, we cannot yet have a unified link here
+        raise NotImplementedError
+
+    @property
+    def snapshot_path(self) -> Optional[str]:
+        """Path to the studio snapshot used to create the job within the distributed teamspace filesystem."""
+        # TODO: Since grouping for those is not done yet on the BE, we cannot yet have a unified link here
+        raise NotImplementedError
+
+    @property
+    def machine(self) -> Union["Machine", str]:
+        """Returns the machine type this job is running on."""
+        return self._job_api._get_job_machine_from_spec(
+            self._guaranteed_job.spec,
+            self.teamspace.id,
+            _get_org_id(self.teamspace),
+        )
+
+    def _update_internal_job(self) -> None:
+        if getattr(self, "_job", None) is None:
+            self._job = self._job_api.get_job_by_name(name=self._name, teamspace_id=self._teamspace.id)
+            return
+
+        self._job = self._job_api.get_job(job_id=self._job.id, teamspace_id=self._teamspace.id)
+
+    @property
+    def name(self) -> str:
+        """The job's name."""
+        return self._name
+
+    @property
+    def teamspace(self) -> "Teamspace":
+        """The teamspace the job is part of."""
+        return self._teamspace
+
+    @property
+    def link(self) -> str:
+        # TODO: MMT env with studio -> go to studio plugin
+        return f"{_get_cloud_url()}/{self.teamspace.owner.name}/{self.teamspace.name}/jobs/{self.name}?app_id=mmt"
+
+    @property
+    def image(self) -> Optional[str]:
+        """The image used to submit the job."""
+        return self._job_api.get_image_name(self._guaranteed_job)
+
+    @property
+    def studio(self) -> Optional["Studio"]:
+        """The studio used to submit the job."""
+        from lightning_sdk.studio import Studio
+
+        studio_name = self._job_api.get_studio_name(self._guaranteed_job)
+
+        # if job was submitted with image, studio will be None
+        if not studio_name:
+            return None
+        return Studio(studio_name, teamspace=self.teamspace)
+
+    @property
+    def command(self) -> str:
+        """The command the job is running."""
+        return self._job_api.get_command(self._guaranteed_job)
+
+    @property
+    def num_machines(self) -> int:
+        """Returns the number of machines assigned to this multi-machine job."""
+        return self._job_api.get_num_machines(self._guaranteed_job)

lightning_sdk/models.py

@@ -0,0 +1,208 @@
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+from lightning_sdk.api import OrgApi, TeamspaceApi, UserApi
+from lightning_sdk.api.utils import AccessibleResource, raise_access_error_if_not_allowed
+from lightning_sdk.lightning_cloud.openapi import V1ModelVersionArchive
+from lightning_sdk.lightning_cloud.openapi.models import V1Membership, V1OwnerType
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.user import User
+from lightning_sdk.utils.resolve import _get_authed_user, _resolve_teamspace
+
+if TYPE_CHECKING:
+    from lightning_sdk.teamspace import Teamspace
+
+
+# TODO: Maybe just have a `Model` object?
+@dataclass
+class UploadedModelInfo:
+    name: str
+    version: str
+    teamspace: str
+    cloud_account: str
+
+
+def _get_teamspace_and_path(
+    ts: V1Membership, org_api: OrgApi, user_api: UserApi, authed_user: User
+) -> Tuple[str, Dict[str, Any]]:
+    if ts.owner_type == V1OwnerType.ORGANIZATION:
+        org = org_api._get_org_by_id(ts.owner_id)
+        return f"{org.name}/{ts.name}", {"name": ts.name, "org": org.name}
+
+    if ts.owner_type == V1OwnerType.USER and ts.owner_id != authed_user.id:
+        user = user_api._get_user_by_id(ts.owner_id)  # todo: check also the name
+        return f"{user.username}/{ts.name}", {"name": ts.name, "user": User(name=user.username)}
+
+    if ts.owner_type == V1OwnerType.USER:
+        return f"{authed_user.name}/{ts.name}", {"name": ts.name, "user": authed_user}
+
+    raise RuntimeError(f"Unknown organization type {ts.owner_type}")
+
+
+def _list_teamspaces() -> List[str]:
+    org_api = OrgApi()
+    user_api = UserApi()
+    authed_user = _get_authed_user()
+
+    return [
+        _get_teamspace_and_path(ts, org_api, user_api, authed_user)[0]
+        for ts in user_api._get_all_teamspace_memberships("")
+    ]
+
+
+def _get_teamspace(name: str, organization: str) -> "Teamspace":
+    """Get a Teamspace object from the SDK."""
+    from lightning_sdk.teamspace import Teamspace
+
+    org_api = OrgApi()
+    user_api = UserApi()
+    authed_user = _get_authed_user()
+
+    requested_teamspace = f"{organization}/{name}".lower()
+
+    for ts in user_api._get_all_teamspace_memberships(""):
+        if ts.name != name:
+            continue
+
+        teamspace_path, teamspace = _get_teamspace_and_path(ts, org_api, user_api, authed_user)
+        if requested_teamspace == teamspace_path:
+            return Teamspace(**teamspace)
+
+    options = f"{os.linesep}\t".join(_list_teamspaces())
+    raise RuntimeError(f"Teamspace `{requested_teamspace}` not found. Available teamspaces: {os.linesep}\t{options}")
+
+
+def _extend_model_name_with_teamspace(name: str) -> str:
+    """Extend the model name with the teamspace if it can be determined from env. variables."""
+    if "/" in name:
+        return name
+    # do some magic if you run studio
+    teamspace = _resolve_teamspace(None, None, None)
+    if not teamspace:
+        raise ValueError(
+            f"Model name must be in the format `organization/teamspace/model_name` but you provided '{name}'."
+        )
+    return f"{teamspace.owner.name}/{teamspace.name}/{name}"
+
+
+def _parse_org_teamspace_model_version(name: str) -> Tuple[str, str, str, Optional[str]]:
+    """Parse the name argument into its components."""
+    try:
+        org_name, teamspace_name, model_name = name.split("/")
+    except ValueError as err:
+        raise ValueError(
+            f"Model name must be in the format `organization/teamspace/model_name` but you provided '{name}'."
+        ) from err
+    parts = model_name.split(":")
+    if len(parts) == 1:
+        return org_name, teamspace_name, parts[0], None
+    if len(parts) == 2:
+        return org_name, teamspace_name, parts[0], parts[1]
+    # The rest of the validation for name and version happens in the backend
+    raise ValueError(
+        "Model version is expected to be in the format `organization/teamspace/model_name:version`"
+        f" separated by a single colon, but got: {name}"
+    )
+
+
+def download_model(
+    name: str,
+    download_dir: Union[Path, str] = ".",
+    progress_bar: bool = True,
+) -> List[str]:
+    """Download a Model.
+
+    Args:
+        name: The name of the Model you want to download.
+        This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>.
+        download_dir: The directory where the Model should be downloaded.
+        progress_bar: Whether to show a progress bar when downloading.
+    """
+    name = _extend_model_name_with_teamspace(name)
+    teamspace_owner_name, teamspace_name, model_name, version = _parse_org_teamspace_model_version(name)
+
+    api = TeamspaceApi()
+
+    try:
+        return api.download_model_files(
+            name=model_name,
+            version=version,
+            download_dir=Path(download_dir).expanduser().resolve(),
+            teamspace_name=teamspace_name,
+            teamspace_owner_name=teamspace_owner_name,
+            progress_bar=progress_bar,
+        )
+    except ApiException as e:
+        if e.status == 404:
+            raise RuntimeError(
+                f"Model '{name}' not found. Either the model doesn't exist or you don't have access to it."
+            ) from None
+        raise RuntimeError(f"Error downloading model. Status code: {e.status}.") from None
+
+
+def upload_model(
+    name: str,
+    path: Union[str, Path, List[Union[str, Path]]] = ".",
+    cloud_account: Optional[str] = None,
+    progress_bar: bool = True,
+    metadata: Optional[Dict[str, Any]] = None,
+) -> UploadedModelInfo:
+    """Upload a Model.
+
+    Args:
+        name: The name of the Model you want to upload.
+            This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>.
+        path: The path to the file or directory you want to upload. Defaults to the current directory.
+        cloud_account: The name of the cloud account to store the Model in.
+            If not provided, the default cloud account for the Teamspace will be used.
+        progress_bar: Whether to show a progress bar for the upload.
+        metadata: Metadata to attach to the uploaded model.
+            If not provided, an empty dictionary will be used.
+    """
+    name = _extend_model_name_with_teamspace(name)
+    org_name, teamspace_name, model_name, version = _parse_org_teamspace_model_version(name)
+    teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
+    raise_access_error_if_not_allowed(AccessibleResource.Models, teamspace.id)
+    return teamspace.upload_model(
+        path=path,
+        name=model_name,
+        version=version,
+        cloud_account=cloud_account,
+        progress_bar=progress_bar,
+        metadata=metadata,
+    )
+
+
+def delete_model(
+    name: str,
+) -> None:
+    """Delete a model or a version of model.
+
+    Args:
+        name: The name of the model you want to delete or with specified version it deltes only that version.
+            This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME> for full model deletion
+            or <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>:<VERSION> for version deletion.
+    """
+    name = _extend_model_name_with_teamspace(name)
+    org_name, teamspace_name, model_name, version = _parse_org_teamspace_model_version(name)
+    teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
+    raise_access_error_if_not_allowed(AccessibleResource.Models, teamspace.id)
+    teamspace.delete_model(name=f"{model_name}:{version}" if version else model_name)
+
+
+def list_model_versions(
+    name: str,
+) -> List[V1ModelVersionArchive]:
+    """List all versions of a model.
+
+    Args:
+        name: The name of the model you want to list versions for.
+            This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>.
+    """
+    name = _extend_model_name_with_teamspace(name)
+    org_name, teamspace_name, model_name, _ = _parse_org_teamspace_model_version(name)
+    teamspace = _get_teamspace(name=teamspace_name, organization=org_name)
+    raise_access_error_if_not_allowed(AccessibleResource.Models, teamspace.id)
+    return teamspace.list_model_versions(name=model_name)

lightning_sdk/organization.py

@@ -0,0 +1,49 @@
+from typing import Optional
+
+from lightning_sdk.api import OrgApi
+from lightning_sdk.owner import Owner
+from lightning_sdk.utils.resolve import _resolve_org_name
+
+
+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()
+        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]:
+        return self._org.preferred_cluster or None
+
+    def __repr__(self) -> str:
+        """Returns reader friendly representation."""
+        return f"Organization(name={self.name})"

lightning_sdk/owner.py

@@ -0,0 +1,50 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, List
+
+from lightning_sdk.api import TeamspaceApi
+from lightning_sdk.utils.logging import TrackCallsABCMeta
+
+if TYPE_CHECKING:
+    from lightning_sdk.teamspace import Teamspace
+
+
+class Owner(ABC, metaclass=TrackCallsABCMeta):
+    """Represents an owner of teamspaces and studios."""
+
+    def __init__(self) -> None:
+        self._teamspace_api = TeamspaceApi()
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """The owner's name."""
+
+    @property
+    @abstractmethod
+    def id(self) -> str:
+        """The owner's ID."""
+
+    @property
+    def teamspaces(self) -> List["Teamspace"]:
+        """All teamspaces by this owner."""
+        from lightning_sdk.teamspace import Teamspace
+        from lightning_sdk.user import User
+
+        is_user = isinstance(self, User)
+        if is_user:
+            user = self
+            org = None
+        else:
+            user = None
+            org = self
+
+        _teamspaces = self._teamspace_api.list_teamspaces(owner_id=self.id, name=None)
+        return [Teamspace(name=t.name, user=user, org=org) for t in _teamspaces]
+
+    def __eq__(self, o: "Owner") -> bool:
+        """Checks for equality with provided object."""
+        return type(o) is type(self) and self.id == o.id and self.name == o.name
+
+    def __str__(self) -> str:
+        """Returns reader friendly representation."""
+        return repr(self)

lightning_sdk/pipeline/__init__.py

@@ -0,0 +1,14 @@
+from lightning_sdk import Studio
+from lightning_sdk.pipeline.pipeline import Pipeline
+from lightning_sdk.pipeline.schedule import Schedule
+from lightning_sdk.pipeline.steps import DeploymentReleaseStep, DeploymentStep, JobStep, MMTStep
+
+__all__ = [
+    "Pipeline",
+    "JobStep",
+    "MMTStep",
+    "DeploymentStep",
+    "Schedule",
+    "Studio",
+    "DeploymentReleaseStep",
+]