lightning-sdk 0.1.54__py3-none-any.whl → 0.1.55__py3-none-any.whl

lightning_sdk/__init__.py

@@ -29,5 +29,5 @@ __all__ = [
     "AIHub",
 ]
 
-__version__ = "0.1.54"
+__version__ = "0.1.55"
 _check_version_and_prompt_upgrade(__version__)

lightning_sdk/api/lit_container_api.py

@@ -1,6 +1,7 @@
-from typing import Generator, List
+from typing import Any, Callable, Generator, List
 
 import docker
+import requests
 
 from lightning_sdk.api.utils import _get_registry_url
 from lightning_sdk.lightning_cloud.env import LIGHTNING_CLOUD_URL
@@ -11,7 +12,22 @@ from lightning_sdk.teamspace import Teamspace
 
 class LCRAuthFailedError(Exception):
     def __init__(self) -> None:
-        super().__init__("Failed to authenticate with Lightning Container Registry")
+        super().__init__(
+            "Failed to authenticate with Lightning Container Registry. Please login manually "
+            "using the following command:\n  "
+            "echo $LIGHTNING_API_KEY | docker login litcr.io --username=LIGHTNING_USERNAME --password-stdin"
+        )
+
+
+def retry_on_lcr_auth_failure(func: Callable) -> Callable:
+    def wrapper(self: "LitContainerApi", *args: Any, **kwargs: Any) -> Callable:
+        try:
+            return func(self, *args, **kwargs)
+        except LCRAuthFailedError:
+            self.authenticate()
+            return func(self, *args, **kwargs)
+
+    return wrapper
 
 
 class LitContainerApi:
@@ -38,10 +54,10 @@ class LitContainerApi:
     def delete_container(self, project_id: str, container: str) -> V1DeleteLitRepositoryResponse:
         try:
             return self._client.lit_registry_service_delete_lit_repository(project_id, container)
-        except Exception as ex:
-            raise ValueError(f"Could not delete container {container} from project {project_id}") from ex
+        except Exception as e:
+            raise ValueError(f"Could not delete container {container} from project {project_id}: {e!s}") from e
 
-    def upload_container(self, container: str, teamspace: Teamspace, tag: str) -> Generator[str, None, None]:
+    def upload_container(self, container: str, teamspace: Teamspace, tag: str) -> Generator[dict, None, None]:
         try:
             self._docker_client.images.get(container)
         except docker.errors.ImageNotFound:
@@ -55,19 +71,23 @@ class LitContainerApi:
             raise ValueError(f"Could not tag container {container} with {repository}:{tag}")
         lines = self._docker_client.api.push(repository, stream=True, decode=True)
         for line in lines:
-            if "errorDetail" in line and "authorization failed" in line["error"]:
+            if "errorDetail" in line and ("authorization failed" in line["error"] or "unauth" in line["error"]):
                 raise LCRAuthFailedError()
             yield line
         yield {
             "finish": True,
-            "url": f"{LIGHTNING_CLOUD_URL}/{teamspace.owner.name}/{teamspace.name}/containers/{container}",
+            "url": f"{LIGHTNING_CLOUD_URL}/{teamspace.owner.name}/{teamspace.name}/containers/{container_basename}",
         }
 
+    @retry_on_lcr_auth_failure
     def download_container(self, container: str, teamspace: Teamspace, tag: str) -> Generator[str, None, None]:
         registry_url = _get_registry_url()
         repository = f"{registry_url}/lit-container/{teamspace.owner.name}/{teamspace.name}/{container}"
         try:
             self._docker_client.images.pull(repository, tag=tag)
+        except requests.exceptions.HTTPError as e:
+            if "unauthorized" in e.response.text:
+                raise LCRAuthFailedError() from e
         except docker.errors.APIError as e:
             raise ValueError(f"Could not pull container {container} from {repository}:{tag}") from e
         return self._docker_client.api.tag(repository, container, tag)

lightning_sdk/cli/delete.py

@@ -4,6 +4,7 @@ from lightning_sdk.cli.exceptions import StudioCliError
 from lightning_sdk.cli.job_and_mmt_action import _JobAndMMTAction
 from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
 from lightning_sdk.lit_container import LitContainer
+from lightning_sdk.studio import Studio
 
 
 class _Delete(_JobAndMMTAction, _TeamspacesMenu):
@@ -56,3 +57,29 @@ class _Delete(_JobAndMMTAction, _TeamspacesMenu):
 
         mmt.delete()
         print(f"Successfully deleted {mmt.name}!")
