vlmparse 0.1.8__py3-none-any.whl → 0.1.9__py3-none-any.whl

vlmparse/servers/docker_run_deployment.py

@@ -0,0 +1,226 @@
+import getpass
+import time
+from contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import docker
+from loguru import logger
+
+if TYPE_CHECKING:
+    from .base_server import BaseServerConfig
+
+
+def _ensure_image_exists(
+    client: docker.DockerClient,
+    image: str,
+    dockerfile_path: Path,
+):
+    """Check if image exists, build it if not."""
+    try:
+        client.images.get(image)
+        logger.info(f"Docker image {image} found")
+        return
+    except docker.errors.ImageNotFound:
+        logger.info(f"Docker image {image} not found, building...")
+
+        if not dockerfile_path.exists():
+            raise FileNotFoundError(
+                f"Dockerfile directory not found at {dockerfile_path}"
+            ) from None
+
+        logger.info(f"Building image from {dockerfile_path}")
+
+        # Use low-level API for real-time streaming
+        api_client = docker.APIClient(base_url="unix://var/run/docker.sock")
+
+        # Build the image with streaming
+        build_stream = api_client.build(
+            path=str(dockerfile_path),
+            tag=image,
+            rm=True,
+            decode=True,  # Automatically decode JSON responses to dict
+        )
+
+        # Stream build logs in real-time
+        for chunk in build_stream:
+            if "stream" in chunk:
+                for line in chunk["stream"].splitlines():
+                    logger.info(line)
+            elif "error" in chunk:
+                logger.error(chunk["error"])
+                raise docker.errors.BuildError(chunk["error"], build_stream) from None
+            elif "status" in chunk:
+                # Handle status updates (e.g., downloading layers)
+                logger.debug(chunk["status"])
+
+        logger.info(f"Successfully built image {image}")
+
+
+@contextmanager
+def docker_server(
+    config: "BaseServerConfig",
+    timeout: int = 1000,
+    cleanup: bool = True,
+):
+    """Generic context manager for Docker server deployment.
+
+    Args:
+        config: BaseServerConfig (can be any subclass like DockerServerConfig or VLLMDockerServerConfig)
+        timeout: Timeout in seconds to wait for server to be ready
+        cleanup: If True, stop and remove container on exit. If False, leave container running
+
+    Yields:
+        tuple: (base_url, container) - The base URL of the server and the Docker container object
+    """
+
+    client = docker.from_env()
+    container = None
+
+    try:
+        # Ensure image exists
+        logger.info(f"Checking for Docker image {config.docker_image}...")
+
+        if config.dockerfile_dir is not None:
+            _ensure_image_exists(
+                client, config.docker_image, Path(config.dockerfile_dir)
+            )
+        else:
+            # Pull pre-built image
+            try:
+                client.images.get(config.docker_image)
+                logger.info(f"Docker image {config.docker_image} found locally")
+            except docker.errors.ImageNotFound:
+                logger.info(
+                    f"Docker image {config.docker_image} not found locally, pulling..."
+                )
+                client.images.pull(config.docker_image)
+                logger.info(f"Successfully pulled {config.docker_image}")
+
+        logger.info(
+            f"Starting Docker container for {config.model_name} on port {config.docker_port}"
+        )
+
+        # Configure GPU access
+        device_requests = None
+
+        if config.gpu_device_ids is None:
+            # Default: Try to use all GPUs if available
+            device_requests = [
+                docker.types.DeviceRequest(count=-1, capabilities=[["gpu"]])
+            ]
+        elif len(config.gpu_device_ids) > 0 and config.gpu_device_ids[0] != "":
+            # Use specific GPU devices
+            device_requests = [
+                docker.types.DeviceRequest(
+                    device_ids=config.gpu_device_ids, capabilities=[["gpu"]]
+                )
+            ]
+        else:
+            # Empty list means CPU-only, no GPU
+            device_requests = None
+
+        # Use generic methods from config
+        command = config.get_command()
+        volumes = config.get_volumes()
+        environment = config.get_environment()
+        container_port = config.container_port
+        log_prefix = config.model_name
+
+        # Construct URI for label
+        uri = f"http://localhost:{config.docker_port}{config.get_base_url_suffix()}"
+
+        # Determine GPU label
+        if config.gpu_device_ids is None:
+            gpu_label = "0"
+        elif len(config.gpu_device_ids) == 0 or (
+            len(config.gpu_device_ids) == 1 and config.gpu_device_ids[0] == ""
+        ):
+            gpu_label = "cpu"
+        else:
+            gpu_label = ",".join(config.gpu_device_ids)
+
+        # Start container
+        container_kwargs = {
+            "image": config.docker_image,
+            "ports": {f"{container_port}/tcp": config.docker_port},
+            "detach": True,
+            "remove": True,
+            "name": f"vlmparse-{config.model_name.replace('/', '-')}-{getpass.getuser()}",
+            "labels": {
+                "vlmparse_model_name": config.model_name,
+                "vlmparse_uri": uri,
+                "vlmparse_gpus": gpu_label,
+            },
+        }
+
+        if device_requests is not None:
+            container_kwargs["device_requests"] = device_requests
+        if command:
+            container_kwargs["command"] = command
+        if environment:
+            container_kwargs["environment"] = environment
+        if volumes:
+            container_kwargs["volumes"] = volumes
+        if config.entrypoint:
+            container_kwargs["entrypoint"] = config.entrypoint
+
+        container = client.containers.run(**container_kwargs)
+
+        logger.info(
+            f"Container {container.short_id} started, waiting for server to be ready..."
+        )
+
+        # Wait for server to be ready
+        start_time = time.time()
+        server_ready = False
+        last_log_position = 0
+
+        while time.time() - start_time < timeout:
+            try:
+                container.reload()
+            except docker.errors.NotFound as e:
+                logger.error("Container stopped unexpectedly during startup")
+                raise RuntimeError(
+                    "Container crashed during initialization. Check Docker logs for details."
+                ) from e
+
+            if container.status == "running":
+                # Get all logs and display new ones
+                all_logs = container.logs().decode("utf-8")
+
+                # Display new log lines
+                if len(all_logs) > last_log_position:
+                    new_logs = all_logs[last_log_position:]
+                    for line in new_logs.splitlines():
+                        if line.strip():  # Only print non-empty lines
+                            logger.info(f"[{log_prefix}] {line}")
+                    last_log_position = len(all_logs)
+
+                # Check if server is ready
+                for indicator in config.server_ready_indicators:
+                    if indicator in all_logs:
+                        server_ready = True
+                if server_ready:
+                    logger.info(f"Server ready indicator '{indicator}' found in logs")
+                    break
+
+            time.sleep(2)
+
+        if not server_ready:
+            raise TimeoutError(f"Server did not become ready within {timeout} seconds")
+
+        # Build base URL using config's suffix method
+        base_url = (
+            f"http://localhost:{config.docker_port}{config.get_base_url_suffix()}"
+        )
+
+        logger.info(f"{log_prefix} server ready at {base_url}")
+
+        yield base_url, container
+
+    finally:
+        if cleanup and container:
+            logger.info(f"Stopping container {container.short_id}")
+            container.stop(timeout=10)
+            logger.info("Container stopped")

