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 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 update
RUN apt install -y \ RUN apt install -y \
apt-transport-https ca-certificates \ 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 add-apt-repository "deb [arch=amd64] https://download.docker.com/linux/ubuntu bionic stable"
RUN apt update RUN apt update
RUN apt install -y docker-ce 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"] CMD ["python3", "-m", "sia"]

3
run.sh
View File

@@ -21,8 +21,7 @@ docker run \
--privileged \ --privileged \
-v /$(pwd)/model/:/root/model/ \ -v /$(pwd)/model/:/root/model/ \
$TAG \ $TAG \
-c "/bin/bash" -c "bash"
#/etc/init.d/docker start
# Clean up image # Clean up image
[ ! -z "$TAG" ] && docker rmi $TAG [ ! -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 from typing import Dict, List, Optional, Tuple
import docker import docker
from docker.models.containers import Container
import logging import logging
import subprocess
import time
class DockerModule: class DockerModule:
""" """
Module for handling Docker container operations in SIA. Module for handling Docker container operations in SIA.
Manages container lifecycle, I/O operations, and status monitoring. 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): 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.containers: Dict[str, Container] = {}
self.container_sockets: Dict[str, docker.APIClient.AttachSocket] = {}
self.logger = logging.getLogger(__name__) self.logger = logging.getLogger(__name__)
def start_container( def start_container(
self, self,
image: str, image: str,
@@ -69,7 +80,9 @@ class DockerModule:
port_bindings[container_port] = host_port port_bindings[container_port] = host_port
try: 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, image,
command=cmd if cmd else None, command=cmd if cmd else None,
name=name, name=name,
@@ -77,21 +90,36 @@ class DockerModule:
volumes=volume_binds, volumes=volume_binds,
ports=port_bindings, ports=port_bindings,
environment=environment, 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: if timeout >= 0:
# For short-lived containers, wait for completion # For short-lived containers, wait for completion
try: try:
container.wait(timeout=timeout/1000) # Convert ms to seconds container.wait(timeout=timeout/1000) # Convert ms to seconds
output = container.logs().decode('utf-8') output = container.logs().decode('utf-8')
socket.close()
container.remove(force=True) container.remove(force=True)
return output return output
except docker.errors.NotFound: except docker.errors.NotFound:
self.logger.warning(f"Container was removed before timeout") self.logger.warning(f"Container was removed before timeout")
return "" return ""
else: else:
# For long-running containers, store reference # For long-running containers, store references
self.containers[name] = container self.containers[name] = container
self.container_sockets[name] = socket
return name return name
except Exception as e: except Exception as e:
@@ -112,6 +140,12 @@ class DockerModule:
raise KeyError(f"Container {name} not found") raise KeyError(f"Container {name} not found")
try: 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 = self.containers[name]
container.stop() container.stop()
container.remove() container.remove()
@@ -131,14 +165,12 @@ class DockerModule:
KeyError: If container name not found KeyError: If container name not found
docker.errors.APIError: If write operation fails docker.errors.APIError: If write operation fails
""" """
if name not in self.containers: if name not in self.containers or name not in self.container_sockets:
raise KeyError(f"Container {name} not found") raise KeyError(f"Container {name} not found or socket not initialized")
try: try:
container = self.containers[name] socket = self.container_sockets[name]
socket = container.attach_socket(params={'stdin': 1, 'stream': 1})
socket._sock.send(data.encode('utf-8')) socket._sock.send(data.encode('utf-8'))
socket.close()
except Exception as e: except Exception as e:
self.logger.error(f"Error writing to container {name} stdin: {str(e)}") self.logger.error(f"Error writing to container {name} stdin: {str(e)}")
raise raise
@@ -251,4 +283,4 @@ class DockerModule:
} }
except Exception as e: except Exception as e:
self.logger.error(f"Error getting status for container {name}: {str(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 \ docker run \
--rm \ --rm \
--gpus=all \ --gpus=all \
--privileged \
-v /$(pwd)/model/:/root/model/ \ -v /$(pwd)/model/:/root/model/ \
$TAG $TAG
# Clean up image # Clean up image
[ ! -z "$TAG" ] && docker rmi $TAG [ ! -z "$TAG" ] && docker rmi $TAG

167
test/docker_module_test.py
View File

@@ -1,6 +1,5 @@
import unittest import unittest
import docker import time
from docker.models.containers import Container
from sia.docker_module import DockerModule from sia.docker_module import DockerModule
class DockerModuleTest(unittest.TestCase): class DockerModuleTest(unittest.TestCase):
@@ -14,7 +13,7 @@ class DockerModuleTest(unittest.TestCase):
docker_module = DockerModule() docker_module = DockerModule()
# Test parameters # Test parameters
image = "sia:latest" image = "busybox:latest"
timeout = 1000 timeout = 1000
command = "echo" command = "echo"
arguments = ["Hello World"] arguments = ["Hello World"]
@@ -33,8 +32,8 @@ class DockerModuleTest(unittest.TestCase):
"""Test starting a long-running container.""" """Test starting a long-running container."""
docker_module = DockerModule() docker_module = DockerModule()
name = "test-container" name = "test_start_container_long_running"
image = "sia:latest" image = "busybox:latest"
container_name = docker_module.start_container( container_name = docker_module.start_container(
image=image, image=image,
@@ -44,23 +43,24 @@ class DockerModuleTest(unittest.TestCase):
# Verify container was started and stored # Verify container was started and stored
self.assertEqual(container_name, name) self.assertEqual(container_name, name)
self.assertIn(name, docker_module.containers) self.assertIn(name, docker_module.containers)
self.assertEqual(docker_module.containers[name], self.mock_container)
def test_start_container_with_volumes(self): def test_start_container_with_volumes(self):
"""Test starting a container with volume mappings.""" """Test starting a container with volume mappings."""
docker_module = DockerModule() 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( output = docker_module.start_container(
image="sia:latest", image="busybox:latest",
timeout=1000, timeout=1000,
arguments=arguments,
volumes=volumes volumes=volumes
) )
print(f"output: {output}", flush=True)
# Verify volume bindings were configured correctly # Verify output was written to volume
expected_binds = { with open("/tmp/test_start_container_with_volumes", "r") as f:
"./pdf": {"bind": "/pdf", "mode": "rw"} self.assertEqual(f.read(), "Hello World\n")
}
def test_invalid_container_start(self): def test_invalid_container_start(self):
"""Test error handling for invalid container start.""" """Test error handling for invalid container start."""
@@ -68,16 +68,16 @@ class DockerModuleTest(unittest.TestCase):
# Test missing required parameters # Test missing required parameters
with self.assertRaises(ValueError): with self.assertRaises(ValueError):
docker_module.start_container(image="test:latest") docker_module.start_container(image="busybox:latest")
def test_stop_container(self): def test_stop_container(self):
"""Test stopping a container.""" """Test stopping a container."""
docker_module = DockerModule() docker_module = DockerModule()
# Start a container first # Start a container first
name = "test-container" name = "test_stop_container"
docker_module.start_container( docker_module.start_container(
image="test:latest", image="busybox:latest",
name=name name=name
) )
@@ -85,8 +85,6 @@ class DockerModuleTest(unittest.TestCase):
docker_module.stop_container(name) docker_module.stop_container(name)
# Verify container was stopped and removed # 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) self.assertNotIn(name, docker_module.containers)
def test_stop_nonexistent_container(self): def test_stop_nonexistent_container(self):
@@ -99,88 +97,83 @@ class DockerModuleTest(unittest.TestCase):
def test_container_io_operations(self): def test_container_io_operations(self):
"""Test container I/O operations.""" """Test container I/O operations."""
docker_module = DockerModule() docker_module = DockerModule()
arguments = [
"sh",
"-c",
"read input; echo \"$input\" | tr '[:lower:]' '[:upper:]'; echo \"$input\" | tr '[:upper:]' '[:lower:]' >&2"
]
# Start a container # Start a container
name = "test-container" name = "test_container_io_operations"
docker_module.start_container( docker_module.start_container(
image="test:latest", image="busybox:latest",
name=name name=name,
arguments=arguments
) )
# Test stdin write # Test stdin write
test_input = "test input" test_input = "Test Input\n"
mock_socket = Mock()
mock_socket._sock = Mock()
self.mock_container.attach_socket.return_value = mock_socket
docker_module.write_container_stdin(name, test_input) 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 # Test stdout read
stdout = docker_module.read_container_stdout(name) stdout = docker_module.read_container_stdout(name)
self.mock_container.logs.assert_called_with( print(f"stdout: {stdout}", flush=True)
stdout=True, self.assertEqual(stdout, "test input\n")
stderr=False,
tail='all'
)
self.assertEqual(stdout, "Test output")
# Test stderr read # Test stderr read
stderr = docker_module.read_container_stderr(name) stderr = docker_module.read_container_stderr(name)
self.mock_container.logs.assert_called_with( print(f"stderr: {stderr}", flush=True)
stdout=False, self.assertEqual(stderr, "TEST INPUT\n")
stderr=True,
tail='all'
)
self.assertEqual(stderr, "Test output")
def test_wait_container(self): # def test_wait_container(self):
"""Test waiting for container completion.""" # """Test waiting for container completion."""
docker_module = DockerModule() # docker_module = DockerModule()
#
# Start a container # # Start a container
name = "test-container" # name = "test-container"
docker_module.start_container( # docker_module.start_container(
image="test:latest", # image="test:latest",
name=name # name=name
) # )
#
# Test wait operation # # Test wait operation
exit_code, output = docker_module.wait_container(name, timeout=1000) # exit_code, output = docker_module.wait_container(name, timeout=1000)
#
# Verify wait was called with correct timeout # # Verify wait was called with correct timeout
self.mock_container.wait.assert_called_once_with(timeout=1.0) # self.mock_container.wait.assert_called_once_with(timeout=1.0)
self.assertEqual(exit_code, 0) # self.assertEqual(exit_code, 0)
self.assertEqual(output, "Test output") # self.assertEqual(output, "Test output")
#
def test_get_container_status(self): # def test_get_container_status(self):
"""Test getting container status information.""" # """Test getting container status information."""
docker_module = DockerModule() # docker_module = DockerModule()
#
# Start a container # # Start a container
name = "test-container" # name = "test-container"
docker_module.start_container( # docker_module.start_container(
image="test:latest", # image="test:latest",
name=name # name=name
) # )
#
# Configure mock logs sizes # # Configure mock logs sizes
self.mock_container.logs.side_effect = [b"stdout", b"stderr"] # self.mock_container.logs.side_effect = [b"stdout", b"stderr"]
#
# Get status # # Get status
status = docker_module.get_container_status(name) # status = docker_module.get_container_status(name)
#
# Verify status information # # Verify status information
self.assertEqual(status['name'], name) # self.assertEqual(status['name'], name)
self.assertEqual(status['status'], "running") # self.assertEqual(status['status'], "running")
self.assertEqual(status['started_at'], "2024-10-23T10:00:00Z") # self.assertEqual(status['started_at'], "2024-10-23T10:00:00Z")
self.assertEqual(status['exit_code'], 0) # self.assertEqual(status['exit_code'], 0)
self.assertEqual(status['error'], "") # self.assertEqual(status['error'], "")
self.assertEqual(status['stdout_size'], 6) # len(b"stdout") # self.assertEqual(status['stdout_size'], 6) # len(b"stdout")
self.assertEqual(status['stderr_size'], 6) # len(b"stderr") # self.assertEqual(status['stderr_size'], 6) # len(b"stderr")
#
# Verify container was reloaded # # Verify container was reloaded
self.mock_container.reload.assert_called_once() # self.mock_container.reload.assert_called_once()
if __name__ == '__main__': if __name__ == '__main__':
unittest.main() unittest.main()