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

lightning_sdk/api/mmt_api.py

@@ -0,0 +1,147 @@
+import time
+from typing import TYPE_CHECKING, Dict, Optional
+
+from lightning_sdk.api.utils import (
+    _COMPUTE_NAME_TO_MACHINE,
+    _MACHINE_TO_COMPUTE_NAME,
+)
+from lightning_sdk.api.utils import (
+    _get_cloud_url as _cloud_url,
+)
+from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__
+from lightning_sdk.lightning_cloud.openapi import (
+    MultimachinejobsIdBody,
+    ProjectIdMultimachinejobsBody,
+    V1EnvVar,
+    V1JobSpec,
+    V1MultiMachineJob,
+    V1MultiMachineJobState,
+)
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+from lightning_sdk.machine import Machine
+
+if TYPE_CHECKING:
+    from lightning_sdk.status import Status
+
+
+class MMTApi:
+    mmt_state_unspecified = "MultiMachineJob_STATE_UNSPECIFIED"
+    mmt_state_running = "MultiMachineJob_STATE_RUNNING"
+    mmt_state_stopped = "MultiMachineJob_STATE_STOPPED"
+    mmt_state_deleted = "MultiMachineJob_STATE_DELETED"
+    mmt_state_failed = "MultiMachineJob_STATE_FAILED"
+    mmt_state_completed = "MultiMachineJob_STATE_COMPLETED"
+
+    def __init__(self) -> None:
+        self._cloud_url = _cloud_url()
+        self._client = LightningClient(max_tries=7)
+
+    def submit_job(
+        self,
+        name: str,
+        num_machines: int,
+        command: Optional[str],
+        cluster_id: Optional[str],
+        teamspace_id: str,
+        studio_id: Optional[str],
+        image: Optional[str],
+        machine: Machine,
+        interruptible: bool,
+        env: Optional[Dict[str, str]],
+        image_credentials: Optional[str],
+        cluster_auth: bool,
+        artifacts_local: Optional[str],
+        artifacts_remote: Optional[str],
+    ) -> V1MultiMachineJob:
+        env_vars = []
+        if env is not None:
+            for k, v in env.items():
+                env_vars.append(V1EnvVar(name=k, value=v))
+
+        instance_name = _MACHINE_TO_COMPUTE_NAME[machine]
+
+        run_id = __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__[studio_id] if studio_id is not None else ""
+
+        spec = V1JobSpec(
+            cloudspace_id=studio_id or "",
+            cluster_id=cluster_id or "",
+            command=command or "",
+            env=env_vars,
+            image=image or "",
+            instance_name=instance_name,
+            run_id=run_id,
+            spot=interruptible,
+            image_cluster_credentials=cluster_auth,
+            image_secret_ref=image_credentials or "",
+            artifacts_source=artifacts_local or "",
+            artifacts_destination=artifacts_remote or "",
+        )
+        body = ProjectIdMultimachinejobsBody(name=name, spec=spec, cluster_id=cluster_id or "", machines=num_machines)
+
+        job: V1MultiMachineJob = self._client.jobs_service_create_multi_machine_job(project_id=teamspace_id, body=body)
+        return job
+
+    def get_job_by_name(self, name: str, teamspace_id: str) -> V1MultiMachineJob:
+        job: V1MultiMachineJob = self._client.jobs_service_get_multi_machine_job_by_name(
+            project_id=teamspace_id, name=name
+        )
+        return job
+
+    def get_job(self, job_id: str, teamspace_id: str) -> V1MultiMachineJob:
+        job: V1MultiMachineJob = self._client.jobs_service_get_multi_machine_job(project_id=teamspace_id, id=job_id)
+        return job
+
+    def stop_job(self, job_id: str, teamspace_id: str) -> None:
+        from lightning_sdk.status import Status
+
+        current_job = self.get_job(job_id=job_id, teamspace_id=teamspace_id)
+
+        current_state = self._job_state_to_external(current_job.desired_state)
+
+        if current_state in (
+            Status.Stopped,
+            Status.Completed,
+            Status.Failed,
+        ):
+            return
+
+        if current_state != Status.Stopped:
+            update_body = MultimachinejobsIdBody(desired_state=self.mmt_state_stopped)
+            self._client.jobs_service_update_multi_machine_job(body=update_body, project_id=teamspace_id, id=job_id)
+
+        while True:
+            current_job = self.get_job(job_id=job_id, teamspace_id=teamspace_id)
+            if self._job_state_to_external(current_job.desired_state) in (
+                Status.Stopped,
+                Status.Completed,
+                Status.Stopped,
+                Status.Failed,
+            ):
+                break
+            time.sleep(1)
+
+    def delete_job(self, job_id: str, teamspace_id: str) -> None:
+        self._client.jobs_service_delete_multi_machine_job(project_id=teamspace_id, id=job_id)
+
+    def _job_state_to_external(self, state: V1MultiMachineJobState) -> "Status":
+        from lightning_sdk.status import Status
+
+        if str(state) == self.mmt_state_unspecified:
+            return Status.Pending
+        if str(state) == self.mmt_state_running:
+            return Status.Running
+        if str(state) == self.mmt_state_stopped:
+            return Status.Stopped
+        if str(state) == self.mmt_state_completed:
+            return Status.Completed
+        if str(state) == self.mmt_state_failed:
+            return Status.Failed
+        return Status.Pending
+
+    def _get_job_machine_from_spec(self, spec: V1JobSpec) -> "Machine":
+        instance_name = spec.instance_name
+        instance_type = spec.instance_type
+
+        return _COMPUTE_NAME_TO_MACHINE.get(
+            instance_type, _COMPUTE_NAME_TO_MACHINE.get(instance_name, instance_type or instance_name)
+        )

lightning_sdk/api/teamspace_api.py

@@ -1,5 +1,4 @@
 import os
-from dataclasses import dataclass
 from pathlib import Path
 from typing import Dict, List, Optional
 
@@ -273,11 +272,3 @@ class TeamspaceApi:
             download_dir=download_dir,
             progress_bar=progress_bar,
         )
-
-
-@dataclass
-class UploadedModelInfo:
-    name: str
-    version: str
-    teamspace: str
-    cluster: str

lightning_sdk/cli/mmt.py

@@ -0,0 +1,137 @@
+from typing import Dict, Optional
+
+from fire import Fire
+
+from lightning_sdk._mmt import MMT
+from lightning_sdk.api.studio_api import _cloud_url
+from lightning_sdk.lightning_cloud.login import Auth
+from lightning_sdk.machine import Machine
+from lightning_sdk.teamspace import Teamspace
+
+_MACHINE_VALUES = tuple([machine.value for machine in Machine])
+
+
+class MMTCLI:
+    """Command line interface (CLI) to interact with/manage Lightning AI MMT."""
+
+    def __init__(self) -> None:
+        # Need to set the docstring here for f-strings to work.
+        # Sadly this is the only way to really show options as f-strings are not allowed as docstrings directly
+        # and fire does not show values for literals, just that it is a literal.
+        docstr = f"""Run async workloads on multiple machines using a docker image.
+
+        Args:
+            name: The name of the job. Needs to be unique within the teamspace.
+            num_machines: The number of Machines to run on. Defaults to 2 Machines
+            machine: The machine type to run the job on. One of {", ".join(_MACHINE_VALUES)}. Defaults to CPU
+            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.
+            cluster: The cluster to run the job on. Defaults to the studio cluster if running with studio compute env.
+                If not provided will fall back to the teamspaces default cluster.
+            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.
+            cluster_auth: Whether to authenticate with the cluster to pull the image.
+                Required if the registry is part of a cluster 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.
+        """
+        # TODO: the docstrings from artifacts_local and artifacts_remote don't show up completely,
+        # might need to switch to explicit cli definition
+        self.run.__func__.__doc__ = docstr
+
+    def login(self) -> None:
+        """Login to Lightning AI Studios."""
+        auth = Auth()
+        auth.clear()
+
+        try:
+            auth.authenticate()
+        except ConnectionError:
+            raise RuntimeError(f"Unable to connect to {_cloud_url()}. Please check your internet connection.") from None
+
+    def logout(self) -> None:
+        """Logout from Lightning AI Studios."""
+        auth = Auth()
+        auth.clear()
+
+    # TODO: sadly, fire displays both Optional[type] and Union[type, None] as Optional[Optional]
+    # see https://github.com/google/python-fire/pull/513
+    # might need to move to different cli library
+    def run(
+        self,
+        name: Optional[str] = None,
+        num_machines: int = 2,
+        machine: Optional[str] = None,
+        command: Optional[str] = None,
+        studio: Optional[str] = None,
+        image: Optional[str] = None,
+        teamspace: Optional[str] = None,
+        org: Optional[str] = None,
+        user: Optional[str] = None,
+        cluster: Optional[str] = None,
+        env: Optional[Dict[str, str]] = None,
+        interruptible: bool = False,
+        image_credentials: Optional[str] = None,
+        cluster_auth: bool = False,
+        artifacts_local: Optional[str] = None,
+        artifacts_remote: Optional[str] = None,
+    ) -> None:
+        if name is None:
+            from datetime import datetime
+
+            timestr = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+            name = f"mmt-{timestr}"
+
+        if machine is None:
+            # TODO: infer from studio
+            machine = "CPU"
+        machine_enum = Machine(machine.upper())
+
+        teamspace = Teamspace(name=teamspace, org=org, user=user)
+        if cluster is None:
+            cluster = teamspace.default_cluster
+
+        if image is None:
+            raise RuntimeError("Currently only docker images are specified")
+        MMT.run(
+            name=name,
+            num_machines=num_machines,
+            machine=machine_enum,
+            command=command,
+            studio=studio,
+            image=image,
+            teamspace=teamspace,
+            org=org,
+            user=user,
+            cluster=cluster,
+            env=env,
+            interruptible=interruptible,
+            image_credentials=image_credentials,
+            cluster_auth=cluster_auth,
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+        )
+
+
+def main_cli() -> None:
+    """CLI entrypoint."""
+    Fire(MMTCLI(), name="_mmt")
+
+
+if __name__ == "__main__":
+    main_cli()

