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

lightning_sdk/cli/legacy/switch.py

@@ -0,0 +1,63 @@
+from typing import Optional
+
+import click
+
+from lightning_sdk import Machine, Studio
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+
+_MACHINE_VALUES = tuple(
+    [machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine) and machine._include_in_cli]
+)
+
+
+@click.group("switch")
+def switch() -> None:
+    """Switch machines for resources on the Lightning AI platform."""
+
+
+@switch.command("studio")
+@click.argument("name")
+@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(
+    "--machine",
+    default="CPU",
+    show_default=True,
+    type=click.Choice(_MACHINE_VALUES),
+    help="The machine type to switch to.",
+)
+def studio(name: str, teamspace: Optional[str] = None, machine: str = "CPU") -> None:
+    """Switch a studio to a given machine.
+
+    Example:
+      lightning switch studio NAME --machine=CPU
+
+    NAME: the name of the studio to switch machine for.
+    """
+    if teamspace is not None:
+        ts_splits = teamspace.split("/")
+        if len(ts_splits) != 2:
+            raise ValueError(f"Teamspace should be of format <OWNER>/<TEAMSPACE_NAME> but got {teamspace}")
+        owner, teamspace = ts_splits
+    else:
+        owner, teamspace = None, None
+
+    try:
+        studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False)
+    except (RuntimeError, ValueError, ApiException):
+        studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
+
+    try:
+        resolved_machine = getattr(Machine, machine.upper(), Machine(machine, machine))
+    except KeyError:
+        resolved_machine = machine
+
+    Studio.show_progress = True
+    studio.switch_machine(resolved_machine)

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

@@ -1,9 +1,10 @@
 from typing import Dict, List, Optional
 
+from rich.console import Console
 from simple_term_menu import TerminalMenu
 
 from lightning_sdk.api import OrgApi
-from lightning_sdk.cli.exceptions import StudioCliError
+from lightning_sdk.cli.legacy.exceptions import StudioCliError
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.user import User
 from lightning_sdk.utils.resolve import _get_authed_user
@@ -21,12 +22,20 @@ class _TeamspacesMenu:
     def _get_teamspace_from_name(
         self, teamspace: str, possible_teamspaces: Dict[str, Dict[str, str]]
     ) -> Dict[str, str]:
-        owner, name = teamspace.split("/", maxsplit=1)
+        try:
+            owner, name = teamspace.split("/", maxsplit=1)
+        except ValueError as e:
+            raise ValueError(
+                f"Invalid teamspace format: '{teamspace}'. "
+                "Teamspace should be specified as '{teamspace_owner}/{teamspace_name}' "
+                "(e.g., 'my-org/my-teamspace')."
+            ) from e
+
         for _, ts in possible_teamspaces.items():
             if ts["name"] == name and (ts["user"] == owner or ts["org"] == owner):
                 return ts
 
-        print("Could not find Teamspace {teamspace}, please select it from the list:")
+        Console().print(f"Could not find Teamspace {teamspace}, please select it from the list:")
         return self._get_teamspace_from_interactive_menu(possible_teamspaces)
 
     @staticmethod

lightning_sdk/cli/legacy/upload.py

