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

lightning_sdk/cli/serve.py

@@ -1,20 +1,26 @@
-import os
 import subprocess
-import warnings
 from pathlib import Path
 from typing import Optional, Union
 
+import click
 import docker
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 from rich.prompt import Confirm
 
+from lightning_sdk.api.lit_container_api import LitContainerApi
+from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.serve import _LitServeDeployer
 
-class _LitServe:
+
+@click.group("serve")
+def serve() -> None:
     """Serve a LitServe model.
 
     Example:
         lightning serve api server.py  # serve locally
+
+    Example:
         lightning serve api server.py --cloud  # deploy to the cloud
 
     You can deploy the API to the cloud by running `lightning serve api server.py --cloud`.
@@ -22,192 +28,123 @@ class _LitServe:
     Deploying to the cloud requires pre-login to the docker registry.
     """
 
-    def api(
-        self,
-        script_path: Union[str, Path],
-        easy: bool = False,
-        cloud: bool = False,
-        repository: Optional[str] = None,
-        non_interactive: bool = False,
-    ) -> None:
-        """Deploy a LitServe model script.
-
-        Args:
-            script_path: Path to the script to serve
-            easy: If True, generates a client for the model
-            cloud: If True, deploy the model to the Lightning Studio
-            repository: Optional Docker repository name (e.g., 'username/model-name')
-            non_interactive: If True, do not prompt for confirmation
-        Raises:
-            FileNotFoundError: If script_path doesn't exist
-            ImportError: If litserve is not installed
-            subprocess.CalledProcessError: If the script fails to run
-            IOError: If client.py generation fails
-        """
-        console = Console()
-        script_path = Path(script_path)
-        if not script_path.exists():
-            raise FileNotFoundError(f"Script not found: {script_path}")
-        if not script_path.is_file():
-            raise ValueError(f"Path is not a file: {script_path}")
-
-        try:
-            from litserve.python_client import client_template
-        except ImportError:
-            raise ImportError(
-                "litserve is not installed. Please install it with `pip install lightning_sdk[serve]`"
-            ) from None
-
-        if easy:
-            client_path = Path("client.py")
-            if client_path.exists():
-                console.print("Skipping client generation: client.py already exists", style="blue")
-            else:
-                try:
-                    client_path.write_text(client_template)
-                    console.print("✅ Client generated at client.py", style="bold green")
-                except OSError as e:
-                    raise OSError(f"Failed to generate client.py: {e!s}") from None
-
-        if cloud:
-            tag = repository if repository else "litserve-model"
-            return self._handle_cloud(script_path, console, tag=tag, non_interactive=non_interactive)
-
-        try:
-            subprocess.run(
-                ["python", str(script_path)],
-                check=True,
-                text=True,
-            )
-        except subprocess.CalledProcessError as e:
-            error_msg = f"Script execution failed with exit code {e.returncode}\nstdout: {e.stdout}\nstderr: {e.stderr}"
-            raise RuntimeError(error_msg) from None
-
-    def _handle_cloud(
-        self,
-        script_path: Union[str, Path],
-        console: Console,
-        tag: str = "litserve-model",
-        non_interactive: bool = False,
-    ) -> None:
-        try:
-            client = docker.from_env()
-            client.ping()
-        except docker.errors.DockerException as e:
-            raise RuntimeError(f"Failed to connect to Docker daemon: {e!s}. Is Docker running?") from None
-
-        dockerizer = _Docker()
-        path = dockerizer.api(script_path, port=8000, gpu=False, tag=tag)
-
-        console.clear()
-        if non_interactive:
-            console.print("[italic]non-interactive[/italic] mode enabled, skipping confirmation prompts", style="blue")
-
-        console.print(f"\nPlease review the Dockerfile at [u]{path}[/u] and make sure it is correct.", style="bold")
-        correct_dockerfile = True if non_interactive else Confirm.ask("Is the Dockerfile correct?", default=True)
-        if not correct_dockerfile:
-            console.print("Please fix the Dockerfile and try again.", style="red")
-            return
-
-        with Progress(
-            SpinnerColumn(),
-            TextColumn("[progress.description]{task.description}"),
-            TimeElapsedColumn(),
-            console=console,
-            transient=False,
-        ) as progress:
-            build_task = progress.add_task("Building Docker image", total=None)
-            build_status = client.api.build(
-                path=os.path.dirname(path), dockerfile=path, tag=tag, decode=True, quiet=False
-            )
-            for line in build_status:
-                if "error" in line:
-                    progress.stop()
-                    console.print(f"\n[red]{line}[/red]")
-                    return
-                if "stream" in line and line["stream"].strip():
-                    console.print(line["stream"].strip(), style="bright_black")
-                    progress.update(build_task, description="Building Docker image")
-
-            progress.update(build_task, description="[green]Build completed![/green]")
-
-            push_task = progress.add_task("Pushing to registry", total=None)
-            console.print("\nPushing image...", style="bold blue")
-            push_status = client.api.push(tag, stream=True, decode=True)
-            for line in push_status:
-                if "error" in line:
-                    progress.stop()
-                    console.print(f"\n[red]{line}[/red]")
-                    return
-                if "status" in line:
-                    console.print(line["status"], style="bright_black")
-                    progress.update(push_task, description="Pushing to registry")
-
-            progress.update(push_task, description="[green]Push completed![/green]")
-
-        console.print(f"\n✅ Image pushed to {tag}", style="bold green")
-        console.print(
-            "Soon you will be able to deploy this model to the Lightning Studio!",
-        )
-        # TODO: Deploy to the cloud
-
-
-class _Docker:
-    """Generate a Dockerfile for a LitServe model."""
-
-    def api(self, server_filename: str, port: int = 8000, gpu: bool = False, tag: str = "litserve-model") -> str:
-        """Generate a Dockerfile for the given server code.
-
-        Args:
-            server_filename: The path to the server file. Example sever.py or app.py.
-            port: The port to expose in the Docker container.
-            gpu: Whether to use a GPU-enabled Docker image.
-            tag: Docker image tag to use in examples.
-        """
-        import litserve as ls
-        from litserve import docker_builder
-
-        requirements = ""
-        if os.path.exists("requirements.txt"):
-            requirements = "-r requirements.txt"
-        else:
-            warnings.warn(
-                f"requirements.txt not found at {os.getcwd()}. "
-                f"Make sure to install the required packages in the Dockerfile.",
-                UserWarning,
-            )
-
-        current_dir = Path.cwd()
-        if not (current_dir / server_filename).is_file():
-            raise FileNotFoundError(f"Server file `{server_filename}` must be in the current directory: {os.getcwd()}")
-
-        version = ls.__version__
-        if gpu:
-            run_cmd = f"docker run --gpus all -p {port}:{port} {tag}:latest"
-            docker_template = docker_builder.CUDA_DOCKER_TEMPLATE
-        else:
-            run_cmd = f"docker run -p {port}:{port} {tag}:latest"
-            docker_template = docker_builder.DOCKERFILE_TEMPLATE
-        dockerfile_content = docker_template.format(
-            server_filename=server_filename,
-            port=port,
-            version=version,
-            requirements=requirements,
-        )
-        with open("Dockerfile", "w") as f:
-            f.write(dockerfile_content)
 
