lightning-sdk 2025.11.13.post0__py3-none-any.whl

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)

lightning_sdk/cli/legacy/create.py

@@ -0,0 +1,115 @@
+import random
+import string
+from contextlib import suppress
+from typing import Optional
+
+import click
+from rich.console import Console
+
+from lightning_sdk import Machine, Studio
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.machine import CloudProvider
+
+_MACHINE_VALUES = tuple(
+    [machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine) and machine._include_in_cli]
+)
+_PROVIDER_VALUES = tuple([provider.value for provider in CloudProvider])
+
+
+@click.group("create")
+def create() -> None:
+    """Create new resources on the Lightning AI platform."""
+
+
+@create.command("studio")
+@click.argument("name")
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "The teamspace the studio will be 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(
+    "--start",
+    default=None,
+    type=click.Choice(_MACHINE_VALUES),
+    help="If specified, will start the created studio on the given machine.",
+)
+@click.option(
+    "--cloud-account",
+    "--cloud_account",
+    default=None,
+    help=(
+        "The cloud account to create the studio on. "
+        "If not specified, will try to infer from the environment (e.g. when run from within a Studio.) "
+        "or fall back to the teamspace default."
+    ),
+)
+@click.option(
+    "--cloud-provider",
+    default=None,
+    type=click.Choice(_PROVIDER_VALUES),
+    help="The provider to create the studio on. If --cloud-account is specified, this option is prioritized.",
+)
+@click.option(
+    "--provider",
+    default=None,
+    type=click.Choice(_PROVIDER_VALUES),
+    help=(
+        "Deprecated. Use --cloud-provider instead. The provider to create the studio on. "
+        "If --cloud-account is specified, this option is prioritized."
+    ),
+)
+def studio(
+    name: str,
+    teamspace: Optional[str] = None,
+    start: Optional[str] = None,
+    cloud_account: Optional[str] = None,
+    cloud_provider: Optional[str] = None,
+    provider: Optional[str] = None,
+) -> None:
+    """Create a new studio on the Lightning AI platform.
+
+    Example:
+        lightning create studio NAME
+
+    NAME: the name of the studio to create. If already present within teamspace, will add a random suffix.
+    """
+    menu = TeamspacesMenu()
+    teamspace_resolved = menu(teamspace)
+
+    # default cloud account to current studios cloud account if run from studio
+    # else it will fall back to teamspace default in the backend
+    if cloud_account is None:
+        with suppress(ValueError):
+            s = Studio()
+            if s.teamspace.name == teamspace_resolved.name and s.teamspace.owner.name == teamspace_resolved.owner.name:
+                cloud_account = s.cloud_account
+
+    console = Console()
+
+    with suppress(ValueError):
+        Studio(name, teamspace=teamspace_resolved, create_ok=False)
+        new_name = name + "-" + "".join(random.choices(string.ascii_lowercase + string.digits, k=4))
+        console.print(f"Studio with name {name} already exists. Using {new_name} instead.")
+        name = new_name
+
+    studio = Studio(
+        name=name,
+        teamspace=teamspace_resolved,
+        cloud_account=cloud_account,
+        create_ok=True,
+        cloud_provider=cloud_provider,
+        provider=provider,
+    )
+
+    console.print(f"Created Studio {studio.name}.")
+
+    if start is not None:
+        start_machine = getattr(Machine, start, start)
+        Studio.show_progress = True
+        studio.start(start_machine)
+        console.print(f"Started Studio {studio.name} on machine {start}")

lightning_sdk/cli/legacy/delete.py

