Make scheduler OpenAI agent compatible

Author: gufengc

src/backend/main.py

@@ -9,6 +9,7 @@
 from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 
+from backend.server.openai_compat import openai_error_response, openai_models_payload
 from backend.server.request_handler import RequestHandler
 from backend.server.scheduler_manage import SchedulerManage
 from backend.server.server_args import parse_args
@@ -83,6 +84,18 @@ async def model_list():
     )
 
 
+@app.get("/v1/models")
+async def openai_v1_models():
+    model_name = None
+    if scheduler_manage is not None:
+        try:
+            model_name = scheduler_manage.get_model_name()
+        except Exception as e:
+            logger.debug(f"Unable to get scheduler model name: {e}")
+
+    return JSONResponse(content=openai_models_payload(model_name), status_code=200)
+
+
 @app.post("/scheduler/init")
 async def scheduler_init(raw_request: Request):
     request_data = await raw_request.json()
@@ -182,7 +195,23 @@ async def cluster_status_json() -> JSONResponse:
 
 @app.post("/v1/chat/completions")
 async def openai_v1_chat_completions(raw_request: Request):
-    request_data = await raw_request.json()
+    try:
+        request_data = await raw_request.json()
+    except Exception:
+        return openai_error_response(
+            "Invalid request body",
+            status_code=400,
+            err_type="invalid_request_error",
+            code="invalid_request_error",
+        )
+    if not isinstance(request_data, dict):
+        return openai_error_response(
+            "Request body must be a JSON object",
+            status_code=400,
+            err_type="invalid_request_error",
+            code="invalid_request_error",
+        )
+
     request_id = uuid.uuid4()
     received_ts = time.time()
     return await request_handler.v1_chat_completions(request_data, request_id, received_ts)

src/backend/server/openai_compat.py

@@ -0,0 +1,95 @@
+import base64
+import json
+from typing import Any, Dict, Optional, Tuple
+
+from fastapi.responses import JSONResponse
+
+PARALLAX_HTTP_RESPONSE_ENVELOPE = "__parallax_http_response__"
+
+
+def openai_error_payload(
+    message: str,
+    *,
+    err_type: str = "server_error",
+    param: Optional[str] = None,
+    code: Optional[str] = None,
+) -> Dict[str, Dict[str, Any]]:
+    return {
+        "error": {
+            "message": message,
+            "type": err_type,
+            "param": param,
+            "code": code or err_type,
+        }
+    }
+
+
+def openai_error_response(
+    message: str,
+    *,
+    status_code: int,
+    err_type: str = "server_error",
+    param: Optional[str] = None,
+    code: Optional[str] = None,
+) -> JSONResponse:
+    return JSONResponse(
+        content=openai_error_payload(
+            message,
+            err_type=err_type,
+            param=param,
+            code=code,
+        ),
+        status_code=status_code,
+    )
+
+
+def openai_models_payload(model_name: Optional[str]) -> Dict[str, Any]:
+    models = []
+    if model_name:
+        models.append(
+            {
+                "id": model_name,
+                "object": "model",
+                "created": 0,
+                "owned_by": "parallax",
+            }
+        )
+    return {"object": "list", "data": models}
+
+
+def encode_http_response_envelope(
+    *,
+    status_code: int,
+    content_type: Optional[str],
+    body: bytes,
+) -> bytes:
+    envelope = {
+        PARALLAX_HTTP_RESPONSE_ENVELOPE: True,
+        "status_code": int(status_code),
+        "content_type": content_type or "application/json",
+        "body_base64": base64.b64encode(body).decode("ascii"),
+    }
+    return json.dumps(envelope, separators=(",", ":")).encode("utf-8")
+
+
+def decode_http_response_envelope(content: bytes) -> Optional[Tuple[int, str, bytes]]:
+    try:
+        envelope = json.loads(content.decode("utf-8"))
+    except (UnicodeDecodeError, json.JSONDecodeError):
+        return None
+
+    if not isinstance(envelope, dict) or envelope.get(PARALLAX_HTTP_RESPONSE_ENVELOPE) is not True:
+        return None
+
+    body_base64 = envelope.get("body_base64")
+    if not isinstance(body_base64, str):
+        return None
+
+    try:
+        body = base64.b64decode(body_base64)
+    except ValueError:
+        return None
+
+    status_code = int(envelope.get("status_code", 200))
+    content_type = envelope.get("content_type") or "application/json"
+    return status_code, str(content_type), body

src/backend/server/request_handler.py

@@ -3,10 +3,14 @@
 from typing import Dict, List, Optional
 
 import aiohttp
-from fastapi.responses import JSONResponse, Response, StreamingResponse
+from fastapi.responses import Response, StreamingResponse
 from starlette.concurrency import iterate_in_threadpool
 
 from backend.server.constants import NODE_STATUS_AVAILABLE
