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

wip docker module

Browse Source
This commit is contained in:
Niels Geens
2024-10-24 17:26:34 +02:00
parent 4bbe636f0e
commit a2f9e9c1c5
5 changed files with 139 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
Dockerfile
View File

@@ -1,15 +1,4 @@
FROM huggingface/transformers-pytorch-gpu AS requirements
WORKDIR /root
COPY requirements.txt /root/requirements.txt
COPY ./sia/ /root/sia/
RUN pip3 install -r requirements.txt
FROM requirements AS test
COPY ./test/ /root/test/
RUN mkdir -p /root/model
CMD ["python3", "-m", "unittest", "discover", "-v", "-p", "*test.py", "-v"]
FROM requirements
RUN apt update
RUN apt install -y \
apt-transport-https ca-certificates \
@@ -18,4 +7,17 @@ RUN curl -fsSL https://download.docker.com/linux/ubuntu/gpg | apt-key add -
RUN add-apt-repository "deb [arch=amd64] https://download.docker.com/linux/ubuntu bionic stable"
RUN apt update
RUN apt install -y docker-ce
COPY requirements.txt /requirements.txt
RUN pip3 install -r /requirements.txt
RUN rm -rf /requirements.txt
FROM requirements AS test
COPY ./ /root/sia/
WORKDIR /root/sia/
RUN mkdir -p /root/model
CMD ["python3", "-m", "unittest", "discover", "-v", "-p", "*test.py", "-v"]
FROM requirements
COPY ./ /root/sia/
WORKDIR /root/sia/
CMD ["python3", "-m", "sia"]

3
run.sh
View File

@@ -21,8 +21,7 @@ docker run \
--privileged \
-v /$(pwd)/model/:/root/model/ \
$TAG \
-c "/bin/bash"
#/etc/init.d/docker start
-c "bash"
# Clean up image
[ ! -z "$TAG" ] && docker rmi $TAG

54
sia/docker_module.py
View File

