diff --git a/src/core/containers/runtime/__init__.py b/src/core/containers/runtime/__init__.py index a72b5301..5cc6cf49 100644 --- a/src/core/containers/runtime/__init__.py +++ b/src/core/containers/runtime/__init__.py @@ -7,9 +7,11 @@ """Container runtime providers.""" from .providers import ContainerProvider, KubernetesProvider, LocalDockerProvider +from .uv_provider import UVProvider __all__ = [ "ContainerProvider", "LocalDockerProvider", "KubernetesProvider", + "UVProvider", ] \ No newline at end of file diff --git a/src/core/containers/runtime/providers.py b/src/core/containers/runtime/providers.py index a8022ddc..03633516 100644 --- a/src/core/containers/runtime/providers.py +++ b/src/core/containers/runtime/providers.py @@ -118,7 +118,11 @@ def __init__(self): capture_output=True, timeout=5, ) - except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired): + except ( + subprocess.CalledProcessError, + FileNotFoundError, + subprocess.TimeoutExpired, + ): raise RuntimeError( "Docker is not available. Please install Docker Desktop or Docker Engine." ) @@ -154,10 +158,13 @@ def start_container( # Build docker run command cmd = [ - "docker", "run", + "docker", + "run", "-d", # Detached - "--name", self._container_name, - "-p", f"{port}:8000", # Map port + "--name", + self._container_name, + "-p", + f"{port}:8000", # Map port ] # Add environment variables @@ -290,4 +297,68 @@ class KubernetesProvider(ContainerProvider): >>> # Pod running in k8s, accessible via service or port-forward >>> provider.stop_container() """ + pass + + +class RuntimeProvider(ABC): + """ + Abstract base class for runtime providers that are not container providers. + Providers implement this interface to support different runtime platforms: + - UVProvider: Runs environments via `uv run` + + The provider manages a single runtime lifecycle and provides the base URL + for connecting to it. + + Example: + >>> provider = UVProvider() + >>> base_url = provider.start_container("echo-env:latest") + >>> print(base_url) # http://localhost:8000 + >>> # Use the environment via base_url + >>> provider.stop_container() + """ + + @abstractmethod + def start( + self, + port: Optional[int] = None, + env_vars: Optional[Dict[str, str]] = None, + **kwargs: Any, + ) -> str: + """ + Start a runtime from the specified image. + + Args: + image: Runtime image name + port: Port to expose (if None, provider chooses) + env_vars: Environment variables for the runtime + **kwargs: Additional runtime options + """ + + @abstractmethod + def stop(self) -> None: + """ + Stop the runtime. + """ + pass + + @abstractmethod + def wait_for_ready(self, timeout_s: float = 30.0) -> None: + """ + Wait for the runtime to be ready to accept requests. + """ + pass + + def __enter__(self) -> "RuntimeProvider": + """ + Enter the runtime provider. + """ + self.start() + return self + + def __exit__(self, exc_type, exc, tb) -> None: + """ + Exit the runtime provider. + """ + self.stop() + return False \ No newline at end of file diff --git a/src/core/containers/runtime/uv_provider.py b/src/core/containers/runtime/uv_provider.py new file mode 100644 index 00000000..a26def0b --- /dev/null +++ b/src/core/containers/runtime/uv_provider.py @@ -0,0 +1,203 @@ +"""Providers for launching Hugging Face Spaces via ``uv run``.""" + +from __future__ import annotations + +import os +import socket +import subprocess +import time +from typing import Dict, Optional + +import requests + +from .providers import RuntimeProvider + + +def _check_uv_installed() -> None: + try: + subprocess.check_output(["uv", "--version"]) + except FileNotFoundError as exc: + raise RuntimeError( + "`uv` executable not found. Install uv from https://docs.astral.sh and ensure it is on PATH." + ) from exc + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("", 0)) + sock.listen(1) + return sock.getsockname()[1] + + +def _create_uv_command( + repo_id: str, + port: int, + reload: bool, +) -> list[str]: + command = [ + "uv", + "run", + "--isolated", + "--with", + f"git+https://huggingface.co/spaces/{repo_id}", + "--", + "server", + "--host", + "0.0.0.0", + "--port", + str(port), + ] + if reload: + command.append("--reload") + return command + + +def _poll_health(health_url: str, timeout_s: float) -> None: + """Poll a health endpoint until it returns HTTP 200 or times out.""" + + deadline = time.time() + timeout_s + while time.time() < deadline: + try: + timeout = max(0.0001, min(deadline - time.time(), 2.0)) + response = requests.get(health_url, timeout=timeout) + if response.status_code == 200: + return + except requests.RequestException: + continue + + time.sleep(0.5) + + raise TimeoutError(f"Server did not become ready within {timeout_s:.1f} seconds") + + +class UVProvider(RuntimeProvider): + """ + RuntimeProvider implementation backed by ``uv run``. + + Args: + repo_id: The repository ID of the environment to run + reload: Whether to reload the environment on code changes + env_vars: Environment variables to pass to the environment + context_timeout_s: The timeout to wait for the environment to become ready + + Example: + >>> provider = UVProvider(repo_id="burtenshaw/echo-cli") + >>> base_url = provider.start() + >>> print(base_url) # http://localhost:8000 + >>> # Use the environment via base_url + >>> provider.stop() + """ + + def __init__( + self, + repo_id: str, + reload: bool = False, + env_vars: Optional[Dict[str, str]] = None, + context_timeout_s: float = 60.0, + ): + """Initialize the UVProvider.""" + self.repo_id = repo_id + self.reload = reload + self.env_vars = env_vars + self.context_timeout_s = context_timeout_s + _check_uv_installed() + self._process = None + self._base_url = None + + def start( + self, + port: Optional[int] = None, + env_vars: Optional[Dict[str, str]] = None, + **_: Dict[str, str], + ) -> str: + """ + Start the environment via `uv run`. + + Args: + port: The port to bind the environment to + env_vars: Environment variables to pass to the environment + + Returns: + The base URL of the environment + + Raises: + RuntimeError: If the environment is already running + """ + if self._process is not None and self._process.poll() is None: + raise RuntimeError("UVProvider is already running") + + bind_port = port or _find_free_port() + + command = _create_uv_command( + repo_id=self.repo_id, + port=bind_port, + reload=self.reload, + ) + + env = os.environ.copy() + + if self.env_vars: + env.update(self.env_vars) + if env_vars: + env.update(env_vars) + + try: + self._process = subprocess.Popen(command, env=env) + except OSError as exc: + raise RuntimeError(f"Failed to launch `uv run`: {exc}") from exc + + self._base_url = f"http://localhost:{bind_port}" + return self._base_url + + def wait_for_ready(self, timeout_s: float = 60.0) -> None: + """ + Wait for the environment to become ready. + + Args: + timeout_s: The timeout to wait for the environment to become ready + + Raises: + RuntimeError: If the environment is not running + TimeoutError: If the environment does not become ready within the timeout + """ + if self._process and self._process.poll() is not None: + code = self._process.returncode + raise RuntimeError(f"uv process exited prematurely with code {code}") + + _poll_health(f"{self._base_url}/health", timeout_s=timeout_s) + + def stop(self) -> None: + """ + Stop the environment. + + Raises: + RuntimeError: If the environment is not running + """ + if self._process is None: + return + + if self._process.poll() is None: + self._process.terminate() + try: + self._process.wait(timeout=10.0) + except subprocess.TimeoutExpired: + self._process.kill() + self._process.wait(timeout=5.0) + + self._process = None + self._base_url = None + + @property + def base_url(self) -> str: + """ + The base URL of the environment. + + Returns: + The base URL of the environment + + Raises: + RuntimeError: If the environment is not running + """ + if self._base_url is None: + raise RuntimeError("UVProvider has not been started") + return self._base_url diff --git a/src/core/http_env_client.py b/src/core/http_env_client.py index 16bbfa5d..25faab46 100644 --- a/src/core/http_env_client.py +++ b/src/core/http_env_client.py @@ -17,10 +17,10 @@ import requests from .client_types import StepResult -from .containers.runtime import LocalDockerProvider +from .containers.runtime import LocalDockerProvider, UVProvider if TYPE_CHECKING: - from .containers.runtime import ContainerProvider + from .containers.runtime import ContainerProvider, RuntimeProvider ActT = TypeVar("ActT") ObsT = TypeVar("ObsT") @@ -106,22 +106,35 @@ def from_docker_image( return cls(base_url=base_url, provider=provider) @classmethod - def from_hub(cls: Type[EnvClientT], repo_id: str, provider: Optional["ContainerProvider"] = None, **kwargs: Any) -> EnvClientT: - """ - Create an environment client by pulling from a Hugging Face model hub. + def from_hub( + cls: Type[EnvClientT], + repo_id: str, + *, + use_docker: bool = True, + provider: Optional["ContainerProvider" | "RuntimeProvider"] = None, + **provider_kwargs: Any, + ) -> EnvClientT: + """Create a client from a Hugging Face Space. + + Set ``use_docker=True`` to launch the registry image with a container + provider. The default ``use_docker=False`` runs the Space locally using + ``uv run`` through :class:`UVProvider`. """ - - if provider is None: - provider = LocalDockerProvider() - - if "tag" in kwargs: - tag = kwargs["tag"] + + if use_docker: + tag = provider_kwargs.pop("tag", "latest") + image = f"registry.hf.space/{repo_id.replace('/', '-')}:{tag}" + return cls.from_docker_image(image, provider=provider, **provider_kwargs) else: - tag = "latest" - - base_url = f"registry.hf.space/{repo_id.replace('/', '-')}:{tag}" - - return cls.from_docker_image(image=base_url, provider=provider) + provider: RuntimeProvider = UVProvider( + repo_id=repo_id, + **provider_kwargs, + ) + base_url = provider.start() + timeout_s = provider_kwargs.pop("timeout_s", 60.0) + provider.wait_for_ready(base_url=provider.base_url, timeout_s=timeout_s) + + return cls(base_url=base_url, provider=provider) @abstractmethod def _step_payload(self, action: ActT) -> dict: