lightning-sdk 2025.11.13.post0__py3-none-any.whl

lightning_sdk/cli/license/list.py

@@ -0,0 +1,45 @@
+"""License list command."""
+
+from typing import Mapping
+
+import click
+from rich.table import Table
+
+from lightning_sdk.cli.utils.richt_print import rich_to_str
+from lightning_sdk.utils.config import _DEFAULT_CONFIG_FILE_PATH, Config, DefaultConfigKeys
+
+
+@click.command("list")
+@click.option("--include-key", help="Print the key as well", is_flag=True)
+@click.option("--config-file", help="Path to the config file")
+def list_licenses(include_key: bool, config_file: str = _DEFAULT_CONFIG_FILE_PATH) -> None:
+    """List configured licenses.
+
+    Example:
+        lightning license list --include-key
+
+    """
+    return list_impl(include_key=include_key, config_path=config_file)
+
+
+def list_impl(include_key: bool, config_path: str) -> None:
+    cfg = Config(config_file=config_path)
+
+    license_cfg = cfg.get_sub_config(DefaultConfigKeys.license)
+
+    if isinstance(license_cfg, Mapping):
+        table = Table(
+            pad_edge=True,
+        )
+
+        table.add_column("Product")
+        table.add_column("License Key")
+
+        # sort by product_name
+        for product_name, license_key in sorted(license_cfg.items(), key=lambda x: x[0]):
+            table.add_row(product_name, license_key if include_key else "********")
+
+        click.echo(rich_to_str(table), color=True)
+
+    else:
+        click.echo("No licenses configured!")

lightning_sdk/cli/license/set.py

@@ -0,0 +1,13 @@
+import click
+
+from lightning_sdk.utils.config import _DEFAULT_CONFIG_FILE_PATH, Config, DefaultConfigKeys
+
+
+@click.command("set")
+@click.argument("product_name")
+@click.argument("license_key")
+@click.option("--config-file", help="Path to the config file")
+def set_license(product_name: str, license_key: str, config_file: str = _DEFAULT_CONFIG_FILE_PATH) -> None:
+    """Set a license key for a given product."""
+    cfg = Config(config_file)
+    cfg.set(f"{DefaultConfigKeys.license}.{product_name}", license_key)

lightning_sdk/cli/mmt/__init__.py

@@ -0,0 +1,7 @@
+"""MMT CLI commands."""
+
+import click
+
+
+def register_commands(group: click.Group) -> None:
+    """Register MMT commands with the given group."""

lightning_sdk/cli/studio/__init__.py

@@ -0,0 +1,26 @@
+"""Studio CLI commands."""
+
+import click
+
+
+def register_commands(group: click.Group) -> None:
+    """Register studio commands with the given group."""
+    from lightning_sdk.cli.studio.connect import connect_studio
+    from lightning_sdk.cli.studio.cp import cp_studio_file
+    from lightning_sdk.cli.studio.create import create_studio
+    from lightning_sdk.cli.studio.delete import delete_studio
+    from lightning_sdk.cli.studio.list import list_studios
+    from lightning_sdk.cli.studio.ssh import ssh_studio
+    from lightning_sdk.cli.studio.start import start_studio
+    from lightning_sdk.cli.studio.stop import stop_studio
+    from lightning_sdk.cli.studio.switch import switch_studio
+
+    group.add_command(delete_studio)
+    group.add_command(create_studio)
+    group.add_command(list_studios)
+    group.add_command(ssh_studio)
+    group.add_command(start_studio)
+    group.add_command(stop_studio)
+    group.add_command(switch_studio)
+    group.add_command(connect_studio)
+    group.add_command(cp_studio_file)

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

