From 4bbe636f0ebe3f9f10dc9b10e57ce2eaa0333722 Mon Sep 17 00:00:00 2001 From: geens Date: Thu, 24 Oct 2024 08:39:31 +0200 Subject: [PATCH] Start work on docker module --- Dockerfile | 8 ++ requirements.txt | 5 +- run.sh | 28 ++++ sia/docker_module.py | 254 +++++++++++++++++++++++++++++++++++++ sia/llm_engine.py | 34 ++--- sia/util.py | 30 ++++- test/docker_module_test.py | 186 +++++++++++++++++++++++++++ test/llm_engine_test.py | 17 +-- 8 files changed, 532 insertions(+), 30 deletions(-) create mode 100755 run.sh create mode 100644 sia/docker_module.py create mode 100644 test/docker_module_test.py diff --git a/Dockerfile b/Dockerfile index ab047eb..fa0b466 100644 --- a/Dockerfile +++ b/Dockerfile @@ -10,4 +10,12 @@ 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 \ + curl software-properties-common +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 CMD ["python3", "-m", "sia"] \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 4803a9a..93b42c2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ -transformers -torch \ No newline at end of file +docker +torch +transformers \ No newline at end of file diff --git a/run.sh b/run.sh new file mode 100755 index 0000000..fc24a3f --- /dev/null +++ b/run.sh @@ -0,0 +1,28 @@ +#!/bin/bash + +# Continue on error +set -e + +# Build with progress output and capture the tag +TAG=$( \ +docker build \ + . \ + 2>&1 | tee /dev/tty | grep "writing image" | cut -d' ' -f4 \ +) + +# Exit if tag is empty +[ -z "$TAG" ] && exit 1 + +# Run tests +docker run \ + --rm \ + -ti \ + --gpus=all \ + --privileged \ + -v /$(pwd)/model/:/root/model/ \ + $TAG \ + -c "/bin/bash" +#/etc/init.d/docker start + +# Clean up image +[ ! -z "$TAG" ] && docker rmi $TAG \ No newline at end of file diff --git a/sia/docker_module.py b/sia/docker_module.py new file mode 100644 index 0000000..5fa59b1 --- /dev/null +++ b/sia/docker_module.py @@ -0,0 +1,254 @@ +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 \ No newline at end of file diff --git a/sia/llm_engine.py b/sia/llm_engine.py index 30c727d..fd47143 100644 --- a/sia/llm_engine.py +++ b/sia/llm_engine.py @@ -1,8 +1,9 @@ -from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextStreamer +from threading import Thread + +from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer import torch from . import util -from .inference_result import InferenceResult class LlmEngine: def __init__(self, model_path: str): @@ -34,33 +35,26 @@ class LlmEngine: self.tokenizer.pad_token_id = self.tokenizer.eos_token_id if model.config.pad_token_id is None: model.config.pad_token_id = model.config.eos_token_id - streamer = TextStreamer( - self.tokenizer, - skip_prompt=True - ) self.pipeline = pipeline( "text-generation", model=model, tokenizer=self.tokenizer, torch_dtype=torch.bfloat16, device_map="auto", - streamer=streamer, return_full_text=False, ) - def infer(self, system_prompt: str, main_context: str, action_schema: str) -> InferenceResult: + def infer(self, system_prompt: str, main_context: str) -> TextIteratorStreamer: """ Run inference using the system prompt and main context, while validating actions against the provided XML schema. Args: system_prompt: The system prompt string main_context: The main context string after templating - action_schema: XML schema to validate the generated actions Returns: - InferenceResult: Tuple containing reasoning and actions that validate against the schema + TextIteratorStreamer: An iterator that yields the generated text. """ - valid_elements = util.get_valid_root_elements(action_schema) messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": main_context} @@ -68,11 +62,19 @@ class LlmEngine: prompt = self.tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) - outputs = self.pipeline(prompt, max_new_tokens=120, do_sample=True) - generated_text = outputs[0]["generated_text"] - #response = generated_text.split("<|start_header_id|>assistant<|end_header_id|>",1)[1].strip() - result = util.split_response(generated_text, valid_elements) - return result + streamer = TextIteratorStreamer( + self.tokenizer, + skip_prompt=True + ) + pipeline_kwargs = dict( + text_inputs=prompt, + do_sample=True, + max_new_tokens=1024, + streamer=streamer + ) + thread = Thread(target=self.pipeline, kwargs=pipeline_kwargs) + thread.start() + return util.stop_before_value(streamer, '<|eot_id|>') def finetune(self, dataset_paths: list, output_dir: str): """ diff --git a/sia/util.py b/sia/util.py index 0704f16..182cd1d 100644 --- a/sia/util.py +++ b/sia/util.py @@ -1,7 +1,8 @@ -from .inference_result import InferenceResult - -import xml.etree.ElementTree as ET +from typing import Iterator, TypeVar import re +import xml.etree.ElementTree as ET + +from .inference_result import InferenceResult def get_valid_root_elements(schema: str) -> set: """ @@ -44,4 +45,25 @@ def split_response(response: str, valid_elements: set) -> InferenceResult: split_point = last_match.start() reasoning = response[:split_point].strip() actions = response[split_point:].strip() - return InferenceResult(reasoning, actions) \ No newline at end of file + return InferenceResult(reasoning, actions) + +def stop_before_value(iterator: Iterator[str], stop_value: str) -> Iterator[str]: + """ + Creates an iterator that yields values from the input iterator + until it encounters the stop_value (exclusive). + + Args: + iterator: The source iterator + stop_value: The value to stop before + + Yields: + Values from the iterator until stop_value is encountered + If stop_value is part of an item, yields the part before stop_value + """ + for item in iterator: + if stop_value in item: + split_point = item.index(stop_value) + if split_point > 0: + yield item[:split_point] + break + yield item \ No newline at end of file diff --git a/test/docker_module_test.py b/test/docker_module_test.py new file mode 100644 index 0000000..4662388 --- /dev/null +++ b/test/docker_module_test.py @@ -0,0 +1,186 @@ +import unittest +import docker +from docker.models.containers import Container +from sia.docker_module import DockerModule + +class DockerModuleTest(unittest.TestCase): + def test_initialization(self): + """Test DockerModule initialization.""" + docker_module = DockerModule() + self.assertIsInstance(docker_module, DockerModule) + + def test_start_container_short_lived(self): + """Test starting a short-lived container.""" + docker_module = DockerModule() + + # Test parameters + image = "sia: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() + + name = "test-container" + image = "sia: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) + 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"} + output = docker_module.start_container( + image="sia:latest", + timeout=1000, + volumes=volumes + ) + + # Verify volume bindings were configured correctly + expected_binds = { + "./pdf": {"bind": "/pdf", "mode": "rw"} + } + + def test_invalid_container_start(self): + """Test error handling for invalid container start.""" + docker_module = DockerModule() + + # Test missing required parameters + with self.assertRaises(ValueError): + docker_module.start_container(image="test:latest") + + def test_stop_container(self): + """Test stopping a container.""" + docker_module = DockerModule() + + # Start a container first + name = "test-container" + docker_module.start_container( + image="test:latest", + name=name + ) + + # Stop the container + 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): + """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() + + # Start a container + name = "test-container" + docker_module.start_container( + image="test:latest", + name=name + ) + + # Test stdin write + test_input = "test input" + mock_socket = Mock() + mock_socket._sock = Mock() + self.mock_container.attach_socket.return_value = mock_socket + + docker_module.write_container_stdin(name, test_input) + mock_socket._sock.send.assert_called_once_with(test_input.encode('utf-8')) + + # 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") + + # 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") + + 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() diff --git a/test/llm_engine_test.py b/test/llm_engine_test.py index f7943ff..93b89c6 100644 --- a/test/llm_engine_test.py +++ b/test/llm_engine_test.py @@ -1,14 +1,15 @@ import unittest +from itertools import tee + from . import test_data from . import test_util -from sia.llm_engine import LlmEngine, InferenceResult +from sia.llm_engine import LlmEngine class LlmEngineTest(unittest.TestCase): def setUp(self): self.model_path = "/root/model" - self.llm_engine = LlmEngine(self.model_path) def test_initialization(self): llm_engine = LlmEngine(self.model_path) @@ -17,9 +18,9 @@ class LlmEngineTest(unittest.TestCase): def test_infer(self): main_context = "This is a test" llm_engine = LlmEngine(self.model_path) - result = llm_engine.infer(test_data.echo_system_prompt, main_context, test_data.echo_action_schema) - self.assertIsInstance(result, InferenceResult) - self.assertIsInstance(result.reasoning, str) - self.assertIsInstance(result.actions, str) - self.assertEqual(result.reasoning, main_context) - self.assertEqual(result.actions, f"{main_context}") \ No newline at end of file + tokens = llm_engine.infer(test_data.echo_system_prompt, main_context) + print_tokens, result_tokens = tee(tokens) + for token in print_tokens: + print(token, end="", flush=True) + result = ''.join(result_tokens) + self.assertEqual(result, f"{main_context}{main_context}") \ No newline at end of file