llama-deploy-appserver 0.2.7a1__py3-none-any.whl → 0.3.0__py3-none-any.whl

llama_deploy/appserver/interrupts.py

@@ -0,0 +1,55 @@
+import asyncio
+import signal
+from asyncio import Event
+from contextlib import suppress
+from typing import Awaitable, TypeVar
+
+shutdown_event = Event()
+
+
+def setup_interrupts() -> None:
+    loop = asyncio.get_running_loop()
+    for sig in (signal.SIGINT, signal.SIGTERM):
+        loop.add_signal_handler(sig, shutdown_event.set)
+
+
+class OperationAborted(Exception):
+    """Raised when an operation is aborted due to shutdown/interrupt."""
+
+
+T = TypeVar("T")
+
+
+async def wait_or_abort(
+    awaitable: Awaitable[T], shutdown_event: asyncio.Event = shutdown_event
+) -> T:
+    """Await an operation, aborting early if shutdown is requested.
+
+    If the shutdown event is set before the awaitable completes, cancel the
+    awaitable and raise OperationAborted. Otherwise, return the awaitable's result.
+    """
+    event = shutdown_event
+    if event.is_set():
+        raise OperationAborted()
+
+    op_task = asyncio.create_task(awaitable)
+    stop_task = asyncio.create_task(event.wait())
+    try:
+        done, _ = await asyncio.wait(
+            {op_task, stop_task}, return_when=asyncio.FIRST_COMPLETED
+        )
+        if stop_task in done:
+            op_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await op_task
+            raise OperationAborted()
+        # Operation finished first
+        stop_task.cancel()
+        with suppress(asyncio.CancelledError):
+            await stop_task
+        return await op_task
+    finally:
+        # Ensure no leaked tasks if an exception propagates
+        for t in (op_task, stop_task):
+            if not t.done():
+                t.cancel()

llama_deploy/appserver/process_utils.py

