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

lightning_sdk/__init__.py

@@ -2,7 +2,7 @@ from lightning_sdk.agents import Agent
 from lightning_sdk.ai_hub import AIHub
 from lightning_sdk.constants import __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__  # noqa: F401
 from lightning_sdk.deployment import Deployment
-from lightning_sdk.helpers import _check_version_and_prompt_upgrade, _set_tqdm_envvars_noninteractive
+from lightning_sdk.helpers import VersionChecker, _set_tqdm_envvars_noninteractive
 from lightning_sdk.job import Job
 from lightning_sdk.machine import CloudProvider, Machine
 from lightning_sdk.mmt import MMT
@@ -35,6 +35,9 @@ __all__ = [
     "VM",
 ]
 
-__version__ = "2025.10.14"
-_check_version_and_prompt_upgrade(__version__)
+__version__ = "2025.10.23"
+
+_version_checker = VersionChecker()
+_version_checker.check_and_prompt_upgrade(__version__)
+
 _set_tqdm_envvars_noninteractive()

lightning_sdk/api/base_studio_api.py

@@ -15,24 +15,28 @@ class BaseStudioApi:
     def __init__(self) -> None:
         self._client = LightningClient(retry=False, max_tries=0)
 
-    def get_base_studio(self, base_studio_id: str, org_id: str) -> V1CloudSpaceEnvironmentTemplate:
+    def get_base_studio(self, base_studio_id: str, org_id: Optional[str] = None) -> V1CloudSpaceEnvironmentTemplate:
         """Retrieve the base studio by its ID."""
         try:
             return self._client.cloud_space_environment_template_service_get_cloud_space_environment_template(
-                base_studio_id, org_id=org_id
+                base_studio_id, org_id=org_id or ""
             )
         except ValueError as e:
             raise ValueError(f"Base studio {base_studio_id} does not exist") from e
 
-    def get_all_base_studios(self, org_id: str, managed: bool = True) -> V1ListCloudSpaceEnvironmentTemplatesResponse:
+    def get_all_base_studios(self, org_id: Optional[str]) -> V1ListCloudSpaceEnvironmentTemplatesResponse:
         """Retrieve all base studios for a given organization."""
-        if managed:
-            return self._client.cloud_space_environment_template_service_list_managed_cloud_space_environment_templates(
-                org_id=org_id
-            )
-        return self._client.cloud_space_environment_template_service_list_cloud_space_environment_templates(
-            org_id=org_id
+        result = self._client.cloud_space_environment_template_service_list_managed_cloud_space_environment_templates(
+            org_id=org_id or ""
         )
+        if org_id is not None:
+            org_templates = (
+                self._client.cloud_space_environment_template_service_list_cloud_space_environment_templates(
+                    org_id=org_id
+                )
+            )
+            result.templates = result.templates + org_templates.templates
+        return result
 
     def update_base_studio(
         self,

lightning_sdk/api/job_api.py

@@ -250,6 +250,7 @@ class JobApiV2:
         artifacts_local: Optional[str],  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str],  # deprecated in favor of path_mappings
         max_runtime: Optional[int] = None,
+        reuse_snapshot: bool = True,
     ) -> V1Job:
         body = self._create_job_body(
             name=name,
@@ -267,6 +268,7 @@ class JobApiV2:
             artifacts_local=artifacts_local,
             artifacts_remote=artifacts_remote,
             max_runtime=max_runtime,
+            reuse_snapshot=reuse_snapshot,
         )
 
         job: V1Job = self._client.jobs_service_create_job(project_id=teamspace_id, body=body)
@@ -288,6 +290,7 @@ class JobApiV2:
         path_mappings: Optional[Dict[str, str]],
         artifacts_local: Optional[str],  # deprecated in favor of path_mappings
         artifacts_remote: Optional[str],  # deprecated in favor of path_mappings)
+        reuse_snapshot: bool,
         max_runtime: Optional[int] = None,
         machine_image_version: Optional[str] = None,
     ) -> ProjectIdJobsBody:
@@ -298,7 +301,7 @@ class JobApiV2:
 
         instance_name = _machine_to_compute_name(machine)
 
