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

lightning_sdk/__init__.py

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

lightning_sdk/ai_hub.py

@@ -8,6 +8,7 @@ from lightning_sdk.utils.resolve import _resolve_teamspace
 
 if TYPE_CHECKING:
     from lightning_sdk import Organization, Teamspace
+    from lightning_sdk.machine import Machine
 
 
 class AIHub:
@@ -107,6 +108,7 @@ class AIHub:
         teamspace: Optional[Union[str, "Teamspace"]] = None,
         org: Optional[Union[str, "Organization"]] = None,
         user: Optional[Union[str, "User"]] = None,
+        machine: Optional[Union[str, "Machine"]] = None,
     ) -> Dict[str, Union[str, bool]]:
         """Deploy an API from the AI Hub.
 
@@ -128,6 +130,7 @@ class AIHub:
             teamspace: The team or group for deployment. Defaults to None.
             org: The organization for deployment. Don't pass user with this. Defaults to None.
             user: The user for deployment. Don't pass org with this. Defaults to None.
+            machine: The machine to run the deployment on. Defaults to the first option set in the AI Hub template.
 
         Returns:
             A dictionary containing the name of the deployed API,
@@ -153,6 +156,7 @@ class AIHub:
             project_id=teamspace_id,
             name=name,
             api_arguments=api_arguments,
+            machine=machine,
         )
 
         url = (
@@ -171,4 +175,22 @@ class AIHub:
             "deployment_url": url,
             "api_endpoint": deployment.status.urls[0],
             "interruptible": deployment.spec.spot,
+            "teamspace id": teamspace_id,
         }
+
+    def delete_deployment(self, deployment: Dict[str, Union[str, bool]]) -> None:
+        """Delete a deployment from the AI Hub.
+
+        Example:
+            from lightning_sdk import AIHub
+            hub = AIHub()
+            deployment = hub.run("temp_xxxx")
+            hub.delete_deployment(deployment)
+
+        Args:
+            deployment: The deployment dictionary returned by the run method.
+        """
+        if "teamspace id" not in deployment or "id" not in deployment:
+            raise ValueError("Deployment dictionary must contain 'teamspace id' and 'id' keys.")
+
+        self._api.delete_api(deployment["id"], deployment["teamspace id"])

lightning_sdk/api/ai_hub_api.py

@@ -1,8 +1,10 @@
 import traceback
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 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,
     V1Deployment,
@@ -16,6 +18,7 @@ from lightning_sdk.lightning_cloud.openapi.models.v1_deployment_template_gallery
     V1DeploymentTemplateGalleryResponse,
 )
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
+from lightning_sdk.machine import Machine
 
 
 class AIHubApi:
@@ -110,7 +113,13 @@ class AIHubApi:
         return job
 
     def run_api(
-        self, template_id: str, project_id: str, cloud_account: str, name: Optional[str], api_arguments: Dict[str, str]
+        self,
+        template_id: str,
+        project_id: str,
+        cloud_account: str,
+        name: Optional[str],
+        api_arguments: Dict[str, str],
+        machine: Optional[Union[str, Machine]],
     ) -> V1Deployment:
         template = self._client.deployment_templates_service_get_deployment_template(template_id)
         name = name or template.name
@@ -123,6 +132,13 @@ class AIHubApi:
             template.spec_v2.autoscaling.enabled = True
 
         AIHubApi._set_parameters(template.spec_v2.job, template.parameter_spec.parameters, api_arguments)
+        if machine and isinstance(machine, Machine):
+            apply_change(template.spec_v2.job, "instance_name", _machine_to_compute_name(machine))
+            apply_change(template.spec_v2.job, "instance_type", _machine_to_compute_name(machine))
+        elif machine and isinstance(machine, str):
+            apply_change(template.spec_v2.job, "instance_name", machine)
+            apply_change(template.spec_v2.job, "instance_type", machine)
+
         return self._client.jobs_service_create_deployment(
             project_id=project_id,
             body=CreateDeploymentRequestDefinesASpecForTheJobThatAllowsForAutoscalingJobs(
@@ -134,3 +150,6 @@ class AIHubApi:
                 spec=template.spec_v2.job,
             ),
         )
+
+    def delete_api(self, deployment_id: str, teamspace_id: str) -> None:
+        self._client.jobs_service_delete_deployment(project_id=teamspace_id, id=deployment_id)

lightning_sdk/api/deployment_api.py

@@ -1,7 +1,7 @@
 from time import sleep
 from typing import Any, List, Literal, Optional, Union
 
-from lightning_sdk.api.utils import _MACHINE_TO_COMPUTE_NAME
+from lightning_sdk.api.utils import _machine_to_compute_name
 from lightning_sdk.lightning_cloud.openapi import (
     CreateDeploymentRequestDefinesASpecForTheJobThatAllowsForAutoscalingJobs,
     V1AutoscalingSpec,
@@ -270,7 +270,8 @@ class DeploymentApi:
 
         # Any updates to the Job Spec triggers a new release
         if machine:
-            apply_change(deployment.spec, "instance_name", _MACHINE_TO_COMPUTE_NAME[machine])
+            apply_change(deployment.spec, "instance_name", _machine_to_compute_name(machine))
+            apply_change(deployment.spec, "instance_type", _machine_to_compute_name(machine))
 
         requires_release = False
         requires_release |= apply_change(deployment.spec, "image", environment)
@@ -554,7 +555,7 @@ def to_spec(
         env=to_env(env),
         image=environment,
         spot=spot,
-        instance_name=_MACHINE_TO_COMPUTE_NAME[machine],
+        instance_name=_machine_to_compute_name(machine),
         readiness_probe=to_health_check(health_check),
     )
 

lightning_sdk/api/job_api.py

@@ -3,7 +3,6 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Union
 from urllib.request import urlopen
 
 from lightning_sdk.api.utils import (
-    _COMPUTE_NAME_TO_MACHINE,
     _create_app,
     _machine_to_compute_name,
     remove_datetime_prefix,
@@ -98,13 +97,11 @@ class JobApiV1:
     def get_machine_from_work(self, work: Externalv1Lightningwork) -> Machine:
         spec: V1LightningworkSpec = work.spec
         # prefer user-requested config if specified
-        compute_config: V1UserRequestedComputeConfig = spec.user_requested_compute_config
-        compute: str = compute_config.name
-        if compute:
-            return _COMPUTE_NAME_TO_MACHINE[compute]
+        user_requested_compute_config: V1UserRequestedComputeConfig = spec.user_requested_compute_config
+        if user_requested_compute_config.name:
+            return Machine(user_requested_compute_config.name, user_requested_compute_config.name)
         compute_config: V1ComputeConfig = spec.compute_config
-        compute: str = compute_config.instance_type
-        return _COMPUTE_NAME_TO_MACHINE[compute]
+        return Machine(compute_config.instance_type, compute_config.instance_type)
 
     def get_studio_name(self, job: Externalv1LightningappInstance) -> str:
         cs: V1CloudSpace = self._client.cloud_space_service_get_cloud_space(
@@ -343,9 +340,7 @@ class JobApiV2:
         instance_name = spec.instance_name
         instance_type = spec.instance_type
 
-        return _COMPUTE_NAME_TO_MACHINE.get(
-            instance_type, _COMPUTE_NAME_TO_MACHINE.get(instance_name, instance_type or instance_name)
-        )
+        return Machine(instance_name, instance_type or instance_name)
 
     def get_total_cost(self, job: V1Job) -> float:
         return job.total_cost

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/api/mmt_api.py

@@ -4,7 +4,6 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 from lightning_sdk.api.job_api import JobApiV1
 from lightning_sdk.api.utils import (
-    _COMPUTE_NAME_TO_MACHINE,
     _create_app,
     _machine_to_compute_name,
     resolve_path_mappings,
@@ -205,9 +204,7 @@ class MMTApiV2:
         instance_name = spec.instance_name
         instance_type = spec.instance_type
 
-        return _COMPUTE_NAME_TO_MACHINE.get(
-            instance_type, _COMPUTE_NAME_TO_MACHINE.get(instance_name, instance_type or instance_name)
-        )
+        return Machine(instance_name, instance_type or instance_name)
 
     def get_total_cost(self, job: V1MultiMachineJob) -> float:
         return job.total_cost

lightning_sdk/api/studio_api.py

@@ -12,8 +12,6 @@ import requests
 from tqdm import tqdm
 
 from lightning_sdk.api.utils import (
-    _COMPUTE_NAME_TO_MACHINE,
-    _MACHINE_TO_COMPUTE_NAME,
     _create_app,
     _DummyBody,
     _DummyResponse,
@@ -257,7 +255,7 @@ class StudioApi:
         response: V1CloudSpaceInstanceConfig = self._client.cloud_space_service_get_cloud_space_instance_config(
             project_id=teamspace_id, id=studio_id
         )
-        return _COMPUTE_NAME_TO_MACHINE[response.compute_config.name]
+        return Machine(response.compute_config.name, 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."""
@@ -582,7 +580,7 @@ class StudioApi:
             plugin_type="job",
             entrypoint=entrypoint,
             name=name,
