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

lightning_sdk/__version__.py

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

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,
@@ -64,7 +64,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,
@@ -293,7 +295,7 @@ class JobApiV2:
         reuse_snapshot: bool,
         max_runtime: Optional[int] = None,
         machine_image_version: Optional[str] = None,
-    ) -> ProjectIdJobsBody:
+    ) -> JobsServiceCreateJobBody:
         env_vars = []
         if env is not None:
             for k, v in env.items():
@@ -330,7 +332,7 @@ class JobApiV2:
             machine_image_version=machine_image_version,
             **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 +357,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,
@@ -174,7 +174,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 +209,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,
         )
@@ -376,7 +378,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 +386,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 +484,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 +496,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 +509,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 +532,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,
         )
 

lightning_sdk/api/utils.py

@@ -17,13 +17,16 @@ from tqdm.auto import tqdm
 from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__, _LIGHTNING_DEBUG
 from lightning_sdk.lightning_cloud.openapi import (
     CloudSpaceServiceApi,
+    CloudSpaceServiceCreateCloudSpaceAppInstanceBody,
     Externalv1LightningappInstance,
     ModelsStoreApi,
-    ProjectIdStorageBody,
-    StorageCompleteBody,
+    ModelsStoreCompleteMultiPartUploadBody,
+    ModelsStoreCreateMultiPartUploadBody,
+    ModelsStoreGetModelFileUploadUrlsBody,
     StorageServiceApi,
-    UploadIdCompleteBody,
-    UploadIdPartsBody,
+    StorageServiceCompleteUploadProjectArtifactBody,
+    StorageServiceUploadProjectArtifactBody,
+    StorageServiceUploadProjectArtifactPartsBody,
     V1CompletedPart,
     V1CompleteUpload,
     V1PathMapping,
@@ -31,20 +34,8 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1SignedUrl,
     V1UploadProjectArtifactPartsResponse,
     V1UploadProjectArtifactResponse,
-    VersionUploadsBody,
 )
 from lightning_sdk.lightning_cloud.openapi.models.v1_model_version_archive import V1ModelVersionArchive
-
-try:
-    from lightning_sdk.lightning_cloud.openapi import AppsIdBody1 as AppsIdBody
-except ImportError:
-    from lightning_sdk.lightning_cloud.openapi import AppsIdBody
-
-try:
-    from lightning_sdk.lightning_cloud.openapi import UploadsUploadIdBody1 as UploadsUploadIdBody
-except ImportError:
-    from lightning_sdk.lightning_cloud.openapi import UploadsUploadIdBody
-
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.machine import Machine
@@ -120,7 +111,7 @@ class _FileUploader:
 
     def _multipart_upload(self, count: int) -> None:
         """Does a parallel multipart upload."""
-        body = ProjectIdStorageBody(cluster_id=self.cloud_account, filename=self.remote_path)
+        body = StorageServiceUploadProjectArtifactBody(cluster_id=self.cloud_account, filename=self.remote_path)
         resp: V1UploadProjectArtifactResponse = self.client.storage_service_upload_project_artifact(
             body=body, project_id=self.teamspace_id
         )
@@ -135,7 +126,7 @@ class _FileUploader:
             for batch in batched_indices:
                 completed.extend(self._process_upload_batch(executor=p, batch=batch, upload_id=resp.upload_id))
 
-        completed_body = StorageCompleteBody(
+        completed_body = StorageServiceCompleteUploadProjectArtifactBody(
             cluster_id=self.cloud_account, filename=self.remote_path, parts=completed, upload_id=resp.upload_id
         )
         self.client.storage_service_complete_upload_project_artifact(body=completed_body, project_id=self.teamspace_id)
@@ -148,7 +139,9 @@ class _FileUploader:
 
     def _request_urls(self, parts: List[int], upload_id: str) -> List[V1PresignedUrl]:
         """Requests urls for a batch of parts."""
-        body = UploadsUploadIdBody(filename=self.remote_path, parts=parts, cluster_id=self.cloud_account)
+        body = StorageServiceUploadProjectArtifactPartsBody(
+            filename=self.remote_path, parts=parts, cluster_id=self.cloud_account
+        )
         resp: V1UploadProjectArtifactPartsResponse = self.client.storage_service_upload_project_artifact_parts(
             body, self.teamspace_id, upload_id
         )
@@ -244,7 +237,7 @@ class _ModelFileUploader:
 
     def _multipart_upload(self, count: int) -> None:
         """Does a parallel multipart upload."""
-        body = VersionUploadsBody(filepath=self.remote_path)
+        body = ModelsStoreCreateMultiPartUploadBody(filepath=self.remote_path)
         resp = self.api.models_store_create_multi_part_upload(
             body,
             project_id=self.teamspace_id,
@@ -262,7 +255,7 @@ class _ModelFileUploader:
             for batch in batched_indices:
                 completed.extend(self._process_upload_batch(executor=p, batch=batch, upload_id=resp.upload_id))
 
-        completed_body = UploadIdCompleteBody(filepath=self.remote_path, parts=completed)
+        completed_body = ModelsStoreCompleteMultiPartUploadBody(filepath=self.remote_path, parts=completed)
         self.api.models_store_complete_multi_part_upload(
             completed_body,
             project_id=self.teamspace_id,
@@ -279,7 +272,7 @@ class _ModelFileUploader:
 
     def _request_urls(self, parts: List[int], upload_id: str) -> List[V1SignedUrl]:
         """Requests urls for a batch of parts."""
-        body = UploadIdPartsBody(filepath=self.remote_path, parts=parts)
+        body = ModelsStoreGetModelFileUploadUrlsBody(filepath=self.remote_path, parts=parts)
         resp = self.api.models_store_get_model_file_upload_urls(
             body,
             project_id=self.teamspace_id,
@@ -363,7 +356,7 @@ def _sanitize_studio_remote_path(path: str, studio_id: str) -> str:
 
 
 def _resolve_teamspace_remote_path(path: str) -> str:
-    return f"/Uploads/{path.replace('/teamspace/studios/this_studio/', '')}"
+    return f"/Uploads/{path.replace('/teamspace/uploads/', '')}"
 
 
 _DOWNLOAD_REQUEST_CHUNK_SIZE = 10 * _BYTES_PER_MB
@@ -676,7 +669,7 @@ def _create_app(
         other_arguments["spot"] = str(other_arguments["interruptible"]).lower()
         del other_arguments["interruptible"]
 
-    body = AppsIdBody(
+    body = CloudSpaceServiceCreateCloudSpaceAppInstanceBody(
         cluster_id=cloud_account,
         plugin_arguments=other_arguments,
         service_id=os.getenv(_LIGHTNING_SERVICE_EXECUTION_ID_KEY),