lightning-sdk 2025.8.14__py3-none-any.whl → 2025.8.18__py3-none-any.whl

lightning_sdk/cli/{mmts_menu.py → legacy/mmts_menu.py}

@@ -3,7 +3,7 @@ from typing import Dict, List, Optional
 from rich.console import Console
 from simple_term_menu import TerminalMenu
 
-from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.cli.legacy.exceptions import StudioCliError
 from lightning_sdk.mmt import MMT
 from lightning_sdk.teamspace import Teamspace
 

lightning_sdk/cli/{open.py → legacy/open.py}

@@ -6,8 +6,8 @@ from typing import Optional
 import click
 from rich.console import Console
 
-from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
-from lightning_sdk.cli.upload import _upload_folder
+from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.legacy.upload import _upload_folder
 from lightning_sdk.studio import Studio
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.utils.resolve import _get_studio_url

lightning_sdk/cli/{stop.py → legacy/stop.py}

@@ -3,7 +3,7 @@ from typing import Optional
 import click
 from rich.console import Console
 
-from lightning_sdk.cli.job_and_mmt_action import _JobAndMMTAction
+from lightning_sdk.cli.legacy.job_and_mmt_action import _JobAndMMTAction
 from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.studio import Studio
 

lightning_sdk/cli/{teamspace_menu.py → legacy/teamspace_menu.py}

@@ -4,7 +4,7 @@ from rich.console import Console
 from simple_term_menu import TerminalMenu
 
 from lightning_sdk.api import OrgApi
-from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.cli.legacy.exceptions import StudioCliError
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
 from lightning_sdk.utils.resolve import _get_authed_user

lightning_sdk/cli/{upload.py → legacy/upload.py}

@@ -14,9 +14,9 @@ from tqdm import tqdm
 
 from lightning_sdk.api.lit_container_api import DockerNotRunningError, LCRAuthFailedError, LitContainerApi
 from lightning_sdk.api.utils import _get_cloud_url
-from lightning_sdk.cli.exceptions import StudioCliError
-from lightning_sdk.cli.studios_menu import _StudiosMenu
-from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.legacy.exceptions import StudioCliError
+from lightning_sdk.cli.legacy.studios_menu import _StudiosMenu
+from lightning_sdk.cli.legacy.teamspace_menu import _TeamspacesMenu
 from lightning_sdk.constants import _LIGHTNING_DEBUG
 from lightning_sdk.models import upload_model as _upload_model
 from lightning_sdk.studio import Studio

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,22 @@
+"""Studio CLI commands."""
+
+import click
+
+
+def register_commands(group: click.Group) -> None:
+    """Register studio commands with the given group."""
+    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)

lightning_sdk/cli/studio/create.py

@@ -0,0 +1,53 @@
+"""Studio start command."""
+
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.machine import CloudProvider
+from lightning_sdk.studio import Studio
+
+
+@click.command("create")
+@click.argument("studio_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)),
+)
+def create_studio(
+    studio_name: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    cloud_provider: Optional[str] = None,
+) -> None:
+    """Create a new Studio.
+
+    Example:
+        lightning studio create [STUDIO_NAME]
+
+    STUDIO_NAME: the name of the studio to create.
+
+    If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
+    """
+    if teamspace is not None:
+        resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
+        if resolved_teamspace is None:
+            raise ValueError(
+                f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
+                "Does the teamspace exist?"
+            )
+    else:
+        resolved_teamspace = None
+
+    try:
+        studio = Studio(studio_name, teamspace=resolved_teamspace, create_ok=True, cloud_provider=cloud_provider)
+    except (RuntimeError, ValueError, ApiException) as e:
+        print(e)
+        if studio_name:
+            raise ValueError(f"Could not create Studio: '{studio_name}'. Does the Studio exist?") from None
+        raise ValueError(f"Could not create Studio: '{studio_name}'. Please provide a Studio name") from None
+
+    click.echo(f"Studio '{studio.name}' created successfully")

lightning_sdk/cli/studio/delete.py

