From c8d5706956380ad2bc51b8d79ac6c7e151d8b857 Mon Sep 17 00:00:00 2001 From: fabolous005 Date: Tue, 27 Oct 2026 13:52:04 +0100 Subject: [PATCH] prain: skeleton --- .gitignore | 5 +++ prain/__init__.py | 1 + pyproject.toml | 19 +++++++++ scripts/greedy.py | 28 ++++++++++++++ tests/__init__.py | 1 + tests/helpers.py | 98 +++++++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 152 insertions(+) create mode 100644 .gitignore create mode 100644 prain/__init__.py create mode 100644 pyproject.toml create mode 100644 scripts/greedy.py create mode 100644 tests/__init__.py create mode 100644 tests/helpers.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..30ff7c4 --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +__pycache__/ +*.egg-info/ +/models/ +/build/ +/.venv/ diff --git a/prain/__init__.py b/prain/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/prain/__init__.py @@ -0,0 +1 @@ + diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..3beb54b --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,19 @@ +[project] +name = "prain" +version = "0.1.0" +description = "Per-expert VRAM/RAM/disk placement for MoE models served by vLLM" +requires-python = ">=3.11" +dependencies = [] # torch and vllm come from the gfx1030 vLLM image + +[project.scripts] +prain = "prain.cli:main" + +[project.entry-points."vllm.general_plugins"] +prain = "prain.plugin:register" + +[build-system] +requires = ["setuptools>=64"] +build-backend = "setuptools.build_meta" + +[tool.setuptools] +packages = ["prain"] diff --git a/scripts/greedy.py b/scripts/greedy.py new file mode 100644 index 0000000..756e5da --- /dev/null +++ b/scripts/greedy.py @@ -0,0 +1,28 @@ +"""Print greedy answers to fixed prompts from an OpenAI-compatible server, as JSON. + +usage: python3 scripts/greedy.py http://127.0.0.1:8000 MODEL > answers.json +""" + +import json +import sys +import urllib.request + +PROMPTS = [ + "Explain what a mixture-of-experts language model is in two sentences.", + "Write a Python function that returns the n-th Fibonacci number.", + "Name three prime numbers greater than 100.", +] + + +def complete(base_url, model, prompt): + body = json.dumps({"model": model, "messages": [{"role": "user", "content": prompt}], + "temperature": 0, "max_tokens": 64}).encode() + request = urllib.request.Request(f"{base_url}/v1/chat/completions", body, + {"Content-Type": "application/json"}) + with urllib.request.urlopen(request, timeout=1800) as response: + return json.load(response)["choices"][0]["message"]["content"] + + +if __name__ == "__main__": + base_url, model = sys.argv[1], sys.argv[2] + print(json.dumps([complete(base_url, model, p) for p in PROMPTS], indent=1)) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/helpers.py b/tests/helpers.py new file mode 100644 index 0000000..ae3e3b5 --- /dev/null +++ b/tests/helpers.py @@ -0,0 +1,98 @@ +import json +import os +import tempfile +import threading +import time +import unittest +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + + +class TmpCase(unittest.TestCase): + """A test case with a fresh temporary directory in self.dir.""" + + def setUp(self): + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + self.dir = tmp.name + + def write(self, name, text): + path = os.path.join(self.dir, name) + with open(path, "w") as f: + f.write(text) + return path + + +def wait_for(predicate, timeout=5.0): + end = time.monotonic() + timeout + while not predicate(): + if time.monotonic() > end: + raise AssertionError("timed out waiting") + time.sleep(0.01) + + +class StubVllm: + """A fake OpenAI-compatible server on a free port. Records every POST body. Streamed + chats answer "Hel", "lo", "!"; others answer "ok". Set .delay (seconds, or a function of + the request body) to make answers slow, .pause (the same) to stream the pieces that far + apart, and .fail to an HTTP code to reject chats.""" + + def __init__(self): + self.bodies = [] + self.delay = 0.0 + self.pause = 0.0 + self.fail = None + stub = self + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + self.send(200, "application/json", {"data": [{"id": "stub-model"}]}) + + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + stub.bodies.append(body) + time.sleep(stub.delay(body) if callable(stub.delay) else stub.delay) + if stub.fail: + return self.send(stub.fail, "application/json", {"error": "rejected"}) + if not body.get("stream"): + return self.send(200, "application/json", + {"choices": [{"message": {"content": "ok"}}]}) + events = [{"choices": [{"delta": {"content": p}}]} for p in ("Hel", "lo", "!")] + pause = stub.pause(body) if callable(stub.pause) else stub.pause + if pause: # stream piece by piece; HTTP/1.0 ends the body by closing + return self.trickle([f"data: {json.dumps(e)}\n\n" for e in events], pause) + text = "".join(f"data: {json.dumps(e)}\n\n" for e in events) + "data: [DONE]\n\n" + self.send(200, "text/event-stream", text) + + def trickle(self, chunks, pause): + try: + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.end_headers() + for i, chunk in enumerate(chunks + ["data: [DONE]\n\n"]): + time.sleep(pause if i else 0) + self.wfile.write(chunk.encode()) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + pass # the client aborted + + def send(self, code, kind, payload): + data = (payload if isinstance(payload, str) else json.dumps(payload)).encode() + try: + self.send_response(code) + self.send_header("Content-Type", kind) + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + except (BrokenPipeError, ConnectionResetError): + pass # the client aborted + + def log_message(self, *args): + pass + + self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self.url = f"http://127.0.0.1:{self.server.server_port}" + threading.Thread(target=self.server.serve_forever, daemon=True).start() + + def close(self): + self.server.shutdown() + self.server.server_close()