lightning-sdk 0.1.3__py3-none-any.whl → 0.1.46__py3-none-any.whl

lightning_sdk/api/studio_api.py

@@ -1,20 +1,36 @@
+import json
 import os
 import tempfile
 import time
+import warnings
 import zipfile
 from threading import Event, Thread
-from typing import Any, Dict, Mapping, Optional, Tuple
-from uuid import uuid4
+from typing import Any, Dict, Mapping, Optional, Tuple, Union
 
+import backoff
 import requests
+from tqdm import tqdm
 
+from lightning_sdk.api.utils import (
+    _COMPUTE_NAME_TO_MACHINE,
+    _MACHINE_TO_COMPUTE_NAME,
+    _create_app,
+    _DummyBody,
+    _DummyResponse,
+    _FileUploader,
+    _machine_to_compute_name,
+    _sanitize_studio_remote_path,
+)
+from lightning_sdk.api.utils import (
+    _get_cloud_url as _cloud_url,
+)
 from lightning_sdk.constants import _LIGHTNING_DEBUG
 from lightning_sdk.lightning_cloud.login import Auth
 from lightning_sdk.lightning_cloud.openapi import (
     CloudspaceIdRunsBody,
     Externalv1LightningappInstance,
     IdCodeconfigBody,
-    IdExecuteBody,
+    IdExecuteBody1,
     IdForkBody1,
     IdStartBody,
     ProjectIdCloudspacesBody,
@@ -29,27 +45,6 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1PluginsListResponse,
     V1UserRequestedComputeConfig,
 )
-
-try:
-    from lightning_sdk.lightning_cloud.openapi import AppsIdBody1 as AppsIdBody
-except ImportError:
-    from lightning_sdk.lightning_cloud.openapi import AppsIdBody
-
-import json
-
-from tqdm import tqdm
-
-from lightning_sdk.api.utils import (
-    _COMPUTE_NAME_TO_MACHINE,
-    _MACHINE_TO_COMPUTE_NAME,
-    _DummyBody,
-    _DummyResponse,
-    _FileUploader,
-    _sanitize_remote_path,
-)
-from lightning_sdk.api.utils import (
-    _get_cloud_url as _cloud_url,
-)
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.machine import Machine
 
@@ -58,10 +53,8 @@ class StudioApi:
     """Internal API client for Studio requests (mainly http requests)."""
 
     def __init__(self) -> None:
-        super().__init__()
-
         self._cloud_url = _cloud_url()
-        self._client = LightningClient(max_tries=3)
+        self._client = LightningClient(max_tries=7)
         self._keep_alive_threads: Mapping[str, Thread] = {}
         self._keep_alive_events: Mapping[str, Event] = {}
 
@@ -115,11 +108,11 @@ class StudioApi:
         self,
         name: str,
         teamspace_id: str,
-        cluster: Optional[str] = None,
+        cloud_account: Optional[str] = None,
     ) -> V1CloudSpace:
-        """Create a Studio with a given name in a given Teamspace on a possibly given cluster."""
+        """Create a Studio with a given name in a given Teamspace on a possibly given cloud_account."""
         body = ProjectIdCloudspacesBody(
-            cluster_id=cluster,
+            cluster_id=cloud_account,
             name=name,
             display_name=name,
             seed_files=[V1CloudSpaceSeedFile(path="main.py", contents="print('Hello, Lightning World!')\n")],
@@ -143,20 +136,42 @@ class StudioApi:
             id=studio_id,
         )
 
-    def start_studio(self, studio_id: str, teamspace_id: str, machine: Machine) -> None:
+    @backoff.on_exception(backoff.expo, AttributeError, max_tries=10)
+    def _check_code_status_top_up_restore_finished(self, studio_id: str, teamspace_id: str) -> bool:
+        """Retries checking the top_up_restore_finished value of the code status when there's an AttributeError."""
+        startup_status = self.get_studio_status(studio_id, teamspace_id).in_use.startup_status
+        return startup_status and startup_status.top_up_restore_finished
+
+    @backoff.on_exception(backoff.expo, AttributeError, max_tries=10)
+    def _check_code_status_sync_in_progress(self, studio_id: str, teamspace_id: str) -> bool:
+        """Retries checking the sync_in_progress value of the code status when there's an AttributeError."""
+        return self.get_studio_status(studio_id, teamspace_id).in_use.sync_in_progress
+
+    def start_studio(
+        self, studio_id: str, teamspace_id: str, machine: Union[Machine, str], interruptible: False
+    ) -> None:
         """Start an existing Studio."""