lightning_sdk/job/base.py

@@ -114,7 +114,7 @@ class _BaseJob(ABC):
                 )
 
         inst = cls(name=name, teamspace=teamspace, org=org, user=user, _fetch_job=False)
-        inst._submit(
+        return inst._submit(
             machine=machine,
             cluster=cluster,
             command=command,
@@ -127,7 +127,6 @@ class _BaseJob(ABC):
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
         )
-        return inst
 
     @abstractmethod
     def _submit(
@@ -143,7 +142,7 @@ class _BaseJob(ABC):
         cluster_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
-    ) -> None:
+    ) -> "_BaseJob":
         """Submits a job and updates the internal _job attribute as well as the _name attribute."""
 
     @abstractmethod

lightning_sdk/job/v1.py

@@ -71,7 +71,7 @@ class _JobV1(_BaseJob):
         cluster_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
-    ) -> None:
+    ) -> "_JobV1":
         if studio is None:
             raise ValueError("Studio is required for submitting jobs")
 
@@ -100,6 +100,7 @@ class _JobV1(_BaseJob):
         )
         self._name = _submitted.name
         self._job = _submitted
+        return self
 
     def _update_internal_job(self) -> None:
         try:

lightning_sdk/job/v2.py

@@ -38,7 +38,7 @@ class _JobV2(_BaseJob):
         cluster_auth: bool = False,
         artifacts_local: Optional[str] = None,
         artifacts_remote: Optional[str] = None,
