lightning-sdk 0.1.58__py3-none-any.whl → 0.2.0__py3-none-any.whl

lightning_sdk/cli/list.py

@@ -1,6 +1,7 @@
 from contextlib import suppress
 from typing import Callable, Optional
 
+import click
 from rich.console import Console
 from rich.table import Table
 from typing_extensions import Literal
@@ -12,233 +13,291 @@ from lightning_sdk.lit_container import LitContainer
 from lightning_sdk.utils.resolve import _get_authed_user
 
 
-class _List(_TeamspacesMenu):
+@click.group(name="list")
+def list_cli() -> None:
     """List resources on the Lightning AI platform."""
 
-    @staticmethod
-    def _sort_studios_key(sort_by: str) -> Callable[[Studio], str]:
-        """Return a key function to sort studios by a given attribute."""
-        sort_key_map = {
-            "name": lambda s: str(s.name or ""),
-            "teamspace": lambda s: str(s.teamspace.name or ""),
-            "status": lambda s: str(s.status or ""),
-            "machine": lambda s: str(s.machine or ""),
-            "cloud-account": lambda s: str(s.cloud_account or ""),
-        }
-        return sort_key_map.get(sort_by, lambda s: s.name)
-
-    def studios(
-        self,
-        teamspace: Optional[str] = None,
-        all: bool = False,  # noqa: A002
-        sort_by: Optional[Literal["name", "teamspace", "status", "machine", "cloud-account"]] = None,
-    ) -> None:
-        """List studios for a given teamspace.
-
-        Args:
-            teamspace: the teamspace to list studios from. Should be specified as {owner}/{name}
-                If not provided, can be selected in an interactive menu.
-            all: if teamspece is not provided, list all studios in all teamspaces.
-            sort_by: the attribute to sort the studios by.
-                Can be one of "name", "teamspace", "status", "machine", "cloud-account".
-
-        """
-        studios = []
-        if all and not teamspace:
-            user = _get_authed_user()
-            possible_teamspaces = self._get_possible_teamspaces(user)
-            for ts in possible_teamspaces.values():
-                teamspace = Teamspace(**ts)
-                studios.extend(teamspace.studios)
-        else:
-            resolved_teamspace = self._resolve_teamspace(teamspace=teamspace)
-            studios = resolved_teamspace.studios
-
-        table = Table(
-            pad_edge=True,
+
+@list_cli.command(name="studios")
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "the teamspace to list studios from. Should be specified as {owner}/{name}"
+        "If not provided, can be selected in an interactive menu."
+    ),
+)
+@click.option(
+    "--all",
+    is_flag=True,
+    flag_value=True,
+    default=False,
+    help="if teamspace is not provided, list all studios in all teamspaces.",
+)
+@click.option(
+    "--sort-by",
+    default=None,
+    type=click.Choice(["name", "teamspace", "status", "machine", "cloud-account"], case_sensitive=False),
+    help="the attribute to sort the studios by.",
+)
+def studios(
+    teamspace: Optional[str] = None,
+    all: bool = False,  # noqa: A002
+    sort_by: Optional[Literal["name", "teamspace", "status", "machine", "cloud-account"]] = None,
+) -> None:
+    """List studios for a given teamspace."""
+    studios = []
+    menu = _TeamspacesMenu()
+    if all and not teamspace:
+        user = _get_authed_user()
+        possible_teamspaces = menu._get_possible_teamspaces(user)
+        for ts in possible_teamspaces.values():
+            teamspace = Teamspace(**ts)
+            studios.extend(teamspace.studios)
+    else:
+        resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+        studios = resolved_teamspace.studios
+
+    table = Table(
+        pad_edge=True,
+    )
+    table.add_column("Name")
+    table.add_column("Teamspace")
+    table.add_column("Status")
+    table.add_column("Machine")
+    table.add_column("Cloud account")
+    for studio in sorted(studios, key=_sort_studios_key(sort_by)):
+        table.add_row(
+            studio.name,
+            f"{studio.teamspace.owner.name}/{studio.teamspace.name}",
+            str(studio.status),
+            str(studio.machine) if studio.machine is not None else None,  # when None the cell is empty
+            str(studio.cloud_account),
         )