@@ -0,0 +1,382 @@
+import concurrent.futures
+import json
+import os
+import webbrowser
+from pathlib import Path
+from typing import Dict, Generator, List, Optional
+
+import click
+import rich
+from rich.console import Console
+from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
+from simple_term_menu import TerminalMenu
+from tqdm import tqdm
+
+from lightning_sdk.api.lit_container_api import DockerNotRunningError, LCRAuthFailedError, LitContainerApi
+from lightning_sdk.api.utils import _get_cloud_url
+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.constants import _LIGHTNING_DEBUG
+from lightning_sdk.models import upload_model as _upload_model
+from lightning_sdk.studio import Studio
+from lightning_sdk.utils.resolve import _get_authed_user
+
+_STUDIO_UPLOAD_STATUS_PATH = "~/.lightning/studios/uploads"
+
+
+@click.group("upload")
+def upload() -> None:
+    """Upload assets to Lightning AI."""
+
+
+@upload.command("model")
+@click.argument("name")
+@click.option(
+    "--path",
+    default=".",
+    help="The path to the file or directory you want to upload. Defaults to the current directory.",
+)
+@click.option(
+    "--cloud-account",
+    "--cloud_account",
+    default=None,
+    help="The name of the cloud account to store the Model in.",
+)
+def model(name: str, path: str = ".", cloud_account: Optional[str] = None) -> None:
+    """Upload a model a teamspace.
+
+    Example:
+        lightning upload model NAME
+
+    NAME: the name of the model to upload (Should be of format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>).
+    """
+    _upload_model(name, path, cloud_account=cloud_account)
+
+
+@upload.command("folder")
+@click.argument("path", type=click.Path(exists=True))
+@click.option(
+    "--studio",
+    default=None,
+    help=(
+        "The name of the studio to upload to. "
+        "Will show a menu for selection if not specified. "
+        "If provided, should be in the form of <TEAMSPACE-NAME>/<STUDIO-NAME>"
+    ),
+)
+@click.option(
+    "--remote-path",
+    "--remote_path",
+    default=None,
+    help=(
+        "The path where the uploaded file should appear on your Studio. "
+        "Has to be within your Studio's home directory and will be relative to that. "
+        "If not specified, will use the name of the folder you want to upload and place it in your home directory."
+    ),
+)
+def folder(path: str, studio: Optional[str], remote_path: Optional[str]) -> None:
+    """Upload a folder to a Studio."""
+    _folder(path=path, studio=studio, remote_path=remote_path)
+
+
+@upload.command("file")
+@click.argument("path", type=click.Path(exists=True))
+@click.option(
+    "--studio",
+    default=None,
+    help=(
+        "The name of the studio to upload to. "
+        "Will show a menu for selection if not specified. "
+        "If provided, should be in the form of <TEAMSPACE-NAME>/<STUDIO-NAME>"
+    ),
+)
+@click.option(
+    "--remote-path",
+    "--remote_path",
+    default=None,
+    help=(
+        "The path where the uploaded file should appear on your Studio. "
+        "Has to be within your Studio's home directory and will be relative to that. "
+        "If not specified, will use the name of the file you want to upload and place it in your home directory."
+    ),
+)
+def file(path: str, studio: Optional[str] = None, remote_path: Optional[str] = None) -> None:
+    """Upload a file to a Studio."""
+    _file(path=path, studio=studio, remote_path=remote_path)
+
+
+@upload.command("container")
+@click.argument("container")
+@click.option("--tag", default="latest", help="The tag of the container to upload.")
+@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(
+    "--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 store the Container in.",
+)
+@click.option(
+    "--platform",
+    default="linux/amd64",
+    help="This is the platform the container pulled and push to Lightning AI will run on.",
+)
+def upload_container(
+    container: str,
+    tag: str = "latest",
+    teamspace: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+    platform: Optional[str] = "linux/amd64",
+) -> None:
+    """Upload a container to Lightning AI's container registry."""
+    menu = TeamspacesMenu()
+    teamspace = menu(teamspace)
+    console = Console()
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        TimeElapsedColumn(),
+        console=console,
+        transient=False,
+    ) as progress:
+        try:
+            if platform != "linux/amd64":
+                console.print(
+                    "[yellow]Warning: The platform you selected (" + platform + ") may not deploy successfully[/yellow]"
+                )
+            api = LitContainerApi()
+            push_task = progress.add_task("Pushing Docker image", total=None)
+            console.print("Authenticating with Lightning Container Registry...")
+            try:
+                api.authenticate()
+                console.print("Authenticated with Lightning Container Registry", style="green")
+            except Exception:
+                # let the push with retry take control of auth moving forward
+                pass
+
+            lines = api.upload_container(container, teamspace, tag, cloud_account, platform, return_final_dict=True)
+            _print_docker_push(lines, console, progress, push_task)
+        except DockerNotRunningError as e:
+            e.print_help()
+            return
+        except LCRAuthFailedError:
+            console.print("Re-authenticating with Lightning Container Registry...")
+            if not api.authenticate(reauth=True):
+                raise StudioCliError("Failed to authenticate with Lightning Container Registry") from None
+            console.print("Authenticated with Lightning Container Registry", style="green")
+            lines = api.upload_container(container, teamspace, tag, cloud_account, platform)
+            _print_docker_push(lines, console, progress, push_task)
+        except Exception as e:
+            if _LIGHTNING_DEBUG:
+                print(e)
+                if e.__cause__:
+                    print(e.__cause__)
+
+            progress.update(push_task, description=f"[bold red]Error: {e!s}[/]")
+            progress.stop()
+            return
+        progress.update(push_task, description="[green]Container pushed![/green]")
+
+
+def _folder(path: str, studio: Optional[str] = None, remote_path: Optional[str] = None) -> None:
+    """Upload a folder to a Studio."""
+    console = Console()
+    if remote_path is None:
+        remote_path = os.path.basename(path)
+
+    if not Path(path).exists():
+        raise FileNotFoundError(f"The provided path does not exist: {path}.")
+    if not Path(path).is_dir():
+        raise StudioCliError(f"The provided path is not a folder: {path}. Use `lightning upload file` instead.")
+
+    selected_studio = _resolve_studio(studio)
+
+    console.print(f"Uploading to {selected_studio.teamspace.name}/{selected_studio.name}")
+
+    _upload_folder(path, remote_path, selected_studio)
+
+    studio_url = (
+        _get_cloud_url().replace(":443", "")
+        + "/"
+        + selected_studio.owner.name
+        + "/"
+        + selected_studio.teamspace.name
+        + "/studios/"
+        + selected_studio.name
+    )
+    console.print(f"See your files at {studio_url}")
+
+
+def _upload_folder(path: str, remote_path: str, studio: Studio) -> None:
+    pairs = {}
+    for root, _, files in os.walk(path):
+        rel_root = os.path.relpath(root, path)
+        for f in files:
+            pairs[os.path.join(root, f)] = os.path.join(remote_path, rel_root, f)
+
+    upload_state = _resolve_previous_upload_state(studio, remote_path, pairs)
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
+        futures = _start_parallel_upload(executor, studio, upload_state)
+
+        update_fn = tqdm(total=len(upload_state)).update if _global_upload_progress(upload_state) else lambda x: None
+
+        for f in concurrent.futures.as_completed(futures):
+            upload_state.pop(f.result())
+            _dump_current_upload_state(studio, remote_path, upload_state)
+            update_fn(1)
+
+
+def _file(path: str, studio: Optional[str] = None, remote_path: Optional[str] = None) -> None:
+    """Upload a file to a Studio."""
+    console = Console()
+    if remote_path is None:
+        remote_path = os.path.basename(path)
+
+    if Path(path).is_dir():
+        raise StudioCliError(f"The provided path is a folder: {path}. Use `lightning upload folder` instead.")
+    if not Path(path).exists():
+        raise FileNotFoundError(f"The provided path does not exist: {path}.")
+
+    selected_studio = _resolve_studio(studio)
+
+    console.print(f"Uploading to {selected_studio.teamspace.name}/{selected_studio.name}")
+
+    _single_file_upload(selected_studio, path, remote_path, True)
+
+    studio_url = (
+        _get_cloud_url().replace(":443", "")
+        + "/"
+        + selected_studio.owner.name
+        + "/"
+        + selected_studio.teamspace.name
+        + "/studios/"
+        + selected_studio.name
+    )
+    console.print(f"See your file at {studio_url}")
+
+
+def _resolve_studio(studio: Optional[str]) -> Studio:
+    user = _get_authed_user()
+    menu = _StudiosMenu()
+    possible_studios = menu._get_possible_studios(user)
+
+    try:
+        if studio is None:
+            selected_studio = menu._get_studio_from_interactive_menu(possible_studios)
+        else:
+            selected_studio = menu._get_studio_from_name(studio, 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 upload files to. "
+            "Please contact Lightning AI directly to resolve this issue."
+        ) from e
+
+    return Studio(**selected_studio)
+
+
+def _print_docker_push(lines: Generator, console: Console, progress: Progress, push_task: rich.progress.TaskID) -> None:
+    for line in lines:
+        if "status" in line:
+            console.print(line["status"], style="bright_black")
+            progress.update(push_task, description="Pushing Docker image")
+        elif "aux" in line:
+            console.print(line["aux"], style="bright_black")
+        elif "error" in line:
+            progress.stop()
+            console.print(f"\n[red]{line}[/red]")
+            return
+        elif "finish" in line:
+            if "url" in line:
+                webbrowser.open(line["url"])
+            console.print(f"Container available at [i]{line['url']}[/i]")
+            return
+        else:
+            console.print(line, style="bright_black")
+
+
+def _start_parallel_upload(
+    executor: concurrent.futures.ThreadPoolExecutor, studio: Studio, upload_state: Dict[str, str]
+) -> List[concurrent.futures.Future]:
+    # only add progress bar on individual uploads with less than 10 files
+    progress_bar = not _global_upload_progress(upload_state)
+
+    futures = []
+    for k, v in upload_state.items():
+        futures.append(
+            executor.submit(_single_file_upload, studio=studio, local_path=k, remote_path=v, progress_bar=progress_bar)
+        )
+
+    return futures
+
+
+def _single_file_upload(studio: Studio, local_path: str, remote_path: str, progress_bar: bool) -> str:
+    studio.upload_file(local_path, remote_path, progress_bar)
+    return local_path
+
+
+def _dump_current_upload_state(studio: Studio, remote_path: str, state_dict: Dict[str, str]) -> None:
+    """Dumps the current upload state so that we can safely resume later."""
+    curr_path = os.path.abspath(
+        os.path.expandvars(
+            os.path.expanduser(os.path.join(_STUDIO_UPLOAD_STATUS_PATH, studio._studio.id, remote_path + ".json"))
+        )
+    )
+
+    dirpath = os.path.dirname(curr_path)
+    if state_dict:
+        os.makedirs(os.path.dirname(curr_path), exist_ok=True)
+        with open(curr_path, "w") as f:
+            json.dump(state_dict, f, indent=4)
+        return
+
+    if os.path.exists(curr_path):
+        os.remove(curr_path)
+    if os.path.exists(dirpath):
+        os.removedirs(dirpath)
+
+
+def _resolve_previous_upload_state(studio: Studio, remote_path: str, state_dict: Dict[str, str]) -> Dict[str, str]:
+    """Resolves potential previous uploads to continue if possible."""
+    curr_path = os.path.abspath(
+        os.path.expandvars(
+            os.path.expanduser(os.path.join(_STUDIO_UPLOAD_STATUS_PATH, studio._studio.id, remote_path + ".json"))
+        )
+    )
+
+    # no previous download exists
+    if not os.path.isfile(curr_path):
+        return state_dict
+
+    menu = TerminalMenu(
+        [
+            "no, I accept that this may cause overwriting existing files",
+            "yes, continue previous upload",
+        ],
+        title=f"Found an incomplete upload for {studio.teamspace.name}/{studio.name}:{remote_path}. "
+        "Should we resume the previous upload?",
+    )
+    index = menu.show()
+    if index == 0:  # selected to start new upload
+        return state_dict
+
+    # at this point we know we want to resume the previous upload
+    with open(curr_path) as f:
+        return json.load(f)
+
+
+def _global_upload_progress(upload_state: Dict[str, str]) -> bool:
+    return len(upload_state) > 10

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)

lightning_sdk/cli/mmt/__init__.py

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

lightning_sdk/cli/studio/__init__.py

@@ -0,0 +1,24 @@
+"""Studio CLI commands."""
+
+import click
+
+
+def register_commands(group: click.Group) -> None:
+    """Register studio commands with the given group."""
+    from lightning_sdk.cli.studio.connect import connect_studio
+    from lightning_sdk.cli.studio.create import create_studio
+    from lightning_sdk.cli.studio.delete import delete_studio
+    from lightning_sdk.cli.studio.list import list_studios
+    from lightning_sdk.cli.studio.ssh import ssh_studio
+    from lightning_sdk.cli.studio.start import start_studio
+    from lightning_sdk.cli.studio.stop import stop_studio
+    from lightning_sdk.cli.studio.switch import switch_studio
+
+    group.add_command(delete_studio)
+    group.add_command(create_studio)
+    group.add_command(list_studios)
+    group.add_command(ssh_studio)
+    group.add_command(start_studio)
+    group.add_command(stop_studio)
+    group.add_command(switch_studio)
+    group.add_command(connect_studio)