-            compute=_MACHINE_TO_COMPUTE_NAME[machine],
+            compute=_machine_to_compute_name(machine),
             interruptible=interruptible,
         )
 
@@ -600,7 +598,7 @@ class StudioApi:
     ) -> Externalv1LightningappInstance:
         """Creates a multi-machine job with given commands."""
         distributed_args = {
-            "cloud_compute": _MACHINE_TO_COMPUTE_NAME[machine],
+            "cloud_compute": _machine_to_compute_name(machine),
             "num_instances": num_instances,
             "strategy": strategy,
         }
@@ -628,7 +626,7 @@ class StudioApi:
     ) -> Externalv1LightningappInstance:
         """Creates a multi-machine job with given commands."""
         data_prep_args = {
-            "cloud_compute": _MACHINE_TO_COMPUTE_NAME[machine],
+            "cloud_compute": _machine_to_compute_name(machine),
             "num_instances": num_instances,
         }
         return self._create_app(
@@ -665,7 +663,7 @@ class StudioApi:
             teamspace_id=teamspace_id,
             cloud_account=cloud_account,
             plugin_type="inference_plugin",
-            compute=_MACHINE_TO_COMPUTE_NAME[machine],
+            compute=_machine_to_compute_name(machine),
             entrypoint=entrypoint,
             name=name,
             min_replicas=min_replicas,

lightning_sdk/api/teamspace_api.py

@@ -14,6 +14,7 @@ from lightning_sdk.lightning_cloud.openapi import (
     ProjectIdModelsBody,
     V1Assistant,
     V1CloudSpace,
+    V1ClusterAccelerator,
     V1Endpoint,
     V1Job,
     V1ModelVersionArchive,
@@ -295,3 +296,9 @@ class TeamspaceApi:
         ).lightningapps
         jobs = self._client.jobs_service_list_multi_machine_jobs(project_id=teamspace_id).multi_machine_jobs
         return apps, jobs
+
+    def list_machines(self, teamspace_id: str, cloud_account: str) -> List[V1ClusterAccelerator]:
+        response = self._client.cluster_service_list_project_cluster_accelerators(
+            project_id=teamspace_id, id=cloud_account
+        )
+        return response.accelerator

lightning_sdk/api/utils.py

@@ -322,38 +322,12 @@ class _DummyResponse:
         self.data = data
 
 
-# TODO: This should really come from some kind of metadata service
-_MACHINE_TO_COMPUTE_NAME: Dict[Machine, str] = {
-    Machine.CPU_SMALL: "m3.medium",
-    Machine.CPU: "cpu-4",
-    Machine.DATA_PREP: "data-large",
-    Machine.DATA_PREP_MAX: "data-max",
-    Machine.DATA_PREP_ULTRA: "data-ultra",
-    Machine.T4: "g4dn.2xlarge",
-    Machine.T4_X_4: "g4dn.12xlarge",
-    Machine.L4: "g6.4xlarge",
-    Machine.L4_X_4: "g6.12xlarge",
-    Machine.L4_X_8: "g6.48xlarge",
-    Machine.A10G: "g5.8xlarge",
-    Machine.A10G_X_4: "g5.12xlarge",
-    Machine.A10G_X_8: "g5.48xlarge",
-    Machine.L40S: "g6e.4xlarge",
-    Machine.L40S_X_4: "g6e.12xlarge",
-    Machine.L40S_X_8: "g6e.48xlarge",
-    Machine.A100_X_8: "p4d.24xlarge",
-    Machine.H100_X_8: "p5.48xlarge",
-    Machine.H200_X_8: "p5e.48xlarge",
-}
-
-
 def _machine_to_compute_name(machine: Union[Machine, str]) -> str:
     if isinstance(machine, Machine):
-        return _MACHINE_TO_COMPUTE_NAME[machine]
+        return machine.instance_type
     return machine
 
 
-_COMPUTE_NAME_TO_MACHINE: Dict[str, Machine] = {v: k for k, v in _MACHINE_TO_COMPUTE_NAME.items()}
-
 _DEFAULT_CLOUD_URL = "https://lightning.ai"
 _DEFAULT_REGISTRY_URL = "litcr.io"
 

lightning_sdk/cli/configure.py

@@ -0,0 +1,92 @@
+import platform
+import uuid
+from pathlib import Path
+from typing import Optional, Union
+
+from rich.console import Console
+
+from lightning_sdk.cli.generate import _Generate
+from lightning_sdk.lightning_cloud.login import Auth
+
+
+def _download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Optional[int] = None) -> None:
+    """Download a file from a URL."""
+    import requests
+
+    if local_path.exists() and not overwrite:
+        raise FileExistsError(f"The file {local_path} already exists and overwrite is set to False.")
+
+    response = requests.get(url, stream=True)
+    response.raise_for_status()
+
+    with open(local_path, "wb") as file:
+        for chunk in response.iter_content(chunk_size=8192):
+            file.write(chunk)
+    if chmod is not None:
+        local_path.chmod(0o600)
+
+
+class _Configure(_Generate):
+    """Configure lightning products."""
+
+    @staticmethod
+    def _download_ssh_keys(
+        api_key: str,
+        key_id: str = "",
+        ssh_home: Union[str, Path] = "",
+        ssh_key_name: str = "lightning_rsa",
+        overwrite: bool = False,
+    ) -> None:
+        if not ssh_home:
+            ssh_home = Path.home() / ".ssh"
+        elif isinstance(ssh_home, str):
+            ssh_home = Path(ssh_home)
+        if not key_id:
+            key_id = str(uuid.uuid4())
+
+        path_key = ssh_home / ssh_key_name
+        path_pub = ssh_home / f"{ssh_key_name}.pub"
+
+        # todo: consider hitting the API to get the key pair directly instead of using wget
+        _download_file(
+            f"https://lightning.ai/setup/ssh-gen?t={api_key}&id={key_id}&machineName={platform.node()}",
+            path_key,
+            overwrite=overwrite,
+            chmod=0o600,
+        )
+        _download_file(f"https://lightning.ai/setup/ssh-public?t={api_key}&id={key_id}", path_pub, overwrite=overwrite)
+
+    def ssh(self, overwrite: bool = False, ssh_key_name: str = "lightning_rsa") -> None:
+        """Get SSH config entry for a studio.
+
+        Args:
+            overwrite: Whether to overwrite the SSH key and config if they already exist.
+            ssh_key_name: The name of the SSH key to generate
+        """
+        auth = Auth()
+        auth.authenticate()
+        console = Console()
+        ssh_dir = Path.home() / ".ssh"
+        ssh_dir.mkdir(parents=True, exist_ok=True)
+
+        key_path = ssh_dir / ssh_key_name
+        config_path = ssh_dir / "config"
+
+        # Check if the SSH key already exists
+        if key_path.exists() and (key_path.with_suffix(".pub")).exists() and not overwrite:
+            console.print(f"SSH key already exists at {key_path}")
+        else:
+            self._download_ssh_keys(auth.api_key, ssh_home=ssh_dir, ssh_key_name=ssh_key_name, overwrite=overwrite)
+            console.print(f"SSH key generated and saved to {key_path}")
+
+        # Check if the SSH config already contains the required configuration
+        config_content = self._generate_ssh_config(str(key_path))
+        if config_path.exists():
+            with config_path.open("r") as config_file:
+                if config_content.strip() in config_file.read():
+                    console.print("SSH config already contains the required configuration.")
+                    return
+
+        with config_path.open("a") as config_file:
+            config_file.write(config_content)
+            console.print(f"SSH config updated at {config_path}")

