inference_endpoint (#2)

- added endpoint scripts (1259b94c50e3fa4f31bce132de4f3ec7596aaa5b)


Co-authored-by: Ivan Moshkov <[email protected]>

Dockerfile

@@ -0,0 +1,26 @@
+FROM igitman/nemo-skills-vllm:0.6.0 as base
+
+# Install NeMo-Skills and dependencies
+RUN git clone https://github.com/NVIDIA/NeMo-Skills \
+    && cd NeMo-Skills \
+    && pip install --ignore-installed blinker \
+    && pip install -e . \
+    && pip install -r requirements/code_execution.txt
+
+# Ensure python is available
+RUN ln -s /usr/bin/python3 /usr/bin/python
+
+# Copy our custom files
+COPY handler.py server.py /usr/local/endpoint/
+
+# Expose port 80
+EXPOSE 80
+
+# Copy and set up entrypoint script
+COPY entrypoint.sh /usr/local/endpoint/
+RUN chmod +x /usr/local/endpoint/entrypoint.sh
+
+# Set working directory
+WORKDIR /usr/local/endpoint
+
+ENTRYPOINT ["/usr/local/endpoint/entrypoint.sh"]

entrypoint.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+set -e
+
+# Default environment variables
+export MODEL_PATH=${MODEL_PATH:-"/repository"}
+
+echo "Starting NeMo Skills inference endpoint..."
+echo "Model path: $MODEL_PATH"
+
+# Function to handle cleanup on exit
+cleanup() {
+    echo "Cleaning up processes..."
+    kill $(jobs -p) 2>/dev/null || true
+    wait
+}
+trap cleanup EXIT
+
+# Start the model server in the background
+echo "Starting model server..."
+ns start_server \
+   --model="$MODEL_PATH" \
+   --server_gpus=2 \
+   --server_type=vllm \
+   --with_sandbox &
+
+# Start the HTTP endpoint
+echo "Starting HTTP endpoint on port 80..."
+python /usr/local/endpoint/server.py &
+
+# Wait for both processes
+echo "Both servers started. Waiting..."
+wait

handler.py

@@ -0,0 +1,139 @@
+import os
+import logging
+import traceback
+from typing import Dict, List, Any
+
+from nemo_skills.inference.server.code_execution_model import get_code_execution_model
+from nemo_skills.code_execution.sandbox import get_sandbox
+from nemo_skills.prompt.utils import get_prompt
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class EndpointHandler:
+    """Custom endpoint handler for NeMo Skills code execution inference."""
+    
+    def __init__(self):
+        """
+        Initialize the handler with the model and prompt configurations.
+        """
+        self.model = None
+        self.prompt = None
+        self.initialized = False
+        
+        # Configuration
+        self.prompt_config_path = os.getenv("PROMPT_CONFIG_PATH", "generic/math")
+        self.prompt_template_path = os.getenv("PROMPT_TEMPLATE_PATH", "openmath-instruct")
+        
+    def _initialize_components(self):
+        """Initialize the model, sandbox, and prompt components lazily."""
+        if self.initialized:
+            return
+            
+        try:
+            logger.info("Initializing sandbox...")
+            sandbox = get_sandbox(sandbox_type="local")
+            
+            logger.info("Initializing code execution model...")
+            self.model = get_code_execution_model(
+                server_type="vllm",
+                sandbox=sandbox,
+                host="127.0.0.1",
+                port=5000
+            )
+            
+            logger.info("Initializing prompt...")
+            if self.prompt_config_path:
+                self.prompt = get_prompt(
+                    prompt_config=self.prompt_config_path,
+                    prompt_template=self.prompt_template_path
+                )
+            
+            self.initialized = True
+            logger.info("All components initialized successfully")
+            
+        except Exception as e:
+            logger.warning(f"Failed to initialize the model")
+    
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Process inference requests.
+        
+        Args:
+            data: Dictionary containing the request data
+                Expected keys:
+                - inputs: str or list of str - the input prompts/problems
+                - parameters: dict (optional) - generation parameters
+                
+        Returns:
+            List of dictionaries containing the generated responses
+        """
+        try:
+            # Initialize components if not already done
+            self._initialize_components()
+            
+            # Extract inputs and parameters
+            inputs = data.get("inputs", "")
+            parameters = data.get("parameters", {})
+            
+            # Handle both single string and list of strings
+            if isinstance(inputs, str):
+                prompts = [inputs]
+            elif isinstance(inputs, list):
+                prompts = inputs
+            else:
+                raise ValueError("inputs must be a string or list of strings")
+            
+            # If we have a prompt template configured, format the inputs
+            if self.prompt is not None:
+                formatted_prompts = []
+                for prompt_text in prompts:
+                    formatted_prompt = self.prompt.fill({"problem": prompt_text, "total_code_executions": 8})
+                    formatted_prompts.append(formatted_prompt)
+                prompts = formatted_prompts
+            
+            # Get code execution arguments from prompt if available
+            extra_generate_params = {}
+            if self.prompt is not None:
+                extra_generate_params = self.prompt.get_code_execution_args()
+            
+            # Set default generation parameters
+            generation_params = {
+                "tokens_to_generate": 12000,
+                "temperature": 0.0,
+                "top_p": 0.95,
+                "top_k": 0,
+                "repetition_penalty": 1.0,
+                "random_seed": 0,
+            }
+            
+            # Update with provided parameters
+            generation_params.update(parameters)
+            generation_params.update(extra_generate_params)
+            
+            logger.info(f"Processing {len(prompts)} prompt(s)")
+            
+            # Generate responses
+            outputs = self.model.generate(
+                prompts=prompts,
+                **generation_params
+            )
+            
+            # Format outputs
+            results = []
+            for output in outputs:
+                result = {
+                    "generated_text": output.get("generation", ""),
+                    "code_rounds_executed": output.get("code_rounds_executed", 0),
+                }
+                results.append(result)
+            
+            logger.info(f"Successfully processed {len(results)} request(s)")
+            return results
+            
+        except Exception as e:
+            logger.error(f"Error processing request: {str(e)}")
+            logger.error(traceback.format_exc())
+            return [{"error": str(e), "generated_text": ""}]

server.py

@@ -0,0 +1,77 @@
+import json
+import logging
+from http.server import HTTPServer, BaseHTTPRequestHandler
+
+from handler import EndpointHandler
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Initialize the handler
+handler = EndpointHandler()
+
+
+class RequestHandler(BaseHTTPRequestHandler):
+    def do_POST(self):
+        try:
+            content_length = int(self.headers['Content-Length'])
+            post_data = self.rfile.read(content_length)
+            data = json.loads(post_data.decode('utf-8'))
+            
+            logger.info(f'Received request with {len(data.get("inputs", []))} inputs')
+            
+            # Process the request
+            result = handler(data)
+            
+            # Send response
+            self.send_response(200)
+            self.send_header('Content-Type', 'application/json')
+            self.end_headers()
+            self.wfile.write(json.dumps(result).encode('utf-8'))
+            
+        except Exception as e:
+            logger.error(f'Error processing request: {str(e)}')
+            self.send_response(500)
+            self.send_header('Content-Type', 'application/json')
+            self.end_headers()
+            error_response = [{'error': str(e), 'generated_text': ''}]
+            self.wfile.write(json.dumps(error_response).encode('utf-8'))
+    
+    def do_GET(self):
+        if self.path == '/health':
+            # Trigger initialisation if needed but don't block.
+            if not handler.initialized:
+                try:
+                    handler._initialize_components()
+                except Exception as e:
+                    logger.error(f'Initialization failed during health check: {str(e)}')
+
+            is_ready = handler.initialized
+            health_response = {
+                'status': 'healthy' if is_ready else 'unhealthy',
+                'model_ready': is_ready
+            }
+
+            try:
+                self.send_response(200 if is_ready else 503)
+                self.send_header('Content-Type', 'application/json')
+                self.end_headers()
+                self.wfile.write(json.dumps(health_response).encode('utf-8'))
+            except BrokenPipeError:
+                # Client disconnected before we replied – safe to ignore.
+                pass
+            return
+        else:
+            self.send_response(404)
+            self.end_headers()
+
+    def log_message(self, format, *args):
+        # Suppress default HTTP server logs to keep output clean
+        pass
+
+
+if __name__ == "__main__":
+    server = HTTPServer(('0.0.0.0', 80), RequestHandler)
+    logger.info('HTTP server started on port 80')
+    server.serve_forever()