@@ -1,19 +1,30 @@
from docker.models.containers import Container
from typing import Dict, List, Optional, Tuple
import docker
from docker.models.containers import Container
import logging
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):
self.client = docker.from_env()
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.logger = logging.getLogger(__name__)
def start_container(
self,
image: str,
@@ -69,7 +80,9 @@ class DockerModule:
port_bindings[container_port] = host_port
try:
container = self.client.containers.run(
self.client.images.pull('busybox:latest')
# Create container with tty and open stdin
container = self.client.containers.create(
image,
command=cmd if cmd else None,
name=name,
@@ -77,21 +90,36 @@ class DockerModule:
volumes=volume_binds,
ports=port_bindings,
environment=environment,
tty=True,
stdin_open=True
)
# Attach socket before starting container
socket = container.attach_socket(params={
'stdin': 1,
'stdout': 1,
'stderr': 1,
'stream': 1
})
# 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 ""
else:
# For long-running containers, store reference
# For long-running containers, store references
self.containers[name] = container
self.container_sockets[name] = socket
return name
except Exception as e:
@@ -112,6 +140,12 @@ class DockerModule:
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()
@@ -131,14 +165,12 @@ class DockerModule:
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")
if name not in self.containers or name not in self.container_sockets:
raise KeyError(f"Container {name} not found or socket not initialized")
try:
container = self.containers[name]
socket = container.attach_socket(params={'stdin': 1, 'stream': 1})
socket = self.container_sockets[name]
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
@@ -251,4 +283,4 @@ class DockerModule:
}
except Exception as e:
self.logger.error(f"Error getting status for container {name}: {str(e)}")
raise
raise

2
test.sh
View File

@@ -18,8 +18,10 @@ docker build \
docker run \
--rm \
--gpus=all \
--privileged \
-v /$(pwd)/model/:/root/model/ \
$TAG
# Clean up image
[ ! -z "$TAG" ] && docker rmi $TAG

167
test/docker_module_test.py
View File

@@ -1,6 +1,5 @@
import unittest
import docker
from docker.models.containers import Container
import time
from sia.docker_module import DockerModule
class DockerModuleTest(unittest.TestCase):
@@ -14,7 +13,7 @@ class DockerModuleTest(unittest.TestCase):
docker_module = DockerModule()
# Test parameters
image = "sia:latest"
image = "busybox:latest"
timeout = 1000
command = "echo"
arguments = ["Hello World"]
@@ -33,8 +32,8 @@ class DockerModuleTest(unittest.TestCase):
"""Test starting a long-running container."""
docker_module = DockerModule()
name = "test-container"
image = "sia:latest"
name = "test_start_container_long_running"
image = "busybox:latest"
container_name = docker_module.start_container(
image=image,
@@ -44,23 +43,24 @@ class DockerModuleTest(unittest.TestCase):
# Verify container was started and stored
self.assertEqual(container_name, name)
self.assertIn(name, docker_module.containers)
self.assertEqual(docker_module.containers[name], self.mock_container)
def test_start_container_with_volumes(self):
"""Test starting a container with volume mappings."""
docker_module = DockerModule()
volumes = {"./pdf": "/pdf"}
arguments = ["sh", "-c", "echo 'Hello World' > /write_here/test_start_container_with_volumes"]
volumes = {"/tmp": "/write_here"}
output = docker_module.start_container(
image="sia:latest",
image="busybox:latest",
timeout=1000,
arguments=arguments,
volumes=volumes
)
# Verify volume bindings were configured correctly
expected_binds = {
"./pdf": {"bind": "/pdf", "mode": "rw"}
}
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."""
@@ -68,16 +68,16 @@ class DockerModuleTest(unittest.TestCase):
# Test missing required parameters
with self.assertRaises(ValueError):
docker_module.start_container(image="test:latest")
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-container"
name = "test_stop_container"
docker_module.start_container(
image="test:latest",
image="busybox:latest",
name=name
)
@@ -85,8 +85,6 @@ class DockerModuleTest(unittest.TestCase):
docker_module.stop_container(name)
# Verify container was stopped and removed
self.mock_container.stop.assert_called_once()
self.mock_container.remove.assert_called_once()
self.assertNotIn(name, docker_module.containers)
def test_stop_nonexistent_container(self):
@@ -99,88 +97,83 @@ class DockerModuleTest(unittest.TestCase):
def test_container_io_operations(self):
"""Test container I/O operations."""
docker_module = DockerModule()
arguments = [
"sh",
"-c",
"read input; echo \"$input\" | tr '[:lower:]' '[:upper:]'; echo \"$input\" | tr '[:upper:]' '[:lower:]' >&2"
]
# Start a container
name = "test-container"
name = "test_container_io_operations"
docker_module.start_container(
image="test:latest",
name=name
image="busybox:latest",
name=name,
arguments=arguments
)
# Test stdin write
test_input = "test input"
mock_socket = Mock()
mock_socket._sock = Mock()
self.mock_container.attach_socket.return_value = mock_socket
test_input = "Test Input\n"
docker_module.write_container_stdin(name, test_input)
mock_socket._sock.send.assert_called_once_with(test_input.encode('utf-8'))
time.sleep(1) # for processing
# Test stdout read
stdout = docker_module.read_container_stdout(name)
self.mock_container.logs.assert_called_with(
stdout=True,
stderr=False,
tail='all'
)
self.assertEqual(stdout, "Test output")
print(f"stdout: {stdout}", flush=True)
self.assertEqual(stdout, "test input\n")
# Test stderr read
stderr = docker_module.read_container_stderr(name)
self.mock_container.logs.assert_called_with(
stdout=False,
stderr=True,
tail='all'
)
self.assertEqual(stderr, "Test output")
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()
# 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()
if __name__ == '__main__':
unittest.main()