-        run_id = __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__[studio_id] if studio_id is not None else ""
+        run_id = __GLOBAL_LIGHTNING_UNIQUE_IDS_STORE__[studio_id] if (studio_id is not None and reuse_snapshot) else ""
 
         path_mappings_list = resolve_path_mappings(
             mappings=path_mappings or {},

lightning_sdk/api/license_api.py

@@ -1,70 +1,37 @@
-import os
-from typing import Optional
-from urllib.parse import urlencode
-
-from lightning_sdk.lightning_cloud import env
-from lightning_sdk.lightning_cloud.rest_client import LightningClient
-
-LICENSE_CODE = os.environ.get("LICENSE_CODE", "d9s79g79ss")
-# https://lightning.ai/home?settings=licenses
-LICENSE_SIGNING_URL = f"{env.LIGHTNING_CLOUD_URL}?settings=licenses"
-
-
-def generate_url_user_settings(name: str, redirect_to: str = LICENSE_SIGNING_URL) -> str:
-    params = urlencode({"redirectTo": redirect_to, "okbhrt": LICENSE_CODE, "licenseName": name})
-    return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}"
+from lightning_sdk.api.utils import _get_cloud_url as _cloud_url
+from lightning_sdk.lightning_cloud.login import Auth
+from lightning_sdk.lightning_cloud.openapi import LicenseKeyValidateBody, ProductLicenseServiceApi
 
 
 class LicenseApi:
-    _client_authenticated: LightningClient = None
-    _client_public: LightningClient = None
-
-    @property
-    def client_public(self) -> LightningClient:
-        if not self._client_public:
-            self._client_public = LightningClient(retry=False, max_tries=0, with_auth=False)
-        return self._client_public
-
-    @property
-    def client_authenticated(self) -> LightningClient:
-        if not self._client_authenticated:
-            self._client_authenticated = LightningClient(retry=True, max_tries=3, with_auth=True)
-        return self._client_authenticated
+    def __init__(self, login_token: str) -> None:
+        self._cloud_url = _cloud_url()
+        self._auth = Auth()
+        self._auth.token_login(login_token, save_token=True)
+        self._client = self._auth.create_api_client()
+        self._api = ProductLicenseServiceApi(self._client)
 
-    def valid_license(
-        self,
-        license_key: str,
-        product_name: str,
-        product_version: Optional[str] = None,
-        product_type: str = "package",
-    ) -> bool:
-        """Check if the license key is valid.
+    def validate_license(self, license_key: str, product_id: str) -> bool:
+        """Validate a license key for a specific product.
 
         Args:
-            license_key: The license key to check.
-            product_name: The name of the product.
-            product_version: The version of the product.
-            product_type: The type of the product. Default is "package".
+            license_key: The license key to validate
+            product_id: The product ID
 
         Returns:
-            True if the license key is valid, False otherwise.
-        """
-        response = self.client_public.product_license_service_validate_product_license(
-            license_key=license_key,
-            product_name=product_name,
-            product_version=product_version,
-            product_type=product_type,
-        )
-        return response.valid
+            bool: True if license is valid, False otherwise
 
-    def list_user_licenses(self, user_id: str) -> list:
-        """List all licenses for a user.
+        Raises:
+            Exception: If license validation fails
+        """
+        try:
+            response = self._api.product_license_service_validate_license(
+                body=LicenseKeyValidateBody(product_id=product_id), license_key=license_key
+            )
+            return response.is_valid
+        except Exception:
+            raise InvalidLicenseError(f"Invalid license key {license_key} for product {product_id}") from None
 
-        Args:
-            user_id: The ID of the user.
 
-        Returns:
-            A list of licenses for the user.
-        """
-        response = self.client_authenticated.product_license_service_list_user_licenses(user_id=user_id)
-        return response.licenses
+class InvalidLicenseError(Exception):
+    pass

lightning_sdk/api/studio_api.py

@@ -589,7 +589,12 @@ class StudioApi:
         )
 
     def duplicate_studio(
-        self, studio_id: str, teamspace_id: str, target_teamspace_id: str, machine: Machine = Machine.CPU
+        self,
+        studio_id: str,
+        teamspace_id: str,
+        target_teamspace_id: str,
+        machine: Machine = Machine.CPU,
+        new_name: Optional[str] = None,
     ) -> Dict[str, Any]:
         """Duplicates the given Studio from a given Teamspace into a given target Teamspace."""
         target_teamspace = self._client.projects_service_get_project(target_teamspace_id)
@@ -604,7 +609,7 @@ class StudioApi:
             init_kwargs["org"] = OrgApi()._get_org_by_id(target_teamspace.owner_id).name
 
         new_cloudspace = self._client.cloud_space_service_fork_cloud_space(
-            IdForkBody1(target_project_id=target_teamspace_id), project_id=teamspace_id, id=studio_id
+            IdForkBody1(target_project_id=target_teamspace_id, new_name=new_name), project_id=teamspace_id, id=studio_id
         )
 
         while self.get_studio_by_id(new_cloudspace.id, target_teamspace_id).state != V1CloudSpaceState.READY:

