lightning-sdk 0.2.11__py3-none-any.whl → 0.2.13__py3-none-any.whl

lightning_sdk/__init__.py

@@ -31,6 +31,6 @@ __all__ = [
     "User",
 ]
 
-__version__ = "0.2.11"
+__version__ = "0.2.13"
 _check_version_and_prompt_upgrade(__version__)
 _set_tqdm_envvars_noninteractive()

lightning_sdk/api/deployment_api.py

@@ -460,7 +460,7 @@ def to_autoscaling(
     if target_metrics is None and (threshold < 0 or threshold > 100):
         raise ValueError("The autoscaling threshold should be defined between 0 and 100.")
 
-    if target_metrics is not None and len(target_metrics) == 0:
+    if target_metrics is not None and len(target_metrics) == 0 and metric is None:
         raise ValueError("The target_metrics must be provided.")
 
     if target_metrics is not None:
@@ -524,8 +524,14 @@ def to_endpoint(
 def to_health_check(
     health_check: Optional[Union[HttpHealthCheck, ExecHealthCheck]] = None
 ) -> Optional[V1JobHealthCheckConfig]:
+    # Use Default health check if none is provided
     if not health_check:
-        return None
+        return V1JobHealthCheckConfig(
+            failure_threshold=600,
+            initial_delay_seconds=0,
+            interval_seconds=1,
+            timeout_seconds=600,
+        )
 
     health_check_config = V1JobHealthCheckConfig(
         failure_threshold=health_check.failure_threshold,
@@ -555,6 +561,7 @@ def to_spec(
     health_check: Optional[Union[HttpHealthCheck, ExecHealthCheck]] = None,
     quantity: Optional[int] = None,
     include_credentials: Optional[bool] = None,
+    cloudspace_id: Optional[None] = None,
 ) -> V1JobSpec:
     if cloud_account is None:
         raise ValueError("The cloud account should be defined.")
@@ -562,9 +569,15 @@ def to_spec(
     if machine is None:
         raise ValueError("The machine should be defined.")
 
-    if image is None:
+    if image is None and cloudspace_id is None:
         raise ValueError("The image should be defined.")
 
+    if entrypoint is not None and cloudspace_id is not None:
+        raise ValueError("The entrypoint shouldn't be defined when a Studio is provided.")
+
+    if command is None and cloudspace_id is not None:
+        raise ValueError("The command should be defined.")
+
     return V1JobSpec(
         cluster_id=cloud_account,
         command=command,
@@ -576,6 +589,7 @@ def to_spec(
         readiness_probe=to_health_check(health_check),
         quantity=quantity,
         include_credentials=include_credentials,
+        cloudspace_id=cloudspace_id,
     )
 
 
@@ -600,3 +614,21 @@ def apply_change(spec: Any, key: str, value: Any) -> bool:
         return True
 
     return False
+
+
+def compose_commands(commands: List[str]) -> str:
+    composite_command = []
+
+    for command in commands:
+        command = command.strip()
+
+        # Check if the command already has '&'
+        if command.endswith("&"):
+            # It's a background command, add it as a subshell without further adjustment
+            composite_command.append(f"( {command} )")
+        else:
+            # Sequential execution, add as-is and use `&&` to connect if followed by another command
+            composite_command.append(command)
+
+    # Joining commands, using `&&` between sequential parts and respecting subshell backgrounds
+    return " && ".join(composite_command)

lightning_sdk/api/lit_container_api.py

@@ -125,6 +125,17 @@ class LitContainerApi:
         except Exception as e:
             raise ValueError(f"Could not delete container {container} from project {project_id}: {e!s}") from e
 
+    def get_container_url(
+        self, repository: str, tag: str, teamspace: Teamspace, cloud_account: Optional[str] = None
+    ) -> str:
+        """Docker container will be pushed to the URL returned from this function."""
+        registry_url = _get_registry_url()
+        container_basename = repository.split("/")[-1]
+        return (
+            f"{registry_url}/lit-container{f'-{cloud_account}' if cloud_account is not None else ''}/"
+            f"{teamspace.owner.name}/{teamspace.name}/{container_basename}"
+        )
+
     @retry_on_lcr_auth_failure
     def upload_container(
         self,
@@ -147,7 +158,6 @@ class LitContainerApi:
             Named cloud-account in the CLI options.
         :param platform: If empty will be linux/amd64. This is important because our entire deployment infra runs on
             linux/amd64. Will show user a warning otherwise.
-        :return_final_dict: Controls whether we respond with the dictionary containing metadata about container upload
         :return: Generator[dict, None, dict]
         """
         try:
@@ -163,18 +173,14 @@ class LitContainerApi:
             except Exception as e:
                 raise ValueError(f"Unable to upload {container}:{tag}") from e
 
-        registry_url = _get_registry_url()
-        container_basename = container.split("/")[-1]
-        repository = (
-            f"{registry_url}/lit-container{f'-{cloud_account}' if cloud_account is not None else ''}/"
-            f"{teamspace.owner.name}/{teamspace.name}/{container_basename}"
-        )
+        repository = self.get_container_url(container, tag, teamspace, cloud_account)
         tagged = self._docker_client.api.tag(f"{container}:{tag}", repository, tag)
         if not tagged:
             raise ValueError(f"Could not tag container {container}:{tag} with {repository}:{tag}")
         yield from self._push_with_retry(repository, tag=tag)
 
         if return_final_dict:
+            container_basename = repository.split("/")[-1]
             yield {
                 "finish": True,
                 "url": f"{LIGHTNING_CLOUD_URL}/{teamspace.owner.name}/{teamspace.name}/containers/"

lightning_sdk/api/llm_api.py

@@ -0,0 +1,46 @@
+from typing import List, Optional
+
+from lightning_sdk.lightning_cloud.openapi.models.v1_conversation_response_chunk import V1ConversationResponseChunk
+from lightning_sdk.lightning_cloud.rest_client import LightningClient
+
+
+class LLMApi:
+    def __init__(self) -> None:
+        self._client = LightningClient(retry=False, max_tries=0)
+
+    def get_public_models(self) -> List[str]:
+        result = self._client.assistants_service_list_assistants(published=True)
+        return result.assistants
+
+    def get_org_models(self, org_id: str) -> List[str]:
+        result = self._client.assistants_service_list_assistants(org_id=org_id)
+        return result.assistants
+
+    def get_user_models(self, user_id: str) -> List[str]:
+        result = self._client.assistants_service_list_assistants(user_id=user_id)
+        return result.assistants
+
+    def start_conversation(
+        self,
+        prompt: str,
+        system_prompt: Optional[str],
+        max_completion_tokens: Optional[int],
+        assistant_id: str,
+        conversation_id: Optional[str],
+    ) -> V1ConversationResponseChunk:
+        body = {
+            "message": {
+                "author": {"role": "user"},
+                "content": [
+                    {
+                        "contentType": "text",
+                        "parts": [prompt],
+                    }
+                ],
+            },
+            "max_completion_tokens": max_completion_tokens,
+        }
+        if conversation_id:
+            body["conversation_id"] = conversation_id
+        result = self._client.assistants_service_start_conversation(body, assistant_id)
+        return result.result

lightning_sdk/api/studio_api.py

@@ -36,13 +36,16 @@ from lightning_sdk.lightning_cloud.openapi import (
     V1CloudSpaceInstanceConfig,
     V1CloudSpaceSeedFile,
     V1CloudSpaceState,
+    V1EndpointType,
     V1GetCloudSpaceInstanceStatusResponse,
     V1GetLongRunningCommandInCloudSpaceResponse,
     V1LoginRequest,
     V1Plugin,
     V1PluginsListResponse,
+    V1UpstreamCloudSpace,
     V1UserRequestedComputeConfig,
 )
+from lightning_sdk.lightning_cloud.openapi.models import ProjectIdEndpointsBody
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.machine import Machine
 
@@ -668,6 +671,20 @@ class StudioApi:
             interruptible=interruptible,
         )
 
+    def start_new_port(self, teamspace_id: str, studio_id: str, name: str, port: int, auto_start: bool = False) -> str:
+        """Starts a new port to the given Studio."""
+        endpoint = self._client.endpoint_service_create_endpoint(
+            project_id=teamspace_id,
+            body=ProjectIdEndpointsBody(
+                name=name,
+                ports=[str(port)],
+                cloudspace=V1UpstreamCloudSpace(
+                    cloudspace_id=studio_id, port=str(port), type=V1EndpointType.PLUGIN_PORT, auto_start=auto_start
+                ),
+            ),
+        )
+        return endpoint.urls[0]
+
     def _create_app(
         self, studio_id: str, teamspace_id: str, cloud_account: str, plugin_type: str, **other_arguments: Any
     ) -> Externalv1LightningappInstance:

lightning_sdk/cli/entrypoint.py

@@ -14,7 +14,7 @@ from lightning_sdk.cli.configure import configure
 from lightning_sdk.cli.connect import connect
 from lightning_sdk.cli.create import create
 from lightning_sdk.cli.delete import delete
-from lightning_sdk.cli.docker import dockerize
+from lightning_sdk.cli.docker_cli import dockerize
 from lightning_sdk.cli.download import download
 from lightning_sdk.cli.generate import generate
 from lightning_sdk.cli.inspect import inspect

lightning_sdk/cli/serve.py

@@ -1,3 +1,4 @@
+import concurrent.futures
 import os
 import socket
 import subprocess
@@ -6,7 +7,8 @@ import webbrowser
 from datetime import datetime
 from enum import Enum
 from pathlib import Path
-from typing import List, Optional, TypedDict, Union
+from threading import Thread
+from typing import Dict, List, Optional, TypedDict, Union
 from urllib.parse import urlencode
 
 import click
@@ -17,16 +19,24 @@ from rich.prompt import Confirm
 from lightning_sdk import Machine, Teamspace
 from lightning_sdk.api import UserApi
 from lightning_sdk.api.lit_container_api import LitContainerApi
+from lightning_sdk.api.utils import _get_registry_url
 from lightning_sdk.cli.teamspace_menu import _TeamspacesMenu
+from lightning_sdk.cli.upload import (
+    _dump_current_upload_state,
+    _resolve_previous_upload_state,
+    _start_parallel_upload,
+)
 from lightning_sdk.lightning_cloud import env
 from lightning_sdk.lightning_cloud.login import Auth, AuthServer
 from lightning_sdk.lightning_cloud.openapi import V1CloudSpace
 from lightning_sdk.lightning_cloud.rest_client import LightningClient
 from lightning_sdk.serve import _LitServeDeployer
-from lightning_sdk.utils.resolve import _get_authed_user, _resolve_teamspace
+from lightning_sdk.studio import Studio
+from lightning_sdk.utils.resolve import _get_authed_user, _get_studio_url, _resolve_teamspace
 
 _MACHINE_VALUES = tuple([machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine)])
 _POLL_TIMEOUT = 600
+LITSERVE_CODE = os.environ.get("LITSERVE_CODE", "j39bzk903h")
 
 
 class _ServeGroup(click.Group):
@@ -83,7 +93,14 @@ def deploy() -> None:
     default="CPU",
     show_default=True,
     type=click.Choice(_MACHINE_VALUES),
-    help="The machine type to deploy the API on.",
+    help="Machine type to deploy the API on. Defaults to CPU.",
+)
+@click.option(
+    "--devbox",
+    default=None,
+    show_default=True,
+    type=click.Choice(_MACHINE_VALUES),
+    help="Machine type to build the API on. Setting this argument will open the server in a Studio.",
 )
 @click.option(
     "--interruptible",
@@ -131,7 +148,8 @@ def api(
     local: bool,
     name: Optional[str],
     non_interactive: bool,
-    machine: str,
+    machine: Optional[str],
+    devbox: Optional[str],
     interruptible: bool,
     teamspace: Optional[str],
     org: Optional[str],
@@ -148,9 +166,10 @@ def api(
         script_path=script_path,
         easy=easy,
         local=local,
-        repository=name,
+        name=name,
         non_interactive=non_interactive,
         machine=machine,
+        devbox=devbox,
         interruptible=interruptible,
         teamspace=teamspace,
         org=org,
@@ -168,10 +187,11 @@ def api_impl(
     script_path: Union[str, Path],
     easy: bool = False,
     local: bool = False,
-    repository: [str] = None,
+    name: Optional[str] = None,
     tag: Optional[str] = None,
     non_interactive: bool = False,
     machine: str = "CPU",
+    devbox: Optional[str] = None,
     interruptible: bool = False,
     teamspace: Optional[str] = None,
     org: Optional[str] = None,
@@ -193,47 +213,50 @@ def api_impl(
 
     _LitServeDeployer.generate_client() if easy else None
 
-    if not repository:
+    if not name:
         timestr = datetime.now().strftime("%b-%d-%H_%M")
-        repository = f"litserve-{timestr}".lower()
-
-    if not local:
-        repository = repository or "litserve-model"
-        machine = Machine.from_str(machine)
-        return _handle_cloud(
-            script_path,
-            console,
-            repository=repository,
-            tag=tag,
-            non_interactive=non_interactive,
-            machine=machine,
-            interruptible=interruptible,
-            teamspace=teamspace,
-            org=org,
-            user=user,
-            cloud_account=cloud_account,
-            port=port,
-            min_replica=min_replica,
-            max_replica=max_replica,
-            replicas=replicas,
-            include_credentials=include_credentials,
-        )
+        name = f"litserve-{timestr}".lower()
 
-    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
+    if local:
+        try:
+            subprocess.run(
+                ["python", str(script_path)],
+                check=True,
+                text=True,
+            )
+            return None
+        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
+
+    if devbox:
+        return _handle_devbox(name, script_path, console, non_interactive, devbox, interruptible, teamspace, org, user)
+
+    machine = Machine.from_str(machine)
+    return _handle_cloud(
+        script_path,
+        console,
+        repository=name,
+        tag=tag,
+        non_interactive=non_interactive,
+        machine=machine,
+        interruptible=interruptible,
+        teamspace=teamspace,
+        org=org,
+        user=user,
+        cloud_account=cloud_account,
+        port=port,
+        min_replica=min_replica,
+        max_replica=max_replica,
+        replicas=replicas,
+        include_credentials=include_credentials,
+    )
 
 
 class _AuthServer(AuthServer):
     def get_auth_url(self, port: int) -> str:
         redirect_uri = f"http://localhost:{port}/login-complete"
-        params = urlencode({"redirectTo": redirect_uri, "inviteCode": "litserve"})
+        params = urlencode({"redirectTo": redirect_uri, "okbhrt": LITSERVE_CODE})
         return f"{env.LIGHTNING_CLOUD_URL}/sign-in?{params}"
 
 
@@ -398,6 +421,128 @@ def is_connected(host: str = "8.8.8.8", port: int = 53, timeout: int = 10) -> bo
         return False
 
 
+def _upload_container(
+    console: Console,
+    ls_deployer: _LitServeDeployer,
+    repository: str,
+    tag: str,
+    resolved_teamspace: Teamspace,
+    lit_cr: LitContainerApi,
+    cloud_account: Optional[str],
+) -> bool:
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        TimeElapsedColumn(),
+        console=console,
+        transient=True,
+    ) as progress:
+        try:
+            push_task = progress.add_task("Uploading container to Lightning registry", total=None)
+            for line in ls_deployer.push_container(
+                repository, tag, resolved_teamspace, lit_cr, cloud_account=cloud_account
+            ):
+                progress.update(push_task, advance=1)
+                if not ("Pushing" in line["status"] or "Waiting" in line["status"]):
+                    console.print(line["status"])
+            progress.update(push_task, description="[green]Push completed![/green]")
+        except Exception as e:
+            console.print(f"❌ Deployment failed: {e}", style="red")
+            return False
+    console.print(f"\n✅ Image pushed to {repository}:{tag}")
+    return True
+
+
+# TODO: Move the rest of the devbox logic here
+class _LitServeDevbox:
+    """Build LitServe API in a Studio."""
+
+    def resolve_previous_upload(self, studio: Studio, folder: str) -> Dict[str, str]:
+        remote_path = "."
+        pairs = {}
+        for root, _, files in os.walk(folder):
+            rel_root = os.path.relpath(root, folder)
+            for f in files:
+                pairs[os.path.join(root, f)] = os.path.join(remote_path, rel_root, f)
+        return _resolve_previous_upload_state(studio, remote_path, pairs)
+
+    def upload_folder(self, studio: Studio, folder: str, upload_state: Dict[str, str]) -> None:
+        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
+            futures = _start_parallel_upload(executor, studio, upload_state)
+            total_files = len(upload_state)
+
+            with Progress(
+                SpinnerColumn(),
+                TextColumn("[progress.description]{task.description}"),
+                TimeElapsedColumn(),
+                console=Console(),
+                transient=True,
+            ) as progress:
+                upload_task = progress.add_task(f"[cyan]Uploading {total_files} files to Studio...", total=total_files)
+                for f in concurrent.futures.as_completed(futures):
+                    upload_state.pop(f.result())
+                    _dump_current_upload_state(studio, ".", upload_state)
+                    progress.update(upload_task, advance=1)
+
+
+def _handle_devbox(
+    name: str,
+    script_path: Path,
+    console: Console,
+    non_interactive: bool = False,
+    devbox: Machine = "CPU",
+    interruptible: bool = False,
+    teamspace: Optional[str] = None,
+    org: Optional[str] = None,
+    user: Optional[str] = None,
+) -> None:
+    if script_path.suffix != ".py":
+        console.print("❌ Error: Only Python files (.py) are supported for development servers", style="red")
+        return
+
+    resolved_teamspace = select_teamspace(teamspace, org, user)
+    studio = Studio(name=name, teamspace=resolved_teamspace)
+    lit_devbox = _LitServeDevbox()
+
+    studio_url = _get_studio_url(studio, turn_on=True)
+    pathlib_path = Path(script_path).resolve()
+    ok = False
+    studio_path = f"{studio.owner.name}/{studio.teamspace.name}/{studio.name}"
+
+    console.print("\n=== Lightning Studio Setup ===")
+    console.print(f"🔧 [bold]Setting up Studio:[/bold] {studio_path}")
+    console.print(f"📁 [bold]Local project:[/bold] {pathlib_path.parent}")
+
+    upload_state = lit_devbox.resolve_previous_upload(studio, pathlib_path.parent)
+    if non_interactive:
+        console.print(f"🌐 [bold]Opening Studio:[/bold] [link={studio_url}]{studio_url}[/link]")
+        ok = webbrowser.open(studio_url)
+    else:
+        if Confirm.ask("Would you like to open your Studio in the browser?", default=True):
+            console.print(f"🌐 [bold]Opening Studio:[/bold] [link={studio_url}]{studio_url}[/link]")
+            ok = webbrowser.open(studio_url)
+
+    if not ok:
+        console.print(f"🔗 [bold]Access Studio:[/bold] [link={studio_url}]{studio_url}[/link]")
+
+    console.print("\n⚡ Initializing Studio (this typically takes 1-2 minutes)...")
+    studio.start(machine=devbox, interruptible=interruptible)
+    studio.install_plugin("custom-port")
+    console.print("🔌 Configuring server port...")
+    studio.run_plugin("custom-port", port=8000)  # TODO: Remove hardcoded port and fetch from LitServe
+
+    console.print("📤 Syncing project files to Studio...")
+    lit_devbox.upload_folder(studio, pathlib_path.parent, upload_state)
+
+    # Add completion message with next steps
+    console.print("\n✅ Studio ready!")
+    console.print("\n📋 [bold]Next steps:[/bold]")
+    console.print("  [bold]1.[/bold] Server code will be available in the Studio")
+    console.print("  [bold]2.[/bold] The Studio is now running with the specified configuration")
+    console.print("  [bold]3.[/bold] Modify and run your server directly in the Studio")
+    # TODO: Once server running is implemented
+
+
 def _handle_cloud(
     script_path: Union[str, Path],
     console: Console,
@@ -480,29 +625,43 @@ def _handle_cloud(
     lit_cr = LitContainerApi()
     lit_cr.list_containers(resolved_teamspace.id, cloud_account=cloud_account)
 
-    with Progress(
-        SpinnerColumn(),
-        TextColumn("[progress.description]{task.description}"),
-        TimeElapsedColumn(),
-        console=console,
-        transient=True,
-    ) as progress:
-        try:
-            push_task = progress.add_task("Pushing to registry", total=None)
-            push_status = {}
-            for line in ls_deployer.push_container(
-                repository, tag, resolved_teamspace, lit_cr, cloud_account=cloud_account
-            ):
-                push_status = line
-                progress.update(push_task, advance=1)
-                if not ("Pushing" in line["status"] or "Waiting" in line["status"]):
-                    console.print(line["status"])
-            progress.update(push_task, description="[green]Push completed![/green]")
-        except Exception as e:
-            console.print(f"❌ Deployment failed: {e}", style="red")
+    registry_url = _get_registry_url()
+    container_basename = repository.split("/")[-1]
+    image = (
+        f"{registry_url}/lit-container{f'-{cloud_account}' if cloud_account is not None else ''}/"
+        f"{resolved_teamspace.owner.name}/{resolved_teamspace.name}/{container_basename}"
+    )
+
+    if from_onboarding:
+        thread = Thread(
+            target=ls_deployer.run_on_cloud,
+            kwargs={
+                "deployment_name": deployment_name,
+                "image": image,
+                "teamspace": resolved_teamspace,
+                "metric": None,
+                "machine": machine,
+                "spot": interruptible,
+                "cloud_account": cloud_account,
+                "port": port,
+                "min_replica": min_replica,
+                "max_replica": max_replica,
+                "replicas": replicas,
+                "include_credentials": include_credentials,
+                "cloudspace_id": cloudspace_id,
+                "from_onboarding": from_onboarding,
+            },
+        )
+        thread.start()
+        console.print("🚀 Deployment started")
+        if not _upload_container(console, ls_deployer, repository, tag, resolved_teamspace, lit_cr, cloud_account):
+            thread.join()
             return
-    console.print(f"\n✅ Image pushed to {repository}:{tag}")
-    image = push_status.get("image")
+        thread.join()
+        return
+
+    if not _upload_container(console, ls_deployer, repository, tag, resolved_teamspace, lit_cr, cloud_account):
+        return
 
     deployment_status = ls_deployer.run_on_cloud(
         deployment_name=deployment_name,