lightning-sdk 2025.12.5__py3-none-any.whl → 2025.12.16__py3-none-any.whl

lightning_sdk/__version__.py

@@ -1,3 +1,3 @@
 """Version information for lightning_sdk."""
 
-__version__ = "2025.12.05"
+__version__ = "2025.12.16"

lightning_sdk/api/agents_api.py

@@ -1,8 +1,8 @@
 from typing import List, Optional
 
 from lightning_sdk.lightning_cloud.openapi import (
-    AgentsIdBody,
-    EndpointsIdBody,
+    AssistantsServiceUpdateAssistantBody,
+    EndpointServiceUpdateEndpointBody,
     V1Assistant,
     V1Endpoint,
     V1PromptSuggestion,
@@ -46,7 +46,7 @@ class AgentApi:
     ) -> V1Assistant:
         """Update the agent with provided details."""
         agent = self.get_agent(agent_id)
-        body = AgentsIdBody(
+        body = AssistantsServiceUpdateAssistantBody(
             cloudspace_id=agent.cloudspace_id,
             cluster_id=agent.cluster_id,
             description=agent.description,
@@ -92,7 +92,7 @@ class AgentApi:
         """Update the agent endpoint with provided details."""
         endpoint = self._get_agent_endpoint(teampsace_id=teamspace_id, endpoint_id=endpoint_id)
 
-        body = EndpointsIdBody(
+        body = EndpointServiceUpdateEndpointBody(
             name=endpoint.name,
             openai=V1UpstreamOpenAI(api_key=endpoint.openai.api_key, base_url=endpoint.openai.base_url),
         )

lightning_sdk/api/ai_hub_api.py

@@ -7,7 +7,7 @@ import backoff
 from lightning_sdk.api.deployment_api import apply_change
 from lightning_sdk.api.utils import _machine_to_compute_name
 from lightning_sdk.lightning_cloud.openapi.models import (
-    CreateDeploymentRequestDefinesASpecForTheJobThatAllowsForAutoscalingJobs,
+    JobsServiceCreateDeploymentBody,
     V1Deployment,
     V1DeploymentTemplate,
     V1DeploymentTemplateParameter,
@@ -158,7 +158,7 @@ class AIHubApi:
             apply_change(template.spec_v2.job, "cluster_id", cloud_account)
         return self._client.jobs_service_create_deployment(
             project_id=project_id,
-            body=CreateDeploymentRequestDefinesASpecForTheJobThatAllowsForAutoscalingJobs(
+            body=JobsServiceCreateDeploymentBody(
                 autoscaling=template.spec_v2.autoscaling,
                 cluster_id=cloud_account,
                 endpoint=template.spec_v2.endpoint,

lightning_sdk/api/base_studio_api.py

@@ -1,10 +1,12 @@
 from typing import Any, List, Optional
 
+from lightning_sdk.lightning_cloud.openapi import (
+    CloudSpaceEnvironmentTemplateServiceUpdateCloudSpaceEnvironmentTemplateBody as BaseStudioUpdateBody,
+)
 from lightning_sdk.lightning_cloud.openapi import (
     V1CloudSpaceEnvironmentType,
     V1ListCloudSpaceEnvironmentTemplatesResponse,
 )
-from lightning_sdk.lightning_cloud.openapi.models.update import Update as BaseStudioUpdateBody
 from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_environment_template import (
     V1CloudSpaceEnvironmentTemplate,
 )

lightning_sdk/api/deployment_api.py

@@ -3,7 +3,7 @@ from typing import Any, Dict, List, Literal, Optional, Union
 
 from lightning_sdk.api.utils import _machine_to_compute_name, resolve_path_mappings
 from lightning_sdk.lightning_cloud.openapi import (
-    CreateDeploymentRequestDefinesASpecForTheJobThatAllowsForAutoscalingJobs,
+    JobsServiceCreateDeploymentBody,
     V1AutoscalingSpec,
     V1AutoscalingTargetMetric,
     V1Deployment,
@@ -233,7 +233,7 @@ class DeploymentApi:
     ) -> V1Deployment:
         return self._client.jobs_service_create_deployment(
             project_id=deployment.project_id,
-            body=CreateDeploymentRequestDefinesASpecForTheJobThatAllowsForAutoscalingJobs(
+            body=JobsServiceCreateDeploymentBody(
                 cloudspace_id=deployment.cloudspace_id,
                 autoscaling=deployment.autoscaling,
                 cluster_id=deployment.spec.cluster_id,

lightning_sdk/api/job_api.py

@@ -11,11 +11,11 @@ from lightning_sdk.api.utils import (
 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 (
-    AppinstancesIdBody,
     Externalv1LightningappInstance,
     Externalv1Lightningwork,
-    JobsIdBody1,
-    ProjectIdJobsBody,
+    JobsServiceCreateJobBody,
+    JobsServiceUpdateJobBody,
+    LightningappInstanceServiceUpdateLightningappInstanceBody,
     V1CloudSpace,
     V1ClusterAccelerator,
     V1DownloadJobLogsResponse,
@@ -30,6 +30,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1LightningworkState,
     V1ListLightningworkResponse,
     V1UserRequestedComputeConfig,
+    V1Volume,
 )
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.machine import Machine
@@ -64,7 +65,9 @@ class JobApiV1:
         return None
 
     def stop_job(self, job_id: str, teamspace_id: str) -> None:
-        body = AppinstancesIdBody(spec=V1LightningappInstanceSpec(desired_state=V1LightningappInstanceState.STOPPED))
+        body = LightningappInstanceServiceUpdateLightningappInstanceBody(
+            spec=V1LightningappInstanceSpec(desired_state=V1LightningappInstanceState.STOPPED)
+        )
         self._client.lightningapp_instance_service_update_lightningapp_instance(
             project_id=teamspace_id,
             id=job_id,
@@ -251,7 +254,16 @@ class JobApiV2:
         artifacts_remote: Optional[str],  # deprecated in favor of path_mappings
         max_runtime: Optional[int] = None,
         reuse_snapshot: bool = True,
+        scratch_disks: Optional[Dict[str, int]] = None,
     ) -> V1Job:
+        if scratch_disks is not None:
+            sanitized_scratch_disks = {}
+            for k, v in scratch_disks.items():
+                sanitized_k = k if k.startswith("/teamspace/scratch/") else f"/teamspace/scratch/{k}"
+                sanitized_scratch_disks[sanitized_k] = v
+        else:
+            sanitized_scratch_disks = None
+
         body = self._create_job_body(
             name=name,
             command=command,
@@ -269,6 +281,7 @@ class JobApiV2:
             artifacts_remote=artifacts_remote,
             max_runtime=max_runtime,
             reuse_snapshot=reuse_snapshot,
+            scratch_disks=sanitized_scratch_disks,
         )
 
         job: V1Job = self._client.jobs_service_create_job(project_id=teamspace_id, body=body)
@@ -293,7 +306,8 @@ class JobApiV2:
         reuse_snapshot: bool,
         max_runtime: Optional[int] = None,
         machine_image_version: Optional[str] = None,
-    ) -> ProjectIdJobsBody:
+        scratch_disks: Optional[Dict[str, int]] = None,
+    ) -> JobsServiceCreateJobBody:
         env_vars = []
         if env is not None:
             for k, v in env.items():
@@ -314,6 +328,9 @@ class JobApiV2:
         if max_runtime:
             optional_spec_kwargs["requested_run_duration_seconds"] = str(max_runtime)
 
+        # don't do default dicts, as they'll be mutable. Create a fresh one here
+        scratch_disks = scratch_disks or {}
+
         spec = V1JobSpec(
             cloudspace_id=studio_id or "",
             cluster_id=cloud_account or "",
@@ -328,9 +345,10 @@ class JobApiV2:
             image_secret_ref=image_credentials or "",
             path_mappings=path_mappings_list,
             machine_image_version=machine_image_version,
+            volumes=[V1Volume(path=k, size_gb=v, ephemeral=True) for k, v in scratch_disks.items()],
             **optional_spec_kwargs,
         )
-        return ProjectIdJobsBody(name=name, spec=spec)
+        return JobsServiceCreateJobBody(name=name, spec=spec)
 
     def get_job_by_name(self, name: str, teamspace_id: str) -> V1Job:
         job: V1Job = self._client.jobs_service_find_job(project_id=teamspace_id, name=name)
@@ -355,7 +373,7 @@ class JobApiV2:
             return
 
         if current_state != Status.Stopping:
-            update_body = JobsIdBody1(state=self.v2_job_state_stop)
+            update_body = JobsServiceUpdateJobBody(state=self.v2_job_state_stop)
             self._client.jobs_service_update_job(body=update_body, project_id=teamspace_id, id=job_id)
 
         while True:

lightning_sdk/api/license_api.py

@@ -1,6 +1,6 @@
 from lightning_sdk.api.utils import _get_cloud_url as _cloud_url
 from lightning_sdk.lightning_cloud.login import Auth
-from lightning_sdk.lightning_cloud.openapi import LicenseKeyValidateBody, ProductLicenseServiceApi
+from lightning_sdk.lightning_cloud.openapi import ProductLicenseServiceApi, ProductLicenseServiceValidateLicenseBody
 
 
 class LicenseApi:
@@ -26,7 +26,7 @@ class LicenseApi:
         """
         try:
             response = self._api.product_license_service_validate_license(
-                body=LicenseKeyValidateBody(product_id=product_id), license_key=license_key
+                body=ProductLicenseServiceValidateLicenseBody(product_id=product_id), license_key=license_key
             )
             return response.is_valid
         except Exception:

lightning_sdk/api/mmt_api.py

@@ -12,8 +12,8 @@ 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,
-    MultimachinejobsIdBody,
-    ProjectIdMultimachinejobsBody,
+    JobsServiceCreateMultiMachineJobBody,
+    JobsServiceUpdateMultiMachineJobBody,
     V1CloudSpace,
     V1EnvVar,
     V1Job,
@@ -133,7 +133,7 @@ class MMTApiV2:
         reuse_snapshot: bool,
         max_runtime: Optional[int] = None,
         machine_image_version: Optional[str] = None,
-    ) -> ProjectIdMultimachinejobsBody:
+    ) -> JobsServiceCreateMultiMachineJobBody:
         env_vars = []
         if env is not None:
             for k, v in env.items():
@@ -170,7 +170,7 @@ class MMTApiV2:
             machine_image_version=machine_image_version,
             **optional_spec_kwargs,
         )
-        return ProjectIdMultimachinejobsBody(
+        return JobsServiceCreateMultiMachineJobBody(
             name=name, spec=spec, cluster_id=cloud_account or "", machines=num_machines
         )
 
@@ -199,7 +199,7 @@ class MMTApiV2:
             return
 
         if current_state != Status.Stopped:
-            update_body = MultimachinejobsIdBody(desired_state=V1MultiMachineJobState.STOP)
+            update_body = JobsServiceUpdateMultiMachineJobBody(desired_state=V1MultiMachineJobState.STOP)
             self._client.jobs_service_update_multi_machine_job(body=update_body, project_id=teamspace_id, id=job_id)
 
         while True:

lightning_sdk/api/pipeline_api.py

@@ -2,8 +2,8 @@ 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,
+    PipelinesServiceCreatePipelineBody,
+    SchedulesServiceCreateScheduleBody,
     V1DeletePipelineResponse,
     V1Pipeline,
     V1PipelineStep,
@@ -52,7 +52,7 @@ class PipelineApi:
         schedules: List["Schedule"],
         parent_pipeline_id: Optional[str],
     ) -> V1Pipeline:
-        body = ProjectIdPipelinesBody(
+        body = PipelinesServiceCreatePipelineBody(
             name=name,
             steps=steps,
             shared_filesystem=self._prepare_shared_filesystem(shared_filesystem, steps, teamspace),
@@ -69,7 +69,7 @@ class PipelineApi:
 
         if len(schedules):
             for schedule in schedules:
-                body = ProjectIdSchedulesBody(
+                body = SchedulesServiceCreateScheduleBody(
                     cron_expression=schedule.cron_expression,
                     display_name=schedule.name,
                     resource_id=pipeline.id,

lightning_sdk/api/studio_api.py

@@ -25,15 +25,18 @@ from lightning_sdk.api.utils import (
 from lightning_sdk.constants import _LIGHTNING_DEBUG
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
-    CloudspaceIdRunsBody,
-    CloudspacesIdBody,
+    AssistantsServiceCreateAssistantBody,
+    AssistantsServiceCreateAssistantManagedEndpointBody,
+    CloudSpaceServiceCreateCloudSpaceBody,
+    CloudSpaceServiceCreateLightningRunBody,
+    CloudSpaceServiceExecuteCommandInCloudSpaceBody,
+    CloudSpaceServiceForkCloudSpaceBody,
+    CloudSpaceServiceStartCloudSpaceInstanceBody,
+    CloudSpaceServiceUpdateCloudSpaceBody,
+    CloudSpaceServiceUpdateCloudSpaceInstanceConfigBody,
+    CloudSpaceServiceUpdateCloudSpaceSleepConfigBody,
+    EndpointServiceCreateEndpointBody,
     Externalv1LightningappInstance,
-    IdCodeconfigBody,
-    IdExecuteBody1,
-    IdForkBody1,
-    IdSleepconfigBody,
-    IdStartBody,
-    ProjectIdCloudspacesBody,
     V1Assistant,
     V1CloudSpace,
     V1CloudSpaceInstanceConfig,
@@ -41,27 +44,20 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1CloudSpaceSourceType,
     V1CloudSpaceState,
     V1ClusterAccelerator,
+    V1Endpoint,
     V1EndpointType,
     V1EnvVar,
     V1GetCloudSpaceInstanceStatusResponse,
     V1GetLongRunningCommandInCloudSpaceResponse,
     V1LoginRequest,
+    V1ManagedEndpoint,
+    V1ManagedModel,
     V1Plugin,
     V1PluginsListResponse,
     V1UpstreamCloudSpace,
+    V1UpstreamManaged,
     V1UserRequestedComputeConfig,
 )
-from lightning_sdk.lightning_cloud.openapi.models import ProjectIdEndpointsBody
-from lightning_sdk.lightning_cloud.openapi.models.project_id_agentmanagedendpoints_body import (
-    ProjectIdAgentmanagedendpointsBody,
-)
-from lightning_sdk.lightning_cloud.openapi.models.project_id_agents_body import (
-    ProjectIdAgentsBody,
-)
-from lightning_sdk.lightning_cloud.openapi.models.v1_endpoint import V1Endpoint
-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_upstream_managed import V1UpstreamManaged
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.machine import Machine
 
@@ -132,7 +128,7 @@ class StudioApi:
         cloud_space_environment_template_id: Optional[str] = None,
     ) -> V1CloudSpace:
         """Create a Studio with a given name in a given Teamspace on a possibly given cloud_account."""
-        body = ProjectIdCloudspacesBody(
+        body = CloudSpaceServiceCreateCloudSpaceBody(
             cluster_id=cloud_account,
             name=name,
             display_name=name,
@@ -144,7 +140,7 @@ class StudioApi:
         )
         studio = self._client.cloud_space_service_create_cloud_space(body, teamspace_id)
 
-        run_body = CloudspaceIdRunsBody(
+        run_body = CloudSpaceServiceCreateLightningRunBody(
             cluster_id=studio.cluster_id,
             local_source=True,
         )
@@ -187,7 +183,7 @@ class StudioApi:
         if max_runtime is not None:
             optional_kwargs_compute_body["requested_run_duration_seconds"] = str(max_runtime)
         self._client.cloud_space_service_start_cloud_space_instance(
-            IdStartBody(
+            CloudSpaceServiceStartCloudSpaceInstanceBody(
                 compute_config=V1UserRequestedComputeConfig(
                     name=_machine_to_compute_name(machine),
                     spot=interruptible,
@@ -223,7 +219,7 @@ class StudioApi:
         if max_runtime is not None:
             optional_kwargs_compute_body["requested_run_duration_seconds"] = str(max_runtime)
         self._client.cloud_space_service_start_cloud_space_instance(
-            IdStartBody(
+            CloudSpaceServiceStartCloudSpaceInstanceBody(
                 compute_config=V1UserRequestedComputeConfig(
                     name=_machine_to_compute_name(machine),
                     spot=interruptible,
@@ -272,7 +268,9 @@ class StudioApi:
         """Switches given Studio to a new machine type."""
         compute_name = _machine_to_compute_name(machine)
         # TODO: UI sends disk size here, maybe we need to also?
-        body = IdCodeconfigBody(compute_config=V1UserRequestedComputeConfig(name=compute_name, spot=interruptible))
+        body = CloudSpaceServiceUpdateCloudSpaceInstanceConfigBody(
+            compute_config=V1UserRequestedComputeConfig(name=compute_name, spot=interruptible)
+        )
         if cloud_account:
             body.compute_config.cluster_override = cloud_account
         self._client.cloud_space_service_update_cloud_space_instance_config(
@@ -516,7 +514,7 @@ class StudioApi:
             timeout: wait for this many seconds for the command to finish.
         """
         response_submit = self._client.cloud_space_service_execute_command_in_cloud_space(
-            IdExecuteBody1("; ".join(commands), detached=True),
+            CloudSpaceServiceExecuteCommandInCloudSpaceBody("; ".join(commands), detached=True),
             project_id=teamspace_id,
             id=studio_id,
         )
@@ -556,7 +554,7 @@ class StudioApi:
     def run_studio_commands(self, studio_id: str, teamspace_id: str, *commands: str) -> Tuple[str, int]:
         """Run given commands in a given Studio."""
         response_submit = self._client.cloud_space_service_execute_command_in_cloud_space(
-            IdExecuteBody1("; ".join(commands), detached=True),
+            CloudSpaceServiceExecuteCommandInCloudSpaceBody("; ".join(commands), detached=True),
             project_id=teamspace_id,
             id=studio_id,
         )
@@ -598,7 +596,7 @@ class StudioApi:
         idle_shutdown_seconds: Optional[int] = None,
     ) -> V1CloudSpaceInstanceConfig:
         """Update the autoshutdown time and behaviour of the given Studio."""
-        body = IdSleepconfigBody(
+        body = CloudSpaceServiceUpdateCloudSpaceSleepConfigBody(
             disable_auto_shutdown=not enabled if enabled is not None else None,
             idle_shutdown_seconds=idle_shutdown_seconds,
         )
@@ -629,7 +627,9 @@ class StudioApi:
             init_kwargs["org"] = OrgApi()._get_org_by_id(target_teamspace.owner_id).name
 
         new_cloudspace = self._client.cloud_space_service_fork_cloud_space(
-            IdForkBody1(target_project_id=target_teamspace_id, new_name=new_name), project_id=teamspace_id, id=studio_id
+            CloudSpaceServiceForkCloudSpaceBody(target_project_id=target_teamspace_id, new_name=new_name),
+            project_id=teamspace_id,
+            id=studio_id,
         )
 
         while self.get_studio_by_id(new_cloudspace.id, target_teamspace_id).state != V1CloudSpaceState.READY:
@@ -966,7 +966,7 @@ class StudioApi:
         """Starts a new port to the given Studio."""
         endpoint = self._client.endpoint_service_create_endpoint(
             project_id=teamspace_id,
-            body=ProjectIdEndpointsBody(
+            body=EndpointServiceCreateEndpointBody(
                 name=name,
                 ports=[str(port)],
                 cloudspace=V1UpstreamCloudSpace(
@@ -983,7 +983,7 @@ class StudioApi:
             org_id = target_teamspace.owner_id
         endpoint = self._client.endpoint_service_create_endpoint(
             project_id=teamspace_id,
-            body=ProjectIdEndpointsBody(
+            body=EndpointServiceCreateEndpointBody(
                 ports=[str(port)],
                 cloudspace=V1UpstreamCloudSpace(
                     cloudspace_id=studio_id,
@@ -994,7 +994,7 @@ class StudioApi:
         )
         valid_url = endpoint.urls[0]
         managed_endpoint = self._client.assistants_service_create_assistant_managed_endpoint(
-            body=ProjectIdAgentmanagedendpointsBody(
+            body=AssistantsServiceCreateAssistantManagedEndpointBody(
                 endpoint=V1ManagedEndpoint(
                     name=assistant_name,
                     base_url=valid_url + "/v1",
@@ -1009,7 +1009,7 @@ class StudioApi:
             project_id=teamspace_id,
         )
 
-        body = ProjectIdAgentsBody(
+        body = AssistantsServiceCreateAssistantBody(
             endpoint=V1Endpoint(
                 cloudspace=V1UpstreamCloudSpace(cloudspace_id=studio_id),
                 name=assistant_name,
@@ -1040,7 +1040,7 @@ class StudioApi:
         )
 
     def _update_cloudspace(self, studio: V1CloudSpace, teamspace_id: str, key: str, value: Any) -> None:
-        body = CloudspacesIdBody(
+        body = CloudSpaceServiceUpdateCloudSpaceBody(
             code_url=studio.code_url,
             data_connection_mounts=studio.data_connection_mounts,
             description=studio.description,

lightning_sdk/api/teamspace_api.py

@@ -17,14 +17,14 @@ from lightning_sdk.api.utils import (
 )
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
-    Create,
+    AssistantsServiceCreateAssistantBody,
+    DataConnectionServiceCreateDataConnectionBody,
     Externalv1LightningappInstance,
-    ModelIdVersionsBody,
     ModelsStoreApi,
-    ProjectIdAgentsBody,
-    ProjectIdModelsBody,
-    ProjectIdSecretsBody,
-    SecretsIdBody,
+    ModelsStoreCreateModelBody,
+    ModelsStoreCreateModelVersionBody,
+    SecretServiceCreateSecretBody,
+    SecretServiceUpdateSecretBody,
     V1Assistant,
     V1CloudSpace,
     V1ClusterAccelerator,
@@ -47,6 +47,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1UpstreamOpenAI,
 )
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
+from lightning_sdk.machine import Machine
 
 __all__ = ["TeamspaceApi"]
 
@@ -174,7 +175,7 @@ class TeamspaceApi:
 
         ([V1PromptSuggestion(content=suggestion) for suggestion in prompt_suggestions] if prompt_suggestions else None)
 
-        body = ProjectIdAgentsBody(
+        body = AssistantsServiceCreateAssistantBody(
             endpoint=endpoint,
             name=name,
             model=model,
@@ -209,12 +210,14 @@ class TeamspaceApi:
         models = self.models_api.models_store_list_models(project_id=teamspace_id, name=name).models
         if len(models) == 0:
             return self.models_api.models_store_create_model(
-                body=ProjectIdModelsBody(cluster_id=cloud_account, metadata=metadata, name=name, private=private),
+                body=ModelsStoreCreateModelBody(
+                    cluster_id=cloud_account, metadata=metadata, name=name, private=private
+                ),
                 project_id=teamspace_id,
             )
         assert len(models) == 1, "Multiple models with the same name found"
         return self.models_api.models_store_create_model_version(
-            body=ModelIdVersionsBody(cluster_id=cloud_account, version=version),
+            body=ModelsStoreCreateModelVersionBody(cluster_id=cloud_account, version=version),
             project_id=teamspace_id,
             model_id=models[0].id,
         )
@@ -321,11 +324,47 @@ class TeamspaceApi:
         jobs = self._client.jobs_service_list_multi_machine_jobs(project_id=teamspace_id).multi_machine_jobs
         return apps, jobs
 
-    def list_machines(self, teamspace_id: str, cloud_account: str) -> List[V1ClusterAccelerator]:
-        response = self._client.cluster_service_list_project_cluster_accelerators(
-            project_id=teamspace_id, id=cloud_account
-        )
-        return response.accelerator
+    def list_machines(
+        self,
+        teamspace_id: str,
+        cloud_accounts: List[str],
+        machine: Optional[Machine] = None,
+        org_id: Optional[str] = None,
+    ) -> List[V1ClusterAccelerator]:
+        from lightning_sdk.api.cloud_account_api import CloudAccountApi
+
+        cloud_account_api = CloudAccountApi()
+        matched_accelerators = []
+        for ca in cloud_accounts:
+            try:
+                accelerators = cloud_account_api.list_cloud_account_accelerators(
+                    teamspace_id=teamspace_id,
+                    cloud_account_id=ca,
+                    org_id=org_id,
+                )
+                if not accelerators.accelerator:
+                    continue
+
+                if accelerators.accelerator:
+                    for cluster_machine in accelerators.accelerator:
+                        if not machine:
+                            matched_accelerators.append(cluster_machine)
+                            continue
+                        if (
+                            cluster_machine.resources.gpu == machine.accelerator_count
+                            or cluster_machine.resources.cpu == machine.accelerator_count
+                        ) and any(
+                            machine.family.lower() in s
+                            for s in (
+                                cluster_machine.slug,
+                                cluster_machine.slug_multi_cloud,
+                                cluster_machine.instance_id,
+                            )
+                        ):
+                            matched_accelerators.append(cluster_machine)
+            except Exception:
+                pass
+        return matched_accelerators
 
     def get_model(self, teamspace_id: str, model_id: Optional[str] = None, model_name: Optional[str] = None) -> V1Model:
         if model_id:
@@ -376,7 +415,6 @@ class TeamspaceApi:
         path: str,
         target_path: str,
         teamspace_id: str,
-        cloud_account: str,
         progress_bar: bool = True,
     ) -> None:
         """Downloads a given file in Teamspace drive to a target location."""
@@ -385,19 +423,29 @@ class TeamspaceApi:
         auth.authenticate()
         token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
 
-        query_params = {
-            "clusterId": cloud_account,
-            "key": _resolve_teamspace_remote_path(path),
-            "token": token,
-        }
+        cluster_ids = [ca.cluster_id for ca in self.list_cloud_accounts(teamspace_id)]
+
+        found = False
+        for cluster_id in cluster_ids:
+            query_params = {
+                "clusterId": cluster_id,
+                "key": _resolve_teamspace_remote_path(path),
+                "token": token,
+            }
+
+            r = requests.get(
+                f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/download",
+                params=query_params,
+                stream=True,
+            )
+
+            if r.status_code == 200:
+                found = True
+                break
+
+        if not found:
+            raise FileNotFoundError(f"The provided path does not exist in the teamspace: {path}")
 
-        r = requests.get(
-            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/download",
-            params=query_params,
-            stream=True,
-        )
-        if r.status_code == 404:
-            raise FileNotFoundError(f"File {path} not found")
         total_length = int(r.headers.get("content-length"))
 
         if progress_bar:
@@ -473,7 +521,7 @@ class TeamspaceApi:
 
     def _update_secret(self, teamspace_id: str, secret_id: str, value: str) -> None:
         self._client.secret_service_update_secret(
-            body=SecretsIdBody(value=value),
+            body=SecretServiceUpdateSecretBody(value=value),
             project_id=teamspace_id,
             id=secret_id,
         )
@@ -485,7 +533,8 @@ class TeamspaceApi:
         value: str,
     ) -> None:
         self._client.secret_service_create_secret(
-            body=ProjectIdSecretsBody(name=key, value=value, type=V1SecretType.UNSPECIFIED), project_id=teamspace_id
+            body=SecretServiceCreateSecretBody(name=key, value=value, type=V1SecretType.UNSPECIFIED),
+            project_id=teamspace_id,
         )
 
     def verify_secret_name(self, name: str) -> bool:
@@ -497,7 +546,7 @@ class TeamspaceApi:
         return re.match(pattern, name) is not None
 
     def new_folder(self, teamspace_id: str, name: str, cluster: Optional[V1ExternalCluster]) -> None:
-        create_request = Create(
+        create_request = DataConnectionServiceCreateDataConnectionBody(
             name=name,
             create_resources=True,
             force=True,
@@ -520,7 +569,7 @@ class TeamspaceApi:
     def new_connection(
         self, teamspace_id: str, name: str, source: str, cluster: V1ExternalCluster, writable: bool, region: str
     ) -> None:
-        create_request = Create(
+        create_request = DataConnectionServiceCreateDataConnectionBody(
             name=name,
             create_resources=False,
             force=True,

lightning_sdk/api/user_api.py

@@ -3,7 +3,7 @@ from typing import Dict, List, Optional, Union
 
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
-    SecretsIdBody1,
+    SecretServiceUpdateUserSecretBody,
     V1CloudSpace,
     V1CreateUserSecretRequest,
     V1GetUserResponse,
@@ -105,7 +105,7 @@ class UserApi:
 
     def _update_secret(self, secret_id: str, value: str) -> None:
         self._client.secret_service_update_user_secret(
-            body=SecretsIdBody1(value=value),
+            body=SecretServiceUpdateUserSecretBody(value=value),
             id=secret_id,
         )