-        success_msg = f"""[bold]Dockerfile created successfully[/bold]
-Update [underline]{os.path.abspath("Dockerfile")}[/underline] to add any additional dependencies or commands.
-
-[bold]Build the container with:[/bold]
-> [underline]docker build -t {tag} .[/underline]
-
-[bold]To run the Docker container on the machine:[/bold]
-> [underline]{run_cmd}[/underline]
-
-[bold]To push the container to a registry:[/bold]
-> [underline]docker push {tag}[/underline]
-"""
-        Console().print(success_msg)
-        return os.path.abspath("Dockerfile")
+@serve.command("api")
+@click.argument("script-path", type=click.Path(exists=True))
+@click.option(
+    "--easy",
+    is_flag=True,
+    default=False,
+    flag_value=True,
+    help="Generate a client for the model",
+)
+@click.option(
+    "--cloud",
+    is_flag=True,
+    default=False,
+    flag_value=True,
+    help="Deploy the model to the Lightning AI platform",
+)
+@click.option("--gpu", is_flag=True, default=False, flag_value=True, help="Use GPU for serving")
+@click.option("--repository", default=None, help="Docker repository name (e.g., 'username/model-name')")
+@click.option(
+    "--non-interactive",
+    "--non_interactive",
+    is_flag=True,
+    default=False,
+    flag_value=True,
+    help="Do not prompt for confirmation",
+)
+def api(
+    script_path: str,
+    easy: bool,
+    cloud: bool,
+    gpu: bool,
+    repository: str,
+    non_interactive: bool,
+) -> None:
+    """Deploy a LitServe model script."""
+    return api_impl(
+        script_path=script_path, easy=easy, cloud=cloud, gpu=gpu, repository=repository, non_interactive=non_interactive
+    )
+
+
+def api_impl(
+    script_path: Union[str, Path],
+    easy: bool = False,
+    cloud: bool = False,
+    gpu: bool = False,
+    repository: Optional[str] = None,
+    non_interactive: bool = False,
+) -> None:
+    """Deploy a LitServe model script."""
+    console = Console()
+    script_path = Path(script_path)
+    if not script_path.exists():
+        raise FileNotFoundError(f"Script not found: {script_path}")
+    if not script_path.is_file():
+        raise ValueError(f"Path is not a file: {script_path}")
+
+    ls_deployer = _LitServeDeployer()
+    ls_deployer.generate_client() if easy else None
+
+    if cloud:
+        tag = repository if repository else "litserve-model"
+        return _handle_cloud(script_path, console, gpu=gpu, tag=tag, non_interactive=non_interactive)
+
+    try:
+        subprocess.run(
+            ["python", str(script_path)],
+            check=True,
+            text=True,
+        )
+    except subprocess.CalledProcessError as e:
+        error_msg = f"Script execution failed with exit code {e.returncode}\nstdout: {e.stdout}\nstderr: {e.stderr}"
+        raise RuntimeError(error_msg) from None
+
+
+def _handle_cloud(
+    script_path: Union[str, Path],
+    console: Console,
+    gpu: bool,
+    repository: str = "litserve-model",
+    tag: Optional[str] = None,
+    teamspace: Optional[str] = None,
+    non_interactive: bool = False,
+) -> None:
+    try:
+        client = docker.from_env()
+        client.ping()
+    except docker.errors.DockerException as e:
+        raise RuntimeError(f"Failed to connect to Docker daemon: {e!s}. Is Docker running?") from None
+
+    ls_deployer = _LitServeDeployer()
+    path = ls_deployer.dockerize_api(script_path, port=8000, gpu=gpu, tag=tag)
+    console.clear()
+    if non_interactive:
+        console.print("[italic]non-interactive[/italic] mode enabled, skipping confirmation prompts", style="blue")
+
+    console.print(f"\nPlease review the Dockerfile at [u]{path}[/u] and make sure it is correct.", style="bold")
+    correct_dockerfile = True if non_interactive else Confirm.ask("Is the Dockerfile correct?", default=True)
+    if not correct_dockerfile:
+        console.print("Please fix the Dockerfile and try again.", style="red")
+        return
+
+    tag = tag if tag else "latest"
+
+    lit_cr = LitContainerApi()
+    menu = _TeamspacesMenu()
+    teamspace = menu._resolve_teamspace(teamspace)
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        TimeElapsedColumn(),
+        console=console,
+        transient=False,
+    ) as progress:
+        ls_deployer._build_container(path, repository, tag, console, progress)
+        ls_deployer._push_container(repository, tag, teamspace, lit_cr, progress)
+    console.print(f"\n✅ Image pushed to {tag}", style="bold green")
+    console.print(
+        "Soon you will be able to deploy this model to the Lightning Studio!",
+    )

