lightning-sdk 0.1.49__py3-none-any.whl → 2025.11.5__py3-none-any.whl

lightning_sdk/cli/studio/connect.py

@@ -0,0 +1,139 @@
+"""Studio connect command."""
+
+import subprocess
+import sys
+from contextlib import suppress
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.utils.get_base_studio import get_base_studio_id
+from lightning_sdk.cli.utils.handle_machine_and_gpus_args import handle_machine_and_gpus_args
+from lightning_sdk.cli.utils.richt_print import studio_name_link
+from lightning_sdk.cli.utils.save_to_config import save_studio_to_config, save_teamspace_to_config
+from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.machine import CloudProvider, Machine
+from lightning_sdk.studio import Studio
+from lightning_sdk.utils.names import random_unique_name
+
+
+def _parse_args_or_get_from_current_studio(
+    teamspace: Optional[str],
+    cloud_account: Optional[str],
+    studio_type: Optional[str],
+    machine: Optional[str],
+    gpus: Optional[str],
+    cloud_provider: Optional[str],
+    name: Optional[str],
+) -> tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
+    # Parse args provided by user
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace)
+    save_teamspace_to_config(resolved_teamspace, overwrite=False)
+
+    template_id = get_base_studio_id(studio_type)
+
+    if cloud_provider is not None:
+        cloud_provider = CloudProvider(cloud_provider)
+
+    name = name or random_unique_name()
+
+    with suppress(ValueError):
+        # Gets current studio context to use its parameters as defaults
+        s = Studio()
+        if not teamspace:
+            resolved_teamspace = s.teamspace
+            save_teamspace_to_config(resolved_teamspace, overwrite=False)
+        if not cloud_account:
+            cloud_account = s.cloud_account
+        if not template_id:
+            template_id = s._studio.environment_template_id
+        if not machine and not gpus:
+            machine = s.machine
+
+    return resolved_teamspace, cloud_account, template_id, machine, cloud_provider, name
+
+
+@click.command("connect")
+@click.argument("name", required=False)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+@click.option(
+    "--cloud-provider",
+    help="The cloud provider to start the studio on. Defaults to teamspace default.",
+    type=click.Choice(m.name for m in list(CloudProvider)),
+)
+@click.option(
+    "--cloud-account",
+    help="The cloud account to create the studio on. Defaults to teamspace default.",
+    type=click.STRING,
+)
+@click.option(
+    "--machine",
+    help="The machine type to start the studio on. Defaults to CPU-4",
+    type=click.Choice(m.name for m in Machine.__dict__.values() if isinstance(m, Machine) and m._include_in_cli),
+)
+@click.option(
+    "--gpus",
+    help="The number and type of GPUs to start the studio on (format: TYPE:COUNT, e.g. L4:4)",
+    type=click.STRING,
+)
+@click.option(
+    "--studio-type",
+    help="The base studio template name to use for creating the studio. "
+    "Must be lowercase and hyphenated (use '-' instead of spaces). "
+    "Run 'lightning base-studio list' to see all available templates. "
+    "Defaults to the first available template.",
+    type=click.STRING,
+)
+@click.option("--interruptible", is_flag=True, help="Start the studio on an interruptible instance.")
+def connect_studio(
+    name: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    cloud_provider: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+    machine: Optional[str] = None,
+    gpus: Optional[str] = None,
+    studio_type: Optional[str] = None,
+    interruptible: bool = False,
+) -> None:
+    """Connect to a Studio.
+
+    Example:
+        lightning studio connect
+    """
+    teamspace, cloud_account, template_id, machine, cloud_provider, name = _parse_args_or_get_from_current_studio(
+        teamspace, cloud_account, studio_type, machine, gpus, cloud_provider, name
+    )
+
+    try:
+        studio = Studio(
+            name=name,
+            teamspace=teamspace,
+            create_ok=True,
+            cloud_provider=cloud_provider,
+            cloud_account=cloud_account,
+            studio_type=template_id,
+        )
+    except (RuntimeError, ValueError, ApiException):
+        raise ValueError(f"Could not create Studio: '{name}'") from None
+
+    click.echo(f"Connecting to Studio '{studio_name_link(studio)}' ...")
+
+    Studio.show_progress = True
+
+    machine = handle_machine_and_gpus_args(machine, gpus)
+
+    save_studio_to_config(studio)
+    studio.start(machine=machine, interruptible=interruptible)
+
+    ssh_private_key_path = configure_ssh_internal()
+
+    ssh_option = "-o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -o LogLevel=ERROR"
+    try:
+        ssh_command = f"ssh -i {ssh_private_key_path} {ssh_option} s_{studio._studio.id}@ssh.lightning.ai"
+        subprocess.run(ssh_command.split())
+    except Exception as ex:
+        print(f"Failed to establish SSH connection: {ex}")
+        sys.exit(1)

