lightning-sdk 2025.8.6rc2__py3-none-any.whl → 2025.11.5__py3-none-any.whl

lightning_sdk/api/studio_api.py

@@ -1,8 +1,7 @@
 import json
 import os
-import tempfile
 import time
-import zipfile
+from pathlib import Path
 from threading import Event, Thread
 from typing import Any, Dict, Generator, List, Mapping, Optional, Tuple, Union
 
@@ -12,6 +11,7 @@ from tqdm import tqdm
 
 from lightning_sdk.api.utils import (
     _create_app,
+    _download_teamspace_files,
     _DummyBody,
     _DummyResponse,
     _FileUploader,
@@ -25,10 +25,12 @@ from lightning_sdk.constants import _LIGHTNING_DEBUG
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
     CloudspaceIdRunsBody,
+    CloudspacesIdBody,
     Externalv1LightningappInstance,
     IdCodeconfigBody,
     IdExecuteBody1,
     IdForkBody1,
+    IdSleepconfigBody,
     IdStartBody,
     ProjectIdCloudspacesBody,
     V1Assistant,
@@ -39,6 +41,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1CloudSpaceState,
     V1ClusterAccelerator,
     V1EndpointType,
+    V1EnvVar,
     V1GetCloudSpaceInstanceStatusResponse,
     V1GetLongRunningCommandInCloudSpaceResponse,
     V1LoginRequest,
@@ -204,6 +207,32 @@ class StudioApi:
             instance_id = code_status.in_use.cloud_space_instance_id
             print(f"Studio started | {teamspace_id=} {studio_id=} {instance_id=}")
 
+    def start_studio_async(
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        machine: Union[Machine, str],
+        interruptible: bool = False,
+        max_runtime: Optional[int] = None,
+    ) -> None:
+        """Start an existing Studio without blocking."""
+        # need to go via kwargs for typing compatibility since autogenerated apis accept None but aren't typed with None
+        optional_kwargs_compute_body = {}
+
+        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(
+                compute_config=V1UserRequestedComputeConfig(
+                    name=_machine_to_compute_name(machine),
+                    spot=interruptible,
+                    **optional_kwargs_compute_body,
+                )
+            ),
+            teamspace_id,
+            studio_id,
+        )
+
     def stop_studio(self, studio_id: str, teamspace_id: str) -> None:
         """Stop an existing Studio."""
         self.stop_keeping_alive(teamspace_id=teamspace_id, studio_id=studio_id)
