lightning-sdk 0.1.49__py3-none-any.whl → 2025.11.5__py3-none-any.whl

lightning_sdk/cli/config/set.py

@@ -0,0 +1,92 @@
+import click
+
+from lightning_sdk.cli.utils.save_to_config import save_teamspace_to_config
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.machine import CloudProvider
+from lightning_sdk.studio import Studio
+from lightning_sdk.utils.config import Config, DefaultConfigKeys
+from lightning_sdk.utils.resolve import _resolve_org, _resolve_user
+
+
+@click.group("set")
+def set_value() -> None:
+    """Set configuration values."""
+
+
+@set_value.command("user")
+@click.argument("user_name")
+def set_user(user_name: str) -> None:
+    """Set the default user name in the config."""
+    try:
+        _resolve_user(user_name)
+    except Exception:
+        # TODO: make this a generic CLI error
+        raise ValueError(f"Could not resolve user: '{user_name}'. Does the user exist?") from None
+
+    config = Config()
+    setattr(config, DefaultConfigKeys.user, user_name)
+
+
+@set_value.command("org")
+@click.argument("org_name")
+def set_org(org_name: str) -> None:
+    """Set the default organization name in the config."""
+    try:
+        _resolve_org(org_name)
+    except Exception:
+        # TODO: make this a generic CLI error
+        raise ValueError(f"Could not resolve organization: '{org_name}'. Does the organization exist?") from None
+
+    config = Config()
+    setattr(config, DefaultConfigKeys.organization, org_name)
+
+
+@set_value.command("studio")
+@click.argument("studio_name")
+def set_studio(studio_name: str) -> None:
+    """Set the default studio name in the config."""
+    try:
+        studio = Studio(studio_name)
+    except Exception:
+        # TODO: make this a generic CLI error
+        raise ValueError(f"Could not resolve studio: '{studio_name}'. Does the studio exist?") from None
+
+    config = Config()
+    setattr(config, DefaultConfigKeys.studio, studio.name)
+
+
+@set_value.command("teamspace")
+@click.argument("teamspace_name")
+def set_teamspace(teamspace_name: str) -> None:
+    """Set the default teamspace name in the config."""
+    menu = TeamspacesMenu()
+    teamspace_resolved = menu(teamspace=teamspace_name)
+
+    # explicit user action, so overwrite the config
+    save_teamspace_to_config(teamspace_resolved, overwrite=True)
+
+
+@set_value.command("cloud-account")
+@click.argument("cloud_account_name")
+def set_cloud_account(cloud_account_name: str) -> None:
+    """Set the default cloud account name in the config."""
+    config = Config()
+    setattr(config, DefaultConfigKeys.cloud_account, cloud_account_name)
+
+
+@set_value.command("cloud-provider")
+@click.argument("cloud_provider_name")
+def set_cloud_provider(cloud_provider_name: str) -> None:
+    """Set the default cloud provider name in the config."""
+    config = Config()
+
+    try:
+        cloud_provider = CloudProvider(cloud_provider_name)
+    except ValueError:
+        # TODO: make this a generic CLI error
+        raise ValueError(
+            f"Could not resolve cloud provider: '{cloud_provider_name}'. "
+            f"Supported values are: {', '.join(m.name for m in list(CloudProvider))}"
+        ) from None
+
+    setattr(config, DefaultConfigKeys.cloud_provider, cloud_provider.name)

lightning_sdk/cli/config/show.py

@@ -0,0 +1,9 @@
+import click
+
+from lightning_sdk.utils.config import Config
+
+
+@click.command("show")
+def show() -> None:
+    """Show configuration values."""
+    click.echo(Config())

lightning_sdk/cli/entrypoint.py

@@ -1,67 +1,109 @@
+"""New Lightning CLI entrypoint with organized command groups."""
+
+import os
 import sys
-from types import TracebackType
-from typing import Type
 
-from fire import Fire
-from lightning_utilities.core.imports import RequirementCache
+import click
 
+from lightning_sdk import __version__
 from lightning_sdk.api.studio_api import _cloud_url
-from lightning_sdk.cli.ai_hub import _AIHub
-from lightning_sdk.cli.delete import _Delete
-from lightning_sdk.cli.download import _Downloads
-from lightning_sdk.cli.inspect import _Inspect
-from lightning_sdk.cli.legacy import _LegacyLightningCLI
-from lightning_sdk.cli.list import _List
-from lightning_sdk.cli.run import _Run
-from lightning_sdk.cli.serve import _Docker, _LitServe
-from lightning_sdk.cli.stop import _Stop
-from lightning_sdk.cli.upload import _Uploads
-from lightning_sdk.lightning_cloud.login import Auth
-
-_LIGHTNING_AVAILABLE = RequirementCache("lightning")
-
-
-class StudioCLI:
-    """Command line interface (CLI) to interact with/manage Lightning AI Studios."""
-
-    def __init__(self) -> None:
-        self.download = _Downloads()
-        self.upload = _Uploads()
-        self.aihub = _AIHub()
-        self.run = _Run(legacy_run=_LegacyLightningCLI() if _LIGHTNING_AVAILABLE else None)
-        self.serve = _LitServe()
-        self.dockerize = _Docker()
-        self.list = _List()
-        self.delete = _Delete()
-        self.inspect = _Inspect()
-        self.stop = _Stop()
 
-        sys.excepthook = _notify_exception
-
-    def login(self) -> None:
-        """Login to Lightning AI Studios."""
-        auth = Auth()
-        auth.clear()
-
-        try:
-            auth.authenticate()
-        except ConnectionError:
-            raise RuntimeError(f"Unable to connect to {_cloud_url()}. Please check your internet connection.") from None
-
-    def logout(self) -> None:
-        """Logout from Lightning AI Studios."""
-        auth = Auth()
-        auth.clear()
-
-
-def _notify_exception(exception_type: Type[BaseException], value: BaseException, tb: TracebackType) -> None:  # No
-    """CLI won't show tracebacks, just print the exception message."""
-    print(value)
+# Import legacy groups directly from groups.py
+from lightning_sdk.cli.groups import (
+    base_studio,
+    config,
+    # job,
+    license,
+    # mmt,
+    studio,
+    vm,
+)
+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
 
 
+@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:
-    """CLI entrypoint."""
-    Fire(StudioCLI(), name="lightning")
+    sys.excepthook = logging_excepthook
+
+
+main_cli.context_class.formatter_class = CustomHelpFormatter
+
+
+@main_cli.command
+def login() -> None:
+    """Login to Lightning AI Studios."""
+    auth = Auth()
+    auth.clear()
+
+    try:
+        auth.authenticate()
+    except ConnectionError:
+        raise RuntimeError(f"Unable to connect to {_cloud_url()}. Please check your internet connection.") from None
+
+
+@main_cli.command
+def logout() -> None:
+    """Logout from Lightning AI Studios."""
+    auth = Auth()
+    auth.clear()
+
+
+# Add new command groups
+main_cli.add_command(config)
+# main_cli.add_command(job)
+# main_cli.add_command(mmt)
+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.
+    # they've grown pretty wild and provide a very inconsistent UX.
+    from lightning_sdk.cli.legacy.ai_hub import aihub
+    from lightning_sdk.cli.legacy.configure import configure
+    from lightning_sdk.cli.legacy.connect import connect
+    from lightning_sdk.cli.legacy.create import create
+    from lightning_sdk.cli.legacy.delete import delete
+    from lightning_sdk.cli.legacy.deploy.serve import deploy
+    from lightning_sdk.cli.legacy.docker_cli import dockerize
+    from lightning_sdk.cli.legacy.download import download
+    from lightning_sdk.cli.legacy.generate import generate
+    from lightning_sdk.cli.legacy.inspection import inspect
+    from lightning_sdk.cli.legacy.list import list_cli
+    from lightning_sdk.cli.legacy.open import open as open_cmd
+    from lightning_sdk.cli.legacy.run import run
+    from lightning_sdk.cli.legacy.start import start
+    from lightning_sdk.cli.legacy.stop import stop
+    from lightning_sdk.cli.legacy.switch import switch
+    from lightning_sdk.cli.legacy.upload import upload
+
+    # Add old command groups
+    main_cli.add_command(aihub)
+    main_cli.add_command(configure)
+    main_cli.add_command(connect)
+    main_cli.add_command(create)
+    main_cli.add_command(delete)
+    main_cli.add_command(deploy)
+    main_cli.add_command(dockerize)
+    main_cli.add_command(download)
+    main_cli.add_command(generate)
+    main_cli.add_command(inspect)
+    main_cli.add_command(list_cli)
+    main_cli.add_command(open_cmd)
+    main_cli.add_command(run)
+    main_cli.add_command(start)
+    main_cli.add_command(stop)
+    main_cli.add_command(switch)
+    main_cli.add_command(upload)
 
 
 if __name__ == "__main__":

lightning_sdk/cli/groups.py

@@ -0,0 +1,56 @@
+"""CLI groups for organizing Lightning SDK commands."""
+
+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
+
+
+@click.group(name="studio")
+def studio() -> None:
+    """Manage Lightning AI Studios."""
+
+
+@click.group(name="job")
+def job() -> None:
+    """Manage Lightning AI Jobs."""
+
+
+@click.group(name="mmt")
+def mmt() -> None:
+    """Manage Lightning AI Multi-Machine Training (MMT)."""
+
+
+@click.group(name="config")
+def config() -> None:
+    """Manage Lightning SDK and CLI configuration."""
+
+
+@click.group(name="vm")
+def vm() -> None:
+    """Manage Lightning AI VMs."""
+
+
+@click.group(name="base-studio")
+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)
+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/job/__init__.py

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