lightning_sdk/cli/studio/create.py

@@ -0,0 +1,96 @@
+"""Studio start command."""
+
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.utils.get_base_studio import get_base_studio_id
+from lightning_sdk.cli.utils.richt_print import studio_name_link
+from lightning_sdk.cli.utils.save_to_config import save_teamspace_to_config
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.machine import CloudProvider
+from lightning_sdk.studio import VM, Studio
+
+
+@click.command("create")
+@click.option("--name", help="The name of the studio to create. If not provided, a random name will be generated.")
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+@click.option(
+    "--cloud-provider",
+    help="The cloud provider to start the studio on. Defaults to teamspace default.",
+    type=click.Choice(m.name for m in list(CloudProvider)),
+)
+@click.option(
+    "--cloud-account",
+    help="The cloud account to create the studio on. Defaults to teamspace default.",
+    type=click.STRING,
+)
+@click.option(
+    "--studio-type",
+    help="The base studio template name to use for creating the studio. "
+    "Must be lowercase and hyphenated (use '-' instead of spaces). "
+    "Run 'lightning base-studio list' to see all available templates. "
+    "Defaults to the first available template.",
+    type=click.STRING,
+)
+def create_studio(
+    name: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    cloud_provider: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+    studio_type: Optional[str] = None,
+) -> None:
+    """Create a new Studio.
+
+    Example:
+        lightning studio create
+    """
+    create_impl(
+        name=name,
+        teamspace=teamspace,
+        cloud_provider=cloud_provider,
+        cloud_account=cloud_account,
+        vm=False,
+        studio_type=studio_type,
+    )
+
+
+def create_impl(
+    name: Optional[str],
+    teamspace: Optional[str],
+    cloud_provider: Optional[str],
+    cloud_account: Optional[str],
+    vm: bool,
+    studio_type: Optional[str],
+) -> None:
+    menu = TeamspacesMenu()
+
+    resolved_teamspace = menu(teamspace)
+    save_teamspace_to_config(resolved_teamspace, overwrite=False)
+
+    if cloud_provider is not None:
+        cloud_provider = CloudProvider(cloud_provider)
+
+    create_cls = VM if vm else Studio
+    cls_name = create_cls.__qualname__
+
+    # check for available base studios
+    template_id = get_base_studio_id(studio_type)
+
+    try:
+        create_cls = VM if vm else Studio
+        studio = create_cls(
+            name=name,
+            teamspace=resolved_teamspace,
+            create_ok=True,
+            cloud_provider=cloud_provider,
+            cloud_account=cloud_account,
+            studio_type=template_id,
+        )
+    except (RuntimeError, ValueError, ApiException):
+        if name:
+            raise ValueError(f"Could not create {cls_name}: '{name}'. Does the {cls_name} exist?") from None
+        raise ValueError(f"Could not create {cls_name}: '{name}'. Please provide a {cls_name} name") from None
+
+    click.echo(f"{cls_name} {studio_name_link(studio)} created successfully")

lightning_sdk/cli/studio/delete.py

@@ -0,0 +1,49 @@
+"""Studio delete command."""
+
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.utils.studio_selection import StudiosMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+
+
+@click.command("delete")
+@click.option(
+    "--name",
+    help=(
+        "The name of the studio to delete. "
+        "If not provided, will try to infer from environment, "
+        "use the default value from the config or prompt for interactive selection."
+    ),
+)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+def delete_studio(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+    """Delete a Studio.
+
+    Example:
+      lightning studio delete --name my-studio
+
+    """
+    return delete_impl(name=name, teamspace=teamspace, vm=False)
+
+
+def delete_impl(name: Optional[str], teamspace: Optional[str], vm: bool) -> None:
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace=teamspace)
+
+    menu = StudiosMenu(resolved_teamspace, vm=vm)
+    studio = menu(studio=name)
+
+    studio_name = f"{studio.teamspace.owner.name}/{studio.teamspace.name}/{studio.name}"
+    confirmed = click.confirm(
+        f"Are you sure you want to delete {studio._cls_name} '{studio_name}'?",
+        abort=True,
+    )
+    if not confirmed:
+        click.echo(f"{studio._cls_name} deletion cancelled")
+        return
+
+    studio.delete()
+
+    click.echo(f"{studio._cls_name} '{studio.name}' deleted successfully")

lightning_sdk/cli/studio/list.py

@@ -0,0 +1,85 @@
+"""Studio list command."""
+
+from typing import Callable, Optional
+
+import click
+from rich.table import Table
+
+from lightning_sdk.cli.utils.cloud_account_map import cloud_account_to_display_name
+from lightning_sdk.cli.utils.richt_print import rich_to_str, studio_name_link
+from lightning_sdk.cli.utils.save_to_config import save_teamspace_to_config
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.studio import Studio
+from lightning_sdk.utils.resolve import _get_authed_user, prevent_refetch_studio
+
+
+@click.command("list")
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+@click.option(
+    "--all",
+    is_flag=True,
+    flag_value=True,
+    default=False,
+    help="List all studios, not just the ones belonging to the authed user",
+)
+@click.option(
+    "--sort-by",
+    default=None,
+    type=click.Choice(["name", "teamspace", "status", "machine", "cloud-account"], case_sensitive=False),
+    help="the attribute to sort the studios by.",
+)
+def list_studios(teamspace: Optional[str] = None, all: bool = False, sort_by: Optional[str] = None) -> None:  # noqa: A002
+    """List Studios in a teamspace.
+
+    Example:
+        lightning studio list --teamspace owner/teamspace
+
+    """
+    return list_impl(teamspace=teamspace, all=all, sort_by=sort_by, vm=False)
+
+
+def list_impl(teamspace: Optional[str], all: bool, sort_by: Optional[str], vm: bool) -> None:  # noqa: A002
+    menu = TeamspacesMenu()
+    teamspace_resolved = menu(teamspace=teamspace)
+    save_teamspace_to_config(teamspace_resolved, overwrite=False)
+
+    user = _get_authed_user()
+
+    studios = teamspace_resolved.vms if vm else teamspace_resolved.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 sorted(
+        filter(lambda s: all or s._studio.user_id == user.id, studios), key=_sort_studios_key(sort_by)
+    ):
+        with prevent_refetch_studio(studio):
+            table.add_row(
+                # cannot convert to ascii here, as the final rich table has to be converted to ascii
+                # otherwise the lack of support for linking in some terminals causes formatting issues.
+                studio_name_link(studio, to_ascii=False),
+                f"{studio.teamspace.owner.name}/{studio.teamspace.name}",
+                str(studio.status),
+                str(studio.machine) if studio.machine is not None else None,  # when None the cell is empty
+                str(cloud_account_to_display_name(studio.cloud_account, studio.teamspace.id)),
+            )
+
+    click.echo(rich_to_str(table), color=True)
+
+
+def _sort_studios_key(sort_by: str) -> Callable[[Studio], str]:
+    """Return a key function to sort studios by a given attribute."""
+    sort_key_map = {
+        "name": lambda s: str(s.name or ""),
+        "teamspace": lambda s: str(s.teamspace.name or ""),
+        "status": lambda s: str(s.status or ""),
+        "machine": lambda s: str(s.machine or ""),
+        "cloud-account": lambda s: str(cloud_account_to_display_name(s.cloud_account or "", s.teamspace.id)),
+    }
+    return sort_key_map.get(sort_by, lambda s: s.name)

lightning_sdk/cli/studio/ssh.py

