254 lines
8.8 KiB
Python
254 lines
8.8 KiB
Python
from typing import Dict, List, Optional, Tuple
|
|
import docker
|
|
from docker.models.containers import Container
|
|
import logging
|
|
|
|
class DockerModule:
|
|
"""
|
|
Module for handling Docker container operations in SIA.
|
|
Manages container lifecycle, I/O operations, and status monitoring.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.client = docker.from_env()
|
|
self.containers: Dict[str, Container] = {}
|
|
self.logger = logging.getLogger(__name__)
|
|
|
|
def start_container(
|
|
self,
|
|
image: str,
|
|
name: Optional[str] = None,
|
|
timeout: int = -1,
|
|
command: Optional[str] = None,
|
|
arguments: Optional[List[str]] = None,
|
|
volumes: Optional[Dict[str, str]] = None,
|
|
ports: Optional[Dict[str, str]] = None,
|
|
environment: Optional[Dict[str, str]] = None,
|
|
) -> str:
|
|
"""Start a new Docker container with the specified configuration.
|
|
|
|
Args:
|
|
image: Docker image to use
|
|
name: Unique container name for long running containers
|
|
timeout: Timeout in milliseconds for short running containers
|
|
command: Main command to run in container
|
|
arguments: List of command line arguments
|
|
volumes: Dictionary mapping host paths to container paths
|
|
ports: Dictionary mapping host ports to container ports
|
|
environment: Dictionary of environment variables and values
|
|
|
|
Returns:
|
|
For short-lived containers (with timeout): Container output
|
|
For long-running containers: Container name
|
|
|
|
Raises:
|
|
docker.errors.ImageNotFound: If specified image doesn't exist
|
|
docker.errors.APIError: If container creation/start fails
|
|
"""
|
|
# Validate inputs
|
|
if name is None and timeout < 0:
|
|
raise ValueError("Either name or timeout must be provided")
|
|
|
|
# Prepare command with arguments
|
|
cmd = []
|
|
if command:
|
|
cmd.append(command)
|
|
if arguments:
|
|
cmd.extend(arguments)
|
|
|
|
# Prepare volume bindings
|
|
volume_binds = {}
|
|
if volumes:
|
|
for host_path, container_path in volumes.items():
|
|
volume_binds[host_path] = {'bind': container_path, 'mode': 'rw'}
|
|
|
|
# Prepare port bindings
|
|
port_bindings = {}
|
|
if ports:
|
|
for host_port, container_port in ports.items():
|
|
port_bindings[container_port] = host_port
|
|
|
|
try:
|
|
container = self.client.containers.run(
|
|
image,
|
|
command=cmd if cmd else None,
|
|
name=name,
|
|
detach=True,
|
|
volumes=volume_binds,
|
|
ports=port_bindings,
|
|
environment=environment,
|
|
)
|
|
|
|
if timeout >= 0:
|
|
# For short-lived containers, wait for completion
|
|
try:
|
|
container.wait(timeout=timeout/1000) # Convert ms to seconds
|
|
output = container.logs().decode('utf-8')
|
|
container.remove(force=True)
|
|
return output
|
|
except docker.errors.NotFound:
|
|
self.logger.warning(f"Container was removed before timeout")
|
|
return ""
|
|
else:
|
|
# For long-running containers, store reference
|
|
self.containers[name] = container
|
|
return name
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error starting container: {str(e)}")
|
|
raise
|
|
|
|
def stop_container(self, name: str) -> None:
|
|
"""Stop and remove a Docker container by name.
|
|
|
|
Args:
|
|
name: Name of the container to stop
|
|
|
|
Raises:
|
|
KeyError: If container name not found
|
|
docker.errors.APIError: If container stop/removal fails
|
|
"""
|
|
if name not in self.containers:
|
|
raise KeyError(f"Container {name} not found")
|
|
|
|
try:
|
|
container = self.containers[name]
|
|
container.stop()
|
|
container.remove()
|
|
del self.containers[name]
|
|
except Exception as e:
|
|
self.logger.error(f"Error stopping container {name}: {str(e)}")
|
|
raise
|
|
|
|
def write_container_stdin(self, name: str, data: str) -> None:
|
|
"""Write data to a container's standard input.
|
|
|
|
Args:
|
|
name: Name of the target container
|
|
data: Data to write to stdin
|
|
|
|
Raises:
|
|
KeyError: If container name not found
|
|
docker.errors.APIError: If write operation fails
|
|
"""
|
|
if name not in self.containers:
|
|
raise KeyError(f"Container {name} not found")
|
|
|
|
try:
|
|
container = self.containers[name]
|
|
socket = container.attach_socket(params={'stdin': 1, 'stream': 1})
|
|
socket._sock.send(data.encode('utf-8'))
|
|
socket.close()
|
|
except Exception as e:
|
|
self.logger.error(f"Error writing to container {name} stdin: {str(e)}")
|
|
raise
|
|
|
|
def read_container_stdout(self, name: str, n: int = -1) -> str:
|
|
"""Read from a container's standard output buffer.
|
|
|
|
Args:
|
|
name: Name of the container
|
|
n: Number of bytes to read; -1 means read all available
|
|
|
|
Returns:
|
|
Data read from stdout
|
|
|
|
Raises:
|
|
KeyError: If container name not found
|
|
docker.errors.APIError: If read operation fails
|
|
"""
|
|
if name not in self.containers:
|
|
raise KeyError(f"Container {name} not found")
|
|
|
|
try:
|
|
container = self.containers[name]
|
|
logs = container.logs(stdout=True, stderr=False, tail=n if n > 0 else 'all')
|
|
return logs.decode('utf-8')
|
|
except Exception as e:
|
|
self.logger.error(f"Error reading container {name} stdout: {str(e)}")
|
|
raise
|
|
|
|
def read_container_stderr(self, name: str, n: int = -1) -> str:
|
|
"""Read from a container's standard error buffer.
|
|
|
|
Args:
|
|
name: Name of the container
|
|
n: Number of bytes to read; -1 means read all available
|
|
|
|
Returns:
|
|
Data read from stderr
|
|
|
|
Raises:
|
|
KeyError: If container name not found
|
|
docker.errors.APIError: If read operation fails
|
|
"""
|
|
if name not in self.containers:
|
|
raise KeyError(f"Container {name} not found")
|
|
|
|
try:
|
|
container = self.containers[name]
|
|
logs = container.logs(stdout=False, stderr=True, tail=n if n > 0 else 'all')
|
|
return logs.decode('utf-8')
|
|
except Exception as e:
|
|
self.logger.error(f"Error reading container {name} stderr: {str(e)}")
|
|
raise
|
|
|
|
def wait_container(self, name: str, timeout: int) -> Tuple[int, str]:
|
|
"""Wait for a container to finish execution.
|
|
|
|
Args:
|
|
name: Name of the container to wait for
|
|
timeout: Time to wait in milliseconds
|
|
|
|
Returns:
|
|
Tuple of (exit_code, output)
|
|
|
|
Raises:
|
|
KeyError: If container name not found
|
|
docker.errors.APIError: If wait operation fails
|
|
TimeoutError: If container doesn't finish within timeout
|
|
"""
|
|
if name not in self.containers:
|
|
raise KeyError(f"Container {name} not found")
|
|
|
|
try:
|
|
container = self.containers[name]
|
|
result = container.wait(timeout=timeout/1000) # Convert ms to seconds
|
|
logs = container.logs().decode('utf-8')
|
|
return (result['StatusCode'], logs)
|
|
except Exception as e:
|
|
self.logger.error(f"Error waiting for container {name}: {str(e)}")
|
|
raise
|
|
|
|
def get_container_status(self, name: str) -> Dict[str, any]:
|
|
"""Get current status information for a container.
|
|
|
|
Args:
|
|
name: Name of the container
|
|
|
|
Returns:
|
|
Dictionary with container status information
|
|
|
|
Raises:
|
|
KeyError: If container name not found
|
|
docker.errors.APIError: If status check fails
|
|
"""
|
|
if name not in self.containers:
|
|
raise KeyError(f"Container {name} not found")
|
|
|
|
try:
|
|
container = self.containers[name]
|
|
container.reload() # Refresh container info
|
|
|
|
return {
|
|
'name': name,
|
|
'status': container.status,
|
|
'started_at': container.attrs['State']['StartedAt'],
|
|
'exit_code': container.attrs['State']['ExitCode'],
|
|
'error': container.attrs['State']['Error'],
|
|
'stdout_size': len(container.logs(stdout=True, stderr=False)),
|
|
'stderr_size': len(container.logs(stdout=False, stderr=True))
|
|
}
|
|
except Exception as e:
|
|
self.logger.error(f"Error getting status for container {name}: {str(e)}")
|
|
raise |