+
+    def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+        """Delete an existing studio.
+
+        Args:
+            name: The name of the studio to delete.
+                If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
+                Note: This could delete your current studio if run without arguments.
+            teamspace: The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.
+                If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
+        """
+        if teamspace is not None:
+            ts_splits = teamspace.split("/")
+            if len(ts_splits) != 2:
+                raise ValueError(f"Teamspace should be of format <OWNER>/<TEAMSPACE_NAME> but got {teamspace}")
+            owner, teamspace = ts_splits
+        else:
+            owner, teamspace = None, None
+
+        try:
+            studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False)
+        except (RuntimeError, ValueError):
+            studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
+
+        studio.delete()
+        print("Studio successfully deleted")

lightning_sdk/cli/entrypoint.py

@@ -9,12 +9,15 @@ from lightning_sdk.api.studio_api import _cloud_url
 from lightning_sdk.cli.ai_hub import _AIHub
 from lightning_sdk.cli.delete import _Delete
 from lightning_sdk.cli.download import _Downloads
+from lightning_sdk.cli.generate import _Generate
 from lightning_sdk.cli.inspect import _Inspect
 from lightning_sdk.cli.legacy import _LegacyLightningCLI
 from lightning_sdk.cli.list import _List
 from lightning_sdk.cli.run import _Run
 from lightning_sdk.cli.serve import _Docker, _LitServe
+from lightning_sdk.cli.start import _Start
 from lightning_sdk.cli.stop import _Stop
+from lightning_sdk.cli.switch import _Switch
 from lightning_sdk.cli.upload import _Uploads
 from lightning_sdk.lightning_cloud.login import Auth
 
@@ -35,6 +38,9 @@ class StudioCLI:
         self.delete = _Delete()
         self.inspect = _Inspect()
         self.stop = _Stop()
+        self.start = _Start()
+        self.switch = _Switch()
+        self.generate = _Generate()
 
         sys.excepthook = _notify_exception
 

lightning_sdk/cli/generate.py

@@ -0,0 +1,58 @@
+from typing import Optional
+
+from rich.console import Console
+
+from lightning_sdk import Studio
+
+
+class _Generate:
+    """Generate configs (such as ssh for studio) and print them to commandline."""
+
+    console = Console()
+
+    def _generate_ssh_config(self, name: str, studio_id: str) -> str:
+        """Generate SSH config entry for the studio.
+
+        Args:
+            name: Studio name
+            studio_id: Studio space ID
+
+        Returns:
+            str: SSH config entry
+        """
+        return f"""# ssh s_{studio_id}@ssh.lightning.ai
+
+Host {name}
+  User s_{studio_id}
+  Hostname ssh.lightning.ai
+  IdentityFile ~/.ssh/lightning_rsa
+  IdentitiesOnly yes
+  ServerAliveInterval 15
+  ServerAliveCountMax 4
+  StrictHostKeyChecking no
+  UserKnownHostsFile=/dev/null"""
+
+    def ssh(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+        """Get SSH config entry for a studio. Will start the studio if needed.
+
+        Args:
+            name: The name of the studio to stop.
+                If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
+            teamspace: The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.
+                If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
+        """
+        if teamspace:
+            ts_splits = teamspace.split("/")
+            if len(ts_splits) != 2:
+                raise ValueError(f"Teamspace should be of format <OWNER>/<TEAMSPACE_NAME> but got {teamspace}")
+            owner, teamspace = ts_splits
+        else:
+            owner, teamspace = None, None
+
+        try:
+            studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False)
+        except (RuntimeError, ValueError):
+            studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
+
+        # Print the SSH config
+        self.console.print(self._generate_ssh_config(name, studio._studio.id))

lightning_sdk/cli/list.py

@@ -3,6 +3,7 @@ from typing import Optional
 from rich.console import Console
 from rich.table import Table
 
+from lightning_sdk import Machine
 from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
 from lightning_sdk.lit_container import LitContainer
 
@@ -10,6 +11,38 @@ from lightning_sdk.lit_container import LitContainer
 class _List(_TeamspacesMenu):
     """List resources on the Lightning AI platform."""
 
