lightning-sdk 0.1.41__py3-none-any.whl → 0.1.43__py3-none-any.whl

lightning_sdk/mmt/v1.py

@@ -0,0 +1,181 @@
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
+
+from lightning_sdk.api.mmt_api import MMTApiV1
+from lightning_sdk.job.v1 import _internal_status_to_external_status
+from lightning_sdk.job.work import Work
+
+if TYPE_CHECKING:
+    from lightning_sdk.machine import Machine
+    from lightning_sdk.organization import Organization
+    from lightning_sdk.status import Status
+    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 _MMTV1(_BaseMMT):
+    """V1 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 = MMTApiV1()
+        super().__init__(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=_fetch_job)
+
+    def _submit(
+        self,
+        num_machines: int,
+        machine: "Machine",
+        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,
+        image_credentials: Optional[str] = None,
+        cloud_account_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
+    ) -> "_MMTV1":
+        """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 will fall back to the teamspaces 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).
+            artifacts_local: The path of inside the docker container, you want to persist images from.
+                CAUTION: When setting this to "/", it will effectively erase your container.
+                Only supported for jobs with a docker image compute environment.
+            artifacts_remote: The remote storage to persist your artifacts to.
+                Should be of format <CONNECTION_TYPE>:<CONNECTION_NAME>:<PATH_WITHIN_CONNECTION>.
+                PATH_WITHIN_CONNECTION hereby is a path relative to the connection's root.
+                E.g. efs:data:some-path would result in an EFS connection named `data` and to the path `some-path`
+                within it.
+                Note that the connection needs to be added to the teamspace already in order for it to be found.
+                Only supported for jobs with a docker image compute environment.
+        """
+        if studio is None:
+            raise ValueError("Studio is required for submitting jobs")
+        if image is not None or image_credentials is not None or cloud_account_auth:
+            raise ValueError("Image is not supported for submitting jobs")
+
+        if artifacts_local is not None or artifacts_remote is not None:
+            raise ValueError("Specifying how to persist artifacts is not yet supported with jobs")
+
+        if env is not None:
+            raise ValueError("Environment variables are not supported for submitting jobs")
+        if command is None:
+            raise ValueError("Command is required for submitting multi-machine jobs")
+
+        _submitted = self._job_api.submit_job(
+            name=self._name,
+            num_machines=num_machines,
+            command=command,
+            studio_id=studio._studio.id,
+            teamspace_id=self._teamspace.id,
+            cloud_account=cloud_account or "",
+            machine=machine,
+            interruptible=interruptible,
+            strategy="parallel",
+        )
+
+        self._name = _submitted.name
+        self._job = _submitted
+        return self
+
+    def _update_internal_job(self) -> None:
+        try:
+            self._job = self._job_api.get_job(self._name, self.teamspace.id)
+        except ValueError as e:
+            raise ValueError(f"Job {self._name} does not exist in Teamspace {self.teamspace.name}") from e
+
+    @property
+    def machines(self) -> Tuple["Work", ...]:
+        """Returns the sub-jobs for each individual instance."""
+        works = self._job_api.list_works(self._guaranteed_job.id, self.teamspace.id)
+
+        return tuple(Work(w.id, self, self.teamspace) for w in works)
+
+    def stop(self) -> None:
+        """Stops the job."""
+        self._job_api.stop_job(self._guaranteed_job.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(self._guaranteed_job.id, self.teamspace.id)
+
+    @property
+    def status(self) -> "Status":
+        """The current status of the job."""
+        try:
+            status = self._job_api.get_job_status(self._job.id, self.teamspace.id)
+            return _internal_status_to_external_status(status)
+        except Exception:
+            raise RuntimeError(
+                f"MMT {self._name} does not exist in Teamspace {self.teamspace.name}. Did you delete it?"
+            ) from None
+
+    @property
+    def artifact_path(self) -> Optional[str]:
+        """Path to the artifacts created by the job within the distributed teamspace filesystem."""
+        return f"/teamspace/jobs/{self.name}"
+
+    @property
+    def snapshot_path(self) -> Optional[str]:
+        """Path to the studio snapshot used to create the job within the distributed teamspace filesystem."""
+        return f"/teamspace/jobs/{self.name}/snapshot"
+
+    @property
+    def machine(self) -> "Machine":
+        """Returns the machine type this job is running on."""
+        return self.machines[0].machine
+
+    @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
+
+    # the following and functions are solely to make the Work class function
+    @property
+    def _id(self) -> str:
+        return self._guaranteed_job.id
+
+    def _name_filter(self, name: str) -> str:
+        return name.replace("root.", "")

lightning_sdk/mmt/v2.py

@@ -0,0 +1,188 @@
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
+
+from lightning_sdk.api.mmt_api import MMTApiV2
+
+if TYPE_CHECKING:
+    from lightning_sdk.job.job import Job
+    from lightning_sdk.machine import Machine
+    from lightning_sdk.organization import Organization
+    from lightning_sdk.status import Status
+    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: "Machine",
+        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,
+        image_credentials: Optional[str] = None,
+        cloud_account_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
+    ) -> "_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 will fall back to the teamspaces 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).
+            artifacts_local: The path of inside the docker container, you want to persist images from.
+                CAUTION: When setting this to "/", it will effectively erase your container.
+                Only supported for jobs with a docker image compute environment.
+            artifacts_remote: The remote storage to persist your artifacts to.
+                Should be of format <CONNECTION_TYPE>:<CONNECTION_NAME>:<PATH_WITHIN_CONNECTION>.
+                PATH_WITHIN_CONNECTION hereby is a path relative to the connection's root.
+                E.g. efs:data:some-path would result in an EFS connection named `data` and to the path `some-path`
+                within it.
+                Note that the connection needs to be added to the teamspace already in order for it to be found.
+                Only supported for jobs with a docker image compute environment.
+        """
+        # 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")
+        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,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+        )
+        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."""
+        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 _latest_job(self) -> Any:
+        """Guarantees to fetch the latest version of a job before returning it."""
+        self._update_internal_job()
+        return self._job
+
+    @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) -> "Machine":
+        """Returns the machine type this job is running on."""
+        return self._job_api._get_job_machine_from_spec(self._guaranteed_job.spec)
+
+    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

lightning_sdk/plugin.py

@@ -127,7 +127,15 @@ class JobsPlugin(_Plugin):
         cloud_compute: Optional[Machine] = None,
         interruptible: bool = False,
     ) -> Job:
-        """Launches an asynchronous job."""
+        """Launches an asynchronous job.
+
+        Args:
+            command: The command to be executed.
+            name: The name of the job.
+            machine: The machine to run the job on.
+            interruptible: Whether to run the job on an interruptible machine.
+                These are cheaper but can be preempted at any time.
+        """
         if not name:
             name = _run_name("job")
 
@@ -139,7 +147,7 @@ class JobsPlugin(_Plugin):
             command=command,
             studio=self._studio,
             teamspace=self._studio.teamspace,
-            cluster=self._studio._cluster,
+            cloud_account=self._studio.cloud_account,
             interruptible=interruptible,
         )
 
@@ -161,29 +169,39 @@ class MultiMachineTrainingPlugin(_Plugin):
         machine: Machine = Machine.CPU,
         cloud_compute: Optional[Machine] = None,
         num_instances: int = 2,
-        strategy: str = "parallel",
         interruptible: bool = False,
     ) -> Job:
-        """Launches an asynchronous multi-machine-training."""
+        """Launches an asynchronous multi-machine-training.
+
+        Args:
+            command: The command to be executed.
+            name: The name of the job.
+            machine: The machine to run the job on.
+            num_instances: The number of instances to run the job on.
+            interruptible: Whether to run the job on an interruptible machine.
+                These are cheaper but can be preempted at any time.
+        """
+        from lightning_sdk.mmt import MMT
+
         if not name:
             name = _run_name("dist-run")
 
         machine = _resolve_deprecated_cloud_compute(machine, cloud_compute)
 
-        # TODO: assert num_instances >=2
-        resp = self._studio._studio_api.create_multi_machine_job(
-            entrypoint=command,
+        MMT._force_v1 = True
+
+        mmt = MMT.run(
             name=name,
-            num_instances=num_instances,
+            num_machines=num_instances,
             machine=machine,
-            strategy=strategy,
-            studio_id=self._studio._studio.id,
-            teamspace_id=self._studio._teamspace.id,
-            cluster_id=self._studio._studio.cluster_id,
+            command=command,
+            studio=self._studio,
+            teamspace=self._studio.teamspace,
             interruptible=interruptible,
         )
 
-        return Job(resp.name, self._studio.teamspace)
+        MMT._force_v1 = False
+        return mmt
 
 
 class MultiMachineDataPrepPlugin(_Plugin):
@@ -201,7 +219,16 @@ class MultiMachineDataPrepPlugin(_Plugin):
         num_instances: int = 2,
         interruptible: bool = False,
     ) -> Job:
-        """Launches an asynchronous multi-machine-processing-job."""
+        """Launches an asynchronous multi-machine-data-processing job.
+
+        Args:
+            command: The command to be executed.
+            name: The name of the job.
+            machine: The machine to run the job on.
+            num_instances: The number of instances to run the job on.
+            interruptible: Whether to run the job on an interruptible machine.
+                These are cheaper but can be preempted at any time.
+        """
         if not name:
             name = _run_name("data-prep")
 
@@ -214,7 +241,7 @@ class MultiMachineDataPrepPlugin(_Plugin):
             machine=machine,
             studio_id=self._studio._studio.id,
             teamspace_id=self._studio._teamspace.id,
-            cluster_id=self._studio._studio.cluster_id,
+            cloud_account=self._studio.cloud_account,
             interruptible=interruptible,
         )
 
@@ -261,7 +288,7 @@ class InferenceServerPlugin(_Plugin):
             endpoint=endpoint,
             studio_id=self._studio._studio.id,
             teamspace_id=self._studio._teamspace.id,
-            cluster_id=self._studio._studio.cluster_id,
+            cloud_account=self._studio.cloud_account,
             interruptible=interruptible,
         )
 

lightning_sdk/services/file_endpoint.py

@@ -11,6 +11,7 @@ from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import CommandArgumentCommandArgumentType
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.services.utilities import _get_cluster, _get_project, _get_service_url
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster
 
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
@@ -25,16 +26,19 @@ class Client:
         self,
         name: str,
         teamspace: Optional[str],
-        cluster_id: Optional[str] = None,
+        cloud_account: Optional[str] = None,
+        cluster_id: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> None:
         """Constructor of the Client.
 
         Args:
             name: The name of the Studio File Endpoint Service.
             teamspace: The name of the teamspace you want to attach the upload data and artifacts to be.
-            cluster_id: The name of the cluster on which to upload the data.
+            cloud_account: The name of the cloud account on which to upload the data.
 
         """
+        cloud_account = _resolve_deprecated_cluster(cloud_account, cluster_id)
+
         self._auth = Auth()
 
         try:
@@ -46,7 +50,9 @@ class Client:
         self._teamspace = teamspace
         self._client = LightningClient()
         self._project = _get_project(client=self._client, project_name=teamspace)
-        self._cluster = _get_cluster(client=self._client, project_id=self._project.project_id, cluster_id=cluster_id)
+        self._cloud_account = _get_cluster(
+            client=self._client, project_id=self._project.project_id, cluster_id=cloud_account
+        )
         self._file_endpoint = self._client.endpoint_service_get_file_endpoint_by_name(
             project_id=self._project.project_id, name=self._name
         )
@@ -101,7 +107,7 @@ class Client:
             _FileUploader(
                 client=self._client,
                 teamspace_id=self._project.project_id,
-                cluster_id=self._cluster.cluster_id,
+                cloud_account=self._cloud_account.cluster_id,
                 file_path=argument.value,
                 progress_bar=True,
                 remote_path=_sanitize_uploads_remote_path(argument.value),
@@ -109,7 +115,7 @@ class Client:
 
         json = {
             "teamspace_id": self._project.project_id,
-            "cluster_id": self._cluster.cluster_id,
+            "cluster_id": self._cloud_account.cluster_id,
             "input": {},
         }
         for argument in self._arguments:

lightning_sdk/studio.py

@@ -11,7 +11,7 @@ from lightning_sdk.owner import Owner
 from lightning_sdk.status import Status
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
-from lightning_sdk.utils.resolve import _resolve_teamspace, _setup_logger
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace, _setup_logger
 
 if TYPE_CHECKING:
     from lightning_sdk.plugin import Plugin
@@ -30,10 +30,9 @@ class Studio:
         teamspace: the name of the teamspace the studio is contained by
         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
-        cluster: the name of the cluster, the studio should be created on.
+        cloud_account: the name of the cloud account, the studio should be created on.
             Doesn't matter when the studio already exists.
         create_ok: whether the studio will be created if it does not yet exist. Defaults to True
-
     Note:
         Since a teamspace can either be owned by an org or by a user directly,
         only one of the arguments can be provided.
@@ -49,13 +48,14 @@ class Studio:
         teamspace: Optional[Union[str, Teamspace]] = None,
         org: Optional[Union[str, Organization]] = None,
         user: Optional[Union[str, User]] = None,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         create_ok: bool = True,
+        cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> None:
         self._studio_api = StudioApi()
 
         self._teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
-        self._cluster = cluster
+        self._cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
         self._setup_done = False
 
         self._plugins = {}
@@ -70,7 +70,9 @@ class Studio:
                 self._studio = self._studio_api.get_studio(name, self._teamspace.id)
             except ValueError as e:
                 if create_ok:
-                    self._studio = self._studio_api.create_studio(name, self._teamspace.id, cluster=self._cluster)
+                    self._studio = self._studio_api.create_studio(
+                        name, self._teamspace.id, cloud_account=self._cloud_account
+                    )
                 else:
                     raise ValueError(f"Studio {name} does not exist.") from e
 
@@ -144,6 +146,11 @@ class Studio:
     @property
     def cluster(self) -> str:
         """Returns the cluster the Studio is running on."""
+        warnings.warn("Studio.cluster is deprecated. Use Studio.cloud_account instead", DeprecationWarning)
+        return self.cloud_account
+
+    @property
+    def cloud_account(self) -> str:
         return self._studio.cluster_id
 
     def start(self, machine: Union[Machine, str] = Machine.CPU, interruptible: bool = False) -> None:
@@ -242,7 +249,7 @@ class Studio:
         self._studio_api.upload_file(
             studio_id=self._studio.id,
             teamspace_id=self._teamspace.id,
-            cluster_id=self._studio.cluster_id,
+            cloud_account=self._studio.cluster_id,
             file_path=file_path,
             remote_path=os.path.normpath(remote_path),
             progress_bar=progress_bar,
@@ -258,7 +265,7 @@ class Studio:
             target_path=file_path,
             studio_id=self._studio.id,
             teamspace_id=self._teamspace.id,
-            cluster_id=self._studio.cluster_id,
+            cloud_account=self._studio.cluster_id,
         )
 
     def download_folder(self, remote_path: str, target_path: Optional[str] = None) -> None:
@@ -271,7 +278,7 @@ class Studio:
             target_path=target_path,
             studio_id=self._studio.id,
             teamspace_id=self._teamspace.id,
-            cluster_id=self._studio.cluster_id,
+            cloud_account=self._studio.cluster_id,
         )
 
     @property