lightning-sdk 2025.7.17__py3-none-any.whl → 2025.7.30rc0__py3-none-any.whl

lightning_sdk/__init__.py

@@ -4,7 +4,7 @@ from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__  # noq
 from lightning_sdk.deployment import Deployment
 from lightning_sdk.helpers import _check_version_and_prompt_upgrade, _set_tqdm_envvars_noninteractive
 from lightning_sdk.job import Job
-from lightning_sdk.machine import Machine
+from lightning_sdk.machine import CloudProvider, Machine
 from lightning_sdk.mmt import MMT
 from lightning_sdk.organization import Organization
 from lightning_sdk.plugin import JobsPlugin, MultiMachineTrainingPlugin, Plugin, SlurmJobsPlugin
@@ -16,6 +16,7 @@ from lightning_sdk.user import User
 __all__ = [
     "AIHub",
     "Agent",
+    "CloudProvider",
     "Deployment",
     "Job",
     "JobsPlugin",
@@ -31,6 +32,6 @@ __all__ = [
     "User",
 ]
 
-__version__ = "2025.07.17"
+__version__ = "2025.07.30rc0"
 _check_version_and_prompt_upgrade(__version__)
 _set_tqdm_envvars_noninteractive()

lightning_sdk/api/cloud_account_api.py

@@ -0,0 +1,204 @@
+from functools import lru_cache
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+from lightning_sdk.lightning_cloud.openapi import (
+    Externalv1Cluster,
+    V1CloudProvider,
+    V1ClusterType,
+    V1ExternalCluster,
+    V1ListClusterAcceleratorsResponse,
+    V1ListDefaultClusterAcceleratorsResponse,
+)
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+if TYPE_CHECKING:
+    from lightning_sdk.machine import CloudProvider
+
+
+class CloudAccountApi:
+    """Internal API client for API requests to cluster endpoints."""
+
+    def __init__(self) -> None:
+        self._client = LightningClient(max_tries=7)
+
+    def get_cloud_account(self, cloud_account_id: str, teamspace_id: str, org_id: str) -> Externalv1Cluster:
+        """Gets the cluster from given params cluster_id, project_id and owner.
+
+        Args:
+            cloud_account_id: the cloud account to get
+            teamspace_id: the teamspace the cloud_account is supposed to be associated with
+            org_id: The owning org of this teamspace
+
+        """
+        res = self._client.cluster_service_get_cluster(id=cloud_account_id, org_id=org_id, project_id=teamspace_id)
+        if not res:
+            raise ValueError(f"CloudAccount {cloud_account_id} does not exist")
+        return res
+
+    @lru_cache(maxsize=None)  # noqa: B019
+    def list_cloud_accounts(self, teamspace_id: str) -> List[V1ExternalCluster]:
+        """Lists the cloud accounts for a given teamspace.
+
+        Args:
+            teamspace_id: The teamspace to list cloud accounts for
+
+        Returns:
+            A list of cloud accounts
+        """
+        res_project = self._client.cluster_service_list_project_clusters(
+            project_id=teamspace_id,
+        )
+        res_global = self._client.cluster_service_list_clusters(
+            project_id=teamspace_id,
+        )
+
+        # can't use set here because the cloud_accounts are not hashable
+        cloud_accounts = []
+        cloud_account_ids = []
+        for cloud_account in res_project.clusters + res_global.clusters:
+            if cloud_account.id not in cloud_account_ids:
+                cloud_accounts.append(cloud_account)
+                cloud_account_ids.append(cloud_account.id)
+
+        return cloud_accounts
+
+    def get_cloud_account_non_org(self, teamspace_id: str, cloud_account_id: str) -> Optional[V1ExternalCluster]:
+        for cluster in self.list_cloud_accounts(teamspace_id=teamspace_id):
+            if cluster.id == cloud_account_id:
+                return cluster
+
+        return None
+
+    @lru_cache(maxsize=None)  # noqa: B019
+    def list_cloud_account_accelerators(
+        self,
+        teamspace_id: str,
+        cloud_account_id: str,
+        org_id: str,
+    ) -> Union[V1ListClusterAcceleratorsResponse, V1ListDefaultClusterAcceleratorsResponse]:
+        """Lists the accelerators for a given cloud account.
+
+        Args:
+            cloud_account_id: cluster ID to list accelerators for
+        """
+        # map cloud_account to provider
+        cloud_provider = None
+        is_default = True
+        for cloud_account in self.list_cloud_accounts(teamspace_id=teamspace_id):
+            if cloud_account.id == cloud_account_id:
+                is_default = cloud_account.spec.cluster_type == V1ClusterType.GLOBAL
+                cloud_provider = self._get_cloud_account_provider(cloud_account)
+                break
+
+        if cloud_provider is None:
+            raise ValueError(
+                f"Cloud Account {cloud_account_id} is not a default cloud account. Are you in the correct teamspace?"
+            )
+
+        if is_default:
+            res = self._list_default_cluster_accelerators(teamspace_id=teamspace_id, cloud_provider=cloud_provider)
+        else:
+            res = self._client.cluster_service_list_cluster_accelerators(
+                id=cloud_account_id,
+                org_id=org_id,
+            )
+
+        if not res:
+            raise ValueError(f"CloudAccount {cloud_account_id} does not exist")
+        return res
+
+    def _list_default_cluster_accelerators(
+        self, teamspace_id: str, cloud_provider: Union[str, "CloudProvider"]
+    ) -> V1ListDefaultClusterAcceleratorsResponse:
+        return self._client.cluster_service_list_default_cluster_accelerators(
+            project_id=teamspace_id, cloud_provider=str(cloud_provider)
+        )
+
+    @lru_cache(maxsize=None)  # noqa: B019
+    def list_global_cloud_accounts(self, teamspace_id: str) -> List[V1ExternalCluster]:
+        """Lists the accelerators for a given teamspace.
+
+        Args:
+            teamspace_id: id of the teamspace to get the associated cloud_accounts for
+        """
+        cloud_accounts = self.list_cloud_accounts(teamspace_id=teamspace_id)
+        if not cloud_accounts:
+            raise ValueError(f"Teamspace {teamspace_id} does not exist")
+        filtered_cloud_accounts = filter(lambda x: x.spec.cluster_type == V1ClusterType.GLOBAL, cloud_accounts)
+        # TODO: remove aggregate filter once finished
+        filtered_cloud_accounts = filter(
+            lambda x: x.spec.driver != V1CloudProvider.LIGHTNING_AGGREGATE, filtered_cloud_accounts
+        )
+        return list(filtered_cloud_accounts)
+
+    def get_cloud_account_provider_mapping(self, teamspace_id: str) -> Dict["CloudProvider", str]:
+        """Gets the cloud account <-> provider mapping."""
+        res = self.list_global_cloud_accounts(teamspace_id=teamspace_id)
+        return {self._get_cloud_account_provider(cloud_account): cloud_account.id for cloud_account in res}
+
+    @staticmethod
+    def _get_cloud_account_provider(cloud_account: Optional[V1ExternalCluster]) -> "CloudProvider":
+        """Determines the cloud provider based on the cloud_account configuration.
+
+        Args:
+            cloud_account: An optional Externalv1Cluster object containing cluster specifications
+
+        Returns:
+            CloudProvider: The determined cloud provider, defaults to AWS if no match is found
+        """
+        from lightning_sdk.machine import CloudProvider
+
+        if not cloud_account:
+            return CloudProvider.AWS
+
+        if cloud_account.spec and cloud_account.spec.driver:
+            if cloud_account.spec.driver == V1CloudProvider.LIGHTNING:
+                return CloudProvider.LIGHTNING
+
+            if cloud_account.spec.driver == V1CloudProvider.DGX:
+                return CloudProvider.DGX
+
+        if cloud_account.spec:
+            if cloud_account.spec.aws_v1:
+                return CloudProvider.AWS
+            if cloud_account.spec.google_cloud_v1:
+                return CloudProvider.GCP
+            if cloud_account.spec.lambda_labs_v1:
+                return CloudProvider.LAMBDA_LABS
+            if cloud_account.spec.vultr_v1:
+                return CloudProvider.VULTR
+            if cloud_account.spec.voltage_park_v1:
+                return CloudProvider.VOLTAGE_PARK
+            if cloud_account.spec.nebius_v1:
+                return CloudProvider.NEBIUS
+
+        return CloudProvider.AWS
+
+    def resolve_cloud_account(
+        self,
+        teamspace_id: str,
+        cloud_account: Optional[str],
+        cloud_provider: Optional[Union["CloudProvider", str]],
+        default_cloud_account: Optional[str],
+    ) -> Optional[str]:
+        if cloud_account:
+            if cloud_provider:
+                cloud_account_resp = self.get_cloud_account_non_org(teamspace_id, cloud_account)
+                cloud_provider_resp = self._get_cloud_account_provider(cloud_account_resp)
+                if cloud_provider_resp != cloud_provider:
+                    raise RuntimeError(
+                        f"Specified both cloud_provider ({cloud_provider}) and "
+                        "cloud_account ({cloud_account} has cloud provider {cloud_provider_resp}) which don't match!"
+                    )
+
+            return cloud_account
+
+        if cloud_provider:
+            cloud_account_mapping = self.get_cloud_account_provider_mapping(teamspace_id=teamspace_id)
+            if cloud_provider and cloud_provider in cloud_account_mapping:
+                return cloud_account_mapping[cloud_provider]
+
+        if default_cloud_account:
+            return default_cloud_account
+
+        return None

lightning_sdk/api/deployment_api.py

@@ -264,6 +264,7 @@ class DeploymentApi:
         custom_domain: Optional[str] = None,
         quantity: Optional[int] = None,
         include_credentials: Optional[bool] = None,
+        max_runtime: Optional[int] = None,
     ) -> V1Deployment:
         # Update the deployment in place
 