+    def studios(self, teamspace: Optional[str] = None) -> None:
+        """List studios for a given teamspace.
+
+        Args:
+            teamspace: the teamspace to list studios from. Should be specified as {owner}/{name}
+                If not provided, can be selected in an interactive menu.
+
+        """
+        resolved_teamspace = self._resolve_teamspace(teamspace=teamspace)
+
+        studios = resolved_teamspace.studios
+
+        table = Table(
+            pad_edge=True,
+        )
+        table.add_column("Name")
+        table.add_column("Teamspace")
+        table.add_column("Status")
+        table.add_column("Machine")
+        table.add_column("Cloud account")
+        for studio in studios:
+            table.add_row(
+                studio.name,
+                f"{studio.teamspace.owner.name}/{studio.teamspace.name}",
+                str(studio.status),
+                str(studio.machine) if studio.machine is not None else None,
+                str(studio.cloud_account),
+            )
+
+        console = Console()
+        console.print(table)
+
     def jobs(self, teamspace: Optional[str] = None) -> None:
         """List jobs for a given teamspace.
 
@@ -110,3 +143,18 @@ class _List(_TeamspacesMenu):
             table.add_row(repo["REPOSITORY"], repo["IMAGE ID"], repo["CREATED"])
         console = Console()
         console.print(table)
+
+    def machines(self) -> None:
+        """Display the list of available machines."""
+        table = Table(pad_edge=True)
+        table.add_column("Name")
+
+        # Get all machine types from the enum
+        machine_types = [name for name in dir(Machine) if not name.startswith("_")]
+
+        # Add rows to table
+        for name in sorted(machine_types):
+            table.add_row(name)
+
+        console = Console()
+        console.print(table)

lightning_sdk/cli/start.py

@@ -0,0 +1,43 @@
+from typing import Optional
+
+from lightning_sdk import Machine, Studio
+
+
+class _Start:
+    """Start resources on the Lightning AI platform."""
+
+    def __init__(self) -> None:
+        _machine_values = tuple([machine.value for machine in Machine])
+
+        docstr_studio = f"""Start a studio on a given machine.
+
+        Args:
+            name: The name of the studio to start.
+                If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
+            teamspace: The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.
+                If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
+            machine: The machine type to start the studio on. One of {", ".join(_machine_values)}.
+                Defaults to the CPU Machine.
+        """
+        self.studio.__func__.__doc__ = docstr_studio
+
+    def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None, machine: str = "CPU") -> None:
+        if teamspace is not None:
+            ts_splits = teamspace.split("/")
+            if len(ts_splits) != 2:
+                raise ValueError(f"Teamspace should be of format <OWNER>/<TEAMSPACE_NAME> but got {teamspace}")
+            owner, teamspace = ts_splits
+        else:
+            owner, teamspace = None, None
+
+        try:
+            studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False)
+        except (RuntimeError, ValueError):
+            studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
+
+        try:
+            resolved_machine = Machine[machine.upper()]
+        except KeyError:
+            resolved_machine = machine
+
+        studio.start(resolved_machine)

lightning_sdk/cli/stop.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 from lightning_sdk.cli.job_and_mmt_action import _JobAndMMTAction
+from lightning_sdk.studio import Studio
 
 
 class _Stop(_JobAndMMTAction):
@@ -35,3 +36,28 @@ class _Stop(_JobAndMMTAction):
 
         mmt.stop()
         print(f"Successfully stopped {mmt.name}!")
+
+    def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+        """Stop a running studio.
+
+        Args:
+            name: The name of the studio to stop.
+                If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
+            teamspace: The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.
+                If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
+        """
+        if teamspace is not None:
+            ts_splits = teamspace.split("/")
+            if len(ts_splits) != 2:
+                raise ValueError(f"Teamspace should be of format <OWNER>/<TEAMSPACE_NAME> but got {teamspace}")
+            owner, teamspace = ts_splits
+        else:
+            owner, teamspace = None, None
+
+        try:
+            studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False)
+        except (RuntimeError, ValueError):
+            studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
+
+        studio.stop()
+        print("Studio successfully stopped")

lightning_sdk/cli/switch.py

@@ -0,0 +1,43 @@
+from typing import Optional
+
+from lightning_sdk import Machine, Studio
+
+
+class _Switch:
+    """Switch machines for resources on the Lightning AI platform."""
+
+    def __init__(self) -> None:
+        _machine_values = tuple([machine.value for machine in Machine])
+
+        docstr_studio = f"""Switch a studio to a given machine.
+
+        Args:
+            name: The name of the studio to start.
+                If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
+            teamspace: The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.
+                If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
+            machine: The machine type to switch to. One of {", ".join(_machine_values)}.
+                Defaults to the CPU Machine.
+        """
+        self.studio.__func__.__doc__ = docstr_studio
+
+    def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None, machine: str = "CPU") -> None:
+        if teamspace is not None:
+            ts_splits = teamspace.split("/")
+            if len(ts_splits) != 2:
+                raise ValueError(f"Teamspace should be of format <OWNER>/<TEAMSPACE_NAME> but got {teamspace}")
+            owner, teamspace = ts_splits
+        else:
+            owner, teamspace = None, None
+
+        try:
+            studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False)
+        except (RuntimeError, ValueError):
+            studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
+
+        try:
+            resolved_machine = Machine[machine.upper()]
+        except KeyError:
+            resolved_machine = machine
+
+        studio.switch_machine(resolved_machine)

lightning_sdk/lightning_cloud/openapi/__init__.py

@@ -571,6 +571,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_list_lightningapp_instances
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_lightningwork_events_response import V1ListLightningworkEventsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_lightningwork_response import V1ListLightningworkResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_lit_pages_response import V1ListLitPagesResponse
+from lightning_sdk.lightning_cloud.openapi.models.v1_list_lit_registry_repository_image_artifact_versions_response import V1ListLitRegistryRepositoryImageArtifactVersionsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_logger_artifact_response import V1ListLoggerArtifactResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_managed_endpoints_response import V1ListManagedEndpointsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_memberships_response import V1ListMembershipsResponse
@@ -606,6 +607,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_list_studio_jobs_response i
 from lightning_sdk.lightning_cloud.openapi.models.v1_list_user_slurm_jobs_response import V1ListUserSLURMJobsResponse
 from lightning_sdk.lightning_cloud.openapi.models.v1_lit_page import V1LitPage
 from lightning_sdk.lightning_cloud.openapi.models.v1_lit_page_type import V1LitPageType
+from lightning_sdk.lightning_cloud.openapi.models.v1_lit_registry_artifact import V1LitRegistryArtifact
 from lightning_sdk.lightning_cloud.openapi.models.v1_lit_registry_project import V1LitRegistryProject
 from lightning_sdk.lightning_cloud.openapi.models.v1_lit_repository import V1LitRepository
 from lightning_sdk.lightning_cloud.openapi.models.v1_locked_resource import V1LockedResource

lightning_sdk/lightning_cloud/openapi/api/cluster_service_api.py

@@ -1096,6 +1096,7 @@ class ClusterServiceApi(object):
         :param async_req bool
         :param str cluster_id: (required)
         :param str id: (required)
+        :param str org_id:
         :return: V1DeleteClusterUsageRestrictionResponse
                  If the method is called asynchronously,
                  returns the request thread.
@@ -1118,12 +1119,13 @@ class ClusterServiceApi(object):
         :param async_req bool
         :param str cluster_id: (required)
         :param str id: (required)
+        :param str org_id:
         :return: V1DeleteClusterUsageRestrictionResponse
                  If the method is called asynchronously,
                  returns the request thread.
         """
 
