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

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,71 @@
+from typing import Dict, Optional, Set
+
+import click
+
+from lightning_sdk.machine import Machine
+
+DEFAULT_MACHINE = "CPU"
+
+
+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:

lightning_sdk/lightning_cloud/login.py

@@ -5,7 +5,7 @@ import logging
 import os
 import pathlib
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Literal
 from urllib.parse import urlencode
 
 import webbrowser
@@ -22,14 +22,21 @@ from lightning_sdk.lightning_cloud.openapi.api import \
     AuthServiceApi
 from lightning_sdk.lightning_cloud.openapi.models.v1_guest_login_request import \
     V1GuestLoginRequest
+from lightning_sdk.lightning_cloud.openapi.models.v1_token_login_request import \
+    V1TokenLoginRequest
+from lightning_sdk.lightning_cloud.openapi.models.v1_refresh_request import \
+    V1RefreshRequest
 
 logger = logging.getLogger(__name__)
 
+# Authentication override types
+AuthOverride = Literal["auth_token", "api_key", "guest"]
 
 class Keys(Enum):
     # USERNAME = "LIGHTNING_USERNAME"
     USER_ID = "LIGHTNING_USER_ID"
     API_KEY = "LIGHTNING_API_KEY"
+    AUTH_TOKEN = "LIGHTNING_AUTH_TOKEN"
 
     @property
     def suffix(self):
@@ -41,15 +48,18 @@ class Auth:
     # username: Optional[str] = None
     user_id: Optional[str] = None
     api_key: Optional[str] = None
+    auth_token: Optional[str] = None
 
     secrets_file = pathlib.Path(env.LIGHTNING_CREDENTIAL_PATH)
 
     def __post_init__(self):
         for key in Keys:
-            setattr(self, key.suffix, os.environ.get(key.value, None))
+            # Only set from environment if not already set
+            if getattr(self, key.suffix) is None:
+                setattr(self, key.suffix, os.environ.get(key.value, None))
 
         # used by authenticate method