@@ -0,0 +1,64 @@
+"""Studio SSH command."""
+
+import subprocess
+from typing import List, Optional
+
+import click
+
+from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
+from lightning_sdk.cli.utils.ssh_connection import configure_ssh_internal
+from lightning_sdk.cli.utils.studio_selection import StudiosMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+
+
+@click.command("ssh")
+@click.option(
+    "--name",
+    help=(
+        "The name of the studio to ssh into. "
+        "If not provided, will try to infer from environment, "
+        "use the default value from the config or prompt for interactive selection."
+    ),
+)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)", type=click.STRING)
+@click.option(
+    "--option",
+    "-o",
+    help="Additional options to pass to the SSH command. Can be specified multiple times.",
+    multiple=True,
+    type=click.STRING,
+)
+def ssh_studio(name: Optional[str] = None, teamspace: Optional[str] = None, option: Optional[List[str]] = None) -> None:
+    """SSH into a Studio.
+
+    Example:
+        lightning studio ssh --name my-studio
+    """
+    return ssh_impl(name=name, teamspace=teamspace, option=option, vm=False)
+
+
+def ssh_impl(name: Optional[str], teamspace: Optional[str], option: Optional[List[str]], vm: bool) -> None:
+    ssh_private_key_path = configure_ssh_internal()
+
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace=teamspace)
+
+    menu = StudiosMenu(resolved_teamspace, vm=vm)
+    studio = menu(
+        studio=name,
+    )
+    save_studio_to_config(studio)
+
+    ssh_options = " -o " + " -o ".join(option) if option else ""
+    ssh_command = f"ssh -i {ssh_private_key_path}{ssh_options} s_{studio._studio.id}@ssh.lightning.ai"
+
+    try:
+        subprocess.run(ssh_command.split())
+    except Exception:
+        # redownload the keys to be sure they are up to date
+        ssh_private_key_path = configure_ssh_internal(force_download=True)
+        try:
+            subprocess.run(ssh_command.split())
+        except Exception:
+            # TODO: make this a generic CLI error
+            raise RuntimeError("Failed to establish SSH connection") from None

lightning_sdk/cli/studio/start.py

@@ -0,0 +1,115 @@
+"""Studio start command."""
+
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.utils.handle_machine_and_gpus_args import handle_machine_and_gpus_args
+from lightning_sdk.cli.utils.richt_print import studio_name_link
+from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
+from lightning_sdk.cli.utils.studio_selection import StudiosMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.machine import CloudProvider, Machine
+from lightning_sdk.studio import VM, Studio
+
+
+@click.command("start")
+@click.option(
+    "--name",
+    help=(
+        "The name of the studio to start. "
+        "If not provided, will try to infer from environment, "
+        "use the default value from the config or prompt for interactive selection."
+    ),
+)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+@click.option("--create", is_flag=True, help="Create the studio if it doesn't exist")
+@click.option(
+    "--machine",
+    help="The machine type to start the studio on. Defaults to CPU-4",
+    type=click.Choice(m.name for m in Machine.__dict__.values() if isinstance(m, Machine) and m._include_in_cli),
+)
+@click.option("--interruptible", is_flag=True, help="Start the studio on an interruptible instance.")
+@click.option(
+    "--cloud-provider",
+    help=(
+        "The cloud provider to start the studio on. Defaults to teamspace default. Only used if --create is specified."
+    ),
+    type=click.Choice(m.name for m in list(CloudProvider)),
+)
+@click.option(
+    "--cloud-account",
+    help="The cloud account to start the studio on. Defaults to teamspace default. Only used if --create is specified.",
+    type=click.STRING,
+)
+@click.option(
+    "--gpus",
+    help="The number and type of GPUs to start the studio on (format: TYPE:COUNT, e.g. L4:4)",
+    type=click.STRING,
+)
+def start_studio(
+    name: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    create: bool = False,
+    machine: str = "CPU",
+    gpus: Optional[str] = None,
+    interruptible: bool = False,
+    cloud_provider: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+) -> None:
+    """Start a Studio.
+
+    Example:
+        lightning studio start --name my-studio
+
+    """
+    return start_impl(
+        name=name,
+        teamspace=teamspace,
+        create=create,
+        machine=machine,
+        gpus=gpus,
+        interruptible=interruptible,
+        cloud_provider=cloud_provider,
+        cloud_account=cloud_account,
+        vm=False,
+    )
+
+
+def start_impl(
+    name: Optional[str],
+    teamspace: Optional[str],
+    create: bool,
+    machine: str,
+    gpus: Optional[str],
+    interruptible: bool,
+    cloud_provider: Optional[str],
+    cloud_account: Optional[str],
+    vm: bool,
+) -> None:
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace=teamspace)
+
+    if cloud_provider is not None:
+        cloud_provider = CloudProvider(cloud_provider)
+
+    if not create:
+        menu = StudiosMenu(resolved_teamspace, vm=vm)
+        studio = menu(studio=name)
+    else:
+        create_cls = VM if vm else Studio
+        studio = create_cls(
+            name=name,
+            teamspace=resolved_teamspace,
+            create_ok=create,
+            cloud_provider=cloud_provider,
+            cloud_account=cloud_account,
+        )
+
+    machine = handle_machine_and_gpus_args(machine, gpus)
+
+    save_studio_to_config(studio)
+
+    Studio.show_progress = True
+    studio.start(machine, interruptible=interruptible)
+    click.echo(f"Studio {studio_name_link(studio)} started successfully")

