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

lightning_sdk/cli/legacy/download.py

@@ -0,0 +1,322 @@
+import json
+import os
+import re
+from pathlib import Path
+from typing import Optional
+
+import click
+from rich.console import Console
+
+from lightning_sdk.api.license_api import LicenseApi
+from lightning_sdk.api.lit_container_api import LitContainerApi
+from lightning_sdk.cli.legacy.exceptions import StudioCliError
+from lightning_sdk.cli.legacy.studios_menu import _StudiosMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.models import download_model
+from lightning_sdk.studio import Studio
+from lightning_sdk.utils.resolve import _get_authed_user
+
+
+def _expand_remote_path(path: str) -> str:
+    """Expand and normalize remote CLI paths.
+
+    - Strips leading `~/` or `~`
+    - Expands `~` to the user's home but returns relative to it
+    - Returns an empty string if path is empty or `~`
+    """
+    if not path:
+        return ""
+
+    local_home = os.path.expanduser("~")
+
+    # Expand to absolute path and remove the local home prefix if present
+    path = os.path.expanduser(path)
+    if path.startswith(local_home):
+        path = path[len(local_home) :]
+
+    # Remove any leading "/" or "~" remnants
+    return path.lstrip("/~")
+
+
+@click.group(name="download")
+def download() -> None:
+    """Download resources from Lightning AI."""
+
+
+@download.command(name="model")
+@click.argument("name")
+@click.option(
+    "--download-dir", "--download_dir", default=".", help="The directory where the Model should be downloaded."
+)
+def model(name: str, download_dir: str = ".") -> None:
+    """Download a model from a teamspace.
+
+    Example:
+      lightning download model NAME
+
+    NAME: The name of the model to download in the format of <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>.
+    """
+    download_model(
+        name=name,
+        download_dir=download_dir,
+        progress_bar=True,
+    )
+
+
+@download.command(name="folder")
+@click.argument("path")
+@click.option(
+    "--studio",
+    default=None,
+    help=(
+        "The name of the studio to download from. "
+        "Will show a menu with user's owned studios for selection if not specified. "
+        "If provided, should be in the form of <TEAMSPACE-NAME>/<STUDIO-NAME> where the names are case-sensitive. "
+        "The teamspace and studio names can be regular expressions to match, "
+        "a menu filtered studios will be shown for final selection."
+    ),
+)
+@click.option(
+    "--teamspace",
+    default=None,
+    help="The teamspace the drive folder is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.",
+)
+@click.option(
+    "--local-path",
+    "--local_path",
+    default=".",
+    type=click.Path(file_okay=False, dir_okay=True),
+    help="The path to the directory you want to download the folder to.",
+)
+def folder(
+    path: str = "", studio: Optional[str] = None, teamspace: Optional[str] = None, local_path: str = "."
+) -> None:
+    """Download a folder from a Studio or a Teamspace drive folder.
+
+    Example:
+      lightning download folder PATH
+
+    PATH: The relative path within the Studio or drive folder you want to download.
+    Defaults to the entire Studio or drive folder.
+    """
+    local_path = Path(local_path)
+    if not local_path.is_dir():
+        raise NotADirectoryError(f"'{local_path}' is not a directory")
+
+    if studio and teamspace:
+        raise ValueError("Either --studio or --teamspace must be provided, not both")
+
+    if studio:
+        path = _expand_remote_path(path)
+        resolved_downloader = _resolve_studio(studio)
+    elif teamspace:
+        menu = TeamspacesMenu()
+        resolved_downloader = menu(teamspace)
+    else:
+        raise ValueError("Either --studio or --teamspace must be provided")
+
+    if not path:
+        local_path /= resolved_downloader.name
+        path = ""
+
+    try:
+        if not path and teamspace:
+            raise FileNotFoundError()
+        resolved_downloader.download_folder(remote_path=path, target_path=str(local_path))
+    except Exception as e:
+        raise StudioCliError(
+            f"Could not download the folder from the given Studio {studio} or Teamspace {teamspace}. "
+            "Please contact Lightning AI directly to resolve this issue."
+        ) from e
+
+
+@download.command(name="file")
+@click.argument("path")
+@click.option(
+    "--studio",
+    default=None,
+    help=(
+        "The name of the studio to download from. "
+        "Will show a menu with user's owned studios for selection if not specified. "
+        "If provided, should be in the form of <TEAMSPACE-NAME>/<STUDIO-NAME> where the names are case-sensitive. "
+        "The teamspace and studio names can be regular expressions to match, "
+        "a menu filtered studios will be shown for final selection."
+    ),
+)
+@click.option(
+    "--teamspace",
+    default=None,
+    help="The teamspace the file is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.",
+)
+@click.option(
+    "--local-path",
+    "--local_path",
+    default=".",
+    type=click.Path(file_okay=False, dir_okay=True),
+    help="The path to the directory you want to download the file to.",
+)
+def file(path: str = "", studio: Optional[str] = None, teamspace: Optional[str] = None, local_path: str = ".") -> None:
+    """Download a file from a Studio or Teamspace drive file.
+
+    Example:
+      lightning download file PATH
+
+    PATH: The relative path to the file within the Studio or Teamspace drive file you want to download.
+    """
+    local_path = Path(local_path)
+    if not local_path.is_dir():
+        raise NotADirectoryError(f"'{local_path}' is not a directory")
+
+    if studio and teamspace:
+        raise ValueError("Either --studio or --teamspace must be provided, not both")
+
+    if studio:
+        resolved_downloader = _resolve_studio(studio)
+    elif teamspace:
+        menu = TeamspacesMenu()
+        resolved_downloader = menu(teamspace)
+    else:
+        raise ValueError("Either --studio or --teamspace must be provided")
+
+    if not path:
+        local_path /= resolved_downloader.name
+        path = ""
+
+    try:
+        if not path:
+            raise FileNotFoundError()
+        resolved_downloader.download_file(remote_path=path, file_path=str(local_path / os.path.basename(path)))
+    except Exception as e:
+        raise StudioCliError(
+            f"Could not download the file from the given Studio {studio} or Teamspace {teamspace}. "
+            "Please contact Lightning AI directly to resolve this issue."
+        ) from e
+
+
+@download.command(name="container")
+@click.argument("container")
+@click.option("--teamspace", default=None, help="The name of the teamspace to download the container from")
+@click.option("--tag", default="latest", show_default=True, help="The tag of the container to download.")
+@click.option(
+    "--cloud-account",
+    "--cloud_account",  # The UI will present the above variant, using this as a secondary to be consistent w/ models
+    default=None,
+    help="The name of the cloud account to download the Container from.",
+)
+def download_container(
+    container: str, teamspace: Optional[str] = None, tag: str = "latest", cloud_account: Optional[str] = None
+) -> None:
+    """Download a docker container from a teamspace.
+
+    Example:
+      lightning download container CONTAINER
+
+    CONTAINER: The name of the container to download.
+    """
+    console = Console()
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace)
+    with console.status("Downloading container..."):
+        api = LitContainerApi()
+        api.download_container(container, resolved_teamspace, tag, cloud_account)
+        console.print("Container downloaded successfully", style="green")
+
+
+def _resolve_studio(studio: Optional[str]) -> Studio:
+    user = _get_authed_user()
+    # if no studio specify suggest/filter only user's studios
+    menu = _StudiosMenu()
+    possible_studios = menu._get_possible_studios(user, is_owner=studio is None)
+
+    try:
+        if studio:
+            team_name, studio_name = studio.split("/")
+            options = [st for st in possible_studios if st["teamspace"] == team_name and st["name"] == studio_name]
+            if len(options) == 1:
+                selected_studio = menu._get_studio_from_name(studio, possible_studios)
+            # user can also use the partial studio name as secondary interactive selection
+            else:
+                # filter matching simple reg expressions or start with the team and studio name
+                possible_studios = filter(
+                    lambda st: (re.match(team_name, st["teamspace"]) or team_name in st["teamspace"])
+                    and (re.match(studio_name, st["name"]) or studio_name in st["name"]),
+                    possible_studios,
+                )
+                if not len(possible_studios):
+                    raise ValueError(
+                        f"Could not find Studio like '{studio}', please consider update your filtering pattern."
+                    )
+                selected_studio = menu._get_studio_from_interactive_menu(list(possible_studios))
+        else:
+            selected_studio = menu._get_studio_from_interactive_menu(possible_studios)
+
+    except KeyboardInterrupt:
+        raise KeyboardInterrupt from None
+
+    # give user friendlier error message
+    except Exception as e:
+        raise StudioCliError(
+            f"Could not find the given Studio {studio} to download files from. "
+            "Please contact Lightning AI directly to resolve this issue."
+        ) from e
+
+    return Studio(**selected_studio)
+
+
+@download.command(name="licenses")
+def download_licenses() -> None:
+    """Download licenses for all user's products/packages.
+
+    Example:
+      lightning download licenses
+
+    """
+    user = _get_authed_user()
+    api = LicenseApi()
+    licenses = api.list_user_licenses(user.id)
+
+    user_home = Path.home()
+    lit_dir = user_home / ".lightning"
+    lit_dir.mkdir(parents=True, exist_ok=True)
+    licenses_file = lit_dir / "licenses.json"
+
+    licenses_short = {ll.product_name: ll.license_key for ll in licenses if ll.is_valid}
+    with licenses_file.open("w") as fp:
+        json.dump(licenses_short, fp, indent=4)
+    Console().print(f"Licenses downloaded to {licenses_file}", style="green")
+
+
+@download.command(name="license")
+@click.argument("name")
+def download_license(name: str) -> None:
+    """Download license for specific products/packages.
+
+    Example:
+      lightning download license NAME
+
+    NAME: The name of the product/package to download the license for.
+    """
+    user = _get_authed_user()
+    api = LicenseApi()
+    licenses = api.list_user_licenses(user.id)
+    licenses_short = {ll.product_name: ll.license_key for ll in licenses if ll.is_valid}
+
+    if name not in licenses_short:
+        Console().print(f"Missing valid license for {name}", style="red")
+        return
+
+    user_home = Path.home()
+    lit_dir = user_home / ".lightning"
+    lit_dir.mkdir(parents=True, exist_ok=True)
+    licenses_file = lit_dir / "licenses.json"
+
+    licenses_loaded = {}
+    if licenses_file.exists():
+        with licenses_file.open("r") as fp:
+            licenses_loaded = json.load(fp)
+
+    licenses_loaded[name] = licenses_short[name]
+
+    with licenses_file.open("w") as fp:
+        json.dump(licenses_loaded, fp, indent=4)
+    Console().print(f"Updated license for {name} in {licenses_file}", style="green")

