lightning-sdk 2025.10.14__py3-none-any.whl → 2025.10.23__py3-none-any.whl

lightning_sdk/cli/studio/connect.py

@@ -2,11 +2,13 @@
 
 import subprocess
 import sys
-from typing import Dict, Optional, Set
+from contextlib import suppress
+from typing import Optional
 
 import click
 
-from lightning_sdk.base_studio import BaseStudio
+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
@@ -16,81 +18,42 @@ from lightning_sdk.machine import CloudProvider, Machine
 from lightning_sdk.studio import Studio
 from lightning_sdk.utils.names import random_unique_name
 
-DEFAULT_MACHINE = "CPU"
 
+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)
 
-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
+    template_id = get_base_studio_id(studio_type)
 
+    if cloud_provider is not None:
+        cloud_provider = CloudProvider(cloud_provider)
 
-def _get_base_studio_id(studio_type: Optional[str]) -> Optional[str]:
-    base_studios = BaseStudio()
-    base_studios = base_studios.list()
-    template_id = None
+    name = name or random_unique_name()
 
-    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
+    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 template_id
+    return resolved_teamspace, cloud_account, template_id, machine, cloud_provider, name
 
 
 @click.command("connect")
@@ -124,6 +87,7 @@ def _get_base_studio_id(studio_type: Optional[str]) -> Optional[str]:
     "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,
