Start work on docker module

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