@@ -0,0 +1,131 @@
+from typing import Optional
+
+import click
+from rich.console import Console
+
+from lightning_sdk.cli.legacy.exceptions import StudioCliError
+from lightning_sdk.cli.legacy.job_and_mmt_action import _JobAndMMTAction
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+from lightning_sdk.lit_container import LitContainer
+from lightning_sdk.studio import Studio
+
+
+@click.group()
+def delete() -> None:
+    """Delete resources on the Lightning AI platform."""
+
+
+@delete.command(name="container")
+@click.argument("name")
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "The teamspace to delete the container from. "
+        "Should be specified as {owner}/{name} "
+        "If not provided, can be selected in an interactive menu."
+    ),
+)
+def container(name: str, teamspace: Optional[str] = None) -> None:
+    """Delete the docker container NAME."""
+    api = LitContainer()
+    menu = TeamspacesMenu()
+    resolved_teamspace = menu(teamspace=teamspace)
+    try:
+        api.delete_container(name, resolved_teamspace.name, resolved_teamspace.owner.name)
+        Console().print(f"Container {name} deleted successfully.")
+    except Exception as e:
+        raise StudioCliError(f"Could not delete container {name} from project {resolved_teamspace.name}: {e}") from None
+
+
+@delete.command(name="job")
+@click.argument(
+    "name",
+)
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "The teamspace to delete the job from. "
+        "Should be specified as {owner}/{name} "
+        "If not provided, can be selected in an interactive menu."
+    ),
+)
+def job(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+    """Delete a job.
+
+    Example:
+      lightning delete job NAME
+
+    NAME: the name of the job to delete.
+    """
+    menu = _JobAndMMTAction()
+    job = menu.job(name=name, teamspace=teamspace)
+
+    job.delete()
+    Console().print(f"Successfully deleted {job.name}!")
+
+
+@delete.command(name="mmt")
+@click.argument(
+    "name",
+)
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "The teamspace to delete the job from. "
+        "Should be specified as {owner}/{name} "
+        "If not provided, can be selected in an interactive menu."
+    ),
+)
+def mmt(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+    """Delete a multi-machine job.
+
+    Example:
+      lightning delete mmt NAME
+
+    NAME: the name of the multi-machine job to delete.
+    """
+    menu = _JobAndMMTAction()
+    mmt = menu.mmt(name=name, teamspace=teamspace)
+
+    mmt.delete()
+    Console().print(f"Successfully deleted {mmt.name}!")
+
+
+@delete.command(name="studio")
+@click.argument("name")
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "The teamspace to delete the studio from. "
+        "Should be specified as {owner}/{name} "
+        "If not provided, can be selected in an interactive menu."
+    ),
+)
+def studio(name: str, teamspace: Optional[str] = None) -> None:
+    """Delete an existing studio.
+
+    Example:
+      lightning delete studio NAME
+
+    NAME: the name of the studio to delete
+    """
+    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)
+
+    studio.delete()
+    Console().print("Studio successfully deleted")

lightning_sdk/cli/legacy/deploy/_auth.py

