lightning-sdk 2025.8.6rc2__py3-none-any.whl → 2025.11.5__py3-none-any.whl

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

lightning_sdk/cli/utils/get_base_studio.py

@@ -0,0 +1,24 @@
+from typing import Optional
+
+from lightning_sdk.base_studio import BaseStudio
+
+
+def get_base_studio_id(studio_type: Optional[str]) -> Optional[str]:
+    base_studios = BaseStudio()
+    base_studios = base_studios.list()
+    template_id = None
+
+    if base_studios and len(base_studios):
+        # if not specified by user, use the first existing template studio
+        template_id = base_studios[0].id
+        # else, try to match the provided studio_type to base studio name
+        if studio_type:
+            normalized_studio_type = studio_type.lower().replace(" ", "-")
+            match = next(
+                (s for s in base_studios if s.name.lower().replace(" ", "-") == normalized_studio_type),
+                None,
+            )
+            if match:
+                template_id = match.id
+
+    return template_id

lightning_sdk/cli/utils/handle_machine_and_gpus_args.py

@@ -0,0 +1,69 @@
+from typing import Dict, Optional, Set
+
+import click
+
+from lightning_sdk.machine import DEFAULT_MACHINE, Machine
+
+
+def _split_gpus_spec(gpus: str) -> tuple[str, int]:
+    machine_name, machine_val = gpus.split(":", 1)
+    machine_name = machine_name.strip()
+    machine_val = machine_val.strip()
+
+    if not machine_val.isdigit() or int(machine_val) <= 0:
+        raise ValueError(f"Invalid GPU count '{machine_val}'. Must be a positive integer.")
+
+    machine_num = int(machine_val)
+    return machine_name, machine_num
+
+
+def _construct_available_gpus(machine_options: Dict[str, str]) -> Set[str]:
+    # returns available gpus:count
+    available_gpus = set()
+    for v in machine_options.values():
+        if "_X_" in v:
+            gpu_type_num = v.replace("_X_", ":")
+            available_gpus.add(gpu_type_num)
+        else:
+            available_gpus.add(v)
+    return available_gpus
+
+
+def _get_machine_from_gpus(gpus: str) -> Machine:
+    machine_name = gpus
+    machine_num = 1
+
+    if ":" in gpus:
+        machine_name, machine_num = _split_gpus_spec(gpus)
+
+    machine_options = {
+        m.name.lower(): m.name for m in Machine.__dict__.values() if isinstance(m, Machine) and m._include_in_cli
+    }
+
+    if machine_num == 1:
+        # e.g. gpus=L4 or gpus=L4:1
+        gpu_key = machine_name.lower()
+        try:
+            return machine_options[gpu_key]
+        except KeyError:
+            available = ", ".join(_construct_available_gpus(machine_options))
+            raise ValueError(f"Invalid GPU type '{machine_name}'. Available options: {available}") from None
+
+    # Else: e.g. gpus=L4:4
+    gpu_key = f"{machine_name.lower()}_x_{machine_num}"
+    try:
+        return machine_options[gpu_key]
+    except KeyError:
+        available = ", ".join(_construct_available_gpus(machine_options))
+        raise ValueError(f"Invalid GPU configuration '{gpus}'. Available options: {available}") from None
+
+
+def handle_machine_and_gpus_args(machine: Optional[str], gpus: Optional[str]) -> str:
+    if machine and gpus:
+        raise click.UsageError("Options --machine and --gpus are mutually exclusive. Provide only one.")
+    elif gpus:
+        machine = _get_machine_from_gpus(gpus.strip())
+    elif not machine:
+        machine = DEFAULT_MACHINE
+
+    return machine