lightning_sdk/cli/legacy/ai_hub.py

@@ -0,0 +1,65 @@
+from typing import Optional
+
+import click
+
+from lightning_sdk.ai_hub import AIHub
+
+
+@click.group(name="aihub")
+def aihub() -> None:
+    """Interact with Lightning Studio - AI Hub."""
+
+
+@aihub.command(name="api-info")
+@click.argument("api-id")
+def api_info(api_id: str) -> None:
+    """Get full API template info such as input details.
+
+    Example:
+      lightning aihub api_info API-ID
+
+    API-ID: The ID of the API for which information is requested.
+    """
+    ai_hub = AIHub()
+    ai_hub.api_info(api_id)
+
+
+@aihub.command(name="list-apis")
+@click.option("--search", default=None, help="Search for API templates by name.")
+def list_apis(search: Optional[str]) -> None:
+    """List API templates available in the AI Hub."""
+    ai_hub = AIHub()
+    ai_hub.list_apis(search=search)
+
+
+@aihub.command(name="deploy")
+@click.argument("api-id")
+@click.option(
+    "--cloud-account",
+    "--cloud_account",
+    default=None,
+    help="Cloud Account to deploy the API to. Defaults to user's default cloud account.",
+)
+@click.option("--name", default=None, help="Name of the deployed API. Defaults to the name of the API template.")
+@click.option(
+    "--teamspace",
+    default=None,
+    help="Teamspace to deploy the API to. Defaults to user's default teamspace.",
+)
+@click.option(
+    "--org",
+    default=None,
+    help="Organization to deploy the API to. Defaults to user's default organization.",
+)
+def deploy(
+    api_id: str, cloud_account: Optional[str], name: Optional[str], teamspace: Optional[str], org: Optional[str]
+) -> None:
+    """Deploy an API template from the AI Hub.
+
+    Example:
+      lightning aihub deploy API-ID
+
+    API-ID: The ID of the API which should be deployed.
+    """
+    ai_hub = AIHub()
+    ai_hub.run(api_id, cloud_account=cloud_account, name=name, teamspace=teamspace, org=org)

lightning_sdk/cli/legacy/clusters_menu.py

@@ -0,0 +1,49 @@
+import sys
+from typing import List, Optional
+
+from rich.console import Console
+from simple_term_menu import TerminalMenu
+
+from lightning_sdk import Teamspace
+from lightning_sdk.api.cloud_account_api import CloudAccountApi
+from lightning_sdk.lightning_cloud.openapi import V1ClusterType, V1ProjectClusterBinding
+
+
+class _ClustersMenu:
+    def _get_cluster_from_interactive_menu(self, possible_clusters: List[V1ProjectClusterBinding]) -> str:
+        terminal_menu = self._prepare_terminal_menu_teamspaces([cluster.cluster_id for cluster in possible_clusters])
+        terminal_menu.show()
+
+        return possible_clusters[terminal_menu.chosen_menu_index].cluster_id
+
+    @staticmethod
+    def _prepare_terminal_menu_teamspaces(cluster_ids: List[str]) -> TerminalMenu:
+        title = "Please select a cluster from the following:"
+
+        return TerminalMenu(cluster_ids, title=title, clear_menu_on_exit=True)
+
+    def _resolve_cluster(self, teamspace: Teamspace) -> Optional[str]:
+        selected_cluster_id = None
+        console = Console()
+        try:
+            selected_cluster_id = teamspace.default_cloud_account or self._get_cluster_from_interactive_menu(
+                possible_clusters=teamspace.cloud_account_objs
+            )
+
+            cloud_account_api = CloudAccountApi()
+
+            resolved_cluster_obj = cloud_account_api.get_cloud_account(
+                cloud_account_id=selected_cluster_id, org_id=teamspace.owner.id, teamspace_id=teamspace.id
+            )
+
+            return None if resolved_cluster_obj.spec.cluster_type == V1ClusterType.GLOBAL else resolved_cluster_obj.id
+        except KeyboardInterrupt:
+            console.print("Operation cancelled by user")
+            sys.exit(0)
+
+        except Exception:
+            console.print(
+                f"[red]Could not find the given Cluster:[/red] {selected_cluster_id}. "
+                "Please contact Lightning AI directly to resolve this issue."
+            )
+            sys.exit(1)

lightning_sdk/cli/legacy/configure.py

