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

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."""

lightning_sdk/job/v2.py

@@ -22,6 +22,15 @@ class _JobV2(_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 = JobApiV2()
         super().__init__(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=_fetch_job)
 
@@ -33,12 +42,40 @@ class _JobV2(_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,
     ) -> "_JobV2":
+        """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.
+        """
         # Command is required if Studio is provided to know what to run
         # Image is mutually exclusive with Studio
         # Command is optional for Image
@@ -59,7 +96,7 @@ class _JobV2(_BaseJob):
         submitted = self._job_api.submit_job(
             name=self.name,
             command=command,
-            cluster_id=cluster,
+            cloud_account=cloud_account,
             teamspace_id=self._teamspace.id,
             studio_id=studio_id,
             image=image,
@@ -67,7 +104,7 @@ class _JobV2(_BaseJob):
             interruptible=interruptible,
             env=env,
             image_credentials=image_credentials,
-            cluster_auth=cluster_auth,
+            cloud_account_auth=cloud_account_auth,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
         )
@@ -76,9 +113,14 @@ class _JobV2(_BaseJob):
         return self
 
     def stop(self) -> None:
+        """Stop the job. If the job is already stopped, this is a no-op. This is blocking until the job is stopped."""
         self._job_api.stop_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id)
 
     def delete(self) -> None:
+        """Delete the job.
+
+        Caution: This also deletes all artifacts created by the job.
+        """
         self._job_api.delete_job(
             job_id=self._guaranteed_job.id,
             teamspace_id=self._teamspace.id,
@@ -91,28 +133,20 @@ class _JobV2(_BaseJob):
         self._update_internal_job()
         return self._job
 
-    @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
-
     @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 machine(self) -> "Machine":
+        """The machine type the job is running on."""
         # only fetch the job it it hasn't been fetched yet as machine cannot change over time
         return self._job_api._get_job_machine_from_spec(self._guaranteed_job.spec)
 
     @property
     def artifact_path(self) -> Optional[str]:
+        """The path to the artifacts of the job within the distributed teamspace filesystem."""
         if self._guaranteed_job.spec.image != "":
             if self._guaranteed_job.spec.artifacts_destination != "":
                 splits = self._guaranteed_job.spec.artifacts_destination.split(":")
@@ -123,12 +157,14 @@ class _JobV2(_BaseJob):
 
     @property
     def snapshot_path(self) -> Optional[str]:
+        """The path to the snapshot of the Studio used to create the job within the distributed teamspace filesystem."""
         if self._guaranteed_job.spec.image != "":
             return None
         return f"/teamspace/jobs/{self._guaranteed_job.name}/snapshot"
 
     @property
     def share_path(self) -> Optional[str]:
+        """The path to the share of the job within the distributed teamspace filesystem."""
         raise NotImplementedError("Not implemented yet")
 
     def _update_internal_job(self) -> None:

lightning_sdk/job/work.py

@@ -1,33 +1,62 @@
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional, Protocol
 
 from lightning_sdk.api.job_api import JobApiV1
 
 if TYPE_CHECKING:
-    from lightning_sdk.job.job import Job
+    from lightning_sdk.status import Status
     from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.machine import Machine
 
 
+class _WorkHolder(Protocol):
+    @property
+    def _id(self) -> str:
+        ...
+
+    @property
+    def name(self) -> str:
+        ...
+
+    def _name_filter(self, name: str) -> str:
+        ...
+
+
 class Work:
-    def __init__(self, work_id: str, job: "Job", teamspace: "Teamspace") -> None:
+    def __init__(self, work_id: str, job: _WorkHolder, teamspace: "Teamspace") -> None:
         self._id = work_id
         self._job = job
         self._teamspace = teamspace
         self._job_api = JobApiV1()
-        self._work = self._job_api.get_work(work_id=work_id, job_id=job.id, teamspace_id=teamspace.id)
+        self._work = None
+
+    @property
+    def _latest_work(self) -> Any:
+        self._work = self._job_api.get_work(work_id=self._id, job_id=self._job._id, teamspace_id=self._teamspace.id)
+        return self._work
+
+    @property
+    def _guaranteed_work(self) -> Any:
+        if self._work is None:
+            return self._latest_work
+
+        return self._work
 
     @property
     def id(self) -> str:
-        return self._work.id
+        return self._guaranteed_work.id
 
     @property
     def name(self) -> str:
-        return self._job._name_filter(self._work.name)
+        return self._job._name_filter(self._guaranteed_work.name)
 
     @property
     def machine(self) -> "Machine":
-        return self._job_api.get_machine_from_work(self._work)
+        return self._job_api.get_machine_from_work(self._guaranteed_work)
 
     @property
     def artifact_path(self) -> Optional[str]:
         return f"/teamspace/jobs/{self._job.name}/{self.name}"
+
+    @property
+    def status(self) -> "Status":
+        return self._job_api.get_status_from_work(self._latest_work)

lightning_sdk/lightning_cloud/openapi/__init__.py

@@ -125,6 +125,7 @@ from lightning_sdk.lightning_cloud.openapi.models.id_start_body import IdStartBo
 from lightning_sdk.lightning_cloud.openapi.models.id_storage_body import IdStorageBody
 from lightning_sdk.lightning_cloud.openapi.models.id_uploads_body import IdUploadsBody
 from lightning_sdk.lightning_cloud.openapi.models.id_uploads_body1 import IdUploadsBody1
+from lightning_sdk.lightning_cloud.openapi.models.id_visibility_body import IdVisibilityBody
 from lightning_sdk.lightning_cloud.openapi.models.jobs_id_body import JobsIdBody
 from lightning_sdk.lightning_cloud.openapi.models.jobs_id_body1 import JobsIdBody1
 from lightning_sdk.lightning_cloud.openapi.models.jobs_id_body2 import JobsIdBody2
@@ -445,6 +446,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_get_long_running_command_in
 from lightning_sdk.lightning_cloud.openapi.models.v1_get_model_file_upload_urls_response import V1GetModelFileUploadUrlsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_get_model_file_url_response import V1GetModelFileUrlResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_get_model_files_response import V1GetModelFilesResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_get_model_files_url_response import V1GetModelFilesUrlResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_get_project_artifact_response import V1GetProjectArtifactResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_get_project_balance_response import V1GetProjectBalanceResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_get_project_compute_usage_response import V1GetProjectComputeUsageResponse
@@ -558,6 +560,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_list_memberships_response i
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_metrics_streams_response import V1ListMetricsStreamsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_model_versions_response import V1ListModelVersionsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_models_response import V1ListModelsResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_list_multi_machine_job_events_response import V1ListMultiMachineJobEventsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_multi_machine_jobs_response import V1ListMultiMachineJobsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_new_features_for_user_response import V1ListNewFeaturesForUserResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_org_members_response import V1ListOrgMembersResponse
@@ -612,9 +615,14 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_metrics_stream import V1Met
 from lightning_sdk.lightning_cloud.openapi.models.v1_metrics_tags import V1MetricsTags
 from lightning_sdk.lightning_cloud.openapi.models.v1_metrics_tracker import V1MetricsTracker
 from lightning_sdk.lightning_cloud.openapi.models.v1_model import V1Model
+from lightning_sdk.lightning_cloud.openapi.models.v1_model_file import V1ModelFile
 from lightning_sdk.lightning_cloud.openapi.models.v1_model_version_archive import V1ModelVersionArchive
 from lightning_sdk.lightning_cloud.openapi.models.v1_mount_target import V1MountTarget
 from lightning_sdk.lightning_cloud.openapi.models.v1_multi_machine_job import V1MultiMachineJob
+from lightning_sdk.lightning_cloud.openapi.models.v1_multi_machine_job_event import V1MultiMachineJobEvent
+from lightning_sdk.lightning_cloud.openapi.models.v1_multi_machine_job_event_type import V1MultiMachineJobEventType
+from lightning_sdk.lightning_cloud.openapi.models.v1_multi_machine_job_fault_tolerance import V1MultiMachineJobFaultTolerance
+from lightning_sdk.lightning_cloud.openapi.models.v1_multi_machine_job_fault_tolerance_strategy import V1MultiMachineJobFaultToleranceStrategy
 from lightning_sdk.lightning_cloud.openapi.models.v1_multi_machine_job_state import V1MultiMachineJobState
 from lightning_sdk.lightning_cloud.openapi.models.v1_multi_machine_job_status import V1MultiMachineJobStatus
 from lightning_sdk.lightning_cloud.openapi.models.v1_named_get_logger_metrics import V1NamedGetLoggerMetrics
@@ -724,6 +732,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_system_metrics import V1Sys
 from lightning_sdk.lightning_cloud.openapi.models.v1_system_metrics_list import V1SystemMetricsList
 from lightning_sdk.lightning_cloud.openapi.models.v1_telemetry import V1Telemetry
 from lightning_sdk.lightning_cloud.openapi.models.v1_timestamp_code_telemetry import V1TimestampCodeTelemetry
+from lightning_sdk.lightning_cloud.openapi.models.v1_trainium_system_metrics import V1TrainiumSystemMetrics
 from lightning_sdk.lightning_cloud.openapi.models.v1_transaction import V1Transaction
 from lightning_sdk.lightning_cloud.openapi.models.v1_transfer_org_balance_response import V1TransferOrgBalanceResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_transfer_project_balance_response import V1TransferProjectBalanceResponse
@@ -740,6 +749,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_update_cluster_accelerators
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_cluster_availability_request import V1UpdateClusterAvailabilityRequest
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_index_response import V1UpdateIndexResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_lit_page_response import V1UpdateLitPageResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_update_metrics_stream_visibility_response import V1UpdateMetricsStreamVisibilityResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_model_visibility_response import V1UpdateModelVisibilityResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_project_cluster_accelerators_response import V1UpdateProjectClusterAcceleratorsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_shared_metrics_stream_response import V1UpdateSharedMetricsStreamResponse
@@ -769,6 +779,8 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_user_slurm_job_action_respo
 from lightning_sdk.lightning_cloud.openapi.models.v1_validate_assistant_status_response import V1ValidateAssistantStatusResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_validate_auto_join_domain_response import V1ValidateAutoJoinDomainResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_validate_data_connection_response import V1ValidateDataConnectionResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_validate_deployment_image_request import V1ValidateDeploymentImageRequest
+from lightning_sdk.lightning_cloud.openapi.models.v1_validate_deployment_image_response import V1ValidateDeploymentImageResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_validate_managed_endpoint_request import V1ValidateManagedEndpointRequest
 from lightning_sdk.lightning_cloud.openapi.models.v1_validate_managed_endpoint_response import V1ValidateManagedEndpointResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_validate_managed_model_response import V1ValidateManagedModelResponse