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

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),

lightning_sdk/cli/entrypoint.py

@@ -40,22 +40,26 @@ main_cli.context_class.formatter_class = CustomHelpFormatter
 @main_cli.command
 def login() -> None:
     """Login to Lightning AI Studios."""
-    if in_studio():
-        # if inside a Studio, no need to run login command
+    # try to fetch credentials, if successful (e.g. in a Studio or already logged in), no need to relogin
+    auth = Auth()
+    if (auth.user_id and auth.api_key) or auth.load():
         try:
             auth_user = _get_authed_user()
+            click.echo(f'You are currently logged in as "{auth_user.name}"')
         except Exception:
-            raise RuntimeError("Unable to identify user within a Studio") from None
+            click.echo("You are already logged in")
+        click.echo('"lightning login" is not required within a Studio or when already logged in')
+        return
 
-        click.echo(f'You are currently logged in as "{auth_user.name}"')
-        click.echo('"lightning login" is not required within a Studio')
-    else:
-        auth = Auth()
-        auth.clear()
-        try:
-            auth.authenticate()
-        except ConnectionError:
-            raise RuntimeError(f"Unable to connect to {_cloud_url()}. Please check your internet connection.") from None
+    if in_studio():
+        # this is unexpected, as we automatically auth within a Studio
+        raise RuntimeError("Unable to login within a Studio. Did you change your shell setup?") from None
+
+    auth.clear()
+    try:
+        auth.authenticate()
+    except ConnectionError:
+        raise RuntimeError(f"Unable to connect to {_cloud_url()}. Please check your internet connection.") from None
 
 
 @main_cli.command

lightning_sdk/job/base.py

@@ -92,6 +92,7 @@ class _BaseJob(ABC, metaclass=TrackCallsABCMeta):
         max_runtime: Optional[int] = None,
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
         reuse_snapshot: bool = True,
+        scratch_disks: Optional[Dict[str, int]] = None,
     ) -> "_BaseJob":
         """Run async workloads using a docker image or a compute environment from your studio.
 
@@ -236,6 +237,7 @@ class _BaseJob(ABC, metaclass=TrackCallsABCMeta):
             path_mappings=path_mappings,
             max_runtime=max_runtime,
             reuse_snapshot=reuse_snapshot,
+            scratch_disks=scratch_disks,
         )
 
     @abstractmethod
@@ -257,6 +259,7 @@ class _BaseJob(ABC, metaclass=TrackCallsABCMeta):
         path_mappings: Optional[Dict[str, str]] = None,
         max_runtime: Optional[int] = None,
         reuse_snapshot: bool = True,
+        scratch_disks: Optional[Dict[str, int]] = None,
     ) -> "_BaseJob":
         """Submit a new job to the Lightning AI platform.
 

lightning_sdk/job/job.py

@@ -104,6 +104,7 @@ class Job(_BaseJob):
         artifacts_remote: Optional[str] = None,  # deprecated in terms of path_mappings
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
         reuse_snapshot: bool = True,
+        scratch_disks: Optional[Dict[str, int]] = None,
     ) -> "Job":
         """Run async workloads using a docker image or a compute environment from your studio.
 
@@ -172,6 +173,7 @@ class Job(_BaseJob):
             max_runtime=max_runtime,
             cluster=cluster,
             reuse_snapshot=reuse_snapshot,
+            scratch_disks=scratch_disks,
         )
         # required for typing with "Job"
         assert isinstance(ret_val, cls)
@@ -197,6 +199,7 @@ class Job(_BaseJob):
         artifacts_remote: Optional[str] = None,  # deprecated in terms of path_mappings
         max_runtime: Optional[int] = None,
         reuse_snapshot: bool = True,
+        scratch_disks: Optional[Dict[str, int]] = None,
     ) -> "Job":
         """Submit a new job to the Lightning AI platform.
 
@@ -255,6 +258,7 @@ class Job(_BaseJob):
             artifacts_remote=artifacts_remote,
             max_runtime=max_runtime,
             reuse_snapshot=reuse_snapshot,