@@ -291,6 +292,9 @@ class DeploymentApi:
         requires_release |= apply_change(deployment.spec, "spot", spot)
         requires_release |= apply_change(deployment.spec, "quantity", quantity)
         requires_release |= apply_change(deployment.spec, "include_credentials", include_credentials)
+        requires_release |= apply_change(
+            deployment.spec, "requested_run_duration_seconds", str(max_runtime) if max_runtime is not None else None
+        )
 
         if requires_release:
             if deployment.strategy is None:
@@ -569,6 +573,7 @@ def to_spec(
     quantity: Optional[int] = None,
     include_credentials: Optional[bool] = None,
     cloudspace_id: Optional[None] = None,
+    max_runtime: Optional[int] = None,
 ) -> V1JobSpec:
     if cloud_account is None:
         raise ValueError("The cloud account should be defined.")
@@ -585,6 +590,11 @@ def to_spec(
     if command is None and cloudspace_id is not None:
         raise ValueError("The command should be defined.")
 
+    # 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)
+
     return V1JobSpec(
         cluster_id=cloud_account,
         command=command,
@@ -597,6 +607,7 @@ def to_spec(
         quantity=quantity,
         include_credentials=include_credentials,
         cloudspace_id=cloudspace_id,
+        **optional_spec_kwargs,
     )
 
 

lightning_sdk/api/job_api.py

@@ -17,7 +17,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     JobsIdBody1,
     ProjectIdJobsBody,
     V1CloudSpace,
-    V1ComputeConfig,
+    V1ClusterAccelerator,
     V1DownloadJobLogsResponse,
     V1DownloadLightningappInstanceLogsResponse,
     V1EnvVar,
@@ -94,15 +94,48 @@ class JobApiV1:
     def get_work(self, job_id: str, teamspace_id: str, work_id: str) -> Externalv1Lightningwork:
         return self._client.lightningwork_service_get_lightningwork(project_id=teamspace_id, app_id=job_id, id=work_id)
 
-    def get_machine_from_work(self, work: Externalv1Lightningwork) -> Machine:
+    def get_machine_from_work(self, work: Externalv1Lightningwork, org_id: str) -> Machine:
         spec: V1LightningworkSpec = work.spec
         # prefer user-requested config if specified
         user_requested_compute_config: V1UserRequestedComputeConfig = spec.user_requested_compute_config
-        if user_requested_compute_config.name:
-            return Machine.from_str(user_requested_compute_config.name)
-        compute_config: V1ComputeConfig = spec.compute_config
+        accelerators = self._get_machines_for_cloud_account(
+            teamspace_id=work.project_id,
+            cloud_account_id=spec.cluster_id,
+            org_id=org_id,
+        )
+
+        identifier = None
+
+        if user_requested_compute_config and user_requested_compute_config.name:
+            identifier = user_requested_compute_config.name
+        else:
+            identifier = spec.compute_config.instance_type
+
+        for accelerator in accelerators:
+            if identifier in (
+                accelerator.slug,
+                accelerator.slug_multi_cloud,
+                accelerator.instance_id,
+            ):
+                return Machine.from_str(accelerator.slug_multi_cloud)
+
+        return Machine.from_str(identifier)
 
-        return Machine.from_str(compute_config.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_studio_name(self, job: Externalv1LightningappInstance) -> str:
         cs: V1CloudSpace = self._client.cloud_space_service_get_cloud_space(
@@ -215,6 +248,7 @@ class JobApiV2:
         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] = None,
     ) -> V1Job:
         body = self._create_job_body(
             name=name,
@@ -231,6 +265,7 @@ class JobApiV2:
             path_mappings=path_mappings,
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
+            max_runtime=max_runtime,
         )
 
         job: V1Job = self._client.jobs_service_create_job(project_id=teamspace_id, body=body)
@@ -252,6 +287,7 @@ class JobApiV2:
         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] = None,
     ) -> ProjectIdJobsBody:
         env_vars = []
         if env is not None:
