"""Start both trained adapters on one Modal L40S with a loopback HTTP gateway.""" import argparse import hmac import json import os from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path import modal from dataset import ReviewItem, Verdict from decision_modal import ( RUN as DECISION_RUN, ) from decision_modal import ( SFT_RUN, VOLUME_NAME, answer_probabilities, image, load_base, ) ROOT = Path(__file__).resolve().parent app = modal.App("hiive-research-inference") volume = modal.Volume.from_name(VOLUME_NAME) serving_image = image.add_local_file( ROOT / "decision_modal.py", "/root/decision_modal.py" ) @app.cls( image=serving_image, gpu="L40S", volumes={"/runs": volume}, memory=32768, min_containers=1, max_containers=1, scaledown_window=120, timeout=120, startup_timeout=1800, ) class Models: # Modal handles one input at a time by default. Keep it that way: adapter # selection changes shared model state and must not race another request. @modal.enter() def load(self): from peft import PeftModel self.systems = {} for name, run in (("decision", DECISION_RUN), ("sft", SFT_RUN)): manifest = json.loads((Path("/runs") / run / "manifest.json").read_text()) self.systems[name] = manifest["system_prompt"] self.tokenizer, base = load_base() self.model = PeftModel.from_pretrained( base, f"/runs/{DECISION_RUN}/adapter", adapter_name="decision" ) self.model.load_adapter(f"/runs/{SFT_RUN}/adapter", adapter_name="sft") self.model.cuda().eval() self.model.config.text_config.use_cache = True def prompt(self, stage, snapshot): snapshot = ReviewItem.model_validate(snapshot).model_dump() ids = self.tokenizer.apply_chat_template( [ {"role": "system", "content": self.systems[stage]}, { "role": "user", "content": json.dumps(snapshot, separators=(",", ":")), }, ], tokenize=True, add_generation_prompt=True, enable_thinking=False, return_dict=False, ) if len(ids) > 8192: raise ValueError("Case exceeds 8192 tokens; evidence was not truncated") return ids @modal.method() def ready(self): return {"decision": DECISION_RUN, "sft": SFT_RUN, "gpu": "L40S"} @modal.method() def decision(self, snapshot): ids = self.prompt("decision", snapshot) self.model.set_adapter("decision") p_ok, p_review, mass = answer_probabilities(self.model, ids) return {"p_ok": p_ok, "p_review": p_review, "answer_mass": mass} @modal.method() def sft(self, snapshot): import torch ids = torch.tensor([self.prompt("sft", snapshot)], device="cuda") self.model.set_adapter("sft") with torch.inference_mode(): output = self.model.generate( input_ids=ids, attention_mask=torch.ones_like(ids), max_new_tokens=160, do_sample=False, eos_token_id=self.tokenizer.eos_token_id, pad_token_id=self.tokenizer.pad_token_id, use_cache=True, ) text = self.tokenizer.decode( output[0, ids.shape[1] :], skip_special_tokens=True ) return Verdict.model_validate_json(text).model_dump() def gateway(models, port, token=None): """Private gateway; models can also be a test backend with remote methods.""" class Handler(BaseHTTPRequestHandler): def setup(self): super().setup() self.connection.settimeout(10) def respond(self, status, body): encoded = json.dumps(body, allow_nan=False).encode() self.send_response(status) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(encoded))) self.end_headers() self.wfile.write(encoded) def authorized(self): if token and not hmac.compare_digest( self.headers.get("Authorization", "").encode(), ("Bearer " + token).encode(), ): self.respond(401, {"error": "unauthorized"}) return False return True def do_GET(self): if self.authorized(): if self.path == "/health": self.respond( 200, {"status": "ready", "decision": DECISION_RUN, "sft": SFT_RUN}, ) else: self.respond(404, {"error": "not_found"}) def do_POST(self): if not self.authorized(): return if self.path not in ("/decision", "/sft"): self.respond(404, {"error": "not_found"}) return try: size = int(self.headers.get("Content-Length", "0")) if not 0 < size <= 1_048_576: self.respond(413, {"error": "body_size_out_of_range"}) return snapshot = ReviewItem.model_validate_json( self.rfile.read(size) ).model_dump() except (ValueError, OSError): self.respond(400, {"error": "invalid_case"}) return try: result = getattr(self.server.models, self.path[1:]).remote(snapshot) if self.path == "/sft": result = Verdict.model_validate(result).model_dump() else: if not all( type(result[key]) in (float, int) and 0 <= result[key] <= 1 for key in ("p_review", "answer_mass") ): raise ValueError("Invalid decision scores") self.respond(200, result) except Exception: # noqa: BLE001 -- Keep remote exception payloads private. # Do not return case text, model output, or credentials in errors. self.respond(503, {"error": "model_request_failed"}) server = ThreadingHTTPServer(("127.0.0.1", port), Handler) server.models = models return server def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--port", type=int, default=8767, help="Loopback gateway port (default: 8767)" ) args = parser.parse_args() # Bind before allocating GPU time so a port conflict fails locally. with ( gateway(None, args.port, os.environ.get("MODEL_API_TOKEN")) as server, modal.enable_output(), app.run(), ): models = Models() print("Loading both adapters on one Modal L40S...", flush=True) models.ready.remote() server.models = models print( f"export DECISION_MODEL_URL=http://127.0.0.1:{server.server_port}/decision", flush=True, ) print( f"export SFT_MODEL_URL=http://127.0.0.1:{server.server_port}/sft", flush=True, ) print( "Both models ready. Ctrl-C stops the gateway and the temporary Modal app.", flush=True, ) try: server.serve_forever() except KeyboardInterrupt: pass if __name__ == "__main__": main()