Create inference/server.py

inference/server.py

@@ -0,0 +1,405 @@
+#!/usr/bin/env python3
+"""
+Helion-2.5-Rnd Inference Server
+High-performance inference server with vLLM backend
+"""
+
+import argparse
+import asyncio
+import json
+import logging
+import os
+import time
+from typing import AsyncGenerator, Dict, List, Optional, Union
+
+import torch
+import uvicorn
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, StreamingResponse
+from pydantic import BaseModel, Field
+from vllm import AsyncEngineArgs, AsyncLLMEngine, SamplingParams
+from vllm.utils import random_uuid
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+class ChatMessage(BaseModel):
+    """Chat message format"""
+    role: str = Field(..., description="Role: system, user, or assistant")
+    content: str = Field(..., description="Message content")
+
+
+class ChatCompletionRequest(BaseModel):
+    """Chat completion request format"""
+    model: str = Field(default="DeepXR/Helion-2.5-Rnd")
+    messages: List[ChatMessage]
+    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
+    top_p: float = Field(default=0.9, ge=0.0, le=1.0)
+    top_k: int = Field(default=50, ge=0)
+    max_tokens: int = Field(default=4096, ge=1)
+    stream: bool = Field(default=False)
+    stop: Optional[List[str]] = None
+    presence_penalty: float = Field(default=0.0, ge=-2.0, le=2.0)
+    frequency_penalty: float = Field(default=0.0, ge=-2.0, le=2.0)
+    repetition_penalty: float = Field(default=1.1, ge=1.0, le=2.0)
+    n: int = Field(default=1, ge=1, le=10)
+    logprobs: Optional[int] = None
+    echo: bool = Field(default=False)
+
+
+class CompletionRequest(BaseModel):
+    """Text completion request format"""
+    model: str = Field(default="DeepXR/Helion-2.5-Rnd")
+    prompt: Union[str, List[str]]
+    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
+    top_p: float = Field(default=0.9, ge=0.0, le=1.0)
+    max_tokens: int = Field(default=4096, ge=1)
+    stream: bool = Field(default=False)
+    stop: Optional[List[str]] = None
+    n: int = Field(default=1, ge=1, le=10)
+
+
+class HelionInferenceServer:
+    """Main inference server class"""
+    
+    def __init__(
+        self,
+        model_path: str,
+        tensor_parallel_size: int = 2,
+        max_model_len: int = 131072,
+        gpu_memory_utilization: float = 0.95,
+        dtype: str = "bfloat16"
+    ):
+        self.model_path = model_path
+        self.model_name = "DeepXR/Helion-2.5-Rnd"
+        
+        # Initialize vLLM engine
+        engine_args = AsyncEngineArgs(
+            model=model_path,
+            tensor_parallel_size=tensor_parallel_size,
+            max_model_len=max_model_len,
+            gpu_memory_utilization=gpu_memory_utilization,
+            dtype=dtype,
+            trust_remote_code=True,
+            enforce_eager=False,
+            disable_log_stats=False,
+        )
+        
+        logger.info(f"Initializing Helion-2.5-Rnd from {model_path}")
+        self.engine = AsyncLLMEngine.from_engine_args(engine_args)
+        logger.info("Engine initialized successfully")
+        
+        # Statistics
+        self.request_count = 0
+        self.start_time = time.time()
+    
+    def format_chat_prompt(self, messages: List[ChatMessage]) -> str:
+        """Format chat messages into prompt"""
+        formatted = ""
+        for msg in messages:
+            formatted += f"<|im_start|>{msg.role}\n{msg.content}<|im_end|>\n"
+        formatted += "<|im_start|>assistant\n"
+        return formatted
+    
+    async def generate(
+        self,
+        prompt: str,
+        sampling_params: SamplingParams,
+        request_id: str
+    ) -> AsyncGenerator[str, None]:
+        """Generate text streaming"""
+        results_generator = self.engine.generate(
+            prompt,
+            sampling_params,
+            request_id
+        )
+        
+        async for request_output in results_generator:
+            text = request_output.outputs[0].text
+            yield text
+    
+    async def chat_completion(
+        self,
+        request: ChatCompletionRequest
+    ) -> Union[Dict, AsyncGenerator]:
+        """Handle chat completion request"""
+        request_id = f"helion-{random_uuid()}"
+        self.request_count += 1
+        
+        # Format prompt
+        prompt = self.format_chat_prompt(request.messages)
+        
+        # Create sampling parameters
+        sampling_params = SamplingParams(
+            temperature=request.temperature,
+            top_p=request.top_p,
+            top_k=request.top_k,
+            max_tokens=request.max_tokens,
+            stop=request.stop or ["<|im_end|>", "<|endoftext|>"],
+            presence_penalty=request.presence_penalty,
+            frequency_penalty=request.frequency_penalty,
+            repetition_penalty=request.repetition_penalty,
+            n=request.n,
+            logprobs=request.logprobs,
+        )
+        
+        if request.stream:
+            return self._stream_chat_completion(
+                prompt,
+                sampling_params,
+                request_id,
+                request.model
+            )
+        else:
+            return await self._complete_chat_completion(
+                prompt,
+                sampling_params,
+                request_id,
+                request.model
+            )
+    
+    async def _complete_chat_completion(
+        self,
+        prompt: str,
+        sampling_params: SamplingParams,
+        request_id: str,
+        model: str
+    ) -> Dict:
+        """Non-streaming chat completion"""
+        final_output = None
+        async for request_output in self.engine.generate(
+            prompt, sampling_params, request_id
+        ):
+            final_output = request_output
+        
+        if final_output is None:
+            raise HTTPException(status_code=500, detail="Generation failed")
+        
+        choice = {
+            "index": 0,
+            "message": {
+                "role": "assistant",
+                "content": final_output.outputs[0].text
+            },
+            "finish_reason": final_output.outputs[0].finish_reason
+        }
+        
+        return {
+            "id": request_id,
+            "object": "chat.completion",
+            "created": int(time.time()),
+            "model": model,
+            "choices": [choice],
+            "usage": {
+                "prompt_tokens": len(final_output.prompt_token_ids),
+                "completion_tokens": len(final_output.outputs[0].token_ids),
+                "total_tokens": len(final_output.prompt_token_ids) + len(final_output.outputs[0].token_ids)
+            }
+        }
+    
+    async def _stream_chat_completion(
+        self,
+        prompt: str,
+        sampling_params: SamplingParams,
+        request_id: str,
+        model: str
+    ) -> AsyncGenerator:
+        """Streaming chat completion"""
+        async def generate():
+            previous_text = ""
+            async for request_output in self.engine.generate(
+                prompt, sampling_params, request_id
+            ):
+                text = request_output.outputs[0].text
+                delta = text[len(previous_text):]
+                previous_text = text
+                
+                chunk = {
+                    "id": request_id,
+                    "object": "chat.completion.chunk",
+                    "created": int(time.time()),
+                    "model": model,
+                    "choices": [{
+                        "index": 0,
+                        "delta": {"content": delta},
+                        "finish_reason": None
+                    }]
+                }
+                yield f"data: {json.dumps(chunk)}\n\n"
+            
+            # Final chunk
+            final_chunk = {
+                "id": request_id,
+                "object": "chat.completion.chunk",
+                "created": int(time.time()),
+                "model": model,
+                "choices": [{
+                    "index": 0,
+                    "delta": {},
+                    "finish_reason": "stop"
+                }]
+            }
+            yield f"data: {json.dumps(final_chunk)}\n\n"
+            yield "data: [DONE]\n\n"
+        
+        return generate()
+
+
+# Initialize FastAPI app
+app = FastAPI(
+    title="Helion-2.5-Rnd Inference API",
+    description="Advanced language model inference server",
+    version="2.5.0-rnd"
+)
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Global server instance
+server: Optional[HelionInferenceServer] = None
+
+
+@app.on_event("startup")
+async def startup_event():
+    """Initialize the model on startup"""
+    global server
+    
+    model_path = os.getenv("MODEL_PATH", "/models/helion")
+    tensor_parallel = int(os.getenv("TENSOR_PARALLEL_SIZE", "2"))
+    max_len = int(os.getenv("MAX_MODEL_LEN", "131072"))
+    gpu_util = float(os.getenv("GPU_MEMORY_UTILIZATION", "0.95"))
+    
+    server = HelionInferenceServer(
+        model_path=model_path,
+        tensor_parallel_size=tensor_parallel,
+        max_model_len=max_len,
+        gpu_memory_utilization=gpu_util
+    )
+    logger.info("Helion-2.5-Rnd server started successfully")
+
+
+@app.get("/")
+async def root():
+    """Root endpoint"""
+    return {
+        "model": "DeepXR/Helion-2.5-Rnd",
+        "version": "2.5.0-rnd",
+        "status": "ready",
+        "type": "research"
+    }
+
+
+@app.get("/health")
+async def health():
+    """Health check endpoint"""
+    if server is None:
+        raise HTTPException(status_code=503, detail="Server not initialized")
+    
+    return {
+        "status": "healthy",
+        "model": server.model_name,
+        "requests_served": server.request_count,
+        "uptime_seconds": int(time.time() - server.start_time)
+    }
+
+
+@app.get("/v1/models")
+async def list_models():
+    """List available models"""
+    return {
+        "object": "list",
+        "data": [{
+            "id": "DeepXR/Helion-2.5-Rnd",
+            "object": "model",
+            "created": int(time.time()),
+            "owned_by": "DeepXR"
+        }]
+    }
+
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: ChatCompletionRequest):
+    """Chat completion endpoint"""
+    if server is None:
+        raise HTTPException(status_code=503, detail="Server not initialized")
+    
+    try:
+        result = await server.chat_completion(request)
+        
+        if request.stream:
+            return StreamingResponse(
+                result,
+                media_type="text/event-stream"
+            )
+        else:
+            return JSONResponse(content=result)
+    
+    except Exception as e:
+        logger.error(f"Error in chat completion: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/v1/completions")
+async def completions(request: CompletionRequest):
+    """Text completion endpoint"""
+    if server is None:
+        raise HTTPException(status_code=503, detail="Server not initialized")
+    
+    # Convert to chat format
+    messages = [ChatMessage(role="user", content=request.prompt)]
+    chat_request = ChatCompletionRequest(
+        model=request.model,
+        messages=messages,
+        temperature=request.temperature,
+        top_p=request.top_p,
+        max_tokens=request.max_tokens,
+        stream=request.stream,
+        stop=request.stop,
+        n=request.n
+    )
+    
+    return await chat_completions(chat_request)
+
+
+def main():
+    """Main entry point"""
+    parser = argparse.ArgumentParser(description="Helion-2.5-Rnd Inference Server")
+    parser.add_argument("--model", type=str, default="/models/helion")
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--tensor-parallel-size", type=int, default=2)
+    parser.add_argument("--max-model-len", type=int, default=131072)
+    parser.add_argument("--gpu-memory-utilization", type=float, default=0.95)
+    
+    args = parser.parse_args()
+    
+    # Set environment variables
+    os.environ["MODEL_PATH"] = args.model
+    os.environ["TENSOR_PARALLEL_SIZE"] = str(args.tensor_parallel_size)
+    os.environ["MAX_MODEL_LEN"] = str(args.max_model_len)
+    os.environ["GPU_MEMORY_UTILIZATION"] = str(args.gpu_memory_utilization)
+    
+    # Run server
+    uvicorn.run(
+        app,
+        host=args.host,
+        port=args.port,
+        log_level="info",
+        access_log=True
+    )
+
+
+if __name__ == "__main__":
+    main()