lightning-sdk 0.1.49__py3-none-any.whl → 2025.11.5__py3-none-any.whl

lightning_sdk/api/mmt_api.py

@@ -1,16 +1,14 @@
 import json
 import time
-from typing import TYPE_CHECKING, Dict, List, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
-from lightning_sdk.api.job_api import JobApiV1
+from lightning_sdk.api.job_api import JobApiV1, V1ClusterAccelerator
 from lightning_sdk.api.utils import (
-    _COMPUTE_NAME_TO_MACHINE,
-    _MACHINE_TO_COMPUTE_NAME,
     _create_app,
+    _machine_to_compute_name,
+    resolve_path_mappings,
 )
-from lightning_sdk.api.utils import (
-    _get_cloud_url as _cloud_url,
-)
+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 (
     Externalv1LightningappInstance,
@@ -43,13 +41,13 @@ class MMTApiV1(JobApiV1):
         cloud_account: Optional[str],
         teamspace_id: str,
         studio_id: str,
-        machine: Machine,
+        machine: Union[Machine, str],
         interruptible: bool,
         strategy: str,
     ) -> Externalv1LightningappInstance:
         """Creates a multi-machine job with given commands."""
         distributed_args = {
-            "cloud_compute": _MACHINE_TO_COMPUTE_NAME[machine],
+            "cloud_compute": _machine_to_compute_name(machine),
             "num_instances": num_machines,
             "strategy": strategy,
         }
@@ -80,23 +78,81 @@ class MMTApiV2:
         teamspace_id: str,
         studio_id: Optional[str],
         image: Optional[str],
-        machine: Machine,
+        machine: Union[Machine, str],
         interruptible: bool,
         env: Optional[Dict[str, str]],
         image_credentials: Optional[str],
         cloud_account_auth: bool,
-        artifacts_local: Optional[str],
-        artifacts_remote: Optional[str],
         entrypoint: str,
+        path_mappings: Optional[Dict[str, str]],
+        artifacts_local: Optional[str],  # deprecated in favor of path_mappings
+        artifacts_remote: Optional[str],  # deprecated in favor of path_mappings
+        max_runtime: Optional[int],
+        reuse_snapshot: bool,
     ) -> V1MultiMachineJob:
+        body = self._create_mmt_body(
+            name=name,
+            num_machines=num_machines,
+            command=command,
+            cloud_account=cloud_account,
+            studio_id=studio_id,
+            image=image,
+            machine=machine,
+            interruptible=interruptible,
+            env=env,
+            image_credentials=image_credentials,
+            cloud_account_auth=cloud_account_auth,
+            entrypoint=entrypoint,
+            path_mappings=path_mappings,
+            artifacts_local=artifacts_local,  # deprecated in favor of path_mappings
+            artifacts_remote=artifacts_remote,  # deprecated in favor of path_mappings
+            max_runtime=max_runtime,
+            reuse_snapshot=reuse_snapshot,
+        )
+
+        job: V1MultiMachineJob = self._client.jobs_service_create_multi_machine_job(project_id=teamspace_id, body=body)
+        return job
+
+    @staticmethod
+    def _create_mmt_body(
+        name: str,
+        num_machines: int,
+        command: Optional[str],
+        cloud_account: Optional[str],
+        studio_id: Optional[str],
+        image: Optional[str],
+        machine: Union[Machine, str],
+        interruptible: bool,
+        env: Optional[Dict[str, str]],
+        image_credentials: Optional[str],
+        cloud_account_auth: bool,
+        entrypoint: str,
+        path_mappings: Optional[Dict[str, str]],
+        artifacts_local: Optional[str],  # deprecated in favor of path_mappings
+        artifacts_remote: Optional[str],  # deprecated in favor of path_mappings
+        reuse_snapshot: bool,
+        max_runtime: Optional[int] = None,
+        machine_image_version: Optional[str] = None,
+    ) -> ProjectIdMultimachinejobsBody:
         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]
+        instance_name = _machine_to_compute_name(machine)
+
+        run_id = __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__[studio_id] if (studio_id is not None and reuse_snapshot) else ""
+
+        path_mappings_list = resolve_path_mappings(
+            mappings=path_mappings or {},
+            artifacts_local=artifacts_local,
+            artifacts_remote=artifacts_remote,
+        )
 