@@ -0,0 +1,214 @@
+import functools
+import os
+import platform
+import pty
+import subprocess
+import sys
+import threading
+from typing import Callable, TextIO, cast
+
+
+def run_process(
+    cmd: list[str],
+    *,
+    cwd: os.PathLike | None = None,
+    env: dict[str, str] | None = None,
+    prefix: str | None = None,
+    color_code: str = "36",
+    line_transform: Callable[[str], str | None] | None = None,
+    use_tty: bool | None = None,
+) -> None:
+    """Run a process and stream its output with optional TTY semantics.
+
+    If use_tty is None, a PTY will be used only when the parent's stdout is a TTY
+    and the platform supports PTYs. When a PTY is used, stdout/stderr are merged.
+    """
+    use_pty = _should_use_pty(use_tty)
+    prefixer = _make_prefixer(prefix, color_code, line_transform)
+
+    process, sources, cleanup = _spawn_process(cmd, cwd=cwd, env=env, use_pty=use_pty)
+    threads: list[threading.Thread] = []
+    try:
+        cleanup()
+        _log_command(cmd, prefixer)
+        threads = _start_stream_threads(sources, prefixer)
+        ret = process.wait()
+        if ret != 0:
+            raise subprocess.CalledProcessError(ret, cmd)
+    finally:
+        for t in threads:
+            t.join()
+
+
+def spawn_process(
+    cmd: list[str],
+    *,
+    cwd: os.PathLike | None = None,
+    env: dict[str, str] | None = None,
+    prefix: str | None = None,
+    color_code: str = "36",
+    line_transform: Callable[[str], str | None] | None = None,
+    use_tty: bool | None = None,
+) -> subprocess.Popen:
+    """Spawn a process and stream its output in background threads.
+
+    Returns immediately with the Popen object. Streaming threads are daemons.
+    """
+    use_pty = _should_use_pty(use_tty)
+    prefixer = _make_prefixer(prefix, color_code, line_transform)
+
+    process, sources, cleanup = _spawn_process(cmd, cwd=cwd, env=env, use_pty=use_pty)
+    cleanup()
+    _log_command(cmd, prefixer)
+    _start_stream_threads(sources, prefixer)
+    return process
+
+
+@functools.cache
+def _use_color() -> bool:
+    """Return True if ANSI colors should be emitted to stdout.
+
+    Respects common environment variables and falls back to TTY detection.
+    """
+    force_color = os.environ.get("FORCE_COLOR")
+
+    return sys.stdout.isatty() or force_color is not None and force_color != "0"
+
+
+def _colored_prefix(prefix: str, color_code: str) -> str:
+    return f"\x1b[{color_code}m{prefix}\x1b[0m " if _use_color() else f"{prefix} "
+
+
+def _make_prefixer(
+    prefix: str | None,
+    color_code: str,
+    line_transform: Callable[[str], str | None] | None = None,
+) -> Callable[[str], str | None]:
+    colored = _colored_prefix(prefix, color_code) if prefix else ""
+
+    def _prefixer(line: str) -> str | None:
+        transformed = line_transform(line) if line_transform else line
+        if transformed is None:
+            return None
+        return f"{colored}{transformed}"
+
+    return _prefixer
+
+
+# Unified PTY/Pipe strategy helpers
+
+
+def _should_use_pty(use_tty: bool | None) -> bool:
+    if platform.system() == "Windows":
+        return False
+    if use_tty is None:
+        return sys.stdout.isatty()
+    return use_tty and sys.stdout.isatty() and not os.environ.get("NO_COLOR")
+
+
+def should_use_color() -> bool:
+    return _should_use_pty(None)
+
+
+def _spawn_process(
+    cmd: list[str],
+    *,
+    cwd: os.PathLike | None,
+    env: dict[str, str] | None,
+    use_pty: bool,
+) -> tuple[subprocess.Popen, list[tuple[int | TextIO, TextIO]], Callable[[], None]]:
+    if use_pty:
+        master_fd, slave_fd = pty.openpty()
+        process = subprocess.Popen(
+            cmd,
+            env=env,
+            cwd=cwd,
+            stdin=slave_fd,
+            stdout=slave_fd,
+            stderr=slave_fd,
+            close_fds=True,
+        )
+
+        def cleanup() -> None:
+            try:
+                os.close(slave_fd)
+            except OSError:
+                pass
+
+        sources: list[tuple[int | TextIO, TextIO]] = [
+            (master_fd, cast(TextIO, sys.stdout)),
+        ]
+        return process, sources, cleanup
+
+    process = subprocess.Popen(
+        cmd,
+        env=env,
+        cwd=cwd,
+        stdin=None,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        encoding="utf-8",
+    )
+
+    def cleanup() -> None:
+        return None
+
+    assert process.stdout is not None and process.stderr is not None
+    sources = [
+        (process.stdout, cast(TextIO, sys.stdout)),
+        (process.stderr, cast(TextIO, sys.stderr)),
+    ]
+    return process, sources, cleanup
+
+
+def _stream_source(
+    source: int | TextIO,
+    writer: TextIO,
+    transform: Callable[[str], str | None] | None,
+) -> None:
+    if isinstance(source, int):
+        try:
+            with os.fdopen(
+                source, "r", encoding="utf-8", errors="replace", buffering=1
+            ) as f:
+                for line in f:
+                    out = transform(line) if transform else line
+                    if out is not None:
+                        writer.write(out)
+                        writer.flush()
+        except OSError:
+            # PTY EOF may raise EIO; ignore
+            pass
+    else:
+        for line in iter(source.readline, ""):
+            out = transform(line) if transform else line
+            if out is None:
+                continue
+            writer.write(out)
+            writer.flush()
+        try:
+            source.close()
+        except Exception:
+            pass
+
+
+def _log_command(cmd: list[str], transform: Callable[[str], str | None] | None) -> None:
+    cmd_str = "> " + " ".join(cmd)
+    if transform:
+        cmd_str = transform(cmd_str)
+    sys.stderr.write(cmd_str + "\n")
+
+
+def _start_stream_threads(
+    sources: list[tuple[int | TextIO, TextIO]],
+    transform: Callable[[str], str | None] | None,
+) -> list[threading.Thread]:
+    threads: list[threading.Thread] = []
+    for src, dst in sources:
+        t = threading.Thread(
+            target=_stream_source, args=(src, dst, transform), daemon=True
+        )
+        t.start()
+        threads.append(t)
+    return threads

llama_deploy/appserver/routers/__init__.py

@@ -1,4 +1,5 @@
-from .deployments import deployments_router
-from .status import status_router
+from .deployments import create_deployments_router
+from .status import health_router
+from .ui_proxy import create_ui_proxy_router
 
-__all__ = ["deployments_router", "status_router"]
+__all__ = ["create_deployments_router", "create_ui_proxy_router", "health_router"]