lightning_sdk/cli/start.py

@@ -1,43 +1,62 @@
 from typing import Optional
 
+import click
+
 from lightning_sdk import Machine, Studio
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
+
+_MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
 
 
-class _Start:
+@click.group("start")
+def start() -> None:
     """Start resources on the Lightning AI platform."""
 
-    def __init__(self) -> None:
-        _machine_values = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
-
-        docstr_studio = f"""Start a studio on a given machine.
-
-        Args:
-            name: The name of the studio to start.
-                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.)
-            machine: The machine type to start the studio on. One of {", ".join(_machine_values)}.
-                Defaults to the CPU Machine.
-        """
-        self.studio.__func__.__doc__ = docstr_studio
-
-    def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None, machine: str = "CPU") -> None:
-        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)
-
-        try:
-            resolved_machine = getattr(Machine, machine.upper(), Machine(machine, machine))
-        except KeyError:
-            resolved_machine = machine
-
-        studio.start(resolved_machine)
+
+@start.command("studio")
+@click.argument(
+    "name",
+)
+@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(
+    "--machine",
+    default="CPU",
+    show_default=True,
+    type=click.Choice(_MACHINE_VALUES),
+    help="The machine type to start the studio on.",
+)
+def studio(name: str, teamspace: Optional[str] = None, machine: str = "CPU") -> None:
+    """Start a studio on a given machine.
+
+    Example:
+      lightning start studio NAME
+
+    NAME: the name of the studio to start
+    """
+    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)
+
+    try:
+        resolved_machine = getattr(Machine, machine.upper(), Machine(machine, machine))
+    except KeyError:
+        resolved_machine = machine
+
+    studio.start(resolved_machine)