-        table.add_column("Name")
-        table.add_column("Teamspace")
-        table.add_column("Status")
-        table.add_column("Machine")
-        table.add_column("Cloud account")
-        for studio in sorted(studios, key=self._sort_studios_key(sort_by)):
+
+    Console().print(table)
+
+
+@list_cli.command(name="jobs")
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "the teamspace to list jobs from. Should be specified as {owner}/{name}"
+        "If not provided, can be selected in an interactive menu."
+    ),
+)
+@click.option(
+    "--all",
+    is_flag=True,
+    flag_value=True,
+    default=False,
+    help="if teamspace is not provided, list all jobs in all teamspaces.",
+)
+@click.option(
+    "--sort-by",
+    "--sort_by",
+    default=None,
+    type=click.Choice(
+        ["name", "teamspace", "status", "studio", "machine", "image", "cloud-account"], case_sensitive=False
+    ),
+    help="the attribute to sort the jobs by.",
+)
+def jobs(
+    teamspace: Optional[str] = None,
+    all: bool = False,  # noqa: A002
+    sort_by: Optional[Literal["name", "teamspace", "status", "studio", "machine", "image", "cloud-account"]] = None,
+) -> None:
+    """List jobs for a given teamspace."""
+    jobs = []
+    menu = _TeamspacesMenu()
+    if all and not teamspace:
+        user = _get_authed_user()
+        possible_teamspaces = menu._get_possible_teamspaces(user)
+        for ts in possible_teamspaces.values():
+            teamspace = Teamspace(**ts)
+            jobs.extend(teamspace.jobs)
+    else:
+        resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+        jobs = resolved_teamspace.jobs
+
+    table = Table(pad_edge=True)
+    table.add_column("Name")
+    table.add_column("Teamspace")
+    table.add_column("Studio")
+    table.add_column("Image")
+    table.add_column("Status")
+    table.add_column("Machine")
+    table.add_column("Total Cost")
+    for j in sorted(jobs, key=_sort_jobs_key(sort_by)):
+        # we know we just fetched these, so no need to refetch
+        j._prevent_refetch_latest = True
+        j._internal_job._prevent_refetch_latest = True
+
+        studio = j.studio
+        with suppress(RuntimeError):
             table.add_row(
-                studio.name,
-                f"{studio.teamspace.owner.name}/{studio.teamspace.name}",
-                str(studio.status),
-                str(studio.machine) if studio.machine is not None else None,  # when None the cell is empty
-                str(studio.cloud_account),
+                j.name,
+                f"{j.teamspace.owner.name}/{j.teamspace.name}",
+                studio.name if studio else None,
+                j.image,
+                str(j.status) if j.status is not None else None,
+                str(j.machine),
+                f"{j.total_cost:.3f}",
             )
 