@@ -0,0 +1,138 @@
+"""Studio cp command."""
+
+from pathlib import Path
+from typing import Optional, TypedDict
+
+import click
+from rich.console import Console
+
+from lightning_sdk.api.utils import _get_cloud_url
+from lightning_sdk.cli.legacy.exceptions import StudioCliError
+from lightning_sdk.cli.utils.owner_selection import OwnerMenu
+from lightning_sdk.cli.utils.studio_selection import StudiosMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.studio import Studio
+
+
+@click.command("cp")
+@click.argument("source", nargs=1)
+@click.argument("destination", nargs=1)
+def cp_studio_file(source: str, destination: str, teamspace: Optional[str] = None) -> None:
+    """Copy a Studio file.
+
+    SOURCE: Source file to copy from. For Studio files, use the format lit://<owner>/<my-teamspace>/studios/<my-studio>/<filepath>.
+
+    DESTINATION: Destination file to copy to. For Studio files, use the format lit://<owner>/<my-teamspace>/studios/<my-studio>/<filepath>.
+
+    Example:
+        lightning studio cp source.txt lit://<owner>/<my-teamspace>/studios/<my-studio>/destination.txt
+
+    """
+    return cp_impl(source=source, destination=destination)
+
+
+def cp_impl(source: str, destination: str) -> None:
+    if "lit://" in source and "lit://" in destination:
+        raise ValueError("Both source and destination cannot be Studio files.")
+    elif "lit://" not in source and "lit://" not in destination:
+        raise ValueError("Either source or destination must be a Studio file.")
+    elif "lit://" in source:
+        # Download from Studio to local
+        cp_download(studio_file_path=source, local_file_path=destination)
+    else:
+        # Upload from local to Studio
+        cp_upload(local_file_path=source, studio_file_path=destination)
+
+
+class StudioPathResult(TypedDict):
+    owner: Optional[str]
+    teamspace: Optional[str]
+    studio: Optional[str]
+    destination: Optional[str]
+
+
+def parse_studio_path(studio_path: str) -> StudioPathResult:
+    path_string = studio_path.removeprefix("lit://")
+    if not path_string:
+        raise ValueError("Studio path cannot be empty after prefix")
+
+    result: StudioPathResult = {"owner": None, "teamspace": None, "studio": None, "destination": None}
+
+    if "/studios/" in path_string:
+        prefix_part, suffix_part = path_string.split("/studios/", 1)
+
+        # org and teamspace
+        if prefix_part:
+            org_ts_components = prefix_part.split("/")
+            if len(org_ts_components) == 2:
+                result["owner"], result["teamspace"] = org_ts_components
+            elif len(org_ts_components) == 1:
+                result["teamspace"] = org_ts_components[0]
+            else:
+                raise ValueError(f"Invalid format: '{prefix_part}'")
+
+        # studio and destination
+        path_parts = suffix_part.split("/")
+
+    else:
+        # studio and destination
+        path_parts = path_string.split("/")
+
+    if not path_parts or len(path_parts) < 2:
+        raise ValueError("Invalid: Missing studio name.")
+
+    result["studio"] = path_parts[0]
+    result["destination"] = "/".join(path_parts[1:])
+
+    return result
+
+
+def cp_upload(
+    local_file_path: str,
+    studio_file_path: str,
+) -> None:
+    console = Console()
+    if Path(local_file_path).is_dir():
+        raise StudioCliError(
+            f"The provided path is a folder: {local_file_path}. Use `lightning upload folder` instead."
+        )
+    if not Path(local_file_path).exists():
+        raise FileNotFoundError(f"The provided path does not exist: {local_file_path}.")
+
+    studio_path_result = parse_studio_path(studio_file_path)
+
+    selected_studio = resolve_studio(
+        studio_path_result["studio"], studio_path_result["teamspace"], studio_path_result["owner"]
+    )
+    console.print(f"Uploading to {selected_studio.teamspace.name}/{selected_studio.name}")
+
+    selected_studio.upload_file(local_file_path, studio_file_path)
+
+    studio_url = (
+        _get_cloud_url().replace(":443", "")
+        + "/"
+        + selected_studio.owner.name
+        + "/"
+        + selected_studio.teamspace.name
+        + "/studios/"
+        + selected_studio.name
+    )
+    console.print(f"See your file at {studio_url}")
+
+
+def cp_download(
+    studio_file_path: str,
+    local_file_path: str,
+) -> None:
+    raise NotImplementedError("Download functionality is not implemented yet.")
+
+
+def resolve_studio(studio_name: Optional[str], teamspace: Optional[str], owner: Optional[str]) -> Studio:
+    owner_menu = OwnerMenu()
+    resolved_owner = owner_menu(owner=owner)
+
+    teamspace_menu = TeamspacesMenu(resolved_owner)
+    resolved_teamspace = teamspace_menu(teamspace=teamspace)
+
+    studio_menu = StudiosMenu(resolved_teamspace)
+    return studio_menu(studio=studio_name)

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