lightning_sdk/cli/legacy/entrypoint.py

@@ -0,0 +1,110 @@
+import sys
+import traceback
+from types import TracebackType
+from typing import Type
+
+import click
+from rich.console import Console, 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
+from lightning_sdk.cli.legacy.ai_hub import aihub
+from lightning_sdk.cli.legacy.coloring import CustomHelpFormatter
+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
+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
+from lightning_sdk.constants import _LIGHTNING_DEBUG
+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."""
+    console = Console()
+
+    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"))
+
+    console.print(Panel(Group(*renderables), title="⚡ Lightning CLI Error", border_style="red"))
+
+
+@click.group(name="lightning", help="Command line interface (CLI) to interact with/manage Lightning AI Studios.")
+@click.version_option(__version__, message="Lightning CLI version %(version)s")
+def main_cli() -> None:
+    sys.excepthook = _notify_exception
+
+
+# colorful help messages
+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()
+
+
+# additional commands
+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(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(run)
+main_cli.add_command(deploy)
+main_cli.add_command(start)
+main_cli.add_command(stop)
+main_cli.add_command(switch)
+main_cli.add_command(upload)
+main_cli.add_command(open)
+
+
+if __name__ == "__main__":
+    main_cli()

lightning_sdk/cli/legacy/generate.py

@@ -0,0 +1,52 @@
+from typing import Optional
+
+import click
+from rich.console import Console
+
+from lightning_sdk.cli.legacy.studios_menu import _StudiosMenu
+
+
+@click.group(name="generate")
+def generate() -> None:
+    """Generate configs (such as ssh for studio) and print them to commandline."""
+
+
+@generate.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.)"
+    ),
+)
+def ssh(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+    """Get SSH config entry for a studio."""
+    menu = _StudiosMenu()
+    studio = menu._get_studio(name=name, teamspace=teamspace)
+
+    conf = _generate_ssh_config(key_path="~/.ssh/lightning_rsa", user=f"s_{studio._studio.id}", host=studio.name)
+    # Print the SSH config
+    Console().print(f"# ssh s_{studio._studio.id}@ssh.lightning.ai\n\n" + conf)
+
+
+def _generate_ssh_config(key_path: str, host: str, user: str) -> str:
+    return f"""Host {host}
+      User {user}
+      Hostname ssh.lightning.ai
+      IdentityFile {key_path}
+      IdentitiesOnly yes
+      ServerAliveInterval 15
+      ServerAliveCountMax 4
+      StrictHostKeyChecking no
+      UserKnownHostsFile=/dev/null
+    """

lightning_sdk/cli/legacy/inspection.py

@@ -0,0 +1,45 @@
+from typing import Optional
+
+import click
+from rich.console import Console
+
+from lightning_sdk.cli.legacy.job_and_mmt_action import _JobAndMMTAction
+
+
+@click.group(name="inspect")
+def inspect() -> None:
+    """Inspect resources of the Lightning AI platform to get additional details as JSON."""
+
+
+@inspect.command(name="job")
+@click.option("--name", default=None, help="the name of the job. If not specified can be selected interactively.")
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "the name of the teamspace the job lives in."
+        "Should be specified as {teamspace_owner}/{teamspace_name} (e.g my-org/my-teamspace). "
+        "If not specified can be selected interactively."
+    ),
+)
+def job(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+    """Inspect a job for further details as JSON."""
+    menu = _JobAndMMTAction()
+    Console().print(menu.job(name=name, teamspace=teamspace).json())
+
+
+@inspect.command(name="mmt")
+@click.option("--name", default=None, help="the name of the job. If not specified can be selected interactively.")
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "the name of the teamspace the job lives in."
+        "Should be specified as {teamspace_owner}/{teamspace_name} (e.g my-org/my-teamspace). "
+        "If not specified can be selected interactively."
+    ),
+)
+def mmt(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+    """Inspect a multi-machine job for further details as JSON."""
+    menu = _JobAndMMTAction()
+    Console().print(menu.mmt(name=name, teamspace=teamspace).json())

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

@@ -1,13 +1,13 @@
 from typing import Optional
 
-from lightning_sdk.cli.jobs_menu import _JobsMenu
-from lightning_sdk.cli.mmts_menu import _MMTsMenu
-from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.legacy.jobs_menu import _JobsMenu
+from lightning_sdk.cli.legacy.mmts_menu import _MMTsMenu
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
 from lightning_sdk.job import Job
 from lightning_sdk.mmt import MMT
 
 
-class _JobAndMMTAction(_TeamspacesMenu, _JobsMenu, _MMTsMenu):
+class _JobAndMMTAction(TeamspacesMenu, _JobsMenu, _MMTsMenu):
     """Inspect resources of the Lightning AI platform to get additional details as JSON."""
 
     def job(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> Job:
@@ -20,7 +20,7 @@ class _JobAndMMTAction(_TeamspacesMenu, _JobsMenu, _MMTsMenu):
                 If not specified can be selected interactively.
 
         """
-        resolved_teamspace = self._resolve_teamspace(teamspace)
+        resolved_teamspace = self(teamspace)
         return self._resolve_job(name, teamspace=resolved_teamspace)
 
     def mmt(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> MMT:
@@ -33,5 +33,5 @@ class _JobAndMMTAction(_TeamspacesMenu, _JobsMenu, _MMTsMenu):
                 If not specified can be selected interactively.
 
         """
-        resolved_teamspace = self._resolve_teamspace(teamspace)
+        resolved_teamspace = self(teamspace)
         return self._resolve_mmt(name, teamspace=resolved_teamspace)

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

@@ -1,8 +1,9 @@
 from typing import Dict, List, Optional
 
+from rich.console import Console
 from simple_term_menu import TerminalMenu
 
-from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.cli.legacy.exceptions import StudioCliError
 from lightning_sdk.job import Job
 from lightning_sdk.teamspace import Teamspace
 
@@ -20,7 +21,7 @@ class _JobsMenu:
             if j.name == job:
                 return j
 
-        print("Could not find Job {job}, please select it from the list:")
+        Console().print("Could not find Job {job}, please select it from the list:")
         return self._get_job_from_interactive_menu(possible_jobs)
 
     @staticmethod