lightning_sdk/cli/connect.py

@@ -0,0 +1,31 @@
+import subprocess
+import sys
+from typing import Optional
+
+from lightning_sdk.cli.configure import _Configure
+from lightning_sdk.lightning_cloud.login import Auth
+
+
+class _Connect(_Configure):
+    """Connect to lightning products."""
+
+    def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+        """Connect to a studio via SSH.
+
+        Args:
+            name: The name of the studio to connect to.
+            teamspace: The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.
+        """
+        auth = Auth()
+        auth.authenticate()  # this is maybe not needed
+        studio = self._get_studio(name=name, teamspace=teamspace)
+        host = "ssh.lightning.ai"
+        username = f"s_{studio._studio.id}"
+
+        self.ssh(overwrite=False)
+
+        try:
+            subprocess.run(["ssh", f"{username}@{host}"])
+        except Exception as ex:
+            print(f"Failed to establish SSH connection: {ex}")
+            sys.exit(1)

lightning_sdk/cli/delete.py

@@ -1,9 +1,12 @@
 from typing import Optional
 
+from rich.console import Console
+
 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):
@@ -21,7 +24,7 @@ class _Delete(_JobAndMMTAction, _TeamspacesMenu):
         resolved_teamspace = self._resolve_teamspace(teamspace=teamspace)
         try:
             api.delete_container(container, resolved_teamspace.name, resolved_teamspace.owner.name)