-        Console().print(table)
-
-    @staticmethod
-    def _sort_jobs_key(sort_by: str) -> Callable[[Job], str]:
-        """Return a key function to sort studios by a given attribute."""
-        sort_key_map = {
-            "name": lambda j: str(j.name or ""),
-            "teamspace": lambda j: str(j.teamspace.name or ""),
-            "status": lambda j: str(j.status or ""),
-            "machine": lambda j: str(j.machine or ""),
-            "studio": lambda j: str(j.studio or ""),
-            "image": lambda j: str(j.image or ""),
-            "cloud-account": lambda j: str(j.cloud_account or ""),
-        }
-        return sort_key_map.get(sort_by, lambda j: j.name)
-
-    def jobs(
-        self,
-        teamspace: Optional[str] = None,
-        all: bool = False,  # noqa: A002
-        sort_by: Optional[Literal["name", "teamspace", "status", "studio", "machine", "image", "cloud-account"]] = None,
-    ) -> None:
-        """List jobs for a given teamspace.
-
-        Args:
-            teamspace: the teamspace to list jobs from. Should be specified as {owner}/{name}
-                If not provided, can be selected in an interactive menu.
-            all: if teamspece is not provided, list all jobs in all teamspaces.
-            sort_by: the attribute to sort the jobs by.
-                Can be one of "name", "teamspace", "status", "studio", "machine", "image", "cloud-account".
-
-        """
-        jobs = []
-        if all and not teamspace:
-            user = _get_authed_user()
-            possible_teamspaces = self._get_possible_teamspaces(user)
-            for ts in possible_teamspaces.values():
-                teamspace = Teamspace(**ts)
-                jobs.extend(teamspace.jobs)
-        else:
-            resolved_teamspace = self._resolve_teamspace(teamspace=teamspace)
-            jobs = resolved_teamspace.jobs
-
-        table = Table(pad_edge=True)
-        table.add_column("Name")
-        table.add_column("Teamspace")
-        table.add_column("Studio")
-        table.add_column("Image")
-        table.add_column("Status")
-        table.add_column("Machine")
-        table.add_column("Total Cost")
-        for j in sorted(jobs, key=self._sort_jobs_key(sort_by)):
-            # we know we just fetched these, so no need to refetch
-            j._prevent_refetch_latest = True
+    Console().print(table)
+
+
+@list_cli.command(name="mmts")
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "the teamspace to list multi-machine jobs from. Should be specified as {owner}/{name}"
+        "If not provided, can be selected in an interactive menu."
+    ),
+)
+@click.option(
+    "--all",
+    is_flag=True,
+    flag_value=True,
+    default=False,
+    help="if teamspace is not provided, list all multi-machine jobs in all teamspaces.",
+)
+@click.option(
+    "--sort-by",
+    "--sort_by",
+    default=None,
+    type=click.Choice(
+        ["name", "teamspace", "studio", "image", "status", "machine", "cloud-account"], case_sensitive=False
+    ),
+    help="the attribute to sort the multi-machine jobs by.",
+)
+def mmts(
+    teamspace: Optional[str] = None,
+    all: bool = False,  # noqa: A002
+    sort_by: Optional[Literal["name", "teamspace", "studio", "image", "status", "machine", "cloud-account"]] = None,
+) -> None:
+    """List multi-machine jobs for a given teamspace."""
+    jobs = []
+    menu = _TeamspacesMenu()
+    if all and not teamspace:
+        user = _get_authed_user()
+        possible_teamspaces = menu._get_possible_teamspaces(user)
+        for ts in possible_teamspaces.values():
+            teamspace = Teamspace(**ts)
+            jobs.extend(teamspace.multi_machine_jobs)
+    else:
+        resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+        jobs = resolved_teamspace.multi_machine_jobs
+
+    table = Table(pad_edge=True)
+    table.add_column("Name")
+    table.add_column("Teamspace")
+    table.add_column("Studio")
+    table.add_column("Image")
+    table.add_column("Status")
+    table.add_column("Machine")
+    table.add_column("Num Machines")
+    table.add_column("Total Cost")
+    for j in sorted(jobs, key=_sort_mmts_key(sort_by)):
+        # we know we just fetched these, so no need to refetch
+        j._prevent_refetch_latest = True
+        with suppress(AttributeError):
             j._internal_job._prevent_refetch_latest = True
 
-            studio = j.studio
-            with suppress(RuntimeError):
-                table.add_row(
-                    j.name,
-                    f"{j.teamspace.owner.name}/{j.teamspace.name}",
-                    studio.name if studio else None,
-                    j.image,
-                    str(j.status) if j.status is not None else None,
-                    str(j.machine),
-                    f"{j.total_cost:.3f}",
-                )
-
-        Console().print(table)
-
-    def _sort_mmts_key(self, sort_by: str) -> Callable[[V1MultiMachineJob], str]:
-        """Return a key function to sort multi-machine jobs by a given attribute."""
-        sort_key_map = {
-            "name": lambda j: str(j.name or ""),
-            "teamspace": lambda j: str(j.teamspace.name or ""),
-            "studio": lambda j: str(j.studio.name or ""),
-            "image": lambda j: str(j.image or ""),
-            "status": lambda j: str(j.status or ""),
-            "machine": lambda j: str(j.machine or ""),
-            "cloud-account": lambda j: str(j.cloud_account or ""),
-        }
-        return sort_key_map.get(sort_by, lambda j: j.name)
-
-    def mmts(
-        self,
-        teamspace: Optional[str] = None,
-        all: bool = False,  # noqa: A002
-        sort_by: Optional[Literal["name", "teamspace", "studio", "image", "status", "machine", "cloud-account"]] = None,
-    ) -> None:
-        """List multi-machine jobs for a given teamspace.
-
-        Args:
-            teamspace: the teamspace to list jobs from. Should be specified as {owner}/{name}
-                If not provided, can be selected in an interactive menu.
-            all: if teamspece is not provided, list all multi-machine jobs in all teamspaces.
-            sort_by: the attribute to sort the multi-machine jobs by.
-                Can be one of "name", "teamspace", "studio", "image", "status", "machine", "cloud-account".
-
-        """
-        jobs = []
-        if all and not teamspace:
-            user = _get_authed_user()
-            possible_teamspaces = self._get_possible_teamspaces(user)
-            for ts in possible_teamspaces.values():
-                teamspace = Teamspace(**ts)
-                jobs.extend(teamspace.multi_machine_jobs)
-        else:
-            resolved_teamspace = self._resolve_teamspace(teamspace=teamspace)
-            jobs = resolved_teamspace.multi_machine_jobs
-
-        table = Table(pad_edge=True)
-        table.add_column("Name")
-        table.add_column("Teamspace")
-        table.add_column("Studio")
-        table.add_column("Image")
-        table.add_column("Status")
-        table.add_column("Machine")
-        table.add_column("Num Machines")
-        table.add_column("Total Cost")
-        for j in sorted(jobs, key=self._sort_mmts_key(sort_by)):
-            # we know we just fetched these, so no need to refetch
-            j._prevent_refetch_latest = True
-            with suppress(AttributeError):
-                j._internal_job._prevent_refetch_latest = True
-
-            studio = j.studio
-            with suppress(RuntimeError):
-                table.add_row(
-                    j.name,
-                    f"{j.teamspace.owner.name}/{j.teamspace.name}",
-                    studio.name if studio else None,
-                    j.image,
-                    str(j.status),
-                    str(j.machine),
-                    str(j.num_machines),
-                    str(j.total_cost),
-                )
-
-        Console().print(table)
-
-    def containers(self, teamspace: Optional[str] = None) -> None:
-        """Display the list of available containers.
-
-        Args:
-            teamspace: The teamspace to list containers 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)
-        result = api.list_containers(teamspace=resolved_teamspace.name, org=resolved_teamspace.owner.name)
-        table = Table(pad_edge=True, box=None)
-        table.add_column("REPOSITORY")
-        table.add_column("IMAGE ID")
-        table.add_column("CREATED")
-        for repo in result:
-            table.add_row(repo["REPOSITORY"], repo["IMAGE ID"], repo["CREATED"])
-        Console().print(table)
-
-    def machines(self) -> None:
-        """Display the list of available machines."""
-        table = Table(pad_edge=True)
-        table.add_column("Name")
-
-        # Get all machine types from the enum
-        machine_types = [name for name in dir(Machine) if not name.startswith("_")]
-
-        # Add rows to table
-        for name in sorted(machine_types):
-            table.add_row(name)
-
-        Console().print(table)
+        studio = j.studio
+        with suppress(RuntimeError):
+            table.add_row(
+                j.name,
+                f"{j.teamspace.owner.name}/{j.teamspace.name}",
+                studio.name if studio else None,
+                j.image,
+                str(j.status),
+                str(j.machine),
+                str(j.num_machines),
+                str(j.total_cost),
+            )
+
+    Console().print(table)
+
+
+@list_cli.command(name="containers")
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "the teamspace to list containers from. Should be specified as {owner}/{name}"
+        "If not provided, can be selected in an interactive menu."
+    ),
+)
+def containers(teamspace: Optional[str] = None) -> None:
+    """Display the list of available containers."""
+    api = LitContainer()
+    menu = _TeamspacesMenu()
+    resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+    result = api.list_containers(teamspace=resolved_teamspace.name, org=resolved_teamspace.owner.name)
+    table = Table(pad_edge=True, box=None)
+    table.add_column("REPOSITORY")
+    table.add_column("IMAGE ID")
+    table.add_column("CREATED")
+    for repo in result:
+        table.add_row(repo["REPOSITORY"], repo["IMAGE ID"], repo["CREATED"])
+    Console().print(table)
+
+
+@list_cli.command(name="machines")
+def machines() -> None:
+    """Display the list of available machines."""
+    table = Table(pad_edge=True)
+    table.add_column("Name")
+
+    # Get all machine types from the enum
+    machine_types = [name for name in dir(Machine) if isinstance(getattr(Machine, name), Machine)]
+
+    # Add rows to table
+    for name in sorted(machine_types):
+        table.add_row(name)
+
+    Console().print(table)
+
+
+def _sort_studios_key(sort_by: str) -> Callable[[Studio], str]:
+    """Return a key function to sort studios by a given attribute."""
+    sort_key_map = {
+        "name": lambda s: str(s.name or ""),
+        "teamspace": lambda s: str(s.teamspace.name or ""),
+        "status": lambda s: str(s.status or ""),
+        "machine": lambda s: str(s.machine or ""),
+        "cloud-account": lambda s: str(s.cloud_account or ""),
+    }
+    return sort_key_map.get(sort_by, lambda s: s.name)
+
+
+def _sort_jobs_key(sort_by: str) -> Callable[[Job], str]:
+    """Return a key function to sort studios by a given attribute."""
+    sort_key_map = {
+        "name": lambda j: str(j.name or ""),
+        "teamspace": lambda j: str(j.teamspace.name or ""),
+        "status": lambda j: str(j.status or ""),
+        "machine": lambda j: str(j.machine or ""),
+        "studio": lambda j: str(j.studio or ""),
+        "image": lambda j: str(j.image or ""),
+        "cloud-account": lambda j: str(j.cloud_account or ""),
+    }
+    return sort_key_map.get(sort_by, lambda j: j.name)
+
+
+def _sort_mmts_key(sort_by: str) -> Callable[[V1MultiMachineJob], str]:
+    """Return a key function to sort multi-machine jobs by a given attribute."""
+    sort_key_map = {
+        "name": lambda j: str(j.name or ""),
+        "teamspace": lambda j: str(j.teamspace.name or ""),
+        "studio": lambda j: str(j.studio.name or ""),
+        "image": lambda j: str(j.image or ""),
+        "status": lambda j: str(j.status or ""),
+        "machine": lambda j: str(j.machine or ""),
+        "cloud-account": lambda j: str(j.cloud_account or ""),
+    }
+    return sort_key_map.get(sort_by, lambda j: j.name)

lightning_sdk/cli/mmts_menu.py

@@ -21,7 +21,7 @@ class _MMTsMenu:
             if j.name == mmt:
                 return j
 
-        Console().print("Could not find Multi-Machine Job {mmt}, please select it from the list:")
+        Console().print(f"Could not find Multi-Machine Job {mmt}, please select it from the list:")
         return self._get_mmt_from_interactive_menu(possible_mmts)
 
     @staticmethod

lightning_sdk/cli/open.py

@@ -0,0 +1,62 @@
+import webbrowser
+from pathlib import Path
+from typing import Optional
+
+import click
+from rich.console import Console
+
+from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.upload import _upload_folder
+from lightning_sdk.studio import Studio
+from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.utils.resolve import _get_studio_url
+
+
+@click.command("open")
+@click.argument("path", default=".", type=click.Path(exists=True))
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "The teamspace to create the Studio in. "
+        "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 open(path: str = ".", teamspace: Optional[str] = None) -> None:  # noqa: A001
+    """Open a local file or folder in a Lightning Studio.
+
+    Example:
+        lightning open PATH
+
+    PATH: the path to the file or folder to open. Defaults to the current directory.
+    """
+    console = Console()
+
+    pathlib_path = Path(path).resolve()
+
+    try:
+        resolved_teamspace = Teamspace()
+    except ValueError:
+        menu = _TeamspacesMenu()
+        resolved_teamspace = menu._resolve_teamspace(teamspace=teamspace)
+
+    new_studio = Studio(name=pathlib_path.stem, teamspace=resolved_teamspace)
+
+    console.print(
+        f"[bold]Uploading {path} to {new_studio.owner.name}/{new_studio.teamspace.name}/{new_studio.name}[/bold]"
+    )
+
+    if pathlib_path.is_dir():
+        _upload_folder(path, remote_path=".", studio=new_studio)
+    else:
+        new_studio.upload_file(path)
+
+    studio_url = _get_studio_url(new_studio, turn_on=True)
+
+    console.line()
+    console.print(f"[bold]Opening {new_studio.owner.name}/{new_studio.teamspace.name}/{new_studio.name}[/bold]")
+
+    ok = webbrowser.open(studio_url)
+    if not ok:
+        console.print(f"Open your Studio at: {studio_url}")