@@ -0,0 +1,44 @@
+"""Studio delete command."""
+
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
+from lightning_sdk.studio import Studio
+
+
+@click.command("delete")
+@click.argument("studio_name", required=False)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+def delete_studio(studio_name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+    """Delete a Studio.
+
+    Example:
+      lightning studio delete [STUDIO_NAME]
+
+    STUDIO_NAME: the name of the studio to delete.
+
+    If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
+    """
+    # missing studio_name and teamspace are handled by the studio class
+    if teamspace is not None:
+        resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
+        if resolved_teamspace is None:
+            raise ValueError(
+                f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
+                "Does the teamspace exist?"
+            )
+    else:
+        resolved_teamspace = None
+
+    try:
+        studio = Studio(studio_name, teamspace=resolved_teamspace, create_ok=False)
+        studio.delete()
+    except Exception:
+        # TODO: make this a generic CLI error
+        if studio_name:
+            raise ValueError(f"Could not delete Studio: '{studio_name}'. Does the Studio exist?") from None
+        raise ValueError("No studio name provided. Use 'lightning studio delete <name>' to delete a studio.") from None
+
+    click.echo(f"Studio '{studio.name}' deleted successfully")

lightning_sdk/cli/studio/list.py

@@ -0,0 +1,67 @@
+"""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.resolve import resolve_teamspace_owner_name_format
+from lightning_sdk.cli.utils.richt_print import rich_to_str
+from lightning_sdk.studio import Studio
+
+
+@click.command("list")
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+@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, sort_by: Optional[str] = None) -> None:
+    """List Studios in a teamspace.
+
+    Example:
+        lightning studio list --teamspace owner/teamspace
+
+    """
+    teamspace_resolved = resolve_teamspace_owner_name_format(teamspace)
+
+    if teamspace_resolved is None:
+        # TODO: make this a generic CLI error
+        raise ValueError(f"Could not resolve teamspace: {teamspace}")
+
+    studios = 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(studios, key=_sort_studios_key(sort_by)):
+        table.add_row(
+            studio.name,
+            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,112 @@
+"""Studio SSH command."""
+
+import os
+import platform
+import subprocess
+import uuid
+from pathlib import Path
+from typing import List, Optional
+
+import click
+
+from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
+from lightning_sdk.lightning_cloud.login import Auth
+from lightning_sdk.studio import Studio
+from lightning_sdk.utils.config import _DEFAULT_CONFIG_FILE_PATH
+
+
+@click.command("ssh")
+@click.argument("studio_name", required=False)
+@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(
+    studio_name: Optional[str] = None, teamspace: Optional[str] = None, option: Optional[List[str]] = None
+) -> None:
+    """SSH into a Studio.
+
+    Example:
+        lightning studio ssh [STUDIO_NAME]
+
+    STUDIO_NAME: the name of the studio to SSH into.
+
+    If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
+    """
+    auth = Auth()
+    auth.authenticate()
+    ssh_private_key_path = _download_ssh_keys(auth.api_key, force_download=False)
+
+    if teamspace is not None:
+        resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
+        if resolved_teamspace is None:
+            raise ValueError(
+                f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
+                "Does the teamspace exist?"
+            )
+    else:
+        resolved_teamspace = None
+
+    studio = Studio(studio_name, teamspace=resolved_teamspace)
+
+    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
+        _download_ssh_keys(auth.api_key, 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
+
+
+def _download_ssh_keys(
+    api_key: str,
+    force_download: bool = False,
+    ssh_key_name: str = "lightning_rsa",
+) -> None:
+    """Download the SSH key for a User."""
+    ssh_private_key_path = os.path.join(os.path.expanduser(os.path.dirname(_DEFAULT_CONFIG_FILE_PATH)), ssh_key_name)
+
+    os.makedirs(os.path.dirname(ssh_private_key_path), exist_ok=True)
+
+    if not os.path.isfile(ssh_private_key_path) or force_download:
+        key_id = str(uuid.uuid4())
+        _download_file(
+            f"https://lightning.ai/setup/ssh-gen?t={api_key}&id={key_id}&machineName={platform.node()}",
+            ssh_private_key_path,
+            overwrite=True,
+            chmod=0o600,
+        )
+        _download_file(
+            f"https://lightning.ai/setup/ssh-public?t={api_key}&id={key_id}",
+            ssh_private_key_path + ".pub",
+            overwrite=True,
+        )
+
+    return ssh_private_key_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 os.path.isfile(local_path) 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:
+        os.chmod(local_path, 0o600)

lightning_sdk/cli/studio/start.py

@@ -0,0 +1,63 @@
+"""Studio start command."""
+
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.machine import CloudProvider, Machine
+from lightning_sdk.studio import Studio
+
+
+@click.command("start")
+@click.argument("studio_name", required=False)
+@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.",
+    type=click.Choice(m.name for m in list(CloudProvider)),
+)
+def start_studio(
+    studio_name: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    create: bool = False,
+    machine: Optional[str] = None,
+    interruptible: bool = False,
+    cloud_provider: Optional[str] = None,
+) -> None:
+    """Start a Studio.
+
+    Example:
+        lightning studio start [STUDIO_NAME]
+
+    STUDIO_NAME: the name of the studio to start.
+
+    If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
+    """
+    if teamspace is not None:
+        resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
+        if resolved_teamspace is None:
+            raise ValueError(
+                f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
+                "Does the teamspace exist?"
+            )
+    else:
+        resolved_teamspace = None
+
+    try:
+        studio = Studio(studio_name, teamspace=resolved_teamspace, create_ok=create, cloud_provider=cloud_provider)
+    except (RuntimeError, ValueError, ApiException):
+        if studio_name:
+            raise ValueError(f"Could not start Studio: '{studio_name}'. Does the Studio exist?") from None
+        raise ValueError(f"Could not start Studio: '{studio_name}'. Please provide a Studio name") from None
+
+    studio.start(machine, interruptible=interruptible)
+    click.echo(f"Studio '{studio.name}' started successfully")

lightning_sdk/cli/studio/stop.py

@@ -0,0 +1,44 @@
+"""Studio stop command."""
+
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
+from lightning_sdk.studio import Studio
+
+
+@click.command("stop")
+@click.argument("studio_name", required=False)
+@click.option("--teamspace", help="Override default teamspace (format: owner/teamspace)")
+def stop_studio(studio_name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+    """Stop a Studio.
+
+    Example:
+        lightning studio stop [STUDIO_NAME]
+
+    STUDIO_NAME: the name of the studio to stop.
+
+    If STUDIO_NAME is not provided, will try to infer from environment or use the default value from the config.
+    """
+    # missing studio_name and teamspace are handled by the studio class
+    if teamspace is not None:
+        resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
+        if resolved_teamspace is None:
+            raise ValueError(
+                f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
+                "Does the teamspace exist?"
+            )
+    else:
+        resolved_teamspace = None
+
+    try:
+        studio = Studio(studio_name, teamspace=resolved_teamspace)
+        studio.stop()
+    except Exception:
+        # TODO: make this a generic CLI error
+        if studio_name:
+            raise ValueError(f"Could not stop studio: '{studio_name}'. Does the studio exist?") from None
+        raise ValueError("No studio name provided. Use 'lightning studio stop <name>' to stop a studio.") from None
+
+    click.echo(f"Studio '{studio.name}' stopped successfully")

lightning_sdk/cli/studio/switch.py

@@ -0,0 +1,52 @@
+"""Studio switch command."""
+
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.utils.resolve import resolve_teamspace_owner_name_format
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.machine import Machine
+from lightning_sdk.studio import Studio
+
+
+@click.command("switch")
+@click.argument("studio_name", required=False)
+@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(
+    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."""
+    if teamspace is not None:
+        resolved_teamspace = resolve_teamspace_owner_name_format(teamspace)
+        if resolved_teamspace is None:
+            raise ValueError(
+                f"Could not resolve teamspace: '{teamspace}'. Teamspace should be specified as 'owner/name'. "
+                "Does the teamspace exist?"
+            )
+    else:
+        resolved_teamspace = None
+
+    try:
+        studio = Studio(
+            studio_name,
+            teamspace=resolved_teamspace,
+        )
+    except (RuntimeError, ValueError, ApiException):
+        if studio_name:
+            raise ValueError(f"Could not switch Studio: '{studio_name}'. Does the Studio exist?") from None
+        raise ValueError(f"Could not switch Studio: '{studio_name}'. Please provide a Studio name") from None
+
+    resolved_machine = Machine.from_str(machine)
+    studio.switch_machine(resolved_machine, interruptible=interruptible)
+
+    click.echo(f"Studio '{studio.name}' 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/resolve.py

@@ -0,0 +1,28 @@
+from typing import Optional
+
+from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.utils.resolve import _resolve_teamspace
+
+
+def resolve_teamspace_owner_name_format(teamspace_name: str) -> Optional[Teamspace]:
+    teamspace_resolved = None
+    if teamspace_name is None:
+        return _resolve_teamspace(None, None, None)
+
+    splits = teamspace_name.split("/")
+    if len(splits) == 1:
+        try:
+            teamspace_resolved = _resolve_teamspace(teamspace_name, None, None)
+        except Exception:
+            teamspace_resolved = None
+
+    elif len(splits) == 2:
+        try:
+            try:
+                teamspace_resolved = _resolve_teamspace(splits[1], splits[0], None)
+            except Exception:
+                teamspace_resolved = _resolve_teamspace(splits[1], None, splits[0])
+        except Exception:
+            teamspace_resolved = None
+
+    return teamspace_resolved

lightning_sdk/cli/utils/richt_print.py

@@ -0,0 +1,11 @@
+import os
+from typing import Any
+
+from rich.console import Console
+
+
+def rich_to_str(*renderables: Any) -> str:
+    with open(os.devnull, "w") as f:
+        console = Console(file=f, record=True)
+        console.print(*renderables)
+    return console.export_text(styles=True)

lightning_sdk/lightning_cloud/openapi/api/billing_service_api.py

@@ -663,6 +663,7 @@ class BillingServiceApi(object):
 
         :param async_req bool
         :param str org_id:
+        :param str project_id:
         :return: V1BillingSubscription
                  If the method is called asynchronously,
                  returns the request thread.
@@ -684,12 +685,13 @@ class BillingServiceApi(object):
 
         :param async_req bool
         :param str org_id:
+        :param str project_id:
         :return: V1BillingSubscription
                  If the method is called asynchronously,
                  returns the request thread.
         """
 
-        all_params = ['org_id']  # noqa: E501
+        all_params = ['org_id', 'project_id']  # noqa: E501
         all_params.append('async_req')
         all_params.append('_return_http_data_only')
         all_params.append('_preload_content')
@@ -712,6 +714,8 @@ class BillingServiceApi(object):
         query_params = []
         if 'org_id' in params:
             query_params.append(('orgId', params['org_id']))  # noqa: E501
+        if 'project_id' in params:
+            query_params.append(('projectId', params['project_id']))  # noqa: E501
 
         header_params = {}