Start work on docker module

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"]

requirements.txt

@@ -1,2 +1,3 @@
-transformers
+docker
 torch
+transformers

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

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

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):
         """

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:
     """
@@ -45,3 +46,24 @@ def split_response(response: str, valid_elements: set) -> InferenceResult:
     reasoning = response[:split_point].strip()
     actions = response[split_point:].strip()
     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

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()

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"<test_tag>{main_context}</test_tag>")
+        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}<test_tag>{main_context}</test_tag>")