lightning_sdk/cli/stop.py

@@ -1,65 +1,107 @@
 from typing import Optional
 
+import click
 from rich.console import Console
 
 from lightning_sdk.cli.job_and_mmt_action import _JobAndMMTAction
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 from lightning_sdk.studio import Studio
 
 
-class _Stop(_JobAndMMTAction):
+@click.group("stop")
+def stop() -> None:
     """Stop resources on the Lightning AI platform."""
 
-    def job(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
-        """Stop 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 = super().job(name=name, teamspace=teamspace)
-
-        job.stop()
-        Console().print(f"Successfully stopped {job.name}!")
-
-    def mmt(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
-        """Stop 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 = super().mmt(name=name, teamspace=teamspace)
-
-        mmt.stop()
-        Console().print(f"Successfully stopped {mmt.name}!")
-
-    def studio(self, name: Optional[str] = None, teamspace: Optional[str] = None) -> None:
-        """Stop a running studio.
-
-        Args:
-            name: The name of the studio to stop.
-                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.)
-        """
-        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.stop()
-        Console().print("Studio successfully stopped")
+
+@stop.command("job")
+@click.argument(
+    "name",
+)
+@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: str, teamspace: Optional[str] = None) -> None:
+    """Stop a job.
+
+    Example:
+      lightning stop job NAME
+
+    NAME: the name of the job to stop.
+    """
+    menu = _JobAndMMTAction()
+    job = menu.job(name=name, teamspace=teamspace)
+
+    job.stop()
+    Console().print(f"Successfully stopped {job.name}!")
+
+
+@stop.command("mmt")
+@click.argument(
+    "name",
+)
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "the name of the teamspace the multi-machine 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: str, teamspace: Optional[str] = None) -> None:
+    """Stop a multi-machine job.
+
+    Example:
+      lightning stop mmt NAME
+
+    NAME: the name of the multi-machine job to stop.
+    """
+    menu = _JobAndMMTAction()
+    mmt = menu.mmt(name=name, teamspace=teamspace)
+
+    mmt.stop()
+    Console().print(f"Successfully stopped {mmt.name}!")
+
+
+@stop.command("studio")
+@click.argument(
+    "name",
+)
+@click.option(
+    "--teamspace",
+    default=None,
+    help=(
+        "the name of the teamspace the studio lives in. "
+        "Should be specified as {teamspace_owner}/{teamspace_name} (e.g my-org/my-teamspace). "
+        "If not specified can be selected interactively."
+    ),
+)
+def studio(name: str, teamspace: Optional[str] = None) -> None:
+    """Stop a running studio.
+
+    Example:
+      lightning stop studio NAME
+
+    NAME: the name of the studio to stop.
+    """
+    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.stop()
+    Console().print("Studio successfully stopped")