@@ -0,0 +1,129 @@
+import os
+import platform
+import uuid
+from pathlib import Path
+from typing import Optional, Union
+
+import click
+from rich.console import Console
+
+from lightning_sdk.cli.legacy.generate import _generate_ssh_config
+from lightning_sdk.cli.legacy.studios_menu import _StudiosMenu
+from lightning_sdk.lightning_cloud.login import Auth
+
+
+@click.group(name="configure")
+def configure() -> None:
+    """Configure access to resources on the Lightning AI platform."""
+
+
+def _configure_ssh_internal(
+    name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False
+) -> None:
+    """Internal function to configure SSH without Click decorators."""
+    auth = Auth()
+    auth.authenticate()
+    console = Console()
+    ssh_dir = Path.home() / ".ssh"
+    ssh_dir.mkdir(parents=True, exist_ok=True)
+
+    key_path = ssh_dir / "lightning_rsa"
+    config_path = ssh_dir / "config"
+
+    # Check if the SSH key already exists
+    if key_path.exists() and (key_path.with_suffix(".pub")).exists() and not overwrite:
+        console.print(f"SSH key already exists at {key_path}")
+    else:
+        _download_ssh_keys(auth.api_key, ssh_home=ssh_dir, ssh_key_name="lightning_rsa", overwrite=overwrite)
+        console.print(f"SSH key generated and saved to {key_path}")
+
+    # Check if the SSH config already contains the required configuration
+    menu = _StudiosMenu()
+    studio = menu._get_studio(name=name, teamspace=teamspace)
+    config_content = _generate_ssh_config(key_path=str(key_path), user=f"s_{studio._studio.id}", host=studio.name)
+    if config_path.exists():
+        with config_path.open("r") as config_file:
+            # check if the host already exists in the config
+            if f"Host {studio.name}" in config_file.read():
+                console.print("SSH config already contains the required configuration.")
+                return
+
+    with config_path.open("a") as config_file:
+        config_file.write(os.linesep)
+        config_file.write(config_content)
+        config_file.write(os.linesep)
+        console.print(f"SSH config updated at {config_path}")
+
+
+@configure.command(name="ssh")
+@click.option(
+    "--name",
+    default=None,
+    help=(
+        "The name of the studio to obtain SSH config. "
+        "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
+    ),
+)
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "The teamspace the studio is part of. "
+        "Should be of format <OWNER>/<TEAMSPACE_NAME>. "
+        "If not specified, tries to infer from the environment (e.g. when run from within a Studio.)"
+    ),
+)
+@click.option(
+    "--overwrite",
+    is_flag=True,
+    flag_value=True,
+    default=False,
+    help="Whether to overwrite the SSH key and config if they already exist.",
+)
+def ssh(name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> None:
+    """Get SSH config entry for a studio."""
+    _configure_ssh_internal(name=name, teamspace=teamspace, overwrite=overwrite)
+
+
+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 local_path.exists() 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:
+        local_path.chmod(0o600)
+
+
+def _download_ssh_keys(
+    api_key: str,
+    key_id: str = "",
+    ssh_home: Union[str, Path] = "",
+    ssh_key_name: str = "lightning_rsa",
+    overwrite: bool = False,
+) -> None:
+    if not ssh_home:
+        ssh_home = Path.home() / ".ssh"
+    elif isinstance(ssh_home, str):
+        ssh_home = Path(ssh_home)
+    if not key_id:
+        key_id = str(uuid.uuid4())
+
+    path_key = ssh_home / ssh_key_name
+    path_pub = ssh_home / f"{ssh_key_name}.pub"
+
+    # todo: consider hitting the API to get the key pair directly instead of using wget
+    _download_file(
+        f"https://lightning.ai/setup/ssh-gen?t={api_key}&id={key_id}&machineName={platform.node()}",
+        path_key,
+        overwrite=overwrite,
+        chmod=0o600,
+    )
+    _download_file(f"https://lightning.ai/setup/ssh-public?t={api_key}&id={key_id}", path_pub, overwrite=overwrite)

lightning_sdk/cli/legacy/connect.py

@@ -0,0 +1,34 @@
+import subprocess
+import sys
+from typing import Optional
+
+import click
+
+from lightning_sdk.cli.legacy.configure import _configure_ssh_internal
+from lightning_sdk.cli.legacy.studios_menu import _StudiosMenu
+
+
+@click.group(name="connect")
+def connect() -> None:
+    """Connect to lightning products."""
+
+
+@connect.command(name="studio")
+@click.option("--name", default=None, help="The name of the studio to connect to.")
+@click.option(
+    "--teamspace",
+    default=None,
+    help="The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.",
+)
+def studio(name: Optional[str], teamspace: Optional[str]) -> None:
+    """Connect to a studio via SSH."""
+    _configure_ssh_internal(name=name, teamspace=teamspace, overwrite=False)
+
+    menu = _StudiosMenu()
+    studio = menu._get_studio(name=name, teamspace=teamspace)
+
+    try:
+        subprocess.run(["ssh", studio.name])
+    except Exception as ex:
+        print(f"Failed to establish SSH connection: {ex}")
+        sys.exit(1)