fixed all remaining issues with the session manager
This commit is contained in:
302
session-manager/async_docker_client.py
Normal file
302
session-manager/async_docker_client.py
Normal file
@@ -0,0 +1,302 @@
|
||||
"""
|
||||
Async Docker Operations Wrapper
|
||||
|
||||
Provides async wrappers for Docker operations to eliminate blocking calls
|
||||
in FastAPI async contexts and improve concurrency and scalability.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, Optional, List, Any
|
||||
from contextlib import asynccontextmanager
|
||||
import os
|
||||
|
||||
from aiodeocker import Docker
|
||||
from aiodeocker.containers import DockerContainer
|
||||
from aiodeocker.exceptions import DockerError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AsyncDockerClient:
|
||||
"""Async wrapper for Docker operations using aiodeocker."""
|
||||
|
||||
def __init__(self):
|
||||
self._docker: Optional[Docker] = None
|
||||
self._connected = False
|
||||
|
||||
async def __aenter__(self):
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
await self.disconnect()
|
||||
|
||||
async def connect(self):
|
||||
"""Connect to Docker daemon."""
|
||||
if self._connected:
|
||||
return
|
||||
|
||||
try:
|
||||
# Configure TLS if available
|
||||
tls_config = None
|
||||
if os.getenv("DOCKER_TLS_VERIFY") == "1":
|
||||
from aiodeocker.utils import create_tls_config
|
||||
|
||||
tls_config = create_tls_config(
|
||||
ca_cert=os.getenv("DOCKER_CA_CERT", "/etc/docker/certs/ca.pem"),
|
||||
client_cert=(
|
||||
os.getenv(
|
||||
"DOCKER_CLIENT_CERT", "/etc/docker/certs/client-cert.pem"
|
||||
),
|
||||
os.getenv(
|
||||
"DOCKER_CLIENT_KEY", "/etc/docker/certs/client-key.pem"
|
||||
),
|
||||
),
|
||||
verify=True,
|
||||
)
|
||||
|
||||
docker_host = os.getenv("DOCKER_HOST", "tcp://host.docker.internal:2376")
|
||||
self._docker = Docker(docker_host, tls=tls_config)
|
||||
|
||||
# Test connection
|
||||
await self._docker.ping()
|
||||
self._connected = True
|
||||
logger.info("Async Docker client connected successfully")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to connect to Docker: {e}")
|
||||
raise
|
||||
|
||||
async def disconnect(self):
|
||||
"""Disconnect from Docker daemon."""
|
||||
if self._docker and self._connected:
|
||||
await self._docker.close()
|
||||
self._connected = False
|
||||
logger.info("Async Docker client disconnected")
|
||||
|
||||
async def ping(self) -> bool:
|
||||
"""Test Docker connectivity."""
|
||||
if not self._docker:
|
||||
return False
|
||||
try:
|
||||
await self._docker.ping()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def create_container(
|
||||
self,
|
||||
image: str,
|
||||
name: str,
|
||||
volumes: Optional[Dict[str, Dict[str, str]]] = None,
|
||||
ports: Optional[Dict[str, int]] = None,
|
||||
environment: Optional[Dict[str, str]] = None,
|
||||
network_mode: str = "bridge",
|
||||
mem_limit: Optional[str] = None,
|
||||
cpu_quota: Optional[int] = None,
|
||||
cpu_period: Optional[int] = None,
|
||||
tmpfs: Optional[Dict[str, str]] = None,
|
||||
**kwargs,
|
||||
) -> DockerContainer:
|
||||
"""
|
||||
Create a Docker container asynchronously.
|
||||
|
||||
Args:
|
||||
image: Container image name
|
||||
name: Container name
|
||||
volumes: Volume mounts
|
||||
ports: Port mappings
|
||||
environment: Environment variables
|
||||
network_mode: Network mode
|
||||
mem_limit: Memory limit (e.g., "4g")
|
||||
cpu_quota: CPU quota
|
||||
cpu_period: CPU period
|
||||
tmpfs: tmpfs mounts
|
||||
**kwargs: Additional container configuration
|
||||
|
||||
Returns:
|
||||
DockerContainer: The created container
|
||||
"""
|
||||
if not self._docker:
|
||||
raise RuntimeError("Docker client not connected")
|
||||
|
||||
config = {
|
||||
"Image": image,
|
||||
"name": name,
|
||||
"Volumes": volumes or {},
|
||||
"ExposedPorts": {f"{port}/tcp": {} for port in ports.values()}
|
||||
if ports
|
||||
else {},
|
||||
"Env": [f"{k}={v}" for k, v in (environment or {}).items()],
|
||||
"NetworkMode": network_mode,
|
||||
"HostConfig": {
|
||||
"Binds": [
|
||||
f"{host}:{container['bind']}:{container.get('mode', 'rw')}"
|
||||
for host, container in (volumes or {}).items()
|
||||
],
|
||||
"PortBindings": {
|
||||
f"{container_port}/tcp": [{"HostPort": str(host_port)}]
|
||||
for container_port, host_port in (ports or {}).items()
|
||||
},
|
||||
"Tmpfs": tmpfs or {},
|
||||
},
|
||||
}
|
||||
|
||||
# Add resource limits
|
||||
host_config = config["HostConfig"]
|
||||
if mem_limit:
|
||||
host_config["Memory"] = self._parse_memory_limit(mem_limit)
|
||||
if cpu_quota is not None:
|
||||
host_config["CpuQuota"] = cpu_quota
|
||||
if cpu_period is not None:
|
||||
host_config["CpuPeriod"] = cpu_period
|
||||
|
||||
# Add any additional host config
|
||||
host_config.update(kwargs.get("host_config", {}))
|
||||
|
||||
try:
|
||||
container = await self._docker.containers.create(config)
|
||||
logger.info(f"Container {name} created successfully")
|
||||
return container
|
||||
except DockerError as e:
|
||||
logger.error(f"Failed to create container {name}: {e}")
|
||||
raise
|
||||
|
||||
async def start_container(self, container: DockerContainer) -> None:
|
||||
"""Start a Docker container."""
|
||||
try:
|
||||
await container.start()
|
||||
logger.info(f"Container {container.id} started successfully")
|
||||
except DockerError as e:
|
||||
logger.error(f"Failed to start container {container.id}: {e}")
|
||||
raise
|
||||
|
||||
async def stop_container(
|
||||
self, container: DockerContainer, timeout: int = 10
|
||||
) -> None:
|
||||
"""Stop a Docker container."""
|
||||
try:
|
||||
await container.stop(timeout=timeout)
|
||||
logger.info(f"Container {container.id} stopped successfully")
|
||||
except DockerError as e:
|
||||
logger.error(f"Failed to stop container {container.id}: {e}")
|
||||
raise
|
||||
|
||||
async def remove_container(
|
||||
self, container: DockerContainer, force: bool = False
|
||||
) -> None:
|
||||
"""Remove a Docker container."""
|
||||
try:
|
||||
await container.delete(force=force)
|
||||
logger.info(f"Container {container.id} removed successfully")
|
||||
except DockerError as e:
|
||||
logger.error(f"Failed to remove container {container.id}: {e}")
|
||||
raise
|
||||
|
||||
async def get_container(self, container_id: str) -> Optional[DockerContainer]:
|
||||
"""Get a container by ID or name."""
|
||||
try:
|
||||
return await self._docker.containers.get(container_id)
|
||||
except DockerError:
|
||||
return None
|
||||
|
||||
async def list_containers(
|
||||
self, all: bool = False, filters: Optional[Dict[str, Any]] = None
|
||||
) -> List[DockerContainer]:
|
||||
"""List Docker containers."""
|
||||
try:
|
||||
return await self._docker.containers.list(all=all, filters=filters)
|
||||
except DockerError as e:
|
||||
logger.error(f"Failed to list containers: {e}")
|
||||
return []
|
||||
|
||||
async def get_container_stats(
|
||||
self, container: DockerContainer
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Get container statistics."""
|
||||
try:
|
||||
stats = await container.stats(stream=False)
|
||||
return stats
|
||||
except DockerError as e:
|
||||
logger.error(f"Failed to get stats for container {container.id}: {e}")
|
||||
return None
|
||||
|
||||
async def get_system_info(self) -> Optional[Dict[str, Any]]:
|
||||
"""Get Docker system information."""
|
||||
if not self._docker:
|
||||
return None
|
||||
try:
|
||||
return await self._docker.system.info()
|
||||
except DockerError as e:
|
||||
logger.error(f"Failed to get system info: {e}")
|
||||
return None
|
||||
|
||||
def _parse_memory_limit(self, memory_str: str) -> int:
|
||||
"""Parse memory limit string to bytes."""
|
||||
memory_str = memory_str.lower().strip()
|
||||
if memory_str.endswith("g"):
|
||||
return int(memory_str[:-1]) * (1024**3)
|
||||
elif memory_str.endswith("m"):
|
||||
return int(memory_str[:-1]) * (1024**2)
|
||||
elif memory_str.endswith("k"):
|
||||
return int(memory_str[:-1]) * 1024
|
||||
else:
|
||||
return int(memory_str)
|
||||
|
||||
|
||||
# Global async Docker client instance
|
||||
_async_docker_client = AsyncDockerClient()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def get_async_docker_client():
|
||||
"""Context manager for async Docker client."""
|
||||
async with _async_docker_client as client:
|
||||
yield client
|
||||
|
||||
|
||||
async def async_docker_ping() -> bool:
|
||||
"""Async ping Docker daemon."""
|
||||
async with get_async_docker_client() as client:
|
||||
return await client.ping()
|
||||
|
||||
|
||||
async def async_create_container(**kwargs) -> DockerContainer:
|
||||
"""Async container creation wrapper."""
|
||||
async with get_async_docker_client() as client:
|
||||
return await client.create_container(**kwargs)
|
||||
|
||||
|
||||
async def async_start_container(container: DockerContainer) -> None:
|
||||
"""Async container start wrapper."""
|
||||
async with get_async_docker_client() as client:
|
||||
await client.start_container(container)
|
||||
|
||||
|
||||
async def async_stop_container(container: DockerContainer, timeout: int = 10) -> None:
|
||||
"""Async container stop wrapper."""
|
||||
async with get_async_docker_client() as client:
|
||||
await client.stop_container(container, timeout)
|
||||
|
||||
|
||||
async def async_remove_container(
|
||||
container: DockerContainer, force: bool = False
|
||||
) -> None:
|
||||
"""Async container removal wrapper."""
|
||||
async with get_async_docker_client() as client:
|
||||
await client.remove_container(container, force)
|
||||
|
||||
|
||||
async def async_list_containers(
|
||||
all: bool = False, filters: Optional[Dict[str, Any]] = None
|
||||
) -> List[DockerContainer]:
|
||||
"""Async container listing wrapper."""
|
||||
async with get_async_docker_client() as client:
|
||||
return await client.list_containers(all=all, filters=filters)
|
||||
|
||||
|
||||
async def async_get_container(container_id: str) -> Optional[DockerContainer]:
|
||||
"""Async container retrieval wrapper."""
|
||||
async with get_async_docker_client() as client:
|
||||
return await client.get_container(container_id)
|
||||
Reference in New Issue
Block a user