lightning_sdk/cli/studio/stop.py

@@ -0,0 +1,45 @@
+"""Studio stop command."""
+
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.utils.richt_print import studio_name_link
+from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
+from lightning_sdk.cli.utils.studio_selection import StudiosMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+
+
+@click.command("stop")
+@click.option(
+    "--name",
+    help=(
+        "The name of the studio to stop. "
+        "If not provided, will try to infer from environment, "
+        "use the default value from the config or prompt for interactive selection."
+    ),
+)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+def stop_studio(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+    """Stop a Studio.
+
+    Example:
+        lightning studio stop --name my-studio
+
+    """
+    return stop_impl(name=name, teamspace=teamspace, vm=False)
+
+
+def stop_impl(name: Optional[str], teamspace: Optional[str], vm: bool) -> None:
+    # missing studio_name and teamspace are handled by the studio class
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace=teamspace)
+
+    menu = StudiosMenu(resolved_teamspace, vm=vm)
+    studio = menu(studio=name)
+
+    studio.stop()
+
+    save_studio_to_config(studio)
+
+    click.echo(f"{studio._cls_name} {studio_name_link(studio)} stopped successfully")

lightning_sdk/cli/studio/switch.py

@@ -0,0 +1,66 @@
+"""Studio switch command."""
+
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.utils.richt_print import studio_name_link
+from lightning_sdk.cli.utils.save_to_config import save_studio_to_config
+from lightning_sdk.cli.utils.studio_selection import StudiosMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.machine import Machine
+
+
+@click.command("switch")
+@click.option(
+    "--name",
+    help=(
+        "The name of the studio to switch to a different machine. "
+        "If not provided, will try to infer from environment, "
+        "use the default value from the config or prompt for interactive selection."
+    ),
+)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+@click.option(
+    "--machine",
+    help="The machine type to switch the studio to.",
+    type=click.Choice(m.name for m in Machine.__dict__.values() if isinstance(m, Machine)),
+)
+@click.option("--interruptible", is_flag=True, help="Switch the studio to an interruptible instance.")
+def switch_studio(
+    name: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    machine: Optional[str] = None,
+    interruptible: bool = False,
+) -> None:
+    """Switch a Studio to a different machine type."""
+    return switch_impl(
+        name=name,
+        teamspace=teamspace,
+        machine=machine,
+        interruptible=interruptible,
+        vm=False,
+    )
+
+
+def switch_impl(
+    name: Optional[str],
+    teamspace: Optional[str],
+    machine: Optional[str],
+    interruptible: bool,
+    vm: bool,
+) -> None:
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace=teamspace)
+
+    menu = StudiosMenu(resolved_teamspace, vm=vm)
+    studio = menu(studio=name)
+
+    resolved_machine = Machine.from_str(machine)
+
+    studio.__class__.show_progress = True
+    studio.switch_machine(resolved_machine, interruptible=interruptible)
+
+    save_studio_to_config(studio)
+
+    click.echo(f"{studio._cls_name} {studio_name_link(studio)} switched to machine '{resolved_machine}' successfully")

lightning_sdk/cli/utils/__init__.py

@@ -0,0 +1,7 @@
+from lightning_sdk.cli.utils.coloring import CustomHelpFormatter
+from lightning_sdk.cli.utils.richt_print import rich_to_str
+
+__all__ = [
+    "CustomHelpFormatter",
+    "rich_to_str",
+]

lightning_sdk/cli/utils/cloud_account_map.py

@@ -0,0 +1,10 @@
+from lightning_sdk.api.cloud_account_api import CloudAccountApi
+
+
+def cloud_account_to_display_name(cloud_account: str, teamspace_id: str) -> str:
+    api = CloudAccountApi()
+    cloud_accounts = api.list_global_cloud_accounts(teamspace_id=teamspace_id)
+    for global_cloud_account in cloud_accounts:
+        if global_cloud_account.id == cloud_account:
+            return "Lightning AI"
+    return cloud_account