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

lightning_sdk/__init__.py

@@ -1,7 +1,7 @@
 from lightning_sdk.agents import Agent
 from lightning_sdk.ai_hub import AIHub
 from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__  # noqa: F401
-from lightning_sdk.helpers import _check_version_and_prompt_upgrade
+from lightning_sdk.helpers import _check_version_and_prompt_upgrade, _set_tqdm_envvars_noninteractive
 from lightning_sdk.job import Job
 from lightning_sdk.machine import Machine
 from lightning_sdk.mmt import MMT
@@ -29,5 +29,6 @@ __all__ = [
     "AIHub",
 ]
 
-__version__ = "0.1.55"
+__version__ = "0.1.57"
 _check_version_and_prompt_upgrade(__version__)
+_set_tqdm_envvars_noninteractive()

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/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/ai_hub.py

@@ -1,4 +1,6 @@
-from typing import List, Optional
+from typing import Optional
+
+import click
 
 from lightning_sdk.ai_hub import AIHub
 from lightning_sdk.cli.studios_menu import _StudiosMenu
@@ -7,10 +9,7 @@ from lightning_sdk.cli.studios_menu import _StudiosMenu
 class _AIHub(_StudiosMenu):
     """Interact with Lightning Studio - AI Hub."""
 
-    def __init__(self) -> None:
-        self._hub = AIHub()
-
-    def api_info(self, api_id: str) -> dict:
+    def api_info(self, api_id: str) -> None:
         """Get full API template info such as input details.
 
         Example:
@@ -19,15 +18,15 @@ class _AIHub(_StudiosMenu):
         Args:
           api_id: The ID of the API for which information is requested.
         """
-        return self._hub.api_info(api_id)
+        return api_info(api_id=api_id)
 
-    def list_apis(self, search: Optional[str] = None) -> List[dict]:
+    def list_apis(self, search: Optional[str] = None) -> None:
         """List API templates available in the AI Hub.
 
         Args:
           search: Search for API templates by name.
         """
-        return self._hub.list_apis(search=search)
+        return list_apis(search=search)
 
     def deploy(
         self,
@@ -36,7 +35,7 @@ class _AIHub(_StudiosMenu):
         name: Optional[str] = None,
         teamspace: Optional[str] = None,
         org: Optional[str] = None,
-    ) -> dict:
+    ) -> None:
         """Deploy an API template from the AI Hub.
 
         Args:
@@ -46,4 +45,56 @@ class _AIHub(_StudiosMenu):
           teamspace: Teamspace to deploy the API to. Defaults to user's default teamspace.
           org: Organization to deploy the API to. Defaults to user's default organization.
         """
-        return self._hub.run(api_id, cloud_account=cloud_account, name=name, teamspace=teamspace, org=org)
+        return deploy(api_id=api_id, cloud_account=cloud_account, name=name, teamspace=teamspace, org=org)
+
+
+@click.group(name="aihub")
+def aihub() -> None:
+    """Interact with Lightning Studio - AI Hub."""
+
+
+# @aihub.command(name="api-info")
+# @click.argument("api-id")
+def api_info(api_id: str) -> None:
+    """Get full API template info such as input details.
+
+    Example:
+      lightning aihub api_info API-ID
+
+    API-ID: The ID of the API for which information is requested.
+    """
+    ai_hub = AIHub()
+    ai_hub.api_info(api_id)
+
+
+# @aihub.command(name="list-apis")
+# @click.option("--search", default=None, help="Search for API templates by name.")
+def list_apis(search: Optional[str]) -> None:
+    """List API templates available in the AI Hub."""
+    ai_hub = AIHub()
+    ai_hub.list_apis(search=search)
+
+
+# @aihub.command(name="deploy")
+# @click.argument("api-id")
+# @click.option(
+#     "--cloud-account",
+#     default=None,
+#     help="Cloud Account to deploy the API to. Defaults to user's default cloud account.",
+# )
+# @click.option("--name", default=None, help="Name of the deployed API. Defaults to the name of the API template.")
+# @click.option(
+#     "--teamspace",
+#     default=None,
+#     help="Teamspace to deploy the API to. Defaults to user's default teamspace.",
+# )
+# @click.option(
+#     "--org",
+#     default=None,
+#     help="Organization to deploy the API to. Defaults to user's default organization.",
+# )
+def deploy(
+    api_id: str, cloud_account: Optional[str], name: Optional[str], teamspace: Optional[str], org: Optional[str]
+) -> None:
+    ai_hub = AIHub()
+    ai_hub.run(api_id, cloud_account=cloud_account, name=name, teamspace=teamspace, org=org)