@@ -232,12 +261,19 @@ class StudioApi:
         return getattr(getattr(studio.code_status, "in_use", None), "phase", None)
 
     def _request_switch(
-        self, studio_id: str, teamspace_id: str, machine: Union[Machine, str], interruptible: bool
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        machine: Union[Machine, str],
+        interruptible: bool,
+        cloud_account: Optional[str],
     ) -> None:
         """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))
+        if cloud_account:
+            body.compute_config.cluster_override = cloud_account
         self._client.cloud_space_service_update_cloud_space_instance_config(
             id=studio_id,
             project_id=teamspace_id,
@@ -245,29 +281,136 @@ class StudioApi:
         )
 
     def switch_studio_machine(
-        self, studio_id: str, teamspace_id: str, machine: Union[Machine, str], interruptible: bool
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        machine: Union[Machine, str],
+        interruptible: bool,
+        cloud_account: Optional[str],
     ) -> None:
         """Switches given Studio to a new machine type."""
         self._request_switch(
-            studio_id=studio_id, teamspace_id=teamspace_id, machine=machine, interruptible=interruptible
+            studio_id=studio_id,
+            teamspace_id=teamspace_id,
+            machine=machine,
+            interruptible=interruptible,
+            cloud_account=cloud_account,
+        )
+
+        # Wait until it's time to switch
+        requested_was_found = False
+        startup_status = None
+        while True:
+            status = self.get_studio_status(studio_id, teamspace_id)
+            requested_machine = status.requested
+
+            if requested_machine is not None:
+                requested_was_found = True
+                startup_status = requested_machine.startup_status
+
+            # if the requested machine was found in the past, use the in_use status instead.
+            # it might be that it either was cancelled or it actually is ready.
+            # Either way, since we're actually blocking below for the in use startup status
+            # it's safe to switch at this point
+            elif requested_was_found:
+                in_use_machine = status.in_use
+                if in_use_machine is not None:
+                    startup_status = in_use_machine.startup_status
+
+            if startup_status and startup_status.initial_restore_finished:
+                break
+            time.sleep(1)
+
+        self._client.cloud_space_service_switch_cloud_space_instance(teamspace_id, studio_id)
+
+        # Wait until the new machine is ready to use
+        while True:
+            in_use = self.get_studio_status(studio_id, teamspace_id).in_use
+            if in_use is None:
+                continue
+            startup_status = in_use.startup_status
+            if startup_status and startup_status.top_up_restore_finished:
+                break
+            time.sleep(1)
+
+    def switch_studio_machine_with_progress(
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        machine: Union[Machine, str],
+        interruptible: bool,
+        progress: Any,  # StudioProgressTracker - avoid circular import
+        cloud_account: Optional[str],
+    ) -> None:
+        """Switches given Studio to a new machine type with progress tracking."""
+        progress.update_progress(10, "Requesting machine switch...")
+
+        self._request_switch(
+            studio_id=studio_id,
+            teamspace_id=teamspace_id,
+            machine=machine,
+            interruptible=interruptible,
+            cloud_account=cloud_account,
         )
 
+        progress.update_progress(20, "Waiting for machine allocation...")
+
         # Wait until it's time to switch
+        requested_was_found = False
+        startup_status = None
+        base_progress = 20
+        max_wait_progress = 60
+        wait_counter = 0
+
         while True:
-            startup_status = self.get_studio_status(studio_id, teamspace_id).requested.startup_status
+            status = self.get_studio_status(studio_id, teamspace_id)
+            requested_machine = status.requested
+
+            if requested_machine is not None:
+                requested_was_found = True
+                startup_status = requested_machine.startup_status
+
+            # if the requested machine was found in the past, use the in_use status instead.
+            # it might be that it either was cancelled or it actually is ready.
+            # Either way, since we're actually blocking below for the in use startup status
+            # it's safe to switch at this point
+            elif requested_was_found:
+                in_use_machine = status.in_use
+                if in_use_machine is not None:
+                    startup_status = in_use_machine.startup_status
+
             if startup_status and startup_status.initial_restore_finished:
                 break
+
+            # Update progress gradually while waiting
+            wait_counter += 1
+            current_progress = min(base_progress + (wait_counter * 2), max_wait_progress)
+            progress.update_progress(current_progress, "Allocating new machine...")
             time.sleep(1)
 
+        progress.update_progress(70, "Starting machine switch...")
         self._client.cloud_space_service_switch_cloud_space_instance(teamspace_id, studio_id)
 
+        progress.update_progress(80, "Configuring new machine...")
+
         # Wait until the new machine is ready to use
+        switch_counter = 0
         while True:
-            startup_status = self.get_studio_status(studio_id, teamspace_id).in_use.startup_status
+            in_use = self.get_studio_status(studio_id, teamspace_id).in_use
+            if in_use is None:
+                continue
+            startup_status = in_use.startup_status
             if startup_status and startup_status.top_up_restore_finished:
                 break
+
+            # Update progress while waiting for machine to be ready
+            switch_counter += 1
+            current_progress = min(80 + switch_counter, 95)
+            progress.update_progress(current_progress, "Finalizing machine setup...")
             time.sleep(1)
 
+        progress.complete("Machine switch completed successfully")
+
     def get_machine(self, studio_id: str, teamspace_id: str, cloud_account_id: str, org_id: str) -> Machine:
         """Get the current machine type the given Studio is running on."""
         response: V1CloudSpaceInstanceConfig = self._client.cloud_space_service_get_cloud_space_instance_config(
@@ -283,7 +426,7 @@ class StudioApi:
                 accelerator.slug_multi_cloud,
                 accelerator.instance_id,
             ):
-                return Machine.from_str(accelerator.slug_multi_cloud)
+                return Machine._from_accelerator(accelerator)
 
         return Machine.from_str(response.compute_config.name)
 
@@ -295,6 +438,14 @@ class StudioApi:
 
         return response.compute_config.spot
 
+    def get_public_ip(self, studio_id: str, teamspace_id: str) -> Optional[str]:
+        """Get the public IP address of the Studio."""
+        internal_status = self.get_studio_status(studio_id=studio_id, teamspace_id=teamspace_id).in_use
+        if internal_status is None:
+            return None
+
+        return internal_status.public_ip_address
+
     def _get_machines_for_cloud_account(
         self, teamspace_id: str, cloud_account_id: str, org_id: str
     ) -> List[V1ClusterAccelerator]:
@@ -423,25 +574,28 @@ class StudioApi:
         self,
         studio_id: str,
         teamspace_id: str,
-        studio: V1CloudSpace,
         enabled: Optional[bool] = None,
-        idle_shutdown_seconds: int = 0,
-    ) -> None:
-        """Update the autoshutdown time of the given Studio."""
-        if enabled is None:
-            enabled = not studio.code_config.disable_auto_shutdown
-        body = IdCodeconfigBody(
-            disable_auto_shutdown=not enabled,
+        idle_shutdown_seconds: Optional[int] = None,
+    ) -> V1CloudSpaceInstanceConfig:
+        """Update the autoshutdown time and behaviour of the given Studio."""
+        body = IdSleepconfigBody(
+            disable_auto_shutdown=not enabled if enabled is not None else None,
             idle_shutdown_seconds=idle_shutdown_seconds,
-            compute_config=studio.code_config.compute_config,
         )
-        self._client.cloud_space_service_update_cloud_space_instance_config(
+        return self._client.cloud_space_service_update_cloud_space_sleep_config(
             id=studio_id,
             project_id=teamspace_id,
             body=body,
         )
 
-    def duplicate_studio(self, studio_id: str, teamspace_id: str, target_teamspace_id: str) -> Dict[str, Any]:
+    def duplicate_studio(
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        target_teamspace_id: str,
+        machine: Machine = Machine.CPU,
+        new_name: Optional[str] = None,
+    ) -> Dict[str, Any]:
         """Duplicates the given Studio from a given Teamspace into a given target Teamspace."""
         target_teamspace = self._client.projects_service_get_project(target_teamspace_id)
         init_kwargs = {}
@@ -455,7 +609,7 @@ 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), project_id=teamspace_id, id=studio_id
+            IdForkBody1(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:
@@ -464,7 +618,7 @@ class StudioApi:
         init_kwargs["name"] = new_cloudspace.name
         init_kwargs["teamspace"] = target_teamspace.name
 
-        self.start_studio(new_cloudspace.id, target_teamspace_id, Machine.CPU, False, None)
+        self.start_studio(new_cloudspace.id, target_teamspace_id, machine, False, None)
         return init_kwargs
 
     def delete_studio(self, studio_id: str, teamspace_id: str) -> None:
@@ -550,46 +704,24 @@ class StudioApi:
         progress_bar: bool = True,
     ) -> None:
         """Downloads a given folder from a Studio to a target location."""
-        # TODO: Update this endpoint to permit basic auth
+        # TODO: implement resumable downloads
         auth = Auth()
         auth.authenticate()
-        token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
 
-        query_params = {
-            "clusterId": cloud_account,
-            "prefix": _sanitize_studio_remote_path(path, studio_id),
-            "token": token,
-        }
+        prefix = _sanitize_studio_remote_path(path, studio_id)
+        # ensure we only download as a directory and not the entire prefix
+        if prefix.endswith("/") is False:
+            prefix = prefix + "/"
 
-        r = requests.get(
-            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/download",
-            params=query_params,
-            stream=True,
+        _download_teamspace_files(
+            client=self._client,
+            teamspace_id=teamspace_id,
+            cluster_id=cloud_account,
+            prefix=prefix,
+            download_dir=Path(target_path),
+            progress_bar=progress_bar,
         )
 
-        if progress_bar:
-            pbar = tqdm(
-                desc=f"Downloading {os.path.split(path)[1]}",
-                unit="B",
-                unit_scale=True,
-                unit_divisor=1000,
-            )
-
-            pbar_update = pbar.update
-        else:
-            pbar_update = lambda x: None
-
-        if target_path:
-            os.makedirs(target_path, exist_ok=True)
-
-        with tempfile.TemporaryFile() as f:
-            for chunk in r.iter_content(chunk_size=4096 * 8):
-                f.write(chunk)
-                pbar_update(len(chunk))
-
-            with zipfile.ZipFile(f) as z:
-                z.extractall(target_path)
-
     def install_plugin(self, studio_id: str, teamspace_id: str, plugin_name: str) -> str:
         """Installs the given plugin."""
         resp: V1Plugin = self._client.cloud_space_service_install_plugin(
@@ -844,3 +976,68 @@ class StudioApi:
             plugin_type=plugin_type,
             **other_arguments,
         )
+
+    def _update_cloudspace(self, studio: V1CloudSpace, teamspace_id: str, key: str, value: Any) -> None:
+        body = CloudspacesIdBody(
+            code_url=studio.code_url,
+            data_connection_mounts=studio.data_connection_mounts,
+            description=studio.description,
+            display_name=studio.display_name,
+            env=studio.env,
+            featured=studio.featured,
+            hide_files=studio.hide_files,
+            is_cloudspace_private=studio.is_cloudspace_private,
+            is_code_private=studio.is_code_private,
+            is_favorite=studio.is_favorite,
+            is_published=studio.is_published,
+            license=studio.license,
+            license_url=studio.license_url,
+            message=studio.message,
+            multi_user_edit=studio.multi_user_edit,
+            operating_cost=studio.operating_cost,
+            paper_authors=studio.paper_authors,
+            paper_org=studio.paper_org,
+            paper_org_avatar_url=studio.paper_org_avatar_url,
+            paper_url=studio.paper_url,
+            switch_to_default_machine_on_idle=studio.switch_to_default_machine_on_idle,
+            tags=studio.tags,
+            thumbnail_file_type=studio.thumbnail_file_type,
+            user_metadata=studio.user_metadata,
+        )
+
+        setattr(body, key, value)
+
+        self._client.cloud_space_service_update_cloud_space(
+            id=studio.id,
+            project_id=teamspace_id,
+            body=body,
+        )
+
+    def set_env(
+        self,
+        studio: V1CloudSpace,
+        teamspace_id: str,
+        new_env: Dict[str, str],
+        partial: bool = True,
+    ) -> None:
+        """Set the environment variables for the Studio.
+
+        Args:
+            new_env: The new environment variables to set.
+            partial: Whether to only set the environment variables that are provided.
+                If False, existing environment variables that are not in new_env will be removed.
+                If True, existing environment variables that are not in new_env will be kept.
+        """
+        updated_env_dict = {}
+        if partial:
+            updated_env_dict = {env.name: env.value for env in studio.env}
+            updated_env_dict.update(new_env)
+        else:
+            updated_env_dict = new_env
+
+        updated_env = [V1EnvVar(name=key, value=value) for key, value in updated_env_dict.items()]
+
+        self._update_cloudspace(studio, teamspace_id, "env", updated_env)
+
+    def get_env(self, studio: V1CloudSpace) -> Dict[str, str]:
+        return {env.name: env.value for env in studio.env}

lightning_sdk/api/teamspace_api.py

@@ -1,6 +1,5 @@
 import os
-import tempfile
-import zipfile
+import re
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
@@ -9,6 +8,7 @@ from tqdm.auto import tqdm
 
 from lightning_sdk.api.utils import (
     _download_model_files,
+    _download_teamspace_files,
     _DummyBody,
     _FileUploader,
     _get_model_version,
@@ -17,15 +17,21 @@ from lightning_sdk.api.utils import (
 )
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
+    Create,
     Externalv1LightningappInstance,
     ModelIdVersionsBody,
     ModelsStoreApi,
     ProjectIdAgentsBody,
     ProjectIdModelsBody,
+    ProjectIdSecretsBody,
+    SecretsIdBody,
     V1Assistant,
     V1CloudSpace,
     V1ClusterAccelerator,
+    V1EfsConfig,
     V1Endpoint,
+    V1ExternalCluster,
+    V1GCSFolderDataConnection,
     V1Job,
     V1LoginRequest,
     V1Model,
@@ -34,6 +40,10 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1Project,
     V1ProjectClusterBinding,
     V1PromptSuggestion,
+    V1R2DataConnection,
+    V1S3FolderDataConnection,
+    V1Secret,
+    V1SecretType,
     V1UpstreamOpenAI,
 )
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
@@ -63,7 +73,7 @@ class TeamspaceApi:
     def _get_teamspace_by_id(self, teamspace_id: str) -> V1Project:
         return self._client.projects_service_get_project(teamspace_id)
 
-    def list_teamspaces(self, owner_id: str, name: Optional[str] = None) -> Optional[V1Project]:
+    def list_teamspaces(self, owner_id: str, name: Optional[str] = None) -> Optional[List[V1Project]]:
         """Lists teamspaces from owner.
 
         If name is passed only teamspaces matching that name will be returned
