280 lines
9.8 KiB
Python
280 lines
9.8 KiB
Python
from docker.models.containers import Container
|
|
from typing import Dict, List, Optional, Tuple
|
|
import docker
|
|
import logging
|
|
import requests
|
|
import subprocess
|
|
import time
|
|
|
|
class DockerModule:
|
|
"""
|
|
Module for handling Docker container operations in SIA.
|
|
Manages container lifecycle, I/O operations, and status monitoring.
|
|
"""
|
|
|
|
def _start_docker_engine():
|
|
"""
|
|
Start the Docker engine
|
|
"""
|
|
subprocess.run(['/etc/init.d/docker', 'start'])
|
|
time.sleep(1)
|
|
|
|
def __init__(self):
|
|
DockerModule._start_docker_engine()
|
|
self.client = docker.DockerClient(base_url='unix://var/run/docker.sock')
|
|
self.containers: Dict[str, (Container, docker.APIClient.AttachSocket)] = {}
|
|
self.logger = logging.getLogger(__name__)
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exec_type, exec_value, traceback):
|
|
del self
|
|
|
|
def __del__(self):
|
|
"""Cleanup all resources before shutdown."""
|
|
try:
|
|
# Stop and remove all containers
|
|
for (container, socket) in self.containers.values():
|
|
self._wait_and_stop(container, socket, 0)
|
|
# Close the Docker client
|
|
if hasattr(self, 'client'):
|
|
self.client.close()
|
|
except Exception as e:
|
|
self.logger.error(f"Error during cleanup: {str(e)}")
|
|
|
|
def _wait_and_stop(self, container: Container, socket, timeout: int) -> Tuple[int, str]:
|
|
"""Helper method for waiting on long- and short-lived containers."""
|
|
status_code = None
|
|
if timeout > 0:
|
|
try:
|
|
result = container.wait(timeout=timeout/1000)
|
|
status_code = result.get('StatusCode')
|
|
except requests.exceptions.ReadTimeout as e:
|
|
container.stop()
|
|
status_code = -1
|
|
except requests.exceptions.ConnectionError as e:
|
|
container.stop()
|
|
status_code = -1
|
|
else:
|
|
container.stop()
|
|
status_code = -1
|
|
try:
|
|
socket.close()
|
|
except Exception as e:
|
|
self.logger.error(f"Error closing socket: {str(e)}")
|
|
output = container.logs(stdout=True, stderr=True).decode('utf-8')
|
|
try:
|
|
container.remove(force=True)
|
|
except Exception as e:
|
|
self.logger.error(f"Error removing container: {str(e)}")
|
|
return status_code, output
|
|
|
|
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")
|
|
cmd = []
|
|
if command:
|
|
cmd.append(command)
|
|
if arguments:
|
|
cmd.extend(arguments)
|
|
volume_binds = {}
|
|
if volumes:
|
|
for host_path, container_path in volumes.items():
|
|
volume_binds[host_path] = {'bind': container_path, 'mode': 'rw'}
|
|
port_bindings = {}
|
|
if ports:
|
|
for host_port, container_port in ports.items():
|
|
port_bindings[container_port] = host_port
|
|
self.client.images.pull(image)
|
|
container = self.client.containers.create(
|
|
image,
|
|
command=cmd if cmd else None,
|
|
name=name,
|
|
detach=True,
|
|
volumes=volume_binds,
|
|
ports=port_bindings,
|
|
environment=environment,
|
|
tty=False,
|
|
stdin_open=True
|
|
)
|
|
socket = container.attach_socket(params={
|
|
'stdin': 1,
|
|
'stdout': 1,
|
|
'stderr': 1,
|
|
'stream': 1,
|
|
'demux': True
|
|
})
|
|
container.start()
|
|
if timeout >= 0:
|
|
# For short-lived containers
|
|
return self._wait_and_stop(container, socket, timeout)[1]
|
|
else:
|
|
# For long-running containers
|
|
self.containers[name] = (container, socket)
|
|
return name
|
|
|
|
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:
|
|
_, socket = self.containers[name]
|
|
socket._sock.send(data.encode('utf-8'))
|
|
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, socket = 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, socket = 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")
|
|
|
|
container, socket = self.containers.pop(name)
|
|
return self._wait_and_stop(container, socket, timeout)
|
|
|
|
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, socket = 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
|
|
|
|
def get_all_container_statuses(self) -> List[Dict[str, any]]:
|
|
"""Get current status information for all containers.
|
|
Returns:
|
|
List of dictionaries with container status information
|
|
"""
|
|
statuses = []
|
|
for name, (container, socket) in self.containers.items():
|
|
statuses.append(self.get_container_status(name)) |