lightning_sdk/cli/configure.py

@@ -0,0 +1,137 @@
+import os
+import platform
+import uuid
+from pathlib import Path
+from typing import Optional, Union
+
+import click
+from rich.console import Console
+
+from lightning_sdk.cli.generate import _generate_ssh_config
+from lightning_sdk.cli.studios_menu import _StudiosMenu
+from lightning_sdk.lightning_cloud.login import Auth
+
+
+class _Configure(_StudiosMenu):
+    """Configure lightning products."""
+
+    def ssh(self, name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> 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.)
+            overwrite: Whether to overwrite the SSH key and config if they already exist.
+        """
+        ssh(name=name, teamspace=teamspace, overwrite=overwrite)
+
+
+@click.group(name="configure")
+def configure() -> None:
+    """Configure access to resources on the Lightning AI platform."""
+
+
+# @configure.command(name="ssh")
+# @click.option(
+#     "--name",
+#     default=None,
+#     help=(
+#         "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.)"
+#     ),
+# )
+# @click.option(
+#     "--teamspace",
+#     default=None,
+#     help=(
+#         "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.)"
+#     ),
+# )
+# @click.option(
+#     "--overwrite",
+#     is_flag=True,
+#     flag_value=True,
+#     help="Whether to overwrite the SSH key and config if they already exist.",
+# )
+def ssh(name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> None:
+    """Get SSH config entry for a studio."""
+    auth = Auth()
+    auth.authenticate()
+    console = Console()
+    ssh_dir = Path.home() / ".ssh"
+    ssh_dir.mkdir(parents=True, exist_ok=True)
+
+    key_path = ssh_dir / "lightning_rsa"
+    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:
+        _download_ssh_keys(auth.api_key, ssh_home=ssh_dir, ssh_key_name="lightning_rsa", overwrite=overwrite)
+        console.print(f"SSH key generated and saved to {key_path}")
+
+    # Check if the SSH config already contains the required configuration
+    menu = _StudiosMenu()
+    studio = menu._get_studio(name=name, teamspace=teamspace)
+    config_content = _generate_ssh_config(key_path=str(key_path), user=f"s_{studio._studio.id}", host=studio.name)
+    if config_path.exists():
+        with config_path.open("r") as config_file:
+            # check if the host already exists in the config
+            if f"Host {studio.name}" 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(os.linesep)
+        config_file.write(config_content)
+        config_file.write(os.linesep)
+        console.print(f"SSH config updated at {config_path}")
+
+
+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)
+
+
+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)

lightning_sdk/cli/connect.py

@@ -0,0 +1,47 @@
+import subprocess
+import sys
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.configure import ssh as configure_ssh
+from lightning_sdk.cli.studios_menu import _StudiosMenu
+
+
+class _Connect:
+    """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>.
+        """
+        return studio(name=name, teamspace=teamspace)
+
+
+@click.group(name="connect")
+def connect() -> None:
+    """Connect to lightning products."""
+
+
+# @connect.command(name="studio")
+# @click.option("--name", default=None, help="The name of the studio to connect to.")
+# @click.option(
+#     "--teamspace",
+#     default=None,
+#     help="The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.",
+# )
+def studio(name: Optional[str], teamspace: Optional[str]) -> None:
+    """Connect to a studio via SSH."""
+    configure_ssh(name=name, teamspace=teamspace, overwrite=False)
+
+    menu = _StudiosMenu()
+    studio = menu._get_studio(name=name, teamspace=teamspace)
+
+    try:
+        subprocess.run(["ssh", studio.name])
+    except Exception as ex:
+        print(f"Failed to establish SSH connection: {ex}")
+        sys.exit(1)