-        self._with_env_var = bool(self.api_key)
+        self._with_env_var = bool(self.api_key or self.auth_token)
 
     def load(self) -> bool:
         """Load credentials from disk and update properties with credentials.
@@ -71,6 +81,7 @@ class Auth:
              token: str = "",
              user_id: str = "",
              api_key: str = "",
+             auth_token: str = "",
              # username: str = "",
         ) -> None:
         """save credentials to disk."""
@@ -81,6 +92,7 @@ class Auth:
                     # f"{Keys.USERNAME.suffix}": username,
                     f"{Keys.USER_ID.suffix}": user_id,
                     f"{Keys.API_KEY.suffix}": api_key,
+                    f"{Keys.AUTH_TOKEN.suffix}": auth_token,
                 },
                 f,
             )
@@ -88,6 +100,7 @@ class Auth:
         # self.username = username
         self.user_id = user_id
         self.api_key = api_key
+        self.auth_token = auth_token
         logger.debug("credentials saved successfully")
 
     @classmethod
@@ -99,15 +112,78 @@ class Auth:
             os.environ.pop(key.value, None)
         logger.debug("credentials removed successfully")
 
-    @property
-    def auth_header(self) -> Optional[str]:
-        """authentication header used by lightning-cloud client."""
-        if self.api_key:
+    def get_auth_header(self, override: Optional[AuthOverride] = None) -> Optional[str]:
+        """Get authentication header with optional override.
+        
+        By default, uses the automatic priority selection (auth_token > api_key).
+        You can override this for specific cases where you need a different auth method.
+        
+        Parameters
+        ----------
+        override : AuthOverride, optional
+            Override the default authentication method:
+            - "auth_token": Force use of JWT auth token (Bearer)
+            - "api_key": Force use of API key (Basic)
+            - "guest": Force use of guest credentials (Basic)
+            - None: Use automatic selection (default)
+            
+        Returns
+        -------
+        Optional[str]
+            The authorization header for the specified method.
+            
+        Raises
+        ------
+        ValueError
+            If the specified override is not available or invalid.
+        """
+        if override == "auth_token":
+            if not self.auth_token:
+                raise ValueError(
+                    "Auth token override requested but no JWT token available. "
+                    "Please use token_login() method first."
+                )
+            return f"Bearer {self.auth_token}"
+        
+        elif override == "api_key":
+            if not self.api_key or not self.user_id:
+                raise ValueError(
+                    "API key override requested but no API key or user ID available. "
+                    "Please set LIGHTNING_API_KEY and LIGHTNING_USER_ID environment variables "
+                    "or use authenticate() method."
+                )
+            token = f"{self.user_id}:{self.api_key}"
+            return f"Basic {base64.b64encode(token.encode('ascii')).decode('ascii')}"  # noqa E501
+        
+        elif override == "guest":
+            if not self.api_key or not self.user_id:
+                raise ValueError(
+                    "Guest override requested but no guest credentials available. "
+                    "Please call guest_login() method first."
+                )
             token = f"{self.user_id}:{self.api_key}"
             return f"Basic {base64.b64encode(token.encode('ascii')).decode('ascii')}"  # noqa E501
-        raise AttributeError(
-            "Authentication Failed, no authentication header available. "
-            "This is most likely a bug in the LightningCloud Framework")
+        
+        elif override is None:
+            # Use the original automatic selection logic (default behavior)
+            if self.auth_token:
+                return f"Bearer {self.auth_token}"
+            elif self.api_key:
+                token = f"{self.user_id}:{self.api_key}"
+                return f"Basic {base64.b64encode(token.encode('ascii')).decode('ascii')}"  # noqa E501
+            else:
+                raise ValueError(
+                    "No authentication credentials available. Please authenticate first using "
+                    "token_login(), guest_login(), or authenticate() methods."
+                )
+        
+        else:
+            raise ValueError(f"Invalid authentication override: {override}")
+
+    @property
+    def auth_header(self) -> Optional[str]:
+        """authentication header used by lightning-cloud client (automatic selection)."""
+        return self.get_auth_header()
 
     def _run_server(self) -> None:
         """start a server to complete authentication."""
@@ -130,6 +206,8 @@ class Auth:
             self._run_server()
             return self.auth_header
 
+        elif self.auth_token:
+            return self.auth_header
         elif self.user_id and self.api_key:
             return self.auth_header
 
@@ -192,6 +270,178 @@ class Auth:
 
         return self.auth_header
 
+    def token_login(self, token_key: str, save_token: bool = True) -> Optional[str]:
+        """Performs token-based authentication.
+        
+        This method sends a request to the token login endpoint to authenticate
+        using an auth token key, optionally saves the returned JWT token, and 
+        returns the authorization header.
+        
+        Parameters
+        ----------
+        token_key : str
+            The auth token key to use for authentication.
+        save_token : bool, optional
+            Whether to save the JWT token for future use. Defaults to True.
+            If False, the token is only used for the current session.
+            
+        Returns
+        -------
+        Optional[str]
+            The authorization header to use for subsequent requests.
+            
+        Raises
+        ------
+        RuntimeError
+            If the token login request fails.
+        ValueError
+            If the response from the server is invalid.
+        """
+        config = Configuration()
+        config.host = env.LIGHTNING_CLOUD_URL
+        api_client = ApiClient(configuration=config)
+        auth_api = AuthServiceApi(api_client)
+
+        logger.debug(f"Attempting token login to {config.host}")
+
+        try:
+            body = V1TokenLoginRequest(token_key=token_key)
+            response = auth_api.auth_service_token_login(body)
+
+        except requests.RequestException as e:
+            logger.error(f"Token login request failed: {e}")
+            raise RuntimeError(
+                "Failed to connect to the token login endpoint. "
+                "Please check your network connection and the server status."
+            ) from e
+
+        # Extract the JWT token from the response
+        jwt_token = getattr(response, "token", None)
+
+        if not jwt_token:
+            logger.error(
+                f"No token received from token login response: {response}"
+            )
+            raise ValueError(
+                "The token login response did not contain a valid JWT token."
+            )
+
+        # Set the JWT token in memory
+        self.auth_token = jwt_token
+        
+        # Optionally save the JWT token to disk
+        if save_token:
+            self.save(auth_token=jwt_token)
+            logger.info("Successfully authenticated using auth token and saved to disk.")
+        else:
+            logger.info("Successfully authenticated using auth token (not saved to disk).")
+
+        return self.auth_header
+
+    def refresh_token(self, duration: int = 43200) -> Optional[str]:
+        """Refreshes the current JWT token.
+        
+        This method sends a request to the refresh endpoint to get a new JWT token
+        with the specified duration, saves the new token, and returns the updated
+        authorization header.
+        
+        Parameters
+        ----------
+        duration : int, optional
+            Duration in seconds for the new token. Can range from 900 seconds (15 minutes)
+            up to a maximum of 129,600 seconds (36 hours), with a default of 43,200 seconds (12 hours).
+            
+        Returns
+        -------
+        Optional[str]
+            The updated authorization header with the new JWT token.
+            
+        Raises
+        ------
+        RuntimeError
+            If the refresh request fails.
+        ValueError
+            If no valid JWT token is available to refresh, or if the response is invalid.
+        """
+        if not self.auth_token:
+            raise ValueError(
+                "No JWT token available to refresh. Please authenticate first using "
+                "token_login() or authenticate() methods."
+            )
+        
+        config = Configuration()
+        config.host = env.LIGHTNING_CLOUD_URL
+        # Set the current auth token as the authorization header
+        config.api_key_prefix['Authorization'] = 'Bearer'
+        config.api_key['Authorization'] = self.auth_token
+        
+        api_client = ApiClient(configuration=config)
+        auth_api = AuthServiceApi(api_client)
+
+        logger.debug(f"Attempting to refresh JWT token with duration {duration} seconds")
+
+        try:
+            body = V1RefreshRequest(duration=str(duration))
+            response = auth_api.auth_service_refresh(body)
+
+        except requests.RequestException as e:
+            logger.error(f"Token refresh request failed: {e}")
+            raise RuntimeError(
+                "Failed to connect to the refresh endpoint. "
+                "Please check your network connection and the server status."
+            ) from e
+
+        # Extract the new JWT token from the response
+        new_jwt_token = getattr(response, "token", None)
+
+        if not new_jwt_token:
+            logger.error(
+                f"No token received from refresh response: {response}"
+            )
+            raise ValueError(
+                "The refresh response did not contain a valid JWT token."
+            )
+
+        # Save the new JWT token
+        self.save(auth_token=new_jwt_token)
+        logger.info("Successfully refreshed JWT token.")
+
+        return self.auth_header
+
+    def create_api_client(self, override: Optional[AuthOverride] = None) -> 'ApiClient':
+        """Create an API client with optional authentication override.
+        
+        This is a convenience method for creating API clients that use a specific
+        authentication method instead of the default automatic selection.
+        
+        Parameters
+        ----------
+        override : AuthOverride, optional
+            Override the default authentication method for this API client.
+            See get_auth_header() for available options.
+            
+        Returns
+        -------
+        ApiClient
+            Configured API client with the specified authentication method.
+        """
+        from lightning_sdk.lightning_cloud.openapi import ApiClient, Configuration
+        
+        config = Configuration()
+        config.host = env.LIGHTNING_CLOUD_URL
+        
+        # Get the auth header for the specified override
+        auth_header = self.get_auth_header(override)
+        
+        # Create the API client
+        client = ApiClient(configuration=config)
+        
+        # Set the Authorization header directly in default_headers
+        if auth_header:
+            client.set_default_header('Authorization', auth_header)
+        
+        return client
+
 
 class AuthServer: