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

lightning_sdk/cli/coloring.py

@@ -0,0 +1,60 @@
+from typing import Optional, Sequence, Tuple
+
+from click.formatting import HelpFormatter, term_len, wrap_text
+from click.termui import style
+
+
+class CustomHelpFormatter(HelpFormatter):
+    _primary = "magenta"
+    _secondary = "cyan"
+
+    def write_heading(self, heading: str) -> None:
+        # Style the heading (command name) in cyan
+        return super().write_heading(style(heading, fg=self._primary, bold=True))
+
+    def write_usage(self, prog: str, args: str = "", prefix: Optional[str] = None) -> None:
+        if prefix is None:
+            prefix = "Usage: "
+
+        styled_prefix = style(prefix, fg=self._primary, bold=True)
+        styled_prog = style(prog, fg=self._secondary, bold=True)
+        styled_args = style(args, fg=self._secondary, bold=True)
+
+        # this is required because otherwise the super() call would include the special characters for coloring in the
+        # width and indent computation leading to wrong results since these characters won't be displayed in the end.
+        usage_prefix = f"{prefix:>{self.current_indent}}{prog} "
+        styled_usage_prefix = f"{styled_prefix:>{self.current_indent}}{styled_prog} "
+        text_width = self.width - self.current_indent
+
+        if text_width >= (term_len(usage_prefix) + 20):
+            # The arguments will fit to the right of the prefix.
+            self.write(styled_usage_prefix)
+            indent = " " * term_len(usage_prefix)
+            self.write(
+                wrap_text(
+                    styled_args,
+                    text_width,
+                    initial_indent="",
+                    subsequent_indent=indent,
+                )
+            )
+        else:
+            # The prefix is too long, put the arguments on the next line.
+            self.write(styled_usage_prefix)
+            self.write("\n")
+            indent = " " * (max(self.current_indent, term_len(prefix)) + 4)
+            self.write(wrap_text(styled_args, text_width, initial_indent=indent, subsequent_indent=indent))
+
+        self.write("\n")
+
+    def write_dl(
+        self,
+        rows: Sequence[Tuple[str, str]],
+        col_max: int = 30,
+        col_spacing: int = 2,
+    ) -> None:
+        new_rows = []
+        for row in rows:
+            new_row = (style(row[0], fg=self._secondary, bold=True), row[1])
+            new_rows.append(new_row)
+        super().write_dl(new_rows, col_max=col_max, col_spacing=col_spacing)

lightning_sdk/cli/configure.py

@@ -12,51 +12,36 @@ 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.",
-# )
+@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."""
     auth = Auth()

lightning_sdk/cli/connect.py

@@ -8,31 +8,18 @@ from lightning_sdk.cli.configure import ssh as configure_ssh
 from lightning_sdk.cli.studios_menu import _StudiosMenu
 
 
-class _Connect:
-    """Connect to lightning products."""
-
-    def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
-        """Connect to a studio via SSH.
-
-        Args:
-            name: The name of the studio to connect to.
-            teamspace: The teamspace the studio is part of. Should be of format <OWNER>/<TEAMSPACE_NAME>.
-        """
-        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>.",
-# )
+@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)

lightning_sdk/cli/create.py

@@ -0,0 +1,83 @@
+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.teamspace_menu import _TeamspacesMenu
+
+_MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
+
+
+@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."
+    ),
+)
+def studio(
+    name: str, teamspace: Optional[str] = None, start: Optional[str] = None, cloud_account: 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._resolve_teamspace(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)
+
+    console.print(f"Created Studio {studio.name}.")
+
+    if start is not None:
+        start_machine = getattr(Machine, start, start)
+        studio.start(start_machine)
+        console.print(f"Started Studio {studio.name} on machine {start}")

lightning_sdk/cli/delete.py

@@ -6,70 +6,27 @@ from rich.console import Console
 from lightning_sdk.cli.exceptions import StudioCliError
 from lightning_sdk.cli.job_and_mmt_action import _JobAndMMTAction
 from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.lit_container import LitContainer
 from lightning_sdk.studio import Studio
 
 
-class _Delete(_JobAndMMTAction, _TeamspacesMenu):
-    """Delete resources on the Lightning AI platform."""
-
-    def container(self, container: str, teamspace: Optional[str] = None) -> None:
-        """Delete a docker container.
-
-        Args:
-            container: The name of the container to delete.
-            teamspace: The teamspace to delete the container from. Should be specified as {owner}/{name}
-                If not provided, can be selected in an interactive menu.
-        """
-        delete_container(container=container, teamspace=teamspace)
-
-    def job(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
-        """Delete a job.
-
-        Args:
-            name: the name of the job. If not specified can be selected interactively.
-            teamspace: 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.
-
-        """
-        job(name=name, teamspace=teamspace)
-
-    def mmt(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
-        """Delete a multi-machine job.
-
-        Args:
-            name: the name of the job. If not specified can be selected interactively.
-            teamspace: 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.
-
-        """
-        mmt(name=name, teamspace=teamspace)
-
-    def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
-        """Delete an existing studio.
-
-        Args:
-            name: The name of the studio to delete.
-                If not specified, tries to infer from the environment (e.g. when run from within a Studio.)
-                Note: This could delete your current studio if run without arguments.
-            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(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."),)
+@delete.command(name="container")
+@click.argument("container")
+@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()
@@ -84,13 +41,27 @@ def delete_container(container: str, teamspace: Optional[str] = None) -> None:
         ) 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."),)
+@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."""
+    """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)
 
@@ -98,13 +69,27 @@ def job(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
     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."),)
+@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."""
+    """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)
 
@@ -112,13 +97,25 @@ def mmt(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
     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."""
+@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:
@@ -129,7 +126,7 @@ def studio(name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
 
     try:
         studio = Studio(name=name, teamspace=teamspace, org=owner, user=None, create_ok=False)
-    except (RuntimeError, ValueError):
+    except (RuntimeError, ValueError, ApiException):
         studio = Studio(name=name, teamspace=teamspace, org=None, user=owner, create_ok=False)
 
     studio.delete()

lightning_sdk/cli/docker.py

@@ -0,0 +1,22 @@
+import click
+
+from lightning_sdk.serve import _LitServeDeployer
+
+
+@click.group(name="dockerize")
+def dockerize() -> None:
+    """Generate a Dockerfile for a LitServe model."""
+
+
+@dockerize.command("api")
+@click.argument("server-filename")
+@click.option("--port", type=int, default=8000, help="Port to expose in the Docker container.")
+@click.option("--gpu", is_flag=True, default=False, flag_value=True, help="Use a GPU-enabled Docker image.")
+@click.option("--tag", default="litserve-model", help="Docker image tag to use in examples.")
+def api(server_filename: str, port: int = 8000, gpu: bool = False, tag: str = "litserve-model") -> None:
+    """Generate a Dockerfile for the given server code."""
+    _api(server_filename=server_filename, port=port, gpu=gpu, tag=tag)
+
+
+def _api(server_filename: str, port: int = 8000, gpu: bool = False, tag: str = "litserve-model") -> str:
+    return _LitServeDeployer().dockerize_api(server_filename=server_filename, port=port, gpu=gpu, tag=tag)