-        run_id = __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__[studio_id] if studio_id is not None else ""
+        # need to go via kwargs for typing compatibility since autogenerated apis accept None but aren't typed with None
+        optional_spec_kwargs = {}
+        if max_runtime:
+            optional_spec_kwargs["requested_run_duration_seconds"] = str(max_runtime)
 
         spec = V1JobSpec(
             cloudspace_id=studio_id or "",
@@ -110,16 +166,14 @@ class MMTApiV2:
             spot=interruptible,
             image_cluster_credentials=cloud_account_auth,
             image_secret_ref=image_credentials or "",
-            artifacts_source=artifacts_local or "",
-            artifacts_destination=artifacts_remote or "",
+            path_mappings=path_mappings_list,
+            machine_image_version=machine_image_version,
+            **optional_spec_kwargs,
         )
-        body = ProjectIdMultimachinejobsBody(
+        return ProjectIdMultimachinejobsBody(
             name=name, spec=spec, cluster_id=cloud_account 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
@@ -145,13 +199,13 @@ class MMTApiV2:
             return
 
         if current_state != Status.Stopped:
-            update_body = MultimachinejobsIdBody(desired_state=V1MultiMachineJobState.STOPPED)
+            update_body = MultimachinejobsIdBody(desired_state=V1MultiMachineJobState.STOP)
             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,
+            if self._job_state_to_external(current_job.state) in (
+                Status.Stopping,
                 Status.Completed,
                 Status.Stopped,
                 Status.Failed,
@@ -179,6 +233,8 @@ class MMTApiV2:
             return Status.Completed
         if str(state) == V1MultiMachineJobState.FAILED:
             return Status.Failed
+        if str(state) == V1MultiMachineJobState.STOP:
+            return Status.Stopping
         return Status.Pending
 
     def get_studio_name(self, job: V1MultiMachineJob) -> Optional[str]:
@@ -196,13 +252,41 @@ class MMTApiV2:
     def get_command(self, job: V1MultiMachineJob) -> str:
         return job.spec.command
 
-    def _get_job_machine_from_spec(self, spec: V1JobSpec) -> "Machine":
-        instance_name = spec.instance_name
-        instance_type = spec.instance_type
+    def _get_job_machine_from_spec(self, spec: V1JobSpec, teamspace_id: str, org_id: str) -> "Machine":
+        accelerators = self._get_machines_for_cloud_account(
+            teamspace_id=teamspace_id,
+            cloud_account_id=spec.cluster_id,
+            org_id=org_id,
+        )
+
+        for accelerator in accelerators:
+            possible_identifiers = (
+                accelerator.slug,
+                accelerator.slug_multi_cloud,
+                accelerator.instance_id,
+            )
+            if (spec.instance_name and spec.instance_name in possible_identifiers) or (
+                spec.instance_type and spec.instance_type in possible_identifiers
+            ):
+                return Machine._from_accelerator(accelerator)
 
-        return _COMPUTE_NAME_TO_MACHINE.get(
-            instance_type, _COMPUTE_NAME_TO_MACHINE.get(instance_name, instance_type or instance_name)
+        return Machine.from_str(spec.instance_name or spec.instance_type)
+
+    def _get_machines_for_cloud_account(
+        self, teamspace_id: str, cloud_account_id: str, org_id: str
+    ) -> List[V1ClusterAccelerator]:
+        from lightning_sdk.api.cloud_account_api import CloudAccountApi
+
+        cloud_account_api = CloudAccountApi()
+        accelerators = cloud_account_api.list_cloud_account_accelerators(
+            teamspace_id=teamspace_id,
+            cloud_account_id=cloud_account_id,
+            org_id=org_id,
         )
+        if not accelerators:
+            return []
+
+        return list(filter(lambda acc: acc.enabled, accelerators.accelerator))
 
     def get_total_cost(self, job: V1MultiMachineJob) -> float:
         return job.total_cost

lightning_sdk/api/pipeline_api.py

@@ -0,0 +1,120 @@
+from typing import TYPE_CHECKING, List, Optional, Union
+
+from lightning_sdk.api.cloud_account_api import CloudAccountApi
+from lightning_sdk.lightning_cloud.openapi.models import (
+    ProjectIdPipelinesBody,
+    ProjectIdSchedulesBody,
+    V1DeletePipelineResponse,
+    V1Pipeline,
+    V1PipelineStep,
+    V1ScheduleResourceType,
+    V1SharedFilesystem,
+)
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+from lightning_sdk.teamspace import Teamspace
+
+if TYPE_CHECKING:
+    from lightning_sdk.pipeline.schedule import Schedule
+
+
+class PipelineApi:
+    """Internal API client for Pipeline requests (mainly http requests)."""
+
+    def __init__(self) -> None:
+        self._client = LightningClient(max_tries=0, retry=False)
+        self._cloud_account_api = CloudAccountApi()
+
+    def get_pipeline_by_id(self, project_id: str, pipeline_id_or_name: str) -> Optional[V1Pipeline]:
+        if pipeline_id_or_name.startswith("pip_"):
+            try:
+                return self._client.pipelines_service_get_pipeline(project_id=project_id, id=pipeline_id_or_name)
+            except ApiException as ex:
+                if "not found" in str(ex):
+                    return None
+                raise ex
+        else:
+            try:
+                return self._client.pipelines_service_get_pipeline_by_name(
+                    project_id=project_id, name=pipeline_id_or_name
+                )
+            except ApiException as ex:
+                if "not found" in str(ex):
+                    return None
+                raise ex
+
+    def create_pipeline(
+        self,
+        name: str,
+        teamspace: Teamspace,
+        steps: List["V1PipelineStep"],
+        shared_filesystem: bool,
+        schedules: List["Schedule"],
+        parent_pipeline_id: Optional[str],
+    ) -> V1Pipeline:
+        body = ProjectIdPipelinesBody(
+            name=name,
+            steps=steps,
+            shared_filesystem=self._prepare_shared_filesystem(shared_filesystem, steps, teamspace),
+            parent_pipeline_id=parent_pipeline_id or "",
+        )
+
+        pipeline = self._client.pipelines_service_create_pipeline(body, teamspace.id)
+
+        # Delete the previous schedules
+        if parent_pipeline_id is not None:
+            current_schedules = self._client.schedules_service_list_schedules(teamspace.id).schedules
+            for schedule in current_schedules:
+                self._client.schedules_service_delete_schedule(teamspace.id, schedule.id)
+
+        if len(schedules):
+            for schedule in schedules:
+                body = ProjectIdSchedulesBody(
+                    cron_expression=schedule.cron_expression,
+                    display_name=schedule.name,
+                    resource_id=pipeline.id,
+                    parent_resource_id=parent_pipeline_id or "",
+                    resource_type=V1ScheduleResourceType.PIPELINE,
+                    timezone=schedule.timezone,
+                    parallel_runs=schedule.parallel_runs or False,
+                )
+
+                self._client.schedules_service_create_schedule(body, teamspace.id)
+
+        return pipeline
+
+    def stop(self, pipeline: V1Pipeline) -> V1Pipeline:
+        body = pipeline
+        body.state = "stop"
+        return self._client.pipelines_service_update_pipeline(body)
+
+    def delete(self, project_id: str, pipeline_id: str) -> V1DeletePipelineResponse:
+        return self._client.pipelines_service_delete_pipeline(project_id, pipeline_id)
+
+    def _prepare_shared_filesystem(
+        self, shared_filesystem: Union[bool, V1SharedFilesystem], steps: List["V1PipelineStep"], teamspace: Teamspace
+    ) -> V1SharedFilesystem:
+        if not shared_filesystem:
+            return V1SharedFilesystem(enabled=False)
+
+        from lightning_sdk.pipeline.utils import _get_cloud_account
+
+        clusters = self._cloud_account_api.list_cloud_accounts(teamspace_id=teamspace.id)
+
+        selected_cluster = None
+        selected_cluster_id = _get_cloud_account(steps)
+        for cluster in clusters:
+            if cluster.id == selected_cluster_id:
+                selected_cluster = cluster
+                break
+
+        if selected_cluster is None:
+            raise ValueError(f"Cloud Account {selected_cluster_id} not found")
+
+        if selected_cluster.spec.aws_v1:
+            return V1SharedFilesystem(enabled=True, s3_folder=True)
+
+        if selected_cluster.spec.google_cloud_v1:
+            return V1SharedFilesystem(enabled=True, gcs_folder=True)
+
+        raise NotImplementedError("This cluster isn't support yet")