Team Ai
Apppublic

openenv/echo_env

sourceHugging Faceupdated 1d agoView on Hugging Face
6likes
uv_provider.py310 linesDownload Raw Back to runtime
1"""Providers for launching ASGI applications via ``uv run``."""2 3from __future__ import annotations4 5import os6import shutil7import socket8import subprocess9import tempfile10import time11from typing import Dict, Optional12 13import requests14 15from .providers import RuntimeProvider16 17GIT_URL_PREFIX = "git+"18 19 20def _check_uv_installed() -> None:21    try:22        subprocess.check_output(["uv", "--version"])23    except FileNotFoundError as exc:24        raise RuntimeError(25            "`uv` executable not found. Install uv from https://docs.astral.sh and ensure it is on PATH."26        ) from exc27 28 29def _clone_git_project(git_url: str, timeout_s: float) -> str:30    """Clone a `git+<url>` spec to a temp dir and return its local path.31 32    `uv run --project` only discovers a project in a local directory (per33    `uv run --help`) -- it has no support for remote git sources. Callers34    that accept a `git+...` `project_path` (e.g. `EnvClient.from_env`)35    must therefore clone it themselves before handing a path to `uv run`.36    """37    repo_url = git_url[len(GIT_URL_PREFIX) :]38    clone_dir = tempfile.mkdtemp(prefix="openenv-uv-clone-")39    try:40        subprocess.run(41            ["git", "clone", "--depth", "1", repo_url, clone_dir],42            check=True,43            capture_output=True,44            text=True,45            # A hung/slow remote would otherwise block start() (and the46            # caller's own readiness timeout) indefinitely.47            timeout=timeout_s,48        )49    except FileNotFoundError as exc:50        shutil.rmtree(clone_dir, ignore_errors=True)51        raise RuntimeError(52            f"`git` executable not found; required to clone project_path {git_url!r}."53        ) from exc54    except subprocess.CalledProcessError as exc:55        shutil.rmtree(clone_dir, ignore_errors=True)56        raise RuntimeError(f"Failed to clone {repo_url!r}: {exc.stderr}") from exc57    except subprocess.TimeoutExpired as exc:58        shutil.rmtree(clone_dir, ignore_errors=True)59        raise RuntimeError(60            f"Timed out cloning {repo_url!r} after {timeout_s:.1f}s"61        ) from exc62    return clone_dir63 64 65def _find_free_port() -> int:66    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:67        sock.bind(("", 0))68        sock.listen(1)69        return sock.getsockname()[1]70 71 72def _create_uv_command(73    *,74    host: str,75    port: int,76    reload: bool,77    workers: int,78    app: str,79    project_path: str,80) -> list[str]:81    command: list[str] = ["uv", "run", "--isolated", "--project", project_path]82 83    command.append("--")84    command.extend(85        [86            "uvicorn",87            app,88            "--host",89            host,90            "--port",91            str(port),92            "--workers",93            str(workers),94        ]95    )96 97    if reload:98        command.append("--reload")99 100    return command101 102 103def _poll_health(health_url: str, timeout_s: float) -> None:104    """Poll a health endpoint until it returns HTTP 200 or times out."""105 106    deadline = time.time() + timeout_s107    while time.time() < deadline:108        try:109            timeout = max(0.0001, min(deadline - time.time(), 2.0))110            response = requests.get(health_url, timeout=timeout)111            if response.status_code == 200:112                return113        except requests.RequestException:114            # Server not accepting connections yet. Fall through to the sleep115            # below rather than `continue`-ing: a refused connection returns116            # immediately, so retrying without a pause busy-spins a CPU core117            # for the whole timeout window while the server boots.118            pass119 120        time.sleep(0.5)121 122    raise TimeoutError(f"Server did not become ready within {timeout_s:.1f} seconds")123 124 125class UVProvider(RuntimeProvider):126    """127    RuntimeProvider implementation backed by ``uv run``.128 129    Args:130        project_path (`str`):131            Local path to a uv project (passed to `uv run --project`), or a132            `git+<url>` spec that is cloned to a temp directory on `start()`.133        app (`str`, *optional*, defaults to `"server.app:app"`):134            ASGI application path for uvicorn.135        host (`str`, *optional*, defaults to `"0.0.0.0"`):136            Host interface to bind to.137        reload (`bool`, *optional*, defaults to `False`):138            Whether to enable uvicorn's reload mode.139        env_vars (`dict`, *optional*):140            Environment variables to pass through to the spawned process.141        context_timeout_s (`float`, *optional*, defaults to `60.0`):142            How long to wait for the environment to become ready.143 144    Examples:145 146        ```python147        provider = UVProvider(project_path="/path/to/env")148        base_url = provider.start()149        print(base_url)  # http://localhost:8000150        # Use the environment via base_url151        provider.stop()152        ```153    """154 155    def __init__(156        self,157        *,158        project_path: str,159        app: str = "server.app:app",160        host: str = "0.0.0.0",161        reload: bool = False,162        env_vars: Optional[Dict[str, str]] = None,163        context_timeout_s: float = 60.0,164    ):165        """Initialize the UVProvider."""166        # `os.path.abspath` would mangle a `git+<url>` spec (e.g. collapsing167        # `https://` to `https:/`) and is meaningless for one anyway -- only168        # resolve it for genuine local paths. Git specs are cloned in `start()`.169        self.project_path = (170            project_path171            if project_path.startswith(GIT_URL_PREFIX)172            else os.path.abspath(project_path)173        )174        self.app = app175        self.host = host176        self.reload = reload177        self.env_vars = env_vars178        self.context_timeout_s = context_timeout_s179        _check_uv_installed()180        self._process = None181        self._base_url = None182        self._clone_dir: Optional[str] = None183 184    def start(185        self,186        port: Optional[int] = None,187        env_vars: Optional[Dict[str, str]] = None,188        workers: int = 1,189        **_: Dict[str, str],190    ) -> str:191        """192        Start the environment via `uv run`.193 194        Args:195            port (`int`, *optional*):196                The port to bind the environment to.197            env_vars (`dict`, *optional*):198                Environment variables to pass to the environment.199            workers (`int`, *optional*, defaults to `1`):200                The number of workers to use.201 202        Returns:203            `str`: Base URL of the environment.204 205        Raises:206            RuntimeError: If the environment is already running.207        """208        if self._process is not None and self._process.poll() is None:209            raise RuntimeError("UVProvider is already running")210 211        bind_port = port or _find_free_port()212 213        # A previous start() may have left a clone dir behind if the spawned214        # process later died on its own -- the "already running" guard above215        # only checks liveness, so a restart in that state would otherwise216        # overwrite self._clone_dir and leak the old directory.217        if self._clone_dir is not None:218            shutil.rmtree(self._clone_dir, ignore_errors=True)219            self._clone_dir = None220 221        project_path = self.project_path222        if project_path.startswith(GIT_URL_PREFIX):223            self._clone_dir = _clone_git_project(project_path, self.context_timeout_s)224            project_path = self._clone_dir225 226        command = _create_uv_command(227            host=self.host,228            port=bind_port,229            reload=self.reload,230            workers=workers,231            app=self.app,232            project_path=project_path,233        )234 235        env = os.environ.copy()236 237        if self.env_vars:238            env.update(self.env_vars)239        if env_vars:240            env.update(env_vars)241 242        try:243            self._process = subprocess.Popen(command, env=env)244        except OSError as exc:245            if self._clone_dir is not None:246                shutil.rmtree(self._clone_dir, ignore_errors=True)247                self._clone_dir = None248            raise RuntimeError(f"Failed to launch `uv run`: {exc}") from exc249 250        client_host = "127.0.0.1" if self.host in {"0.0.0.0", "::"} else self.host251        self._base_url = f"http://{client_host}:{bind_port}"252        return self._base_url253 254    def wait_for_ready(self, timeout_s: float | None = None) -> None:255        """256        Wait for the environment to become ready.257 258        Args:259            timeout_s (`float`, *optional*):260                Maximum time in seconds to wait for the environment to become261                ready. Defaults to the provider's `context_timeout_s`.262 263        Raises:264            RuntimeError: If the environment is not running.265            TimeoutError: If the environment does not become ready within the timeout.266        """267        if timeout_s is None:268            timeout_s = self.context_timeout_s269 270        if self._process and self._process.poll() is not None:271            code = self._process.returncode272            raise RuntimeError(f"uv process exited prematurely with code {code}")273 274        _poll_health(f"{self._base_url}/health", timeout_s=timeout_s)275 276    def stop(self) -> None:277        """278        Stop the environment.279        """280        if self._process is not None:281            if self._process.poll() is None:282                self._process.terminate()283                try:284                    self._process.wait(timeout=10.0)285                except subprocess.TimeoutExpired:286                    self._process.kill()287                    self._process.wait(timeout=5.0)288 289            self._process = None290            self._base_url = None291 292        if self._clone_dir is not None:293            shutil.rmtree(self._clone_dir, ignore_errors=True)294            self._clone_dir = None295 296    @property297    def base_url(self) -> str:298        """299        The base URL of the environment.300 301        Returns:302            `str`: Base URL of the running environment.303 304        Raises:305            RuntimeError: If the environment has not been started.306        """307        if self._base_url is None:308            raise RuntimeError("UVProvider has not been started")309        return self._base_url310