-    ) -> None:
+    ) -> "_JobV2":
         # Command is required if Studio is provided to know what to run
         # Image is mutually exclusive with Studio
         # Command is optional for Image
@@ -73,19 +73,16 @@ class _JobV2(_BaseJob):
         )
         self._job = submitted
         self._name = submitted.name
+        return self
 
     def stop(self) -> None:
-        if self._job is None:
-            self._update_internal_job()
-
-        self._job_api.stop_job(job_id=self._job.id, teamspace_id=self._teamspace.id)
+        self._job_api.stop_job(job_id=self._guaranteed_job.id, teamspace_id=self._teamspace.id)
 
     def delete(self) -> None:
-        if self._job is None:
-            self._update_internal_job()
-
         self._job_api.delete_job(
-            job_id=self._job.id, teamspace_id=self._teamspace.id, cloudspace_id=self._job.spec.cloudspace_id
+            job_id=self._guaranteed_job.id,
+            teamspace_id=self._teamspace.id,
+            cloudspace_id=self._guaranteed_job.spec.cloudspace_id,
         )
 
     @property

lightning_sdk/lightning_cloud/openapi/__init__.py

@@ -137,6 +137,7 @@ from lightning_sdk.lightning_cloud.openapi.models.metricsstream_create_body impo
 from lightning_sdk.lightning_cloud.openapi.models.metricsstream_delete_body import MetricsstreamDeleteBody
 from lightning_sdk.lightning_cloud.openapi.models.metricsstream_id_body import MetricsstreamIdBody
 from lightning_sdk.lightning_cloud.openapi.models.model_id_versions_body import ModelIdVersionsBody
+from lightning_sdk.lightning_cloud.openapi.models.model_id_visibility_body import ModelIdVisibilityBody
 from lightning_sdk.lightning_cloud.openapi.models.models_model_id_body import ModelsModelIdBody
 from lightning_sdk.lightning_cloud.openapi.models.multimachinejobs_id_body import MultimachinejobsIdBody
 from lightning_sdk.lightning_cloud.openapi.models.multipartuploads_upload_id_body import MultipartuploadsUploadIdBody
@@ -146,7 +147,6 @@ from lightning_sdk.lightning_cloud.openapi.models.orgs_id_body import OrgsIdBody
 from lightning_sdk.lightning_cloud.openapi.models.profiler_captures_body import ProfilerCapturesBody
 from lightning_sdk.lightning_cloud.openapi.models.profiler_enabled_body import ProfilerEnabledBody
 from lightning_sdk.lightning_cloud.openapi.models.project_id_agentmanagedendpoints_body import ProjectIdAgentmanagedendpointsBody
-from lightning_sdk.lightning_cloud.openapi.models.project_id_agentmanagedmodels_body import ProjectIdAgentmanagedmodelsBody
 from lightning_sdk.lightning_cloud.openapi.models.project_id_agents_body import ProjectIdAgentsBody
 from lightning_sdk.lightning_cloud.openapi.models.project_id_cloudspaces_body import ProjectIdCloudspacesBody
 from lightning_sdk.lightning_cloud.openapi.models.project_id_clusters_body import ProjectIdClustersBody
@@ -222,6 +222,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_billing_feature import V1Bi
 from lightning_sdk.lightning_cloud.openapi.models.v1_billing_period import V1BillingPeriod
 from lightning_sdk.lightning_cloud.openapi.models.v1_billing_subscription import V1BillingSubscription
 from lightning_sdk.lightning_cloud.openapi.models.v1_billing_tier import V1BillingTier
+from lightning_sdk.lightning_cloud.openapi.models.v1_body import V1Body
 from lightning_sdk.lightning_cloud.openapi.models.v1_build_spec import V1BuildSpec
 from lightning_sdk.lightning_cloud.openapi.models.v1_cpu_system_metrics import V1CPUSystemMetrics
 from lightning_sdk.lightning_cloud.openapi.models.v1_cancel_cloud_space_instance_switch_response import V1CancelCloudSpaceInstanceSwitchResponse