+from backend.server.openai_compat import (
+    decode_http_response_envelope,
+    openai_error_response,
+)
 from parallax_utils.logging_config import get_logger
 from parallax_utils.request_metrics import get_request_metrics
 
@@ -100,9 +104,11 @@ async def _forward_request(self, request_data: Dict, request_id: str, received_t
             self.scheduler_manage is None
             or not self.scheduler_manage.get_schedule_status() == NODE_STATUS_AVAILABLE
         ):
-            return JSONResponse(
-                content={"error": "Server is not ready"},
-                status_code=500,
+            return openai_error_response(
+                "Server is not ready",
+                status_code=503,
+                err_type="server_unavailable",
+                code="server_not_ready",
             )
 
         # Try to get a success response
@@ -119,16 +125,20 @@ async def _forward_request(self, request_data: Dict, request_id: str, received_t
                     )
                 except Exception as e:
                     logger.exception(f"get_routing_table error: {e}")
-                    return JSONResponse(
-                        content={"error": "Get routing table error"},
+                    return openai_error_response(
+                        "Get routing table error",
                         status_code=500,
+                        err_type="server_error",
+                        code="routing_table_error",
                     )
 
                 # None -> scheduler has not set yet; treat as hard error (no waiting here)
                 if routing_table is None:
-                    return JSONResponse(
-                        content={"error": "Routing pipelines not ready"},
+                    return openai_error_response(
+                        "Routing pipelines not ready",
                         status_code=503,
+                        err_type="server_unavailable",
+                        code="routing_not_ready",
                     )
 
                 # Non-empty -> proceed
@@ -143,9 +153,11 @@ async def _forward_request(self, request_data: Dict, request_id: str, received_t
 
             # If still empty after retries, return 429 Too Many Requests
             if routing_table is not None and len(routing_table) == 0:
-                return JSONResponse(
-                    content={"error": "All pipelines are busy or not ready. Please retry later."},
+                return openai_error_response(
+                    "All pipelines are busy or not ready. Please retry later.",
                     status_code=429,
+                    err_type="rate_limit_error",
+                    code="rate_limit_exceeded",
                 )
 
             backend_request = self._prepare_backend_request(
@@ -203,19 +215,33 @@ async def stream_generator():
                     return resp
                 else:
                     response = stub.chat_completion(backend_request)
-                    content = (await anext(iterate_in_threadpool(response))).decode()
+                    content = await anext(iterate_in_threadpool(response))
+                    decoded_response = decode_http_response_envelope(content)
+                    if decoded_response is None:
+                        status_code = 200
+                        content_type = "application/json"
+                        body = content
+                    else:
+                        status_code, content_type, body = decoded_response
                     logger.debug(f"Non-stream response completed for {request_id}")
-                    return Response(content=content, media_type="application/json")
+                    return Response(
+                        content=body,
+                        status_code=status_code,
+                        headers={"content-type": content_type},
+                        media_type=None,
+                    )
             except Exception as e:
                 forward_attempts += 1
                 if forward_attempts < self.MAX_FORWARD_RETRY:
                     # small async delay before re-forwarding
                     await asyncio.sleep(self.FORWARD_DELAY_SEC)
                 logger.warning(f"Error in _forward_request: {e}. Retry attemps {forward_attempts}")
 
-        return JSONResponse(
-            content={"error": "Internal server error"},
-            status_code=500,
+        return openai_error_response(
+            "Downstream request failed",
+            status_code=502,
+            err_type="upstream_error",
+            code="upstream_error",
         )
 
     async def v1_chat_completions(self, request_data: Dict, request_id: str, received_ts: int):

src/parallax/p2p/server.py

@@ -22,6 +22,7 @@
 import zmq
 from lattica import ConnectionHandler, Lattica, rpc_method, rpc_stream, rpc_stream_iter
 
+from backend.server.openai_compat import encode_http_response_envelope
 from backend.server.rpc_connection_handler import RPCConnectionHandler
 from parallax.p2p.proto import forward_pb2
 from parallax.p2p.utils import AsyncWorker
@@ -203,10 +204,21 @@ def chat_completion(
                     response = client.post(
                         f"http://localhost:{self.http_port}/v1/chat/completions", json=request
                     )
-                    yield response.content
+                    yield encode_http_response_envelope(
+                        status_code=response.status_code,
+                        content_type=response.headers.get("content-type"),
+                        body=response.content,
+                    )
         except Exception as e:
             logger.exception(f"Error in chat completion: {e}")
-            yield b"internal server error"
+            yield encode_http_response_envelope(
+                status_code=502,
+                content_type="application/json",
+                body=(
+                    b'{"error":{"message":"Internal server error",'
+                    b'"type":"upstream_error","param":null,"code":"upstream_error"}}'
+                ),
+            )
 
 
 def check_and_run_weight_refit(gradient_server, message):