-            print(f"Container {container} deleted successfully.")
+            Console().print(f"Container {container} deleted successfully.")
         except Exception as e:
             raise StudioCliError(
                 f"Could not delete container {container} from project {resolved_teamspace.name}: {e}"
@@ -40,7 +43,7 @@ class _Delete(_JobAndMMTAction, _TeamspacesMenu):
         job = super().job(name=name, teamspace=teamspace)
 
         job.delete()
-        print(f"Successfully deleted {job.name}!")
+        Console().print(f"Successfully deleted {job.name}!")
 
     def mmt(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
         """Delete a multi-machine job.
@@ -55,4 +58,30 @@ class _Delete(_JobAndMMTAction, _TeamspacesMenu):
         mmt = super().mmt(name=name, teamspace=teamspace)
 
         mmt.delete()
-        print(f"Successfully deleted {mmt.name}!")
+        Console().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()
+        Console().print("Studio successfully deleted")

lightning_sdk/cli/download.py

@@ -143,8 +143,8 @@ class _Downloads(_StudiosMenu, _TeamspacesMenu):
           teamspace: The name of the teamspace to download the container from.
           tag: The tag of the container to download.
         """
-        resolved_teamspace = self._resolve_teamspace(teamspace)
         console = Console()
+        resolved_teamspace = self._resolve_teamspace(teamspace)
         with console.status("Downloading container..."):
             api = LitContainerApi()
             api.download_container(container, resolved_teamspace, tag)

lightning_sdk/cli/entrypoint.py

@@ -4,17 +4,24 @@ from typing import Type
 
 from fire import Fire
 from lightning_utilities.core.imports import RequirementCache
+from rich.console import Console
+from rich.panel import Panel
 
 from lightning_sdk.api.studio_api import _cloud_url
 from lightning_sdk.cli.ai_hub import _AIHub
+from lightning_sdk.cli.configure import _Configure
+from lightning_sdk.cli.connect import _Connect
 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 +42,11 @@ class StudioCLI:
         self.delete = _Delete()
         self.inspect = _Inspect()
         self.stop = _Stop()
+        self.start = _Start()
+        self.switch = _Switch()
+        self.generate = _Generate()
+        self.connect = _Connect()
+        self.configure = _Configure()
 
         sys.excepthook = _notify_exception
 
@@ -56,7 +68,8 @@ class StudioCLI:
 
 def _notify_exception(exception_type: Type[BaseException], value: BaseException, tb: TracebackType) -> None:  # No
     """CLI won't show tracebacks, just print the exception message."""
-    print(value)
+    console = Console()
+    console.print(Panel(value))
 
 
 def main_cli() -> None:

lightning_sdk/cli/generate.py

@@ -0,0 +1,35 @@
+from typing import Optional
+
+from rich.console import Console
+
+from lightning_sdk.cli.studios_menu import _StudiosMenu
+
+
+class _Generate(_StudiosMenu):
+    """Generate configs (such as ssh for studio) and print them to commandline."""
+
+    @staticmethod
+    def _generate_ssh_config(key_path: str) -> str:
+        return f"""Host ssh.lightning.ai
+  IdentityFile {key_path}
+  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.
+
+        Args:
+            name: The name of the studio to obtain SSH config.
+                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.)
+        """
+        studio = self._get_studio(name=name, teamspace=teamspace)
+
+        # Print the SSH config
+        conf = f"# ssh s_{studio._studio.id}@ssh.lightning.ai\n\n" + self._generate_ssh_config("~/.ssh/lightning_rsa")
+        Console().print(conf)