@@ -268,6 +304,11 @@ class JobApiV2:
             artifacts_remote=artifacts_remote,
         )
 
+        # 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 "",
             cluster_id=cloud_account or "",
@@ -281,6 +322,7 @@ class JobApiV2:
             image_cluster_credentials=cloud_account_auth,
             image_secret_ref=image_credentials or "",
             path_mappings=path_mappings_list,
+            **optional_spec_kwargs,
         )
         return ProjectIdJobsBody(name=name, spec=spec)
 
@@ -371,11 +413,41 @@ class JobApiV2:
             return Status.Stopping
         return Status.Pending
 
-    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_str(accelerator.slug_multi_cloud)
+
+        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 Machine.from_str(instance_name, instance_type or instance_name)
+        return list(filter(lambda acc: acc.enabled, accelerators.accelerator))
 
     def get_total_cost(self, job: V1Job) -> float:
         return job.total_cost

lightning_sdk/api/llm_api.py

@@ -98,7 +98,7 @@ class LLMApi:
                     {"contentType": "text", "parts": [prompt]},
                 ],
             },
-            "max_completion_tokens": max_completion_tokens,
+            "max_tokens": max_completion_tokens,
             "conversation_id": conversation_id,
             "billing_project_id": billing_project_id,
             "name": name,

