lightning-sdk 0.1.40__py3-none-any.whl → 0.1.42__py3-none-any.whl

lightning_sdk/job/base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
-from lightning_sdk.utils.resolve import _resolve_teamspace
+from lightning_sdk.utils.resolve import _resolve_deprecated_cluster, _resolve_teamspace
 
 if TYPE_CHECKING:
     from lightning_sdk.machine import Machine
@@ -13,6 +13,8 @@ if TYPE_CHECKING:
 
 
 class _BaseJob(ABC):
+    """Base interface to all job types."""
+
     def __init__(
         self,
         name: str,
@@ -22,6 +24,15 @@ class _BaseJob(ABC):
         *,
         _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.
+        """
         _teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
         if _teamspace is None:
             raise ValueError(
@@ -47,22 +58,59 @@ class _BaseJob(ABC):
         teamspace: Union[str, "Teamspace", None] = None,
         org: Union[str, "Organization", None] = None,
         user: Union[str, "User", None] = None,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
-        cluster_auth: bool = False,
+        cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> "_BaseJob":
+        """Run async workloads using a docker image or a compute environment from your studio.
+
+        Args:
+            name: The name of the job. Needs to be unique within the teamspace.
+            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.
+            teamspace: The teamspace the job should be associated with. Defaults to the current teamspace.
+            org: The organization owning the teamspace (if any). Defaults to the current organization.
+            user: The user owning the teamspace (if any). Defaults to the current user.
+            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.
+            env: Environment variables to set inside the job.
+            interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
+            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.
+        """
         from lightning_sdk.studio import Studio
 
+        cloud_account = _resolve_deprecated_cluster(cloud_account, cluster)
+
         if not name:
             raise ValueError("A job needs to have a name!")
 
         if image is None:
             if not isinstance(studio, Studio):
-                studio = Studio(name=studio, teamspace=teamspace, org=org, user=user, cluster=cluster, create_ok=False)
+                studio = Studio(
+                    name=studio, teamspace=teamspace, org=org, user=user, cloud_account=cloud_account, create_ok=False
+                )
 
             # studio is a Studio instance at this point
             if teamspace is None:
@@ -76,20 +124,20 @@ class _BaseJob(ABC):
                         "Can only run jobs with Studio envs in the teamspace of that Studio."
                     )
 
-            if cluster is None:
-                cluster = studio.cluster
+            if cloud_account is None:
+                cloud_account = studio.cloud_account
 
-            if cluster != studio.cluster:
+            if cloud_account != studio.cloud_account:
                 raise ValueError(
-                    "Studio cluster does not match provided cluster. "
-                    "Can only run jobs with Studio envs in the same cluster."
+                    "Studio cloud account does not match provided cloud account. "
+                    "Can only run jobs with Studio envs in the same cloud account."
                 )
 
             if image_credentials is not None:
                 raise ValueError("image_credentials is only supported when using a custom image")
 
-            if cluster_auth:
-                raise ValueError("cluster_auth is only supported when using a custom image")
+            if cloud_account_auth:
+                raise ValueError("cloud_account_auth is only supported when using a custom image")
 
             if artifacts_local is not None or artifacts_remote is not None:
                 raise ValueError(
@@ -116,14 +164,14 @@ class _BaseJob(ABC):
         inst = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False)
         return inst._submit(
             machine=machine,
-            cluster=cluster,
+            cloud_account=cloud_account,
             command=command,
             studio=studio,
             image=image,
             env=env,
             interruptible=interruptible,
             image_credentials=image_credentials,
-            cluster_auth=cluster_auth,
+            cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
         )
@@ -137,46 +185,79 @@ class _BaseJob(ABC):
         image: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         image_credentials: Optional[str] = None,
-        cluster_auth: bool = False,
+        cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
     ) -> "_BaseJob":
-        """Submits a job and updates the internal _job attribute as well as the _name attribute."""
+        """Submit a new job to the Lightning AI platform.
+
+        Args:
+            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.
+        """
 
     @abstractmethod
     def stop(self) -> None:
-        pass
+        """Stops the job.
+
+        This is blocking until the job is stopped.
+        """
 
     @abstractmethod
     def delete(self) -> None:
-        pass
+        """Deletes the job.
+
+        Caution: This also deletes all artifacts and snapshots associated with the job.
+        """
 
     @property
     @abstractmethod
     def status(self) -> "Status":
-        pass
+        """The current status of the job."""
 
     @property
     @abstractmethod
     def machine(self) -> "Machine":
-        pass
+        """The machine type the job is running on."""
 
     @property
     @abstractmethod
     def artifact_path(self) -> Optional[str]:
-        pass
+        """Path to the artifacts created by the job within the distributed teamspace filesystem."""
 
     @property
     @abstractmethod
     def snapshot_path(self) -> Optional[str]:
-        pass
+        """Path to the studio snapshot used to create the job within the distributed teamspace filesystem."""
 
     @property
     @abstractmethod
     def share_path(self) -> Optional[str]:
-        pass
+        """Path to the jobs share path."""
 
     @abstractmethod
     def _update_internal_job(self) -> None:
@@ -184,8 +265,21 @@ class _BaseJob(ABC):
 
     @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 _guaranteed_job(self) -> Any:
+        """Guarantees that the job was fetched at some point before returning it.
+
+        Doesn't guarantee to have the lastest version of the job. Use _latest_job for that.
+        """
+        if getattr(self, "_job", None) is None:
+            self._update_internal_job()
+
+        return self._job

lightning_sdk/job/job.py

@@ -55,14 +55,47 @@ class Job(_BaseJob):
         teamspace: Union[str, "Teamspace", None] = None,
         org: Union[str, "Organization", None] = None,
         user: Union[str, "User", None] = None,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
         image_credentials: Optional[str] = None,
-        cluster_auth: bool = False,
+        cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
+        cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> "Job":
+        """Run async workloads using a docker image or a compute environment from your studio.
+
+        Args:
+        name: The name of the job. Needs to be unique within the teamspace.
+        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.
+        teamspace: The teamspace the job should be associated with. Defaults to the current teamspace.
+        org: The organization owning the teamspace (if any). Defaults to the current organization.
+        user: The user owning the teamspace (if any). Defaults to the current user.
+        cloud_account: The cloud acocunt 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.
+        env: Environment variables to set inside the job.
+        interruptible: Whether the job should run on interruptible instances. They are cheaper but can be preempted.
+        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.
+        """
         ret_val = super().run(
             name=name,
             machine=machine,
@@ -72,13 +105,14 @@ class Job(_BaseJob):
             teamspace=teamspace,
             org=org,
             user=user,
-            cluster=cluster,
+            cloud_account=cloud_account,
             env=env,
             interruptible=interruptible,
             image_credentials=image_credentials,
-            cluster_auth=cluster_auth,
+            cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            cluster=cluster,
         )
         # required for typing with "Job"
         assert isinstance(ret_val, cls)
@@ -92,51 +126,92 @@ class Job(_BaseJob):
         image: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         image_credentials: Optional[str] = None,
-        cluster_auth: bool = False,
+        cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
-    ) -> None:
+    ) -> "Job":
+        """Submit a new job to the Lightning AI platform.
+
+        Args:
+            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.
+        """
         self._job = self._internal_job._submit(
             machine=machine,
-            cluster=cluster,
+            cloud_account=cloud_account,
             command=command,
             studio=studio,
             image=image,
             env=env,
             interruptible=interruptible,
             image_credentials=image_credentials,
-            cluster_auth=cluster_auth,
+            cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
         )
         return self
 
     def stop(self) -> None:
+        """Stops the job.
+
+        This is blocking until the job is stopped.
+        """
         return self._internal_job.stop()
 
     def delete(self) -> None:
+        """Deletes the job.
+
+        Caution: This also deletes all artifacts and snapshots associated with the job.
+        """
         return self._internal_job.delete()
 
     @property
     def status(self) -> "Status":
+        """The current status of the job."""
         return self._internal_job.status
 
     @property
     def machine(self) -> "Machine":
+        """The machine type the job is running on."""
         return self._internal_job.machine
 
     @property
     def artifact_path(self) -> Optional[str]:
+        """Path to the artifacts created by the job within the distributed teamspace filesystem."""
         return self._internal_job.artifact_path
 
     @property
     def snapshot_path(self) -> Optional[str]:
+        """Path to the studio snapshot used to create the job within the distributed teamspace filesystem."""
         return self._internal_job.snapshot_path
 
     @property
     def share_path(self) -> Optional[str]:
+        """Path to the jobs share path."""
         return self._internal_job.share_path
 
     def _update_internal_job(self) -> None:
@@ -144,14 +219,17 @@ class Job(_BaseJob):
 
     @property
     def name(self) -> str:
+        """The job's name."""
         return self._internal_job.name
 
     @property
     def teamspace(self) -> "Teamspace":
+        """The teamspace the job is part of."""
         return self._internal_job._teamspace
 
     @property
     def cluster(self) -> Optional[str]:
+        """The cluster the job is running on."""
         return self._internal_job.cluster
 
     def __getattr__(self, key: str) -> Any:

lightning_sdk/job/v1.py

@@ -17,6 +17,8 @@ from lightning_sdk.job.work import Work
 
 
 class _JobV1(_BaseJob):
+    """Implementation to run async workloads from your Studio."""
+
     def __init__(
         self,
         name: str,
@@ -26,6 +28,15 @@ class _JobV1(_BaseJob):
         *,
         _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 = JobApiV1()
         super().__init__(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=_fetch_job)
 
@@ -39,9 +50,26 @@ class _JobV1(_BaseJob):
         teamspace: Union[str, "Teamspace", None] = None,
         org: Union[str, "Organization", None] = None,
         user: Union[str, "User", None] = None,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         interruptible: bool = False,
+        cluster: Optional[str] = None,  # deprecated in favor of cloud_account
     ) -> "_BaseJob":
+        """Start a new async workload from your studio.
+
+        Args:
+            name: the name of the job
+            machine: the machine to run the workload on
+            command: the command to execute
+            studio: the studio the job belongs to
+            teamspace: the teamspace the job is part of
+            org: the organization owning the teamspace (if applicable)
+            user: the user owning the teamspace (if applicable)
+            cloud_account: the cloud account to run the workload on
+            interruptible: whether the workload can be interrupted
+
+        Returns:
+            the created job
+        """
         return super().run(
             name=name,
             machine=machine,
@@ -51,11 +79,12 @@ class _JobV1(_BaseJob):
             teamspace=teamspace,
             org=org,
             user=user,
-            cluster=cluster,
+            cloud_account=cloud_account,
             env=None,
             interruptible=interruptible,
             image_credentials=None,
-            cluster_auth=False,
+            cloud_account_auth=False,
+            cluster=cluster,
         )
 
     def _submit(
@@ -66,16 +95,34 @@ class _JobV1(_BaseJob):
         image: Optional[str] = None,
         env: Optional[Dict[str, str]] = None,
         interruptible: bool = False,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
         image_credentials: Optional[str] = None,
-        cluster_auth: bool = False,
+        cloud_account_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
     ) -> "_JobV1":
+        """Submit a job to run on a machine.
+
+        Args:
+            machine: The machine to run the job on.
+            command: The command to execute.
+            studio: The studio the job belongs to.
+            image: The image to use for the job (not supported).
+            env: The environment variables for the job (not supported).
+            interruptible: Whether the job can be interrupted.
+            cloud_account: The cloud account to run the job on.
+            image_credentials: The image credentials for the job (not supported).
+            cloud_account_auth: Whether to use cloud account authentication for the job (not supported).
+            artifacts_local: The local path for persisting artifacts (not supported).
+            artifacts_remote: The remote path for persisting artifacts (not supported).
+
+        Returns:
+            The submitted job.
+
+        """
         if studio is None:
             raise ValueError("Studio is required for submitting jobs")
-
-        if image is not None or image_credentials is not None or cluster_auth:
+        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:
@@ -83,18 +130,15 @@ class _JobV1(_BaseJob):
 
         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 jobs")
-
         # TODO: add support for empty names (will give an empty string)
-
         _submitted = self._job_api.submit_job(
             name=self._name,
             command=command,
             studio_id=studio._studio.id,
             teamspace_id=self._teamspace.id,
-            cluster_id=cluster,
+            cloud_account=cloud_account or "",
             machine=machine,
             interruptible=interruptible,
         )
@@ -110,6 +154,7 @@ class _JobV1(_BaseJob):
 
     @property
     def status(self) -> "Status":
+        """Returns the 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)
@@ -119,19 +164,22 @@ class _JobV1(_BaseJob):
             ) from None
 
     def stop(self) -> None:
+        """Stops the job. is blocking until the ob is stopped."""
         if self.status in (Status.Stopped, Status.Failed):
             return None
 
         return self._job_api.stop_job(self._job.id, self.teamspace.id)
 
     def delete(self) -> None:
-        self._job_api.delete_job(self._job.id, self.teamspace.id)
+        """Deletes the job.
 
-    def _name_filter(self, orig_name: str) -> str:
-        return orig_name.replace("root.", "")
+        Caution: this also deletes all artifacts created by the job.
+        """
+        self._job_api.delete_job(self._job.id, self.teamspace.id)
 
     @cached_property
     def work(self) -> Work:
+        """Get the work associated with the job."""
         _work = self._job_api.list_works(self._job.id, self.teamspace.id)
         if len(_work) == 0:
             raise ValueError("No works found for job")
@@ -139,28 +187,37 @@ class _JobV1(_BaseJob):
 
     @property
     def machine(self) -> "Machine":
+        """Get the machine the job is running on."""
         return self.work.machine
 
-    @property
-    def id(self) -> str:
-        return self._job.id
-
     @property
     def name(self) -> str:
+        """The name of the job."""
         return self._job.name
 
     @property
     def artifact_path(self) -> Optional[str]:
+        """The path to the artifacts of the job in the distributed teamspace filesystem."""
         return self.work.artifact_path
 
     @property
     def snapshot_path(self) -> Optional[str]:
+        """The path to the snapshot of the job in the distributed teamspace filesystem."""
         return f"/teamspace/jobs/{self.name}/snapshot"
 
     @property
     def share_path(self) -> Optional[str]:
+        """The path to the share of the job in the distributed teamspace filesystem."""
         return f"/teamspace/jobs/{self.name}/share"
 
+    # 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.", "")
+
 
 def _internal_status_to_external_status(internal_status: str) -> "Status":
     """Converts internal status strings from HTTP requests to external enums."""