lightning-sdk 0.1.56__py3-none-any.whl → 0.1.57__py3-none-any.whl

lightning_sdk/__init__.py

@@ -1,7 +1,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.helpers import _check_version_and_prompt_upgrade
+from lightning_sdk.helpers import _check_version_and_prompt_upgrade, _set_tqdm_envvars_noninteractive
 from lightning_sdk.job import Job
 from lightning_sdk.machine import Machine
 from lightning_sdk.mmt import MMT
@@ -29,5 +29,6 @@ __all__ = [
     "AIHub",
 ]
 
-__version__ = "0.1.56"
+__version__ = "0.1.57"
 _check_version_and_prompt_upgrade(__version__)
+_set_tqdm_envvars_noninteractive()

lightning_sdk/cli/ai_hub.py

@@ -1,4 +1,6 @@
-from typing import List, Optional
+from typing import Optional
+
+import click
 
 from lightning_sdk.ai_hub import AIHub
 from lightning_sdk.cli.studios_menu import _StudiosMenu
@@ -7,10 +9,7 @@ from lightning_sdk.cli.studios_menu import _StudiosMenu
 class _AIHub(_StudiosMenu):
     """Interact with Lightning Studio - AI Hub."""
 
-    def __init__(self) -> None:
-        self._hub = AIHub()
-
-    def api_info(self, api_id: str) -> dict:
+    def api_info(self, api_id: str) -> None:
         """Get full API template info such as input details.
 
         Example:
@@ -19,15 +18,15 @@ class _AIHub(_StudiosMenu):
         Args:
           api_id: The ID of the API for which information is requested.
         """
-        return self._hub.api_info(api_id)
+        return api_info(api_id=api_id)
 
-    def list_apis(self, search: Optional[str] = None) -> List[dict]:
+    def list_apis(self, search: Optional[str] = None) -> None:
         """List API templates available in the AI Hub.
 
         Args:
           search: Search for API templates by name.
         """
-        return self._hub.list_apis(search=search)
+        return list_apis(search=search)
 
     def deploy(
         self,
@@ -36,7 +35,7 @@ class _AIHub(_StudiosMenu):
         name: Optional[str] = None,
         teamspace: Optional[str] = None,
         org: Optional[str] = None,
-    ) -> dict:
+    ) -> None:
         """Deploy an API template from the AI Hub.
 
         Args:
@@ -46,4 +45,56 @@ class _AIHub(_StudiosMenu):
           teamspace: Teamspace to deploy the API to. Defaults to user's default teamspace.
           org: Organization to deploy the API to. Defaults to user's default organization.
         """
-        return self._hub.run(api_id, cloud_account=cloud_account, name=name, teamspace=teamspace, org=org)
+        return deploy(api_id=api_id, cloud_account=cloud_account, name=name, teamspace=teamspace, org=org)
+
+
+@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",
+#     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:
+    ai_hub = AIHub()
+    ai_hub.run(api_id, cloud_account=cloud_account, name=name, teamspace=teamspace, org=org)

lightning_sdk/cli/configure.py

@@ -1,14 +1,98 @@
+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.generate import _Generate
+from lightning_sdk.cli.generate import _generate_ssh_config
+from lightning_sdk.cli.studios_menu import _StudiosMenu
 from lightning_sdk.lightning_cloud.login import Auth
 
 
+class _Configure(_StudiosMenu):
+    """Configure lightning products."""
+
+    def ssh(self, name: Optional[str] = None, teamspace: Optional[str] = None, overwrite: bool = False) -> None:
+        """Get SSH config entry for a studio.
+
+        Args:
+            name: 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.)
+            teamspace: 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.)
+            overwrite: Whether to overwrite the SSH key and config if they already exist.
+        """
+        ssh(name=name, teamspace=teamspace, overwrite=overwrite)
+
+
+@click.group(name="configure")
+def configure() -> None:
+    """Configure access to resources on the Lightning AI platform."""
+
+
+# @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,
+#     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."""
+    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}")
+
+
 def _download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Optional[int] = None) -> None:
     """Download a file from a URL."""
     import requests
@@ -26,67 +110,28 @@ def _download_file(url: str, local_path: Path, overwrite: bool = True, chmod: Op
         local_path.chmod(0o600)
 
 
-class _Configure(_Generate):
-    """Configure lightning products."""
-
-    @staticmethod
-    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)
-
-    def ssh(self, overwrite: bool = False, ssh_key_name: str = "lightning_rsa") -> None:
-        """Get SSH config entry for a studio.
-
-        Args:
-            overwrite: Whether to overwrite the SSH key and config if they already exist.
-            ssh_key_name: The name of the SSH key to generate
-        """
-        auth = Auth()
-        auth.authenticate()
-        console = Console()
-        ssh_dir = Path.home() / ".ssh"
-        ssh_dir.mkdir(parents=True, exist_ok=True)
-
-        key_path = ssh_dir / ssh_key_name
-        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:
-            self._download_ssh_keys(auth.api_key, ssh_home=ssh_dir, ssh_key_name=ssh_key_name, overwrite=overwrite)
-            console.print(f"SSH key generated and saved to {key_path}")
-
-        # Check if the SSH config already contains the required configuration
-        config_content = self._generate_ssh_config(str(key_path))
-        if config_path.exists():
-            with config_path.open("r") as config_file:
-                if config_content.strip() 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(config_content)
-            console.print(f"SSH config updated at {config_path}")
+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/connect.py