-        all_params = ['cluster_id', 'id']  # noqa: E501
+        all_params = ['cluster_id', 'id', 'org_id']  # noqa: E501
         all_params.append('async_req')
         all_params.append('_return_http_data_only')
         all_params.append('_preload_content')
@@ -1156,6 +1158,8 @@ class ClusterServiceApi(object):
             path_params['id'] = params['id']  # noqa: E501
 
         query_params = []
+        if 'org_id' in params:
+            query_params.append(('orgId', params['org_id']))  # noqa: E501
 
         header_params = {}
 
@@ -2311,6 +2315,7 @@ class ClusterServiceApi(object):
 
         :param async_req bool
         :param str cluster_id: (required)
+        :param str org_id:
         :return: V1ListClusterUsageRestrictionsResponse
                  If the method is called asynchronously,
                  returns the request thread.
@@ -2332,12 +2337,13 @@ class ClusterServiceApi(object):
 
         :param async_req bool
         :param str cluster_id: (required)
+        :param str org_id:
         :return: V1ListClusterUsageRestrictionsResponse
                  If the method is called asynchronously,
                  returns the request thread.
         """
 
-        all_params = ['cluster_id']  # noqa: E501
+        all_params = ['cluster_id', 'org_id']  # noqa: E501
         all_params.append('async_req')
         all_params.append('_return_http_data_only')
         all_params.append('_preload_content')
@@ -2364,6 +2370,8 @@ class ClusterServiceApi(object):
             path_params['clusterId'] = params['cluster_id']  # noqa: E501
 
         query_params = []
+        if 'org_id' in params:
+            query_params.append(('orgId', params['org_id']))  # noqa: E501
 
         header_params = {}