llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Docker module basic functionality

Browse Source
This commit is contained in:
Niels Geens
2024-10-25 14:37:56 +02:00
parent a2f9e9c1c5
commit 016187fa21
2 changed files with 176 additions and 249 deletions
Download Patch File Download Diff File Expand all files Collapse all files

141
sia/docker_module.py
View File

@@ -2,6 +2,7 @@ from docker.models.containers import Container
from typing import Dict, List, Optional, Tuple
import docker
import logging
import requests
import subprocess
import time
@@ -21,10 +22,54 @@ class DockerModule:
def __init__(self):
DockerModule._start_docker_engine()
self.client = docker.DockerClient(base_url='unix://var/run/docker.sock')
self.containers: Dict[str, Container] = {}
self.container_sockets: Dict[str, docker.APIClient.AttachSocket] = {}
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,
@@ -59,29 +104,20 @@ class DockerModule:
# 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:
self.client.images.pull('busybox:latest')
# Create container with tty and open stdin
self.client.images.pull(image)
container = self.client.containers.create(
image,
command=cmd if cmd else None,
@@ -90,70 +126,25 @@ class DockerModule:
volumes=volume_binds,
ports=port_bindings,
environment=environment,
tty=True,
tty=False,
stdin_open=True
)
# Attach socket before starting container
socket = container.attach_socket(params={
'stdin': 1,
'stdout': 1,
'stderr': 1,
'stream': 1
'stream': 1,
'demux': True
})
# Start the container
container.start()
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')
socket.close()
container.remove(force=True)
return output
except docker.errors.NotFound:
self.logger.warning(f"Container was removed before timeout")
return ""
# For short-lived containers
return self._wait_and_stop(container, socket, timeout)[1]
else:
# For long-running containers, store references
self.containers[name] = container
self.container_sockets[name] = socket
# For long-running containers
self.containers[name] = (container, socket)
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:
# Close socket if it exists
if name in self.container_sockets:
self.container_sockets[name].close()
del self.container_sockets[name]
# Stop and remove container
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.
@@ -165,11 +156,11 @@ class DockerModule:
KeyError: If container name not found
docker.errors.APIError: If write operation fails
"""
if name not in self.containers or name not in self.container_sockets:
raise KeyError(f"Container {name} not found or socket not initialized")
if name not in self.containers:
raise KeyError(f"Container {name} not found")
try:
socket = self.container_sockets[name]
_, 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)}")
@@ -193,7 +184,7 @@ class DockerModule:
raise KeyError(f"Container {name} not found")
try:
container = self.containers[name]
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:
@@ -218,7 +209,7 @@ class DockerModule:
raise KeyError(f"Container {name} not found")
try:
container = self.containers[name]
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:
@@ -243,14 +234,8 @@ class DockerModule:
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
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.
@@ -269,7 +254,7 @@ class DockerModule:
raise KeyError(f"Container {name} not found")
try:
container = self.containers[name]
container, socket = self.containers[name]
container.reload() # Refresh container info
return {

166
test/docker_module_test.py
View File

@@ -5,175 +5,117 @@ from sia.docker_module import DockerModule
class DockerModuleTest(unittest.TestCase):
def test_initialization(self):
"""Test DockerModule initialization."""
docker_module = DockerModule()
with DockerModule() as docker_module:
self.assertIsInstance(docker_module, DockerModule)
def test_start_container_short_lived(self):
"""Test starting a short-lived container."""
docker_module = DockerModule()
# Test parameters
with DockerModule() as docker_module:
image = "busybox:latest"
timeout = 1000
command = "echo"
arguments = ["Hello World"]
output = docker_module.start_container(
image=image,
timeout=timeout,
command=command,
arguments=arguments
)
# Verify output was returned
self.assertEqual(output, "Hello World\n")
def test_start_container_long_running(self):
"""Test starting a long-running container."""
docker_module = DockerModule()
with DockerModule() as docker_module:
name = "test_start_container_long_running"
image = "busybox:latest"
container_name = docker_module.start_container(
image=image,
name=name
)
# Verify container was started and stored
self.assertEqual(container_name, name)
self.assertIn(name, docker_module.containers)
docker_module.wait_container(container_name, 0)
self.assertNotIn(name, docker_module.containers)
def test_start_container_with_volumes(self):
"""Test starting a container with volume mappings."""
docker_module = DockerModule()
with DockerModule() as docker_module:
arguments = ["sh", "-c", "echo 'Hello World' > /write_here/test_start_container_with_volumes"]
volumes = {"/tmp": "/write_here"}
output = docker_module.start_container(
docker_module.start_container(
image="busybox:latest",
timeout=1000,
arguments=arguments,
volumes=volumes
)
print(f"output: {output}", flush=True)
# Verify output was written to volume
with open("/tmp/test_start_container_with_volumes", "r") as f:
self.assertEqual(f.read(), "Hello World\n")
def test_invalid_container_start(self):
"""Test error handling for invalid container start."""
docker_module = DockerModule()
# Test missing required parameters
with DockerModule() as docker_module:
with self.assertRaises(ValueError):
docker_module.start_container(image="busybox:latest")
def test_stop_container(self):
"""Test stopping a container."""
docker_module = DockerModule()
# Start a container first
name = "test_stop_container"
docker_module.start_container(
image="busybox:latest",
name=name
)
# Stop the container
docker_module.stop_container(name)
# Verify container was stopped and removed
self.assertNotIn(name, docker_module.containers)
def test_stop_nonexistent_container(self):
"""Test stopping a container that doesn't exist."""
docker_module = DockerModule()
with self.assertRaises(KeyError):
docker_module.stop_container("nonexistent")
def test_container_io_operations(self):
"""Test container I/O operations."""
docker_module = DockerModule()
with DockerModule() as docker_module:
arguments = [
"sh",
"-c",
"read input; echo \"$input\" | tr '[:lower:]' '[:upper:]'; echo \"$input\" | tr '[:upper:]' '[:lower:]' >&2"
]
# Start a container
name = "test_container_io_operations"
container_name = docker_module.start_container(
image="busybox:latest",
name=name,
arguments=arguments
)
test_input = "Test Input\n"
docker_module.write_container_stdin(name, test_input)
time.sleep(1) # for processing
stdout = docker_module.read_container_stdout(name)
self.assertEqual(stdout, "TEST INPUT\n")
stderr = docker_module.read_container_stderr(name)
self.assertEqual(stderr, "test input\n")
docker_module.wait_container(container_name, 0)
def test_wait_container(self):
"""Test waiting for container completion."""
with DockerModule() as docker_module:
name = "test_wait_container"
docker_module.start_container(
image="busybox:latest",
name=name
)
exit_code, output = docker_module.wait_container(name, timeout=1000)
self.assertEqual(exit_code, -1)
self.assertEqual(output, "")
def test_get_container_status(self):
"""Test getting container status information."""
with DockerModule() as docker_module:
arguments = [
"sh",
"-c",
"echo 'standard output'; echo 'standard error' >&2; sleep 3"
]
name = "test_get_container_status"
docker_module.start_container(
image="busybox:latest",
name=name,
arguments=arguments
)
# Test stdin write
test_input = "Test Input\n"
docker_module.write_container_stdin(name, test_input)
time.sleep(1) # for processing
# Test stdout read
stdout = docker_module.read_container_stdout(name)
print(f"stdout: {stdout}", flush=True)
self.assertEqual(stdout, "test input\n")
# Test stderr read
stderr = docker_module.read_container_stderr(name)
print(f"stderr: {stderr}", flush=True)
self.assertEqual(stderr, "TEST INPUT\n")
# def test_wait_container(self):
# """Test waiting for container completion."""
# docker_module = DockerModule()
#
# # Start a container
# name = "test-container"
# docker_module.start_container(
# image="test:latest",
# name=name
# )
#
# # Test wait operation
# exit_code, output = docker_module.wait_container(name, timeout=1000)
#
# # Verify wait was called with correct timeout
# self.mock_container.wait.assert_called_once_with(timeout=1.0)
# self.assertEqual(exit_code, 0)
# self.assertEqual(output, "Test output")
#
# def test_get_container_status(self):
# """Test getting container status information."""
# docker_module = DockerModule()
#
# # Start a container
# name = "test-container"
# docker_module.start_container(
# image="test:latest",
# name=name
# )
#
# # Configure mock logs sizes
# self.mock_container.logs.side_effect = [b"stdout", b"stderr"]
#
# # Get status
# status = docker_module.get_container_status(name)
#
# # Verify status information
# self.assertEqual(status['name'], name)
# self.assertEqual(status['status'], "running")
# self.assertEqual(status['started_at'], "2024-10-23T10:00:00Z")
# self.assertEqual(status['exit_code'], 0)
# self.assertEqual(status['error'], "")
# self.assertEqual(status['stdout_size'], 6) # len(b"stdout")
# self.assertEqual(status['stderr_size'], 6) # len(b"stderr")
#
# # Verify container was reloaded
# self.mock_container.reload.assert_called_once()
time.sleep(1)
status = docker_module.get_container_status(name)
self.assertEqual(status['name'], name)
self.assertEqual(status['status'], "running")
self.assertGreater(len(status['started_at']), 0)
self.assertEqual(status['exit_code'], 0)
self.assertEqual(status['error'], "")
self.assertEqual(status['stdout_size'], 16)
self.assertEqual(status['stderr_size'], 15)
docker_module.wait_container(name, 0)
if __name__ == '__main__':
unittest.main()