@@ -339,7 +340,6 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lightningwork_respon
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_lit_page_response import V1DeleteLitPageResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_logger_artifact_response import V1DeleteLoggerArtifactResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_managed_endpoint_response import V1DeleteManagedEndpointResponse
-from lightning_sdk.lightning_cloud.openapi.models.v1_delete_managed_model_response import V1DeleteManagedModelResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_metrics_stream_response import V1DeleteMetricsStreamResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_model_response import V1DeleteModelResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_delete_model_version_response import V1DeleteModelVersionResponse
@@ -365,6 +365,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_delete_user_slurm_job_respo
 from lightning_sdk.lightning_cloud.openapi.models.v1_dependency_cache_state import V1DependencyCacheState
 from lightning_sdk.lightning_cloud.openapi.models.v1_dependency_file_info import V1DependencyFileInfo
 from lightning_sdk.lightning_cloud.openapi.models.v1_deployment import V1Deployment
+from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_api import V1DeploymentAPI
 from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_event import V1DeploymentEvent
 from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_event_type import V1DeploymentEventType
 from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_metrics import V1DeploymentMetrics
@@ -552,7 +553,6 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_list_lightningwork_response
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_lit_pages_response import V1ListLitPagesResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_logger_artifact_response import V1ListLoggerArtifactResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_managed_endpoints_response import V1ListManagedEndpointsResponse
-from lightning_sdk.lightning_cloud.openapi.models.v1_list_managed_models_response import V1ListManagedModelsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_memberships_response import V1ListMembershipsResponse
 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
@@ -598,6 +598,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_magic_link_login_request im
 from lightning_sdk.lightning_cloud.openapi.models.v1_magic_link_login_response import V1MagicLinkLoginResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_managed_endpoint import V1ManagedEndpoint
 from lightning_sdk.lightning_cloud.openapi.models.v1_managed_model import V1ManagedModel
+from lightning_sdk.lightning_cloud.openapi.models.v1_managed_model_abilities import V1ManagedModelAbilities
 from lightning_sdk.lightning_cloud.openapi.models.v1_membership import V1Membership
 from lightning_sdk.lightning_cloud.openapi.models.v1_message import V1Message
 from lightning_sdk.lightning_cloud.openapi.models.v1_message_author import V1MessageAuthor
@@ -653,6 +654,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_publish_cloud_space_respons
 from lightning_sdk.lightning_cloud.openapi.models.v1_published_cloud_space_response import V1PublishedCloudSpaceResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_purchase_capacity_block_response import V1PurchaseCapacityBlockResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_python_dependency_info import V1PythonDependencyInfo
+from lightning_sdk.lightning_cloud.openapi.models.v1_query_param import V1QueryParam
 from lightning_sdk.lightning_cloud.openapi.models.v1_query_result import V1QueryResult
 from lightning_sdk.lightning_cloud.openapi.models.v1_query_result_row import V1QueryResultRow
 from lightning_sdk.lightning_cloud.openapi.models.v1_quest import V1Quest
@@ -671,6 +673,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_request_cluster_access_requ
 from lightning_sdk.lightning_cloud.openapi.models.v1_request_cluster_access_response import V1RequestClusterAccessResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_request_verification_code_response import V1RequestVerificationCodeResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_resource_tag import V1ResourceTag
+from lightning_sdk.lightning_cloud.openapi.models.v1_resource_visibility import V1ResourceVisibility
 from lightning_sdk.lightning_cloud.openapi.models.v1_resources import V1Resources
 from lightning_sdk.lightning_cloud.openapi.models.v1_response_choice import V1ResponseChoice
 from lightning_sdk.lightning_cloud.openapi.models.v1_response_choice_delta import V1ResponseChoiceDelta
@@ -736,6 +739,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_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
 from lightning_sdk.lightning_cloud.openapi.models.v1_update_snowflake_query_response import V1UpdateSnowflakeQueryResponse