lightning_sdk/base_studio.py

@@ -3,11 +3,11 @@ from typing import List, Optional, Union
 
 from lightning_sdk.api.base_studio_api import BaseStudioApi
 from lightning_sdk.api.user_api import UserApi
-from lightning_sdk.lightning_cloud import login
 from lightning_sdk.lightning_cloud.openapi.models.v1_cloud_space_environment_type import V1CloudSpaceEnvironmentType
 from lightning_sdk.organization import Organization
+from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
-from lightning_sdk.utils.resolve import _resolve_org, _resolve_user
+from lightning_sdk.utils.resolve import _resolve_teamspace
 
 
 @dataclass
@@ -24,6 +24,7 @@ class BaseStudio:
     def __init__(
         self,
         name: Optional[str] = None,
+        teamspace: Optional[Union[str, Teamspace]] = None,
         org: Optional[Union[str, Organization]] = None,
         user: Optional[Union[str, User]] = None,
     ) -> None:
@@ -38,26 +39,35 @@ class BaseStudio:
         Raises:
             ConnectionError: If there is an issue with the authentication process.
         """
-        self._auth = login.Auth()
-        self._user = None
+        self._teamspace = None
 
-        try:
-            self._auth.authenticate()
-            if user is None:
-                self._user = User(name=UserApi()._get_user_by_id(self._auth.user_id).username)
-        except ConnectionError as e:
-            raise e
+        _teamspace = _resolve_teamspace(teamspace=teamspace, org=org, user=user)
+        if _teamspace is None:
+            raise ValueError("Couldn't resolve teamspace from the provided name, org, or user")
 
-        self._user = _resolve_user(self._user or user)
-        self._org = _resolve_org(org)
+        self._teamspace = _teamspace
+
+        # self._auth = login.Auth()
+        # self._user = None
+
+        # try:
+        #     self._auth.authenticate()
+        #     if user is None:
+        #         self._user = User(name=UserApi()._get_user_by_id(self._auth.user_id).username)
+        # except ConnectionError as e:
+        #     raise e
+
+        # self._user = _resolve_user(self._user or user)
+        # self._org = _resolve_org(org)
 
         self._base_studio_api = BaseStudioApi()
 
         if name is not None:
-            base_studio = self._base_studio_api.get_base_studio(name, self._org.id)
+            org_id = self._teamspace._org.id if self._teamspace._org is not None else None
+            base_studio = self._base_studio_api.get_base_studio(name, org_id)
 
             if base_studio is None:
-                raise ValueError(f"Base studio with name {name} does not exist in organization {self._org.name}")
+                raise ValueError(f"Base studio with name {name} does not exist")
             self._base_studio = base_studio
 
     def update(
@@ -70,9 +80,11 @@ class BaseStudio:
         machine_image_version: Optional[str] = None,
         setup_script_text: Optional[str] = None,
     ) -> None:
+        org_id = self._teamspace._org.id if self._teamspace._org is not None else None
+        # TODO: if not in an org, can't update them
         self._base_studio = self._base_studio_api.update_base_studio(
             self._base_studio.id,
-            self._org.id,
+            org_id,
             name=name,
             allowed_machines=allowed_machines,
             default_machine=default_machine,
@@ -82,7 +94,7 @@ class BaseStudio:
             disabled=disabled,
         )
 
-    def list(self, managed: bool = True, include_disabled: bool = False) -> List[BaseStudioInfo]:
+    def list(self, include_disabled: bool = False) -> List[BaseStudioInfo]:
         """List all base studios in the organization.
 
         Args:
@@ -92,7 +104,8 @@ class BaseStudio:
         Returns:
             List[BaseStudioInfo]: A list of base studio templates.
         """
-        templates = self._base_studio_api.get_all_base_studios(self._org.id, managed).templates
+        org_id = self._teamspace._org.id if self._teamspace._org is not None else None
+        templates = self._base_studio_api.get_all_base_studios(org_id).templates
 
         return [
             BaseStudioInfo(

lightning_sdk/cli/base_studio/list.py

@@ -21,9 +21,7 @@ def list_base_studios(include_disabled: bool) -> None:
 
 def list_impl(include_disabled: bool) -> None:
     base_studio_cls = BaseStudio()
-    base_studios = base_studio_cls.list(include_disabled=include_disabled) + base_studio_cls.list(
-        managed=False, include_disabled=include_disabled
-    )
+    base_studios = base_studio_cls.list(include_disabled=include_disabled)
 
     table = Table(
         pad_edge=True,

lightning_sdk/cli/entrypoint.py

@@ -2,15 +2,8 @@
 
 import os
 import sys
-import traceback
-from types import TracebackType
-from typing import 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 import __version__
 from lightning_sdk.api.studio_api import _cloud_url
@@ -20,42 +13,24 @@ from lightning_sdk.cli.groups import (
     base_studio,
     config,
     # job,
+    license,
     # mmt,
     studio,
     vm,
 )
-from lightning_sdk.cli.utils import CustomHelpFormatter, rich_to_str
-from lightning_sdk.constants import _LIGHTNING_DEBUG
+from lightning_sdk.cli.utils import CustomHelpFormatter
+from lightning_sdk.cli.utils.logging import CommandLoggingGroup, logging_excepthook
 from lightning_sdk.lightning_cloud.login import Auth
 
 
-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)
-
-
-@click.group(name="lightning", help="Command line interface (CLI) to interact with/manage Lightning AI Studios.")
+@click.group(
+    name="lightning",
+    help="Command line interface (CLI) to interact with/manage Lightning AI Studios.",
+    cls=CommandLoggingGroup,
+)
 @click.version_option(__version__, message="Lightning CLI version %(version)s")
 def main_cli() -> None:
-    sys.excepthook = _notify_exception
+    sys.excepthook = logging_excepthook
 
 
 main_cli.context_class.formatter_class = CustomHelpFormatter
@@ -87,6 +62,8 @@ main_cli.add_command(config)
 main_cli.add_command(studio)
 main_cli.add_command(vm)
 main_cli.add_command(base_studio)
+main_cli.add_command(license)
+
 if os.environ.get("LIGHTNING_EXPERIMENTAL_CLI_ONLY", "0") != "1":
     #### LEGACY COMMANDS ####
     # these commands are currently supported for backwards compatibility, but will potentially be removed in the future.

lightning_sdk/cli/groups.py

@@ -5,6 +5,7 @@ import click
 from lightning_sdk.cli.base_studio import register_commands as register_base_studio_commands
 from lightning_sdk.cli.config import register_commands as register_config_commands
 from lightning_sdk.cli.job import register_commands as register_job_commands
+from lightning_sdk.cli.license import register_commands as register_license_commands
 from lightning_sdk.cli.mmt import register_commands as register_mmt_commands
 from lightning_sdk.cli.studio import register_commands as register_studio_commands
 from lightning_sdk.cli.vm import register_commands as register_vm_commands
@@ -40,6 +41,11 @@ def base_studio() -> None:
     """Manage Lightning AI Base Studios."""
 
 
+@click.group(name="license")
+def license() -> None:  # noqa: A001
+    """Manage Lightning AI Product Licenses."""
+
+
 # Register config commands with the main config group
 register_job_commands(job)
 register_mmt_commands(mmt)
@@ -47,3 +53,4 @@ register_studio_commands(studio)
 register_config_commands(config)
 register_vm_commands(vm)
 register_base_studio_commands(base_studio)
+register_license_commands(license)

lightning_sdk/cli/license/__init__.py

@@ -0,0 +1,14 @@
+"""Base Studio CLI commands."""
+
+import click
+
+
+def register_commands(group: click.Group) -> None:
+    """Register base studio commands with the given group."""
+    from lightning_sdk.cli.license.get import get_license
+    from lightning_sdk.cli.license.list import list_licenses
+    from lightning_sdk.cli.license.set import set_license
+
+    group.add_command(list_licenses)
+    group.add_command(get_license)
+    group.add_command(set_license)

lightning_sdk/cli/license/get.py

@@ -0,0 +1,15 @@
+import click
+
+from lightning_sdk.utils.config import _DEFAULT_CONFIG_FILE_PATH, Config, DefaultConfigKeys
+
+
+@click.command("get")
+@click.argument("product_name")
+@click.option("--config-file", help="Path to the config file")
+def get_license(product_name: str, config_file: str = _DEFAULT_CONFIG_FILE_PATH) -> None:
+    """Get a license key for a given product."""
+    cfg = Config(config_file)
+    license_key = cfg.get(f"{DefaultConfigKeys.license}.{product_name}")
+    if license_key:
+        # echo the license key without any additional output to make parsing simpler
+        click.echo(license_key)

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)