dgx-harness-engineering / scripts /v64 /devstral_chat_proxy_v2.py
koreallmdev's picture
Upload anonymized DGX harness engineering export V1.4 (part 3)
0236599 verified
Raw
History Blame Contribute Delete
5.89 kB
#!/usr/bin/env python3
import json, os, time, uuid, urllib.request, urllib.error
from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler
from transformers import AutoTokenizer
MODEL_DIR=os.environ["DEVSTRAL_MODEL"]
BACKEND=os.environ.get("DEVSTRAL_BACKEND","http://192.0.2.11:8031").rstrip("/")
PUBLIC_MODEL=os.environ.get("DEVSTRAL_PUBLIC_MODEL","devstral-small2-24b")
HOST=os.environ.get("PROXY_HOST","192.0.2.11")
PORT=int(os.environ.get("PROXY_PORT","8025"))
TOKENIZER=AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
def normalize_content(content):
if isinstance(content,str):
return content
if isinstance(content,list):
parts=[]
for item in content:
if isinstance(item,dict) and item.get("type") in ("text","input_text"):
parts.append(str(item.get("text","")))
elif isinstance(item,str):
parts.append(item)
return "\n".join(x for x in parts if x)
if content is None:
return ""
return str(content)
def render_messages(messages):
normalized=[]
for m in messages:
normalized.append({
"role": str(m.get("role","user")),
"content": normalize_content(m.get("content",""))
})
return TOKENIZER.apply_chat_template(
normalized,
tokenize=False,
add_generation_prompt=True
)
def backend_post(path,payload,timeout=600):
req=urllib.request.Request(
BACKEND+path,
data=json.dumps(payload).encode("utf-8"),
headers={"Content-Type":"application/json"},
method="POST"
)
with urllib.request.urlopen(req,timeout=timeout) as r:
return r.status,json.loads(r.read().decode("utf-8"))
def backend_health():
try:
with urllib.request.urlopen(BACKEND+"/health",timeout=3) as r:
return r.status==200
except Exception:
return False
class Handler(BaseHTTPRequestHandler):
server_version="DevstralChatProxyV2/1.0"
def log_message(self,fmt,*args):
print("%s - %s" % (self.address_string(), fmt%args), flush=True)
def send_json(self,status,obj):
raw=json.dumps(obj,ensure_ascii=False).encode("utf-8")
self.send_response(status)
self.send_header("Content-Type","application/json; charset=utf-8")
self.send_header("Content-Length",str(len(raw)))
self.end_headers()
self.wfile.write(raw)
def do_GET(self):
if self.path=="/health":
if backend_health():
self.send_json(200,{"status":"ok","backend":"ok"})
else:
self.send_json(503,{"status":"degraded","backend":"down"})
return
if self.path=="/v1/models":
self.send_json(200,{
"object":"list",
"data":[{
"id":PUBLIC_MODEL,
"object":"model",
"created":int(time.time()),
"owned_by":"local-vllm-proxy",
"root":MODEL_DIR,
"parent":None
}]
})
return
self.send_json(404,{"error":{"message":"not found"}})
def do_POST(self):
try:
length=int(self.headers.get("Content-Length","0"))
body=json.loads(self.rfile.read(length) or b"{}")
except Exception as e:
self.send_json(400,{"error":{"message":"invalid json: "+repr(e)}})
return
if self.path=="/v1/chat/completions":
try:
messages=body.get("messages") or []
prompt=render_messages(messages)
payload={
"model":PUBLIC_MODEL,
"prompt":prompt,
"max_tokens":int(body.get("max_tokens",body.get("max_completion_tokens",256))),
"temperature":float(body.get("temperature",0.0)),
"top_p":float(body.get("top_p",1.0)),
"stream":False
}
if body.get("stop") is not None:
payload["stop"]=body["stop"]
status,out=backend_post("/v1/completions",payload)
text=((out.get("choices") or [{}])[0].get("text") or "")
resp={
"id":"chatcmpl-"+uuid.uuid4().hex[:24],
"object":"chat.completion",
"created":int(time.time()),
"model":PUBLIC_MODEL,
"choices":[{
"index":0,
"message":{"role":"assistant","content":text},
"finish_reason":((out.get("choices") or [{}])[0].get("finish_reason"))
}],
"usage":out.get("usage",{})
}
self.send_json(200 if status==200 else status,resp)
except urllib.error.HTTPError as e:
raw=e.read().decode(errors="replace")
self.send_json(e.code,{"error":{"message":"backend http error","body":raw}})
except Exception as e:
self.send_json(500,{"error":{"message":"proxy render/backend failure","detail":repr(e)}})
return
if self.path=="/v1/completions":
try:
status,out=backend_post("/v1/completions",body)
self.send_json(status,out)
except urllib.error.HTTPError as e:
self.send_json(e.code,{"error":{"message":e.read().decode(errors="replace")}})
except Exception as e:
self.send_json(500,{"error":{"message":repr(e)}})
return
self.send_json(404,{"error":{"message":"not found"}})
print(f"DEVSTRAL_CHAT_PROXY_V2_START host={HOST} port={PORT} backend={BACKEND}",flush=True)
ThreadingHTTPServer((HOST,PORT),Handler).serve_forever()