@@ -386,6 +396,8 @@ class TeamspaceApi:
             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:
@@ -421,39 +433,103 @@ class TeamspaceApi:
         # TODO: Update this endpoint to permit basic auth
         auth = Auth()
         auth.authenticate()
-        token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
 
-        query_params = {
-            "clusterId": cloud_account,
-            "prefix": _resolve_teamspace_remote_path(path),
-            "token": token,
-        }
+        prefix = _resolve_teamspace_remote_path(path)
 
-        r = requests.get(
-            f"{self._client.api_client.configuration.host}/v1/projects/{teamspace_id}/artifacts/download",
-            params=query_params,
-            stream=True,
+        # ensure we only download as a directory and not the entire prefix
+        if prefix.endswith("/") is False:
+            prefix = prefix + "/"
+
+        _download_teamspace_files(
+            client=self._client,
+            teamspace_id=teamspace_id,
+            cluster_id=cloud_account,
+            prefix=prefix,
+            download_dir=Path(target_path),
+            progress_bar=progress_bar,
         )
 
-        if progress_bar:
-            pbar = tqdm(
-                desc=f"Downloading {os.path.split(path)[1]}",
-                unit="B",
-                unit_scale=True,
-                unit_divisor=1000,
-            )
+    def get_secrets(self, teamspace_id: str) -> Dict[str, str]:
+        """Get all secrets for a teamspace."""
+        secrets = self._get_secrets(teamspace_id)
+        # this returns encrypted values for security. It doesn't make sense to show them,
+        # so we just return a placeholder
+        # not a security issue to replace in the client as we get the encrypted values from the server.
+        return {secret.name: "***REDACTED***" for secret in secrets if secret.type == V1SecretType.UNSPECIFIED}
 
-            pbar_update = pbar.update
+    def set_secret(self, teamspace_id: str, key: str, value: str) -> None:
+        """Set a secret for a teamspace.
+
+        This will replace the existing secret if it exists and create a new one if it doesn't.
+        """
+        secrets = self._get_secrets(teamspace_id)
+        for secret in secrets:
+            if secret.name == key:
+                return self._update_secret(teamspace_id, secret.id, value)
+        return self._create_secret(teamspace_id, key, value)
+
+    def _get_secrets(self, teamspace_id: str) -> List[V1Secret]:
+        return self._client.secret_service_list_secrets(project_id=teamspace_id).secrets
+
+    def _update_secret(self, teamspace_id: str, secret_id: str, value: str) -> None:
+        self._client.secret_service_update_secret(
+            body=SecretsIdBody(value=value),
+            project_id=teamspace_id,
+            id=secret_id,
+        )
+
+    def _create_secret(
+        self,
+        teamspace_id: str,
+        key: str,
+        value: str,
+    ) -> None:
+        self._client.secret_service_create_secret(
+            body=ProjectIdSecretsBody(name=key, value=value, type=V1SecretType.UNSPECIFIED), project_id=teamspace_id
+        )
+
+    def verify_secret_name(self, name: str) -> bool:
+        """Verify if a secret name is valid.
+
+        A valid secret name starts with a letter or underscore, followed by letters, digits, or underscores.
+        """
+        pattern = r"^[A-Za-z_][A-Za-z0-9_]*$"
+        return re.match(pattern, name) is not None
+
+    def new_folder(self, teamspace_id: str, name: str, cluster: Optional[V1ExternalCluster]) -> None:
+        create_request = Create(
+            name=name,
+            create_resources=True,
+            force=True,
+            writable=True,
+        )
+
+        if cluster is None:
+            create_request.r2 = V1R2DataConnection(name=name)
         else:
-            pbar_update = lambda x: None
+            create_request.cluster_id = cluster.id
+            create_request.access_cluster_ids = [cluster.id]
 
-        if target_path:
-            os.makedirs(target_path, exist_ok=True)
+            if cluster.spec.aws_v1:
+                create_request.s3_folder = V1S3FolderDataConnection()
+            elif cluster.spec.google_cloud_v1:
+                create_request.gcs_folder = V1GCSFolderDataConnection()
 
-        with tempfile.TemporaryFile() as f:
-            for chunk in r.iter_content(chunk_size=4096 * 8):
-                f.write(chunk)
-                pbar_update(len(chunk))
+        self._client.data_connection_service_create_data_connection(create_request, teamspace_id)
+
+    def new_connection(
+        self, teamspace_id: str, name: str, source: str, cluster: V1ExternalCluster, writable: bool, region: str
+    ) -> None:
+        create_request = Create(
+            name=name,
+            create_resources=False,
+            force=True,
+            writable=writable,
+            cluster_id=cluster.id,
+            access_cluster_ids=[cluster.id],
+        )
+
+        # TODO: Add support for other connection types
+        create_request.efs = V1EfsConfig(file_system_id=source, region=region)
 
-            with zipfile.ZipFile(f) as z:
-                z.extractall(target_path)
+        self._client.data_connection_service_create_data_connection(create_request, teamspace_id)