vlmparse/servers/docker_server.py

@@ -1,52 +1,28 @@
 import os
-from typing import Callable
 
-from loguru import logger
 from pydantic import Field
 
-from .model_identity import ModelIdentityMixin
-from .utils import docker_server
+from .base_server import BaseServer, BaseServerConfig
+from .docker_run_deployment import docker_server
 
 
-class DockerServerConfig(ModelIdentityMixin):
-    """Base configuration for deploying a Docker server.
+class DockerServerConfig(BaseServerConfig):
+    """Configuration for deploying a Docker server.
 
-    Inherits from ModelIdentityMixin which provides:
-    - model_name: str
-    - default_model_name: str | None
-    - aliases: list[str]
-    - _create_client_kwargs(base_url): Helper for creating client configs
-    - get_all_names(): All names this model can be referenced by
+    Inherits from BaseServerConfig which provides common server configuration.
     """
 
     docker_image: str
     dockerfile_dir: str | None = None
     command_args: list[str] = Field(default_factory=list)
-    server_ready_indicators: list[str] = Field(
-        default_factory=lambda: [
-            "Application startup complete",
-            "Uvicorn running",
-            "Starting vLLM API server",
-        ]
-    )
-    docker_port: int = 8056
-    gpu_device_ids: list[str] | None = None
-    container_port: int = 8000
-    environment: dict[str, str] = Field(default_factory=dict)
     volumes: dict[str, dict] | None = None
     entrypoint: str | None = None
 
-    class Config:
-        extra = "allow"
-
     @property
     def client_config(self):
         """Override in subclasses to return appropriate client config."""
         raise NotImplementedError
 
-    def get_client(self, **kwargs):
-        return self.client_config.get_client(**kwargs)
-
     def get_server(self, auto_stop: bool = True):
         return ConverterServer(config=self, auto_stop=auto_stop)
 
@@ -71,14 +47,6 @@ class DockerServerConfig(ModelIdentityMixin):
         """Setup volumes for container. Override in subclasses for specific logic."""
         return self.volumes
 
-    def get_environment(self) -> dict | None:
-        """Setup environment variables. Override in subclasses for specific logic."""
-        return self.environment if self.environment else None
-
-    def get_base_url_suffix(self) -> str:
-        """Return URL suffix (e.g., '/v1' for OpenAI-compatible APIs). Override in subclasses."""
-        return ""
-
 
 DEFAULT_MODEL_NAME = "vllm-model"
 
@@ -144,93 +112,9 @@ class VLLMDockerServerConfig(DockerServerConfig):
         return "/v1"
 
 
-class ConverterServer:
+class ConverterServer(BaseServer):
     """Manages Docker server lifecycle with start/stop methods."""
 
-    def __init__(self, config: DockerServerConfig, auto_stop: bool = True):
-        self.config = config
-        self.auto_stop = auto_stop
-        self._server_context = None
-        self._container = None
-        self.base_url = None
-
-    def start(self):
-        """Start the Docker server."""
-        if self._server_context is not None:
-            logger.warning("Server already started")
-            return self.base_url, self._container
-
-        # Use the generic docker_server for all server types
-        self._server_context = docker_server(config=self.config, cleanup=self.auto_stop)
-
-        self.base_url, self._container = self._server_context.__enter__()
-        logger.info(f"Server started at {self.base_url}")
-        logger.info(f"Container ID: {self._container.id}")
-        logger.info(f"Container name: {self._container.name}")
-        return self.base_url, self._container
-
-    def stop(self):
-        """Stop the Docker server."""
-        if self._server_context is not None:
-            try:
-                self._server_context.__exit__(None, None, None)
-            except Exception as e:
-                logger.warning(f"Error during server cleanup: {e}")
-            finally:
-                self._server_context = None
-                self._container = None
-                self.base_url = None
-            logger.info("Server stopped")
-
-    def __del__(self):
-        """Automatically stop server when object is destroyed if auto_stop is True.
-
-        Note: This is a fallback mechanism. Prefer using the context manager
-        or explicitly calling stop() for reliable cleanup.
-        """
-        try:
-            if self.auto_stop and self._server_context is not None:
-                self.stop()
-        except Exception:
-            pass  # Suppress errors during garbage collection
-
-
-class DockerConfigRegistry:
-    """Registry for mapping model names to their Docker configurations.
-
-    Thread-safe registry that maps model names to their Docker configuration factories.
-    """
-
-    def __init__(self):
-        import threading
-
-        self._registry: dict[str, Callable[[], DockerServerConfig | None]] = {}
-        self._lock = threading.RLock()
-
-    def register(
-        self, model_name: str, config_factory: Callable[[], DockerServerConfig | None]
-    ):
-        """Register a config factory for a model name (thread-safe)."""
-        with self._lock:
-            self._registry[model_name] = config_factory
-
-    def get(self, model_name: str, default=False) -> DockerServerConfig | None:
-        """Get config for a model name (thread-safe). Returns default if not registered."""
-        with self._lock:
-            if model_name not in self._registry:
-                if default:
-                    return VLLMDockerServerConfig(
-                        model_name=model_name, default_model_name=DEFAULT_MODEL_NAME
-                    )
-                return None
-            factory = self._registry[model_name]
-        return factory()
-
-    def list_models(self) -> list[str]:
-        """List all registered model names (thread-safe)."""
-        with self._lock:
-            return list(self._registry.keys())
-
-
-# Global registry instance
-docker_config_registry = DockerConfigRegistry()
+    def _create_server_context(self):
+        """Create the Docker server context."""
+        return docker_server(config=self.config, cleanup=self.auto_stop)

vlmparse/servers/server_registry.py

@@ -0,0 +1,42 @@
+from typing import Callable
+
+from vlmparse.servers.base_server import BaseServerConfig
+
+
+class DockerConfigRegistry:
+    """Registry for mapping model names to their Docker configurations.
+
+    Thread-safe registry that maps model names to their Docker configuration factories.
+    """
+
+    def __init__(self):
+        import threading
+
+        self._registry: dict[str, Callable[[], BaseServerConfig | None]] = {}
+        self._lock = threading.RLock()
+
+    def register(
+        self,
+        model_name: str,
+        config_factory: Callable[[], BaseServerConfig | None],
+    ):
+        """Register a config factory for a model name (thread-safe)."""
+        with self._lock:
+            self._registry[model_name] = config_factory
+
+    def get(self, model_name: str) -> BaseServerConfig | None:
+        """Get config for a model name (thread-safe). Returns None if not registered."""
+        with self._lock:
+            if model_name not in self._registry:
+                return None
+            factory = self._registry[model_name]
+        return factory()
+
+    def list_models(self) -> list[str]:
+        """List all registered model names (thread-safe)."""
+        with self._lock:
+            return list(self._registry.keys())
+
+
+# Global registry instance
+docker_config_registry = DockerConfigRegistry()