@@ -2,11 +2,13 @@ import subprocess
 import sys
 from typing import Optional
 
-from lightning_sdk.cli.configure import _Configure
-from lightning_sdk.lightning_cloud.login import Auth
+import click
 
+from lightning_sdk.cli.configure import ssh as configure_ssh
+from lightning_sdk.cli.studios_menu import _StudiosMenu
 
-class _Connect(_Configure):
+
+class _Connect:
     """Connect to lightning products."""
 
     def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
@@ -16,16 +18,30 @@ class _Connect(_Configure):
             name: The name of the studio to connect to.
             teamspace: The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.
         """
-        auth = Auth()
-        auth.authenticate()  # this is maybe not needed
-        studio = self._get_studio(name=name, teamspace=teamspace)
-        host = "ssh.lightning.ai"
-        username = f"s_{studio._studio.id}"
-
-        self.ssh(overwrite=False)
-
-        try:
-            subprocess.run(["ssh", f"{username}@{host}"])
-        except Exception as ex:
-            print(f"Failed to establish SSH connection: {ex}")
-            sys.exit(1)
+        return studio(name=name, teamspace=teamspace)
+
+
+@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(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/delete.py

@@ -1,5 +1,6 @@
 from typing import Optional
 
+import click
 from rich.console import Console
 
 from lightning_sdk.cli.exceptions import StudioCliError
@@ -20,15 +21,7 @@ class _Delete(_JobAndMMTAction, _TeamspacesMenu):
             teamspace: The teamspace to delete the container from. Should be specified as {owner}/{name}
                 If not provided, can be selected in an interactive menu.
         """
-        api = LitContainer()
-        resolved_teamspace = self._resolve_teamspace(teamspace=teamspace)
-        try:
-            api.delete_container(container, resolved_teamspace.name, resolved_teamspace.owner.name)
-            Console().print(f"Container {container} deleted successfully.")
-        except Exception as e:
-            raise StudioCliError(
-                f"Could not delete container {container} from project {resolved_teamspace.name}: {e}"
-            ) from None
+        delete_container(container=container, teamspace=teamspace)
 
     def job(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
         """Delete a job.
@@ -40,10 +33,7 @@ class _Delete(_JobAndMMTAction, _TeamspacesMenu):
                 If not specified can be selected interactively.
 
         """
-        job = super().job(name=name, teamspace=teamspace)
-
-        job.delete()
-        Console().print(f"Successfully deleted {job.name}!")
+        job(name=name, teamspace=teamspace)
 
     def mmt(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
         """Delete a multi-machine job.
@@ -55,10 +45,7 @@ class _Delete(_JobAndMMTAction, _TeamspacesMenu):
                 If not specified can be selected interactively.
 
         """
-        mmt = super().mmt(name=name, teamspace=teamspace)
-
-        mmt.delete()
-        Console().print(f"Successfully deleted {mmt.name}!")
+        mmt(name=name, teamspace=teamspace)
 
     def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
         """Delete an existing studio.
@@ -70,18 +57,80 @@ class _Delete(_JobAndMMTAction, _TeamspacesMenu):
             teamspace: 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.)
         """
-        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):
-            studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
-
-        studio.delete()
-        Console().print("Studio successfully deleted")
+        studio(name=name, teamspace=teamspace)
+
+
+@click.group()
+def delete() -> None:
+    """Delete resources on the Lightning AI platform."""
+
+
+# @delete.command(name="container")
+# @click.option("--container", help="The name of the container to delete.")
+# @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 delete_container(container: str, teamspace: Optional[str] = None) -> None:
+    """Delete the docker container CONTAINER."""
+    api = LitContainer()
+    menu = _TeamspacesMenu()
+    resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+    try:
+        api.delete_container(container, resolved_teamspace.name, resolved_teamspace.owner.name)
+        Console().print(f"Container {container} deleted successfully.")
+    except Exception as e:
+        raise StudioCliError(
+            f"Could not delete container {container} from project {resolved_teamspace.name}: {e}"
+        ) from None
+
+
+# @delete.command(name="job")
+# @click.option("--name", help="The name of the job to delete.")
+# @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."""
+    menu = _JobAndMMTAction()
+    job = menu.job(name=name, teamspace=teamspace)
+
+    job.delete()
+    Console().print(f"Successfully deleted {job.name}!")
+
+
+# @delete.command(name="mmt")
+# @click.option("--name", help="The name of the multi-machine job to delete.")
+# @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."""
+    menu = _JobAndMMTAction()
+    mmt = menu.mmt(name=name, teamspace=teamspace)
+
+    mmt.delete()
+    Console().print(f"Successfully deleted {mmt.name}!")
+
+
+# @delete.command(name="studio")
+# @click.option("--name", help="The name of the studio to delete.")
+# @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: Optional[str] = None, teamspace: Optional[str] = None) -> None:
+    """Delete an existing studio."""
+    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):
+        studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
+
+    studio.delete()
+    Console().print("Studio successfully deleted")