Team Ai
Apppublic

openenv/coding_env

sourceHugging Faceupdated 3mo agoView on Hugging Face
21likes
uv_provider.py225 linesDownload Raw Back to runtime
1"""Providers for launching ASGI applications via ``uv run``."""2 3from __future__ import annotations4 5import os6import socket7import subprocess8import time9from typing import Dict, Optional10 11import requests12 13from .providers import RuntimeProvider14 15 16def _check_uv_installed() -> None:17    try:18        subprocess.check_output(["uv", "--version"])19    except FileNotFoundError as exc:20        raise RuntimeError(21            "`uv` executable not found. Install uv from https://docs.astral.sh and ensure it is on PATH."22        ) from exc23 24 25def _find_free_port() -> int:26    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:27        sock.bind(("", 0))28        sock.listen(1)29        return sock.getsockname()[1]30 31 32def _create_uv_command(33    *,34    host: str,35    port: int,36    reload: bool,37    workers: int,38    app: str,39    project_path: str,40) -> list[str]:41    command: list[str] = ["uv", "run", "--isolated", "--project", project_path]42 43    command.append("--")44    command.extend(45        [46            "uvicorn",47            app,48            "--host",49            host,50            "--port",51            str(port),52            "--workers",53            str(workers),54        ]55    )56 57    if reload:58        command.append("--reload")59 60    return command61 62 63def _poll_health(health_url: str, timeout_s: float) -> None:64    """Poll a health endpoint until it returns HTTP 200 or times out."""65 66    deadline = time.time() + timeout_s67    while time.time() < deadline:68        try:69            timeout = max(0.0001, min(deadline - time.time(), 2.0))70            response = requests.get(health_url, timeout=timeout)71            if response.status_code == 200:72                return73        except requests.RequestException:74            continue75 76        time.sleep(0.5)77 78    raise TimeoutError(f"Server did not become ready within {timeout_s:.1f} seconds")79 80 81class UVProvider(RuntimeProvider):82    """83    RuntimeProvider implementation backed by ``uv run``.84 85    Args:86        project_path: Local path to a uv project (passed to ``uv run --project``)87        app: ASGI application path for uvicorn (defaults to ``server.app:app``)88        host: Host interface to bind to (defaults to ``0.0.0.0``)89        reload: Whether to enable uvicorn's reload mode90        env_vars: Environment variables to pass through to the spawned process91        context_timeout_s: How long to wait for the environment to become ready92 93    Example:94        >>> provider = UVProvider(project_path="/path/to/env")95        >>> base_url = provider.start()96        >>> print(base_url)  # http://localhost:800097        >>> # Use the environment via base_url98        >>> provider.stop()99    """100 101    def __init__(102        self,103        *,104        project_path: str,105        app: str = "server.app:app",106        host: str = "0.0.0.0",107        reload: bool = False,108        env_vars: Optional[Dict[str, str]] = None,109        context_timeout_s: float = 60.0,110    ):111        """Initialize the UVProvider."""112        self.project_path = os.path.abspath(project_path)113        self.app = app114        self.host = host115        self.reload = reload116        self.env_vars = env_vars117        self.context_timeout_s = context_timeout_s118        _check_uv_installed()119        self._process = None120        self._base_url = None121 122    def start(123        self,124        port: Optional[int] = None,125        env_vars: Optional[Dict[str, str]] = None,126        workers: int = 1,127        **_: Dict[str, str],128    ) -> str:129        """130        Start the environment via `uv run`.131 132        Args:133            port: The port to bind the environment to134            env_vars: Environment variables to pass to the environment135            workers: The number of workers to use136 137        Returns:138            The base URL of the environment139 140        Raises:141            RuntimeError: If the environment is already running142        """143        if self._process is not None and self._process.poll() is None:144            raise RuntimeError("UVProvider is already running")145 146        bind_port = port or _find_free_port()147 148        command = _create_uv_command(149            host=self.host,150            port=bind_port,151            reload=self.reload,152            workers=workers,153            app=self.app,154            project_path=self.project_path,155        )156 157        env = os.environ.copy()158 159        if self.env_vars:160            env.update(self.env_vars)161        if env_vars:162            env.update(env_vars)163 164        try:165            self._process = subprocess.Popen(command, env=env)166        except OSError as exc:167            raise RuntimeError(f"Failed to launch `uv run`: {exc}") from exc168 169        client_host = "127.0.0.1" if self.host in {"0.0.0.0", "::"} else self.host170        self._base_url = f"http://{client_host}:{bind_port}"171        return self._base_url172 173    def wait_for_ready(self, timeout_s: float = 60.0) -> None:174        """175        Wait for the environment to become ready.176 177        Args:178            timeout_s: The timeout to wait for the environment to become ready179 180        Raises:181            RuntimeError: If the environment is not running182            TimeoutError: If the environment does not become ready within the timeout183        """184        if self._process and self._process.poll() is not None:185            code = self._process.returncode186            raise RuntimeError(f"uv process exited prematurely with code {code}")187 188        _poll_health(f"{self._base_url}/health", timeout_s=timeout_s)189 190    def stop(self) -> None:191        """192        Stop the environment.193 194        Raises:195            RuntimeError: If the environment is not running196        """197        if self._process is None:198            return199 200        if self._process.poll() is None:201            self._process.terminate()202            try:203                self._process.wait(timeout=10.0)204            except subprocess.TimeoutExpired:205                self._process.kill()206                self._process.wait(timeout=5.0)207 208        self._process = None209        self._base_url = None210 211    @property212    def base_url(self) -> str:213        """214        The base URL of the environment.215 216        Returns:217            The base URL of the environment218 219        Raises:220            RuntimeError: If the environment is not running221        """222        if self._base_url is None:223            raise RuntimeError("UVProvider has not been started")224        return self._base_url225