lightning_sdk/api/mmt_api.py

@@ -2,7 +2,7 @@ import json
 import time
 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 (
     _create_app,
     _machine_to_compute_name,
@@ -87,6 +87,7 @@ class MMTApiV2:
         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],
     ) -> V1MultiMachineJob:
         body = self._create_mmt_body(
             name=name,
@@ -104,6 +105,7 @@ class MMTApiV2:
             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,
         )
 
         job: V1MultiMachineJob = self._client.jobs_service_create_multi_machine_job(project_id=teamspace_id, body=body)
@@ -126,6 +128,7 @@ class MMTApiV2:
         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] = None,
     ) -> ProjectIdMultimachinejobsBody:
         env_vars = []
         if env is not None:
@@ -142,6 +145,11 @@ class MMTApiV2:
             artifacts_remote=artifacts_remote,
         )
 
+        # 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 "",
             cluster_id=cloud_account or "",
@@ -155,6 +163,7 @@ class MMTApiV2:
             image_cluster_credentials=cloud_account_auth,
             image_secret_ref=image_credentials or "",
             path_mappings=path_mappings_list,
+            **optional_spec_kwargs,
         )
         return ProjectIdMultimachinejobsBody(
             name=name, spec=spec, cluster_id=cloud_account or "", machines=num_machines
@@ -238,11 +247,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_str(accelerator.slug_multi_cloud)
+
+        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 Machine.from_str(instance_name, instance_type or instance_name)
+        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

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, List, Optional, Union
 
-from lightning_sdk.api.cluster_api import ClusterApi
+from lightning_sdk.api.cloud_account_api import CloudAccountApi
 from lightning_sdk.lightning_cloud.openapi.models import (
     ProjectIdPipelinesBody,
     ProjectIdSchedulesBody,
@@ -23,7 +23,7 @@ class PipelineApi:
 
     def __init__(self) -> None:
         self._client = LightningClient(max_tries=0, retry=False)
-        self._cluster_api = ClusterApi()
+        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_"):
@@ -75,6 +75,7 @@ class PipelineApi:
                     resource_id=pipeline.id,
                     parent_resource_id=parent_pipeline_id or "",
                     resource_type=V1ScheduleResourceType.PIPELINE,
+                    timezone=schedule.timezone,
                 )
 
                 self._client.schedules_service_create_schedule(body, teamspace.id)
@@ -97,7 +98,7 @@ class PipelineApi:
 
         from lightning_sdk.pipeline.utils import _get_cloud_account
 
-        clusters = self._cluster_api.list_clusters(project_id=teamspace.id)
+        clusters = self._cloud_account_api.list_cloud_accounts(teamspace_id=teamspace.id)
 
         selected_cluster = None
         selected_cluster_id = _get_cloud_account(steps)