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

lightning_sdk/cli/download.py

@@ -3,6 +3,7 @@ import re
 from pathlib import Path
 from typing import Optional
 
+import click
 from rich.console import Console
 
 from lightning_sdk.api.lit_container_api import LitContainerApi
@@ -25,51 +26,7 @@ class _Downloads(_StudiosMenu, _TeamspacesMenu):
             This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>.
           download_dir: The directory where the Model should be downloaded.
         """
-        download_model(
-            name=name,
-            download_dir=download_dir,
-            progress_bar=True,
-        )
-
-    def _resolve_studio(self, studio: Optional[str]) -> Studio:
-        user = _get_authed_user()
-        # if no studio specify suggest/filter only user's studios
-        possible_studios = self._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 = self._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 possible_studios:
-                        raise ValueError(
-                            f"Could not find Studio like '{studio}', please consider update your filtering pattern."
-                        )
-                    selected_studio = self._get_studio_from_interactive_menu(list(possible_studios))
-            else:
-                selected_studio = self._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 upload files to. "
-                "Please contact Lightning AI directly to resolve this issue."
-            ) from e
-
-        with skip_studio_init():
-            return Studio(**selected_studio)
+        model(name=name, download_dir=download_dir)
 
     def folder(self, path: str = "", studio: Optional[str] = None, local_path: str = ".") -> None:
         """Download a folder from a Studio.
@@ -84,27 +41,9 @@ class _Downloads(_StudiosMenu, _TeamspacesMenu):
             with filtered studios will be shown for final selection.
           local_path: The path to the directory you want to download the folder to.
         """
-        local_path = Path(local_path)
-        if not local_path.is_dir():
-            raise NotADirectoryError(f"'{local_path}' is not a directory")
-
-        resolved_studio = self._resolve_studio(studio)
+        folder(path=path, studio=studio, local_path=local_path)
 
-        if not path:
-            local_path /= resolved_studio.name
-            path = ""
-
-        try:
-            if not path:
-                raise FileNotFoundError()
-            resolved_studio.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}. "
-                "Please contact Lightning AI directly to resolve this issue."
-            ) from e
-
-    def file(self, path: str = "", studio: Optional[str] = None, local_path: str = ".") -> None:
+    def file(self, path: str, studio: Optional[str] = None, local_path: str = ".") -> None:
         """Download a file from a Studio.
 
         Args:
@@ -115,25 +54,7 @@ class _Downloads(_StudiosMenu, _TeamspacesMenu):
             with filtered studios will be shown for final selection.
           local_path: The path to the directory you want to download the file to.
         """
-        local_path = Path(local_path)
-        if not local_path.is_dir():
-            raise NotADirectoryError(f"'{local_path}' is not a directory")
-
-        resolved_studio = self._resolve_studio(studio)
-
-        if not path:
-            local_path /= resolved_studio.name
-            path = ""
-
-        try:
-            if not path:
-                raise FileNotFoundError()
-            resolved_studio.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}. "
-                "Please contact Lightning AI directly to resolve this issue."
-            ) from e
+        file(path=path, studio=studio, local_path=local_path)
 
     def container(self, container: str, teamspace: Optional[str] = None, tag: str = "latest") -> None:
         """Download a docker container from a teamspace.
@@ -143,9 +64,175 @@ class _Downloads(_StudiosMenu, _TeamspacesMenu):
           teamspace: The name of the teamspace to download the container from.
           tag: The tag of the container to download.
         """
-        console = Console()
-        resolved_teamspace = self._resolve_teamspace(teamspace)
-        with console.status("Downloading container..."):
-            api = LitContainerApi()
-            api.download_container(container, resolved_teamspace, tag)
-            console.print("Container downloaded successfully", style="green")
+        download_container(container=container, teamspace=teamspace, tag=tag)
+
+
+@click.group(name="download")
+def download() -> None:
+    """Download resources from Lightning AI."""
+
+
+# @download.command(name="model")
+# @click.option(
+#     "--name",
+#     help=(
+#         "The name of the Model you want to download. "
+#         "This should have the format <ORGANIZATION-NAME>/<TEAMSPACE-NAME>/<MODEL-NAME>."
+#     ),
+# )
+# @click.option("--download-dir", default=".", help="The directory where the Model should be downloaded.")
+def model(name: str, download_dir: str = ".") -> None:
+    """Download a Model."""
+    download_model(
+        name=name,
+        download_dir=download_dir,
+        progress_bar=True,
+    )
+
+
+# @download.command(name="folder")
+# @click.option(
+#     "--path",
+#     default="",
+#     help=(
+#         "The relative path within the Studio you want to download. "
+#         "If you leave it empty it will download whole studio and locally creates a "
+#         "new folder with the same name as the selected studio."
+#     ),
+# )
+# @click.option(
+#     "--studio",
+#     default=None,
+#     help=(
+#         "The name of the studio to upload to. "
+#         "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("--local-path", default=".", help="The path to the directory you want to download the folder to.")
+def folder(path: str = "", studio: Optional[str] = None, local_path: str = ".") -> None:
+    """Download a folder from a Studio."""
+    local_path = Path(local_path)
+    if not local_path.is_dir():
+        raise NotADirectoryError(f"'{local_path}' is not a directory")
+
+    menu = _StudiosMenu()
+    resolved_studio = menu._resolve_studio(studio)
+
+    if not path:
+        local_path /= resolved_studio.name
+        path = ""
+
+    try:
+        if not path:
+            raise FileNotFoundError()
+        resolved_studio.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}. "
+            "Please contact Lightning AI directly to resolve this issue."
+        ) from e
+
+
+# @download.command(name="file")
+# @click.option(
+#     "--path",
+#     default="",
+#     help=(
+#         "The relative path within the Studio you want to download. "
+#         "If you leave it empty it will download whole studio and locally creates a new folder "
+#         "with the same name as the selected studio."
+#     ),
+# )
+# @click.option(
+#     "--studio",
+#     default=None,
+#     help=(
+#         "The name of the studio to upload to. "
+#         "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("--local-path", default=".", help="The path to the directory you want to download the folder to.")
+def file(path: str = "", studio: Optional[str] = None, local_path: str = ".") -> None:
+    """Download a file from a Studio."""
+    local_path = Path(local_path)
+    if not local_path.is_dir():
+        raise NotADirectoryError(f"'{local_path}' is not a directory")
+
+    resolved_studio = _resolve_studio(studio)
+
+    if not path:
+        local_path /= resolved_studio.name
+        path = ""
+
+    try:
+        if not path:
+            raise FileNotFoundError()
+        resolved_studio.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}. "
+            "Please contact Lightning AI directly to resolve this issue."
+        ) from e
+
+
+# @download.command(name="container")
+# @click.argument("container")
+# @click.argument("teamspace", default=None, help="The name of the teamspace to download the container from")
+# @click.argument("tag", default="latest", show_default=True, help="The tag of the container to download.")
+def download_container(container: str, teamspace: Optional[str] = None, tag: str = "latest") -> None:
+    """Download the docker container CONTAINER from a teamspace."""
+    console = Console()
+    menu = _TeamspacesMenu()
+    resolved_teamspace = menu._resolve_teamspace(teamspace)
+    with console.status("Downloading container..."):
+        api = LitContainerApi()
+        api.download_container(container, resolved_teamspace, tag)
+        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 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 upload files to. "
+            "Please contact Lightning AI directly to resolve this issue."
+        ) from e
+
+    with skip_studio_init():
+        return Studio(**selected_studio)

lightning_sdk/cli/entrypoint.py

@@ -2,19 +2,20 @@ import sys
 from types import TracebackType
 from typing import Type
 
+import click
 from fire import Fire
 from lightning_utilities.core.imports import RequirementCache
 from rich.console import Console
 from rich.panel import Panel
 
 from lightning_sdk.api.studio_api import _cloud_url
-from lightning_sdk.cli.ai_hub import _AIHub
-from lightning_sdk.cli.configure import _Configure
-from lightning_sdk.cli.connect import _Connect
-from lightning_sdk.cli.delete import _Delete
-from lightning_sdk.cli.download import _Downloads
-from lightning_sdk.cli.generate import _Generate
-from lightning_sdk.cli.inspect import _Inspect
+from lightning_sdk.cli.ai_hub import _AIHub, aihub
+from lightning_sdk.cli.configure import _Configure, configure
+from lightning_sdk.cli.connect import _Connect, connect
+from lightning_sdk.cli.delete import _Delete, delete
+from lightning_sdk.cli.download import _Downloads, download
+from lightning_sdk.cli.generate import _Generate, generate
+from lightning_sdk.cli.inspect import _Inspect, inspect
 from lightning_sdk.cli.legacy import _LegacyLightningCLI
 from lightning_sdk.cli.list import _List
 from lightning_sdk.cli.run import _Run
@@ -52,18 +53,11 @@ class StudioCLI:
 
     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
+        return login()
 
     def logout(self) -> None:
         """Logout from Lightning AI Studios."""
-        auth = Auth()
-        auth.clear()
+        return logout()
 
 
 def _notify_exception(exception_type: Type[BaseException], value: BaseException, tb: TracebackType) -> None:  # No
@@ -77,5 +71,39 @@ def main_cli() -> None:
     Fire(StudioCLI(), name="lightning")
 
 
+@click.group(name="lightning", help="Command line interface (CLI) to interact with/manage Lightning AI Studios.")
+def main_cli_click() -> None:
+    pass
+
+
+# @main_cli_click.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_click.command
+def logout() -> None:
+    """Logout from Lightning AI Studios."""
+    auth = Auth()
+    auth.clear()
+
+
+# TODO: handle exception hook registration
+main_cli_click.add_command(aihub)
+main_cli_click.add_command(configure)
+main_cli_click.add_command(connect)
+main_cli_click.add_command(delete)
+main_cli_click.add_command(download)
+main_cli_click.add_command(generate)
+main_cli_click.add_command(inspect)
+
+
 if __name__ == "__main__":
     main_cli()

lightning_sdk/cli/generate.py

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

lightning_sdk/cli/inspect.py

@@ -1,11 +1,12 @@
 from typing import Optional
 
+import click
 from rich.console import Console
 
 from lightning_sdk.cli.job_and_mmt_action import _JobAndMMTAction
 
 
-class _Inspect(_JobAndMMTAction):
+class _Inspect:
     """Inspect resources of the Lightning AI platform to get additional details as JSON."""
 
     def job(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
@@ -18,7 +19,7 @@ class _Inspect(_JobAndMMTAction):
                 If not specified can be selected interactively.
 
         """
-        (super().job(name=name, teamspace=teamspace).json())
+        job(name=name, teamspace=teamspace)
 
     def mmt(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
         """Inspect a multi-machine job for further details as JSON.
@@ -30,4 +31,43 @@ class _Inspect(_JobAndMMTAction):
                 If not specified can be selected interactively.
 
         """
-        Console().print(super().mmt(name=name, teamspace=teamspace).json())
+        mmt(name=name, teamspace=teamspace)
+
+
+@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())