prain: skeleton

.gitignore

@@ -0,0 +1,5 @@
+__pycache__/
+*.egg-info/
+/models/
+/build/
+/.venv/

prain/__init__.py

@@ -0,0 +1 @@
+

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"]

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))

tests/__init__.py

@@ -0,0 +1 @@
+

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()