+        if _machine_to_compute_name(machine) == _machine_to_compute_name(Machine.CPU_SMALL):
+            warnings.warn(
+                f"{Machine.CPU_SMALL} is not a valid machine for starting a Studio. "
+                "It is reserved for running jobs only. "
+                "The Studio will be started with a CPU machine instead."
+            )
+            machine = Machine.CPU
+
         self._client.cloud_space_service_start_cloud_space_instance(
-            IdStartBody(compute_config=V1UserRequestedComputeConfig(name=_MACHINE_TO_COMPUTE_NAME[machine])),
+            IdStartBody(
+                compute_config=V1UserRequestedComputeConfig(name=_machine_to_compute_name(machine), spot=interruptible)
+            ),
             teamspace_id,
             studio_id,
         )
 
-        while self.get_studio_status(studio_id, teamspace_id).in_use.sync_in_progress:
+        while self._check_code_status_sync_in_progress(studio_id, teamspace_id):
             time.sleep(1)
 
         while True:
-            startup_status = self.get_studio_status(studio_id, teamspace_id).in_use.startup_status
-            if startup_status and startup_status.top_up_restore_finished:
+            if self._check_code_status_top_up_restore_finished(studio_id, teamspace_id):
                 break
             time.sleep(1)
 
@@ -192,20 +207,34 @@ class StudioApi:
     def _get_studio_instance_status_from_object(self, studio: V1CloudSpace) -> Optional[str]:
         return getattr(getattr(studio.code_status, "in_use", None), "phase", None)
 
-    def _request_switch(self, studio_id: str, teamspace_id: str, machine: Machine) -> None:
+    def _request_switch(
+        self, studio_id: str, teamspace_id: str, machine: Union[Machine, str], interruptible: bool
+    ) -> None:
         """Switches given Studio to a new machine type."""
-        compute_name = _MACHINE_TO_COMPUTE_NAME[machine]
+        if _machine_to_compute_name(machine) == _machine_to_compute_name(Machine.CPU_SMALL):
+            warnings.warn(
+                f"{Machine.CPU_SMALL} is not a valid machine for switching a Studio. "
+                "It is reserved for running jobs only. "
+                "The Studio will be switched to a CPU machine instead."
+            )
+            machine = Machine.CPU
+
+        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))
+        body = IdCodeconfigBody(compute_config=V1UserRequestedComputeConfig(name=compute_name, spot=interruptible))
         self._client.cloud_space_service_update_cloud_space_instance_config(
             id=studio_id,
             project_id=teamspace_id,
             body=body,
         )
 
-    def switch_studio_machine(self, studio_id: str, teamspace_id: str, machine: Machine) -> None:
+    def switch_studio_machine(
+        self, studio_id: str, teamspace_id: str, machine: Union[Machine, str], interruptible: bool
+    ) -> None:
         """Switches given Studio to a new machine type."""
-        self._request_switch(studio_id=studio_id, teamspace_id=teamspace_id, machine=machine)
+        self._request_switch(
+            studio_id=studio_id, teamspace_id=teamspace_id, machine=machine, interruptible=interruptible
+        )
 
         # Wait until it's time to switch
         while True:
@@ -230,6 +259,14 @@ class StudioApi:
         )
         return _COMPUTE_NAME_TO_MACHINE[response.compute_config.name]
 
+    def get_interruptible(self, studio_id: str, teamspace_id: str) -> bool:
+        """Get whether the Studio is running on a interruptible instance."""
+        response: V1CloudSpaceInstanceConfig = self._client.cloud_space_service_get_cloud_space_instance_config(
+            project_id=teamspace_id, id=studio_id
+        )
+
+        return response.compute_config.spot
+
     def _get_detached_command_status(
         self, studio_id: str, teamspace_id: str, session_id: str
     ) -> V1GetLongRunningCommandInCloudSpaceResponse:
@@ -257,9 +294,8 @@ 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."""
-        session_id = str(uuid4())
         response_submit = self._client.cloud_space_service_execute_command_in_cloud_space(
-            IdExecuteBody("; ".join(commands), detached=True, session_name=session_id),
+            IdExecuteBody1("; ".join(commands), detached=True),
             project_id=teamspace_id,
             id=studio_id,
         )
@@ -267,28 +303,54 @@ class StudioApi:
         if not response_submit:
             raise RuntimeError("Unable to submit command")
 
+        if response_submit.session_name == "":
+            raise RuntimeError("The session name should be defined.")
+
         while True:
             output = ""
             exit_code = None
 
             for resp in self._get_detached_command_status(
-                studio_id=studio_id, teamspace_id=teamspace_id, session_id=session_id
+                studio_id=studio_id,
+                teamspace_id=teamspace_id,
+                session_id=response_submit.session_name,
             ):
-                if resp.exit_code != -1:
-                    if exit_code is None:
-                        exit_code = resp.exit_code
-                    elif exit_code != resp.exit_code:
-                        raise RuntimeError("Cannot determine exit code")
-
-                    output += resp.output
-                else:
+                if resp.exit_code == -1:
                     break
+                if exit_code is None:
+                    exit_code = resp.exit_code
+                elif exit_code != resp.exit_code:
+                    raise RuntimeError("Cannot determine exit code")
+
+                output += resp.output
 
             if exit_code is not None:
                 return output, exit_code
 
             time.sleep(1)
 
+    def update_autoshutdown(
+        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=idle_shutdown_seconds,
+            compute_config=studio.code_config.compute_config,
+        )
+        self._client.cloud_space_service_update_cloud_space_instance_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, str]:
         """Duplicates the given Studio from a given Teamspace into a given target Teamspace."""
         target_teamspace = self._client.projects_service_get_project(target_teamspace_id)
@@ -312,7 +374,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)
+        self.start_studio(new_cloudspace.id, target_teamspace_id, Machine.CPU, False)
         return init_kwargs
 
     def delete_studio(self, studio_id: str, teamspace_id: str) -> None:
@@ -321,21 +383,32 @@ class StudioApi:
         self._client.cloud_space_service_delete_cloud_space(project_id=teamspace_id, id=studio_id)
 
     def upload_file(
-        self, studio_id: str, teamspace_id: str, cluster_id: str, file_path: str, remote_path: str, progress_bar: bool
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        cloud_account: str,
+        file_path: str,
+        remote_path: str,
+        progress_bar: bool,
     ) -> None:
         """Uploads file to given remote path on the studio."""
         _FileUploader(
             client=self._client,
-            studio_id=studio_id,
             teamspace_id=teamspace_id,
-            cluster_id=cluster_id,
+            cloud_account=cloud_account,
             file_path=file_path,
-            remote_path=remote_path,
+            remote_path=_sanitize_studio_remote_path(remote_path, studio_id),
             progress_bar=progress_bar,
         )()
 
     def download_file(
-        self, path: str, target_path: str, studio_id: str, teamspace_id: str, cluster_id: str, progress_bar: bool = True
+        self,
+        path: str,
+        target_path: str,
+        studio_id: str,
+        teamspace_id: str,
+        cloud_account: str,
+        progress_bar: bool = True,
     ) -> None:
         """Downloads a given file from a Studio to a target location."""
         # TODO: Update this endpoint to permit basic auth
@@ -344,8 +417,8 @@ class StudioApi:
         token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
 
         query_params = {
-            "clusterId": cluster_id,
-            "key": _sanitize_remote_path(path, studio_id),
+            "clusterId": cloud_account,
+            "key": _sanitize_studio_remote_path(path, studio_id),
             "token": token,
         }
 
@@ -378,7 +451,13 @@ class StudioApi:
                 pbar_update(len(chunk))
 
     def download_folder(
-        self, path: str, target_path: str, studio_id: str, teamspace_id: str, cluster_id: str, progress_bar: bool = True
+        self,
+        path: str,
+        target_path: str,
+        studio_id: str,
+        teamspace_id: str,
+        cloud_account: str,
+        progress_bar: bool = True,
     ) -> None:
         """Downloads a given folder from a Studio to a target location."""
         # TODO: Update this endpoint to permit basic auth
@@ -387,8 +466,8 @@ class StudioApi:
         token = self._client.auth_service_login(V1LoginRequest(auth.api_key)).token
 
         query_params = {
-            "clusterId": cluster_id,
-            "prefix": _sanitize_remote_path(path, studio_id),
+            "clusterId": cloud_account,
+            "prefix": _sanitize_studio_remote_path(path, studio_id),
             "token": token,
         }
 
@@ -486,17 +565,25 @@ class StudioApi:
         return resp.plugins
 
     def create_job(
-        self, entrypoint: str, name: str, machine: Machine, studio_id: str, teamspace_id: str, cluster_id: str
+        self,
+        entrypoint: str,
+        name: str,
+        machine: Machine,
+        studio_id: str,
+        teamspace_id: str,
+        cloud_account: str,
+        interruptible: bool,
     ) -> Externalv1LightningappInstance:
         """Creates a job with given commands."""
         return self._create_app(
             studio_id=studio_id,
             teamspace_id=teamspace_id,
-            cluster_id=cluster_id,
+            cloud_account=cloud_account,
             plugin_type="job",
             entrypoint=entrypoint,
             name=name,
             compute=_MACHINE_TO_COMPUTE_NAME[machine],
+            interruptible=interruptible,
         )
 
     def create_multi_machine_job(
@@ -508,7 +595,8 @@ class StudioApi:
         strategy: str,
         studio_id: str,
         teamspace_id: str,
-        cluster_id: str,
+        cloud_account: str,
+        interruptible: bool,
     ) -> Externalv1LightningappInstance:
         """Creates a multi-machine job with given commands."""
         distributed_args = {
@@ -519,11 +607,12 @@ class StudioApi:
         return self._create_app(
             studio_id=studio_id,
             teamspace_id=teamspace_id,
-            cluster_id=cluster_id,
+            cloud_account=cloud_account,
             plugin_type="distributed_plugin",
             entrypoint=entrypoint,
             name=name,
             distributedArguments=json.dumps(distributed_args),
+            interruptible=interruptible,
         )
 
     def create_data_prep_machine_job(
@@ -534,7 +623,8 @@ class StudioApi:
         machine: Machine,
         studio_id: str,
         teamspace_id: str,
-        cluster_id: str,
+        cloud_account: str,
+        interruptible: bool,
     ) -> Externalv1LightningappInstance:
         """Creates a multi-machine job with given commands."""
         data_prep_args = {
@@ -544,11 +634,12 @@ class StudioApi:
         return self._create_app(
             studio_id=studio_id,
             teamspace_id=teamspace_id,
-            cluster_id=cluster_id,
+            cloud_account=cloud_account,
             plugin_type="litdata",
             entrypoint=entrypoint,
             name=name,
             dataPrepArguments=json.dumps(data_prep_args),
+            interruptible=interruptible,
         )
 
     def create_inference_job(
@@ -565,13 +656,14 @@ class StudioApi:
         endpoint: str,
         studio_id: str,
         teamspace_id: str,
-        cluster_id: str,
+        cloud_account: str,
+        interruptible: bool,
     ) -> Externalv1LightningappInstance:
         """Creates an inference job for given endpoint."""
         return self._create_app(
             studio_id=studio_id,
             teamspace_id=teamspace_id,
-            cluster_id=cluster_id,
+            cloud_account=cloud_account,
             plugin_type="inference_plugin",
             compute=_MACHINE_TO_COMPUTE_NAME[machine],
             entrypoint=entrypoint,
@@ -583,19 +675,18 @@ class StudioApi:
             scale_in_interval=scale_in_interval,
             scale_out_interval=scale_out_interval,
             endpoint=endpoint,
+            interruptible=interruptible,
         )
 
     def _create_app(
-        self, studio_id: str, teamspace_id: str, cluster_id: str, plugin_type: str, **other_arguments: Any
+        self, studio_id: str, teamspace_id: str, cloud_account: str, plugin_type: str, **other_arguments: Any
     ) -> Externalv1LightningappInstance:
         """Creates an arbitrary app."""
-        body = AppsIdBody(cluster_id=cluster_id, plugin_arguments=other_arguments)
-
-        resp = self._client.cloud_space_service_create_cloud_space_app_instance(
-            body=body, project_id=teamspace_id, cloudspace_id=studio_id, id=plugin_type
-        ).lightningappinstance
-
-        if _LIGHTNING_DEBUG:
-            print(f"Create App: {resp.id=} {teamspace_id=} {studio_id=} {cluster_id=}")
-
-        return resp
+        return _create_app(
+            self._client,
+            studio_id=studio_id,
+            teamspace_id=teamspace_id,
+            cloud_account=cloud_account,
+            plugin_type=plugin_type,
+            **other_arguments,
+        )