+            scratch_disks=scratch_disks,
         )
         return self
 

lightning_sdk/job/v1.py

@@ -55,6 +55,7 @@ class _JobV1(_BaseJob):
         interruptible: bool = False,
         cluster: Optional[str] = None,  # deprecated in favor of cloud_account
         reuse_snapshot: bool = True,
+        scratch_disks: Optional[Dict[str, int]] = None,
     ) -> "_BaseJob":
         """Start a new async workload from your studio.
 
@@ -111,6 +112,7 @@ class _JobV1(_BaseJob):
         path_mappings: Optional[Dict[str, str]] = None,
         max_runtime: Optional[int] = None,
         reuse_snapshot: bool = True,
+        scratch_disks: Optional[Dict[str, int]] = None,
     ) -> "_JobV1":
         """Submit a job to run on a machine.
 

lightning_sdk/job/v2.py

@@ -1,3 +1,4 @@
+from pathlib import PurePath
 from typing import TYPE_CHECKING, Dict, Optional, Union
 
 from lightning_sdk.api.job_api import JobApiV2
@@ -54,6 +55,7 @@ class _JobV2(_BaseJob):
         artifacts_local: Optional[str] = None,  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str] = None,  # deprecated in favor of path_mappings
         reuse_snapshot: bool = True,
+        scratch_disks: Optional[Dict[str, int]] = None,
     ) -> "_JobV2":
         """Submit a new job to the Lightning AI platform.
 
@@ -90,6 +92,11 @@ class _JobV2(_BaseJob):
                 Defaults to 3h
             reuse_snapshot: Whether the job should reuse a Studio snapshot when multiple jobs for the same Studio are
                 submitted. Turning this off may result in longer job startup times. Defaults to True.
+            scratch_disks: Dictionary of scratch disks to add to the job. The keys are the path that the disk
+                will be mounted to, relative to the /teamspace/scratch directory. The values are the size of
+                the volume in GiB. For example, { "data": 100 } will add a 100GiB volume available under
+                /teamspace/scratch/data.
+
         """
         # Command is required if Studio is provided to know what to run
         # Image is mutually exclusive with Studio
@@ -115,6 +122,30 @@ class _JobV2(_BaseJob):
             default_cloud_account=self._teamspace.default_cloud_account,
         )
 
+        if scratch_disks is not None and len(scratch_disks) > 0:
+            if studio is None:
+                raise ValueError("scratch_disks are only supported within a studio job")
+
+            if len(scratch_disks) > 5:
+                raise ValueError("scratch_disk may only contain up to 5 elements")
+
+            for path, size in scratch_disks.items():
+                if size > 50000:
+                    raise ValueError("scratch_disk size cannot exceed 50TiB")
+
+                path = PurePath(path)
+
+                if path.is_absolute():
+                    # For compatibility with Python 3.8, which doesn't provide
+                    # pathlib.PurePath.is_relative_to.
+                    try:
+                        path.relative_to("/teamspace/scratch")
+                    except ValueError:
+                        raise ValueError("scratch_disk paths must be relative to /teamspace/scratch") from None
+
+                if ".." in path.parts:
+                    raise ValueError("scratch_disk path cannot contain '..'")
+
         submitted = self._job_api.submit_job(
             name=self.name,
             command=command,
@@ -133,6 +164,7 @@ class _JobV2(_BaseJob):
             path_mappings=path_mappings,
             max_runtime=max_runtime,
             reuse_snapshot=reuse_snapshot,
+            scratch_disks=scratch_disks,
         )
         self._job = submitted
         self._name = submitted.name

lightning_sdk/lightning_cloud/login.py

@@ -268,7 +268,7 @@ class Auth:
         self.save(user_id=user_id, api_key=api_key)
         logger.info("Successfully authenticated as a guest user.")
 
-        return self.auth_header
+        return self.authenticate()
 
     def token_login(self, token_key: str, save_token: bool = True) -> Optional[str]:
         """Performs token-based authentication.