@@ -0,0 +1,196 @@
+import os
+import time
+from datetime import datetime
+from enum import Enum
+from typing import Any, List, Optional, TypedDict
+from urllib.parse import urlencode
+
+from rich.console import Console
+from rich.prompt import Confirm
+
+from lightning_sdk import Teamspace
+from lightning_sdk.api import UserApi
+from lightning_sdk.cli.utils.teamspace_selection import TeamspacesMenu
+from lightning_sdk.lightning_cloud import env
+from lightning_sdk.lightning_cloud.login import Auth, AuthServer
+from lightning_sdk.lightning_cloud.openapi import V1CloudSpace
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+from lightning_sdk.utils.resolve import _get_authed_user, _resolve_teamspace
+
+LITSERVE_CODE = os.environ.get("LITSERVE_CODE", "j39bzk903h")
+_POLL_TIMEOUT = 120
+
+
+class _AuthMode(Enum):
+    DEVBOX = "dev"
+    DEPLOY = "deploy"
+
+
+class _AuthServer(AuthServer):
+    def __init__(self, mode: _AuthMode, *args: Any, **kwargs: Any) -> None:
+        self._mode = mode
+        super().__init__(*args, **kwargs)
+
+    def get_auth_url(self, port: int) -> str:
+        redirect_uri = f"http://localhost:{port}/login-complete"
+        params = urlencode({"redirectTo": redirect_uri, "mode": self._mode.value, "okbhrt": LITSERVE_CODE})
+        return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}"
+
+
+class _AuthLitServe(Auth):
+    def __init__(self, mode: _AuthMode, shall_confirm: bool = False) -> None:
+        super().__init__()
+        self._mode = mode
+        self._shall_confirm = shall_confirm
+
+    def _run_server(self) -> None:
+        if self._shall_confirm:
+            proceed = Confirm.ask(
+                "[bold yellow]LitServe needs to authenticate with Lightning AI to deploy your server.[/bold yellow]\n"
+                "This will open a browser window for login.\n"
+                "Do you want to continue?",
+                default=True,
+            )
+            if not proceed:
+                raise RuntimeError(
+                    "Login cancelled. Please login to Lightning AI to deploy the API. Run `lightning login` to login."
+                )
+        print("Opening browser for authentication...")
+        print("Please come back to the terminal after logging in.")
+        time.sleep(3)
+        _AuthServer(self._mode).login_with_browser(self)
+
+
+def authenticate(mode: _AuthMode, shall_confirm: bool = True) -> None:
+    """Authenticate with Lightning AI.
+
+    This will open a browser window for authentication.
+    If `shall_confirm` is True, it will ask for confirmation before proceeding.
+    """
+    auth = _AuthLitServe(mode, shall_confirm)
+    auth.authenticate()
+
+
+def select_teamspace(teamspace: Optional[str], org: Optional[str], user: Optional[str]) -> Teamspace:
+    if teamspace is None:
+        menu = TeamspacesMenu()
+        possible_teamspaces = menu._get_possible_teamspaces(_get_authed_user())
+        if len(possible_teamspaces) == 1:
+            name = next(iter(possible_teamspaces.values()))["name"]
+            return Teamspace(name=name, org=org, user=user)
+
+        return menu(teamspace)
+
+    return _resolve_teamspace(teamspace=teamspace, org=org, user=user)
+
+
+class _UserStatus(TypedDict):
+    verified: bool
+    onboarded: bool
+
+
+def poll_verified_status(timeout: int = _POLL_TIMEOUT) -> _UserStatus:
+    """Polls the verified status of the user until it is True or a timeout occurs."""
+    user_api = UserApi()
+    user = _get_authed_user()
+    start_time = datetime.now()
+    result = _UserStatus(onboarded=False, verified=False)
+    while True:
+        user_resp = user_api.get_user(name=user.name)
+        result["onboarded"] = user_resp.status.completed_project_onboarding
+        result["verified"] = user_resp.status.verified
+        if user_resp.status.verified:
+            return result
+        if (datetime.now() - start_time).total_seconds() > timeout:
+            break
+        time.sleep(5)
+    return result
+
+
+class _OnboardingStatus(Enum):
+    NOT_VERIFIED = "not_verified"
+    ONBOARDING = "onboarding"
+    ONBOARDED = "onboarded"
+
+
+class _Onboarding:
+    def __init__(self, console: Console) -> None:
+        self.console = console
+        self.user = _get_authed_user()
+        self.user_api = UserApi()
+        self.client = LightningClient(max_tries=7)
+
+    @property
+    def verified(self) -> bool:
+        return self.user_api.get_user(name=self.user.name).status.verified
+
+    @property
+    def is_onboarded(self) -> bool:
+        return self.user_api.get_user(name=self.user.name).status.completed_project_onboarding
+
+    @property
+    def can_join_org(self) -> bool:
+        return len(self.client.organizations_service_list_joinable_organizations().joinable_organizations) > 0
+
+    @property
+    def status(self) -> _OnboardingStatus:
+        if not self.verified:
+            return _OnboardingStatus.NOT_VERIFIED
+        if self.is_onboarded:
+            return _OnboardingStatus.ONBOARDED
+        return _OnboardingStatus.ONBOARDING
+
+    def _wait_user_onboarding(self, timeout: int = _POLL_TIMEOUT) -> None:
+        """Wait for user onboarding if they can join the teamspace otherwise move to select a teamspace."""
+        status = self.status
+        if status == _OnboardingStatus.ONBOARDED:
+            return
+
+        self.console.print("Waiting for account setup. Visit lightning.ai")
+        start_time = datetime.now()
+        while self.status != _OnboardingStatus.ONBOARDED:
+            time.sleep(5)
+            if self.is_onboarded:
+                return
+            if (datetime.now() - start_time).total_seconds() > timeout:
+                break
+
+        raise RuntimeError("Timed out waiting for onboarding status")
+
+    def get_cloudspace_id(self, teamspace: Teamspace) -> Optional[str]:
+        cloudspaces: List[V1CloudSpace] = self.client.cloud_space_service_list_cloud_spaces(teamspace.id).cloudspaces
+        cloudspaces = sorted(cloudspaces, key=lambda cloudspace: cloudspace.created_at, reverse=True)
+        if len(cloudspaces) == 0:
+            raise RuntimeError("Error creating deployment! Finish account setup at lightning.ai first.")
+        # get the first cloudspace
+        cloudspace = cloudspaces[0]
+        if "scratch-studio" in cloudspace.name or "scratch-studio" in cloudspace.display_name:
+            return cloudspace.id
+        return None
+
+    def select_teamspace(self, teamspace: Optional[str], org: Optional[str], user: Optional[str]) -> Teamspace:
+        """Select a teamspace while onboarding.
+
+        If user is being onboarded and can't join any org, the teamspace it will be resolved to the default
+         personal teamspace.
+        If user is being onboarded and can join an org then it will select default teamspace from the org.
+        """
+        if self.is_onboarded:
+            return select_teamspace(teamspace, org, user)
+
+        # Run only when user hasn't completed onboarding yet.
+        menu = TeamspacesMenu()
+        self._wait_user_onboarding()
+        # Onboarding has been completed - user already selected organization if they could
+        possible_teamspaces = menu._get_possible_teamspaces(self.user)
+        if len(possible_teamspaces) == 1:
+            # User didn't select any org
+            value = next(iter(possible_teamspaces.values()))
+            return Teamspace(name=value["name"], org=value["org"], user=value["user"])
+
+        for _, value in possible_teamspaces.items():
+            # User select an org
+            # Onboarding teamspace will be the default teamspace in the selected org
+            if value["org"]:
+                return Teamspace(name=value["name"], org=value["org"], user=value["user"])
+        raise RuntimeError("Unable to select teamspace. Visit lightning.ai")