@@ -132,29 +96,21 @@ def connect_studio(
     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
     """
-    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)
-
-    name = name or random_unique_name()
-
-    # check for available base studios
-    template_id = _get_base_studio_id(studio_type)
+    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=resolved_teamspace,
+            teamspace=teamspace,
             create_ok=True,
             cloud_provider=cloud_provider,
             cloud_account=cloud_account,
@@ -167,16 +123,10 @@ def connect_studio(
 
     Studio.show_progress = True
 
-    if machine and gpus:
-        raise click.UsageError("Options --machine and --gpu are mutually exclusive. Provide only one.")
-    elif gpus:
-        machine = _get_machine_from_gpus(gpus.strip())
-    elif not machine:
-        machine = DEFAULT_MACHINE
+    machine = handle_machine_and_gpus_args(machine, gpus)
 
     save_studio_to_config(studio)
-    # by default, interruptible is False
-    studio.start(machine=machine, interruptible=False)
+    studio.start(machine=machine, interruptible=interruptible)
 
     ssh_private_key_path = configure_ssh_internal()
 

lightning_sdk/cli/studio/create.py

@@ -4,6 +4,7 @@ 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
@@ -25,18 +26,34 @@ from lightning_sdk.studio import VM, Studio
     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)
+    create_impl(
+        name=name,
+        teamspace=teamspace,
+        cloud_provider=cloud_provider,
+        cloud_account=cloud_account,
+        vm=False,
+        studio_type=studio_type,
+    )
 
 
 def create_impl(
@@ -45,6 +62,7 @@ def create_impl(
     cloud_provider: Optional[str],
     cloud_account: Optional[str],
     vm: bool,
+    studio_type: Optional[str],
 ) -> None:
     menu = TeamspacesMenu()
 
@@ -57,6 +75,9 @@ def create_impl(
     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(
@@ -65,6 +86,7 @@ def create_impl(
             create_ok=True,
             cloud_provider=cloud_provider,
             cloud_account=cloud_account,
+            studio_type=template_id,
         )
     except (RuntimeError, ValueError, ApiException):
         if name:

lightning_sdk/cli/studio/start.py

@@ -4,6 +4,7 @@ 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
@@ -32,8 +33,7 @@ from lightning_sdk.studio import VM, Studio
 @click.option(
     "--cloud-provider",
     help=(
-        "The cloud provider to start the studio on. Defaults to teamspace default. "
-        "Only used if --create is specified."
+        "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)),
 )
@@ -42,11 +42,17 @@ from lightning_sdk.studio import VM, Studio
     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,
@@ -62,6 +68,7 @@ def start_studio(
         teamspace=teamspace,
         create=create,
         machine=machine,
+        gpus=gpus,
         interruptible=interruptible,
         cloud_provider=cloud_provider,
         cloud_account=cloud_account,
@@ -74,6 +81,7 @@ def start_impl(
     teamspace: Optional[str],
     create: bool,
     machine: str,
+    gpus: Optional[str],
     interruptible: bool,
     cloud_provider: Optional[str],
     cloud_account: Optional[str],
@@ -98,6 +106,8 @@ def start_impl(
             cloud_account=cloud_account,
         )
 
+    machine = handle_machine_and_gpus_args(machine, gpus)
+
     save_studio_to_config(studio)
 
     Studio.show_progress = True

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

lightning_sdk/cli/utils/logging.py

@@ -0,0 +1,121 @@
+import shlex
+import sys
+import traceback
+from contextlib import suppress
+from time import time
+from types import TracebackType
+from typing import Optional, Type
+
+import click
+from rich.console import Group
+from rich.panel import Panel
+from rich.syntax import Syntax
+from rich.text import Text
+
+from lightning_sdk.cli.utils import rich_to_str
+from lightning_sdk.constants import _LIGHTNING_DEBUG
+from lightning_sdk.lightning_cloud.openapi.models.v1_create_sdk_command_history_request import (
+    V1CreateSDKCommandHistoryRequest,
+)
+from lightning_sdk.lightning_cloud.openapi.models.v1_sdk_command_history_severity import V1SDKCommandHistorySeverity
+from lightning_sdk.lightning_cloud.openapi.models.v1_sdk_command_history_type import V1SDKCommandHistoryType
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+
+def _log_command(message: str = "", duration: int = 0, error: Optional[str] = None) -> None:
+    original_command = " ".join(shlex.quote(arg) for arg in sys.argv)
+    client = LightningClient(retry=False, max_tries=0)
+
+    body = V1CreateSDKCommandHistoryRequest(
+        command=original_command,
+        duration=duration,
+        message=message,
+        project_id=None,
+        severity=V1SDKCommandHistorySeverity.INFO,
+        type=V1SDKCommandHistoryType.CLI,
+    )
+
+    if error:
+        body.severity = V1SDKCommandHistorySeverity.WARNING if error == "0" else V1SDKCommandHistorySeverity.ERROR
+        body.message = body.message + f" | Error: {error}"
+
+    # limit characters
+    body.message = body.message[:1000]
+
+    with suppress(Exception):
+        client.s_dk_command_history_service_create_sdk_command_history(body=body)
+
+
+def _notify_exception(exception_type: Type[BaseException], value: BaseException, tb: TracebackType) -> None:
+    """CLI won't show tracebacks, just print the exception message."""
+    message = str(value.args[0]) if value.args else str(value) or "An unknown error occurred"
+
+    error_text = Text()
+    error_text.append(f"{exception_type.__name__}: ", style="bold red")
+    error_text.append(message, style="white")
+
+    renderables = [error_text]
+
+    if _LIGHTNING_DEBUG:
+        tb_text = "".join(traceback.format_exception(exception_type, value, tb))
+        renderables.append(Text("\n\nFull traceback:\n", style="bold yellow"))
+        renderables.append(Syntax(tb_text, "python", theme="monokai light", line_numbers=False, word_wrap=True))
+    else:
+        renderables.append(Text("\n\n🐞 To view the full traceback, set: LIGHTNING_DEBUG=1"))
+
+    renderables.append(Text("\n📘 Need help? Run: lightning <command> --help", style="cyan"))
+
+    text = rich_to_str(Panel(Group(*renderables), title="⚡ Lightning CLI Error", border_style="red"))
+    click.echo(text, color=True)
+
+
+def logging_excepthook(exception_type: Type[BaseException], value: BaseException, tb: TracebackType) -> None:
+    try:
+        tb_str = "".join(traceback.format_exception(exception_type, value, tb))
+        ctx = click.get_current_context(silent=True)
+        command_context = ctx.command_path if ctx else "outside_command_context"
+
+        message = (
+            f"Command: {command_context} | Type: {exception_type.__name__!s} | Value: {value!s} | Traceback: {tb_str}"
+        )
+        _log_command(message=message)
+    finally:
+        _notify_exception(exception_type, value, tb)
+
+
+class CommandLoggingGroup(click.Group):
+    def _format_ctx(self, ctx: click.Context) -> str:
+        parts = []
+        for k, v in ctx.params.items():
+            if v is True:
+                parts.append(f"--{k}")
+            elif v is False or v is None:
+                continue
+            else:
+                parts.append(f"--{k} {v}")
+        params = " ".join(parts)
+        args = " ".join(ctx.args)
+        return (
+            f"""Commands: {ctx.command_path} | Subcommand: {ctx.invoked_subcommand} | Params: {params} | Args:{args}"""
+        )
+
+    def invoke(self, ctx: click.Context) -> any:
+        """Overrides the default invoke to wrap command execution with tracking."""
+        start_time = time()
+        error_message = None
+
+        try:
+            return super().invoke(ctx)
+        except click.ClickException as e:
+            error_message = str(e)
+            e.show()
+            ctx.exit(e.exit_code)
+        except Exception as e:
+            error_message = str(e)
+            raise
+        finally:
+            _log_command(
+                message=self._format_ctx(ctx),
+                duration=int(time() - start_time),
+                error=error_message,
+            )

lightning_sdk/cli/utils/ssh_connection.py

@@ -16,7 +16,7 @@ def configure_ssh_internal(force_download: bool = False) -> str:
 
 
 def download_ssh_keys(
-    api_key: str | None,
+    api_key: Optional[str],
     force_download: bool = False,
     ssh_key_name: str = "lightning_rsa",
 ) -> str:

lightning_sdk/constants.py

@@ -29,3 +29,4 @@ class Store:
 
 
 __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__ = Store()
+_LIGHTNING_DISABLE_VERSION_CHECK = int(os.getenv("LIGHTNING_DISABLE_VERSION_CHECK", "0"))

lightning_sdk/helpers.py

@@ -1,4 +1,3 @@
-import functools
 import importlib
 import os
 import sys
@@ -10,48 +9,68 @@ import tqdm
 import tqdm.std
 from packaging import version as packaging_version
 
-__package_name__ = "lightning-sdk"
+from lightning_sdk.constants import _LIGHTNING_DISABLE_VERSION_CHECK
 
 
-@functools.lru_cache(maxsize=1)
-def _get_newer_version(curr_version: str) -> Optional[str]:
-    """Check PyPI for newer versions of ``lightning-sdk``.
-
-    Returning the newest version if different from the current or ``None`` otherwise.
+class VersionChecker:
+    """Handles version checking and upgrade prompts for lightning-sdk.
 
+    This class ensures that version check warnings are only shown once per session,
+    preventing duplicate warnings in multithreaded scenarios.
     """
-    if packaging_version.parse(curr_version).is_prerelease:
-        return None
-    try:
-        response = requests.get(f"https://pypi.org/pypi/{__package_name__}/json")
-        response_json = response.json()
-        releases = response_json["releases"]
-        if curr_version not in releases:
-            # Always return None if not installed from PyPI (e.g. dev versions)
+
+    def __init__(self, package_name: str = "lightning-sdk") -> None:
+        self.package_name = package_name
+        self._warning_shown = False
+        self._cached_version: Optional[str] = None
+
+    def _get_newer_version(self, curr_version: str) -> Optional[str]:
+        """Check PyPI for newer versions of ``lightning-sdk``.
+
+        Returning the newest version if different from the current or ``None`` otherwise.
+        """
+        if self._cached_version is not None:
+            return self._cached_version
+
+        if _LIGHTNING_DISABLE_VERSION_CHECK == 1 or packaging_version.parse(curr_version).is_prerelease:
+            self._cached_version = None
             return None
-        latest_version = response_json["info"]["version"]
-        parsed_version = packaging_version.parse(latest_version)
-        is_invalid = response_json["info"]["yanked"] or parsed_version.is_devrelease or parsed_version.is_prerelease
-        return None if curr_version == latest_version or is_invalid else latest_version
-    except requests.exceptions.RequestException:
-        return None
 
+        try:
+            response = requests.get(f"https://pypi.org/pypi/{self.package_name}/json")
+            response_json = response.json()
+            releases = response_json["releases"]
+            if curr_version not in releases:
+                # Always return None if not installed from PyPI (e.g. dev versions)
+                self._cached_version = None
+                return None
+            latest_version = response_json["info"]["version"]
+            parsed_version = packaging_version.parse(latest_version)
+            is_invalid = response_json["info"]["yanked"] or parsed_version.is_devrelease or parsed_version.is_prerelease
+            self._cached_version = None if curr_version == latest_version or is_invalid else latest_version
+            return self._cached_version
+        except requests.exceptions.RequestException:
+            self._cached_version = None
+            return None
 
-def _check_version_and_prompt_upgrade(curr_version: str) -> None:
-    """Checks that the current version of ``lightning-sdk`` is the latest on PyPI.
+    def check_and_prompt_upgrade(self, curr_version: str) -> None:
+        """Checks that the current version of ``lightning-sdk`` is the latest on PyPI.
 
-    If not, warn the user to upgrade ``lightning-sdk``.
+        If not, warn the user to upgrade ``lightning-sdk``.
+        Tracks if the warning has already been shown in this session to avoid duplicate warnings.
+        """
+        if self._warning_shown:
+            return
 
-    """
-    new_version = _get_newer_version(curr_version)
-    if new_version:
-        warnings.warn(
-            f"A newer version of {__package_name__} is available ({new_version}). "
-            f"Please consider upgrading with `pip install -U {__package_name__}`. "
-            "Not all platform functionality can be guaranteed to work with the current version.",
-            UserWarning,
-        )
-    return
+        new_version = self._get_newer_version(curr_version)
+        if new_version:
+            warnings.warn(
+                f"A newer version of {self.package_name} is available ({new_version}). "
+                f"Please consider upgrading with `pip install -U {self.package_name}`. "
+                "Not all platform functionality can be guaranteed to work with the current version.",
+                UserWarning,
+            )
+            self._warning_shown = True
 
 
 def _set_tqdm_envvars_noninteractive() -> None: