beam-boot / scripts /deploy_inweights_modal.py

Publish Beam Boot / Super Beam source

324ecd2 verified 7 days ago

8.2 kB

"""Serve the in-weights CHIP-8 computer (beamboot/chip8_lm.py) on a Modal GPU so the demo page can
run one game frame on the real model, live, from a machine state the visitor created.
  modal deploy scripts/deploy_inweights_modal.py        # prints the public URL
POST /frame  {"mem": base64 of 4096 bytes, "v": [16 ints], "i": int, "pc": int, "dt": int,
              "keys": [16 bools]}   (the machine state at a frame boundary, keys already applied)
-> {"rows": 32 hex strings (the display decoded from the tokens the model accepted), "verified": bool,
    "first_reject": int|null, "prompt_tokens", "frame_tokens", "instructions", "seconds", "gpu",
    "sample": first token records as text, "state": the machine state after the frame, decoded from the
    model's tokens (mem, v, i, pc, dt), so a client can keep playing on the model}
Same as beamboot/chip8_fast.FastSession.advance: checkpoint prompt + emulator draft, one forward pass
of the constructed BeamForCausalLM, every drafted token compared with the model's own argmax. If the
model disagrees anywhere it answers verified=false with the position (no greedy fallback here, to keep
requests short). One GPU container at most (L40S, else A10G, L4, T4), one forward at a time; it scales to zero when idle.
"""
from pathlib import Path

import modal

PKG = Path(__file__).resolve().parent.parent / "beamboot"

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("torch==2.11.0", "transformers==5.22.1", "fastapi", "numpy")
    .add_local_file(str(PKG / "__init__.py"), "/app/beamboot/__init__.py")
    .add_local_file(str(PKG / "chip8.py"), "/app/beamboot/chip8.py")
    .add_local_file(str(PKG / "chip8_lm.py"), "/app/beamboot/chip8_lm.py")
    .add_local_file(str(PKG / "chip8_fast.py"), "/app/beamboot/chip8_fast.py")
)

app = modal.App("superbeam-inweights", image=image)

ORIGINS = ["https://staratnyte1-beam-boot.static.hf.space", "https://huggingface.co",
           "http://localhost:8765"]
SAMPLE_RECORDS = 6


def _crop_to(cache, n):
    """Keep the first n positions of a DynamicCache (this transformers version takes a negative count)."""
    extra = cache.get_seq_length() - n
    if extra > 0:
        cache.crop(-extra)


def _fmt(C, t):
    k, v, a = C.untok(t)
    return {"VW": "V%X=%d" % (a, v), "PCW": "PC=%03X" % v, "IW": "I=%03X" % v}.get(k, "%s %d" % (k, v))


@app.cls(gpu=["L40S", "A10G", "L4", "T4"], min_containers=0, max_containers=1, scaledown_window=300, timeout=300)
@modal.concurrent(max_inputs=12)
class InWeights:
    @modal.enter()
    def load(self):
        import sys
        import time
        import torch
        sys.path.insert(0, "/app")
        from beamboot import chip8_lm as C
        t0 = time.perf_counter()
        self.model = C.get_model(300, "cuda")
        self.model.config._attn_implementation = "sdpa"
        torch.backends.cuda.matmul.allow_tf32 = False
        torch.backends.cudnn.allow_tf32 = False
        self.gpu = torch.cuda.get_device_name(0)
        self.build_seconds = round(time.perf_counter() - t0, 1)
        import threading
        self.lock = threading.Lock()          # one forward at a time on the GPU; drafting overlaps it
        # Drafting only needs the record tokens; skip the per-record state snapshot (a full display
        # dump after every instruction), which was ~90% of the draft time. Tokens are unchanged.
        C.snapshot = lambda c: None
        self.pids, self.pcache = None, None   # prompt-prefix KV cache: consecutive frames share it

    def frame(self, req):
        import base64
        import copy
        import time
        import torch
        from transformers import DynamicCache
        from beamboot import chip8_fast as F
        from beamboot import chip8_lm as C
        from beamboot.chip8 import Chip8

        mem = base64.b64decode(req["mem"])
        if len(mem) != 4096 or len(req["v"]) != 16 or len(req["keys"]) != 16:
            raise ValueError("bad machine state")
        c = Chip8()
        c.mem[:] = mem
        c.v = [int(x) & 0xFF for x in req["v"]]
        c.i, c.pc, c.dt = int(req["i"]) & 0xFFF, int(req["pc"]) & 0xFFF, int(req["dt"]) & 0xFF
        c.keys = [bool(k) for k in req["keys"]]
        s = F.FastSession.__new__(F.FastSession)
        s.rom, s.ipf, s.c = bytes(mem[0x200:]), 300, c
        t0 = time.perf_counter()
        p = F.checkpoint_prompt(c, 300)
        # Draft (as FastSession._draft) on a copy of the state; keep that copy: it is the state after the
        # frame implied by the drafted tokens.
        after, d, n_rec = copy.deepcopy(c), [], 0
        for kind, ts, _ in C._ref_records(s.rom, 300, 0, lambda fr: (), chip=after):
            d += ts
            if kind in ("instr", "tick"):
                n_rec += 1
            if kind == "tick":
                break
        seq = p + d
        td = time.perf_counter()
        with torch.inference_mode(), self.lock:
            t1 = time.perf_counter()              # GPU time is measured after the queue
            # Reuse the KV cache of the longest prompt prefix shared with the previous request (at most
            # len(p) - 1, so the logits for every draft position are recomputed).
            L = 0
            if self.pids is not None:
                n = min(len(self.pids), len(p) - 1)
                while L < n and self.pids[L] == p[L]:
                    L += 1
            if L >= 256:
                cache = copy.deepcopy(self.pcache)
                _crop_to(cache, L)
            else:
                cache, L = DynamicCache(config=self.model.config), 0
            out_ = self.model(input_ids=torch.tensor([seq[L:]], device="cuda"), past_key_values=cache,
                              use_cache=True, logits_to_keep=len(d) + 1)
            am = out_.logits[0, :len(d)].argmax(-1)
            # first position where the model's own next token differs from the draft (one GPU compare)
            bad = torch.nonzero(am != torch.tensor(d, device="cuda")).flatten()
            rej = int(bad[0]) if len(bad) else None
            _crop_to(cache, len(p))
            self.pids, self.pcache = p, cache
            torch.cuda.synchronize()
        t2 = time.perf_counter()
        sec = t2 - t1
        if rej is None:
            # The model accepted every drafted token, so its tokens are exactly the draft and the state
            # they encode is `after` (token replay == reference: tests/test_chip8_fast.py).
            fin = after
        else:
            # The model disagreed: decode what it did accept, token by token, and report where it stopped.
            replay = F._new_replay(c, 300)
            for t in d[:rej]:
                replay.feed(t)
            fin = replay.c
        out = {"rows": fin.display_rows(), "verified": rej is None, "first_reject": rej,
               "prompt_tokens": len(p), "frame_tokens": len(d), "instructions": max(0, n_rec - 1),
               "seconds": round(sec, 3), "gpu": self.gpu, "draft_seconds": round(td - t0, 3), "cached_prefix": L,
               "state": None if rej is not None else {"mem": base64.b64encode(bytes(fin.mem)).decode(),
                        "v": list(fin.v), "i": fin.i, "pc": fin.pc, "dt": fin.dt}}
        out["server_seconds"] = round(time.perf_counter() - t0, 3)
        return out

    @modal.asgi_app()
    def web(self):
        from fastapi import FastAPI, HTTPException, Request
        from fastapi.middleware.cors import CORSMiddleware
        api = FastAPI()
        api.add_middleware(CORSMiddleware, allow_origins=ORIGINS, allow_methods=["POST", "GET"],
                           allow_headers=["content-type"])

        @api.get("/health")
        def health():
            return {"ok": True, "gpu": self.gpu, "model_build_seconds": self.build_seconds}

        @api.post("/frame")
        async def frame(request: Request):
            from starlette.concurrency import run_in_threadpool
            try:
                return await run_in_threadpool(self.frame, await request.json())
            except Exception as e:                      # bad state / unsupported program
                raise HTTPException(status_code=400, detail=str(e)[:200])

        return api

	"""Serve the in-weights CHIP-8 computer (beamboot/chip8_lm.py) on a Modal GPU so the demo page can
	run one game frame on the real model, live, from a machine state the visitor created.

	modal deploy scripts/deploy_inweights_modal.py # prints the public URL

	POST /frame {"mem": base64 of 4096 bytes, "v": [16 ints], "i": int, "pc": int, "dt": int,
	"keys": [16 bools]} (the machine state at a frame boundary, keys already applied)
	-> {"rows": 32 hex strings (the display decoded from the tokens the model accepted), "verified": bool,
	"first_reject": int\|null, "prompt_tokens", "frame_tokens", "instructions", "seconds", "gpu",
	"sample": first token records as text, "state": the machine state after the frame, decoded from the
	model's tokens (mem, v, i, pc, dt), so a client can keep playing on the model}

	Same as beamboot/chip8_fast.FastSession.advance: checkpoint prompt + emulator draft, one forward pass
	of the constructed BeamForCausalLM, every drafted token compared with the model's own argmax. If the
	model disagrees anywhere it answers verified=false with the position (no greedy fallback here, to keep
	requests short). One GPU container at most (L40S, else A10G, L4, T4), one forward at a time; it scales to zero when idle.
	"""
	from pathlib import Path

	import modal

	PKG = Path(__file__).resolve().parent.parent / "beamboot"

	image = (
	modal.Image.debian_slim(python_version="3.12")
	.pip_install("torch==2.11.0", "transformers==5.22.1", "fastapi", "numpy")
	.add_local_file(str(PKG / "__init__.py"), "/app/beamboot/__init__.py")
	.add_local_file(str(PKG / "chip8.py"), "/app/beamboot/chip8.py")
	.add_local_file(str(PKG / "chip8_lm.py"), "/app/beamboot/chip8_lm.py")
	.add_local_file(str(PKG / "chip8_fast.py"), "/app/beamboot/chip8_fast.py")
	)

	app = modal.App("superbeam-inweights", image=image)

	ORIGINS = ["https://staratnyte1-beam-boot.static.hf.space", "https://huggingface.co",
	"http://localhost:8765"]
	SAMPLE_RECORDS = 6


	def _crop_to(cache, n):
	"""Keep the first n positions of a DynamicCache (this transformers version takes a negative count)."""
	extra = cache.get_seq_length() - n
	if extra > 0:
	cache.crop(-extra)


	def _fmt(C, t):
	k, v, a = C.untok(t)
	return {"VW": "V%X=%d" % (a, v), "PCW": "PC=%03X" % v, "IW": "I=%03X" % v}.get(k, "%s %d" % (k, v))


	@app.cls(gpu=["L40S", "A10G", "L4", "T4"], min_containers=0, max_containers=1, scaledown_window=300, timeout=300)
	@modal.concurrent(max_inputs=12)
	class InWeights:
	@modal.enter()
	def load(self):
	import sys
	import time
	import torch
	sys.path.insert(0, "/app")
	from beamboot import chip8_lm as C
	t0 = time.perf_counter()
	self.model = C.get_model(300, "cuda")
	self.model.config._attn_implementation = "sdpa"
	torch.backends.cuda.matmul.allow_tf32 = False
	torch.backends.cudnn.allow_tf32 = False
	self.gpu = torch.cuda.get_device_name(0)
	self.build_seconds = round(time.perf_counter() - t0, 1)
	import threading
	self.lock = threading.Lock() # one forward at a time on the GPU; drafting overlaps it
	# Drafting only needs the record tokens; skip the per-record state snapshot (a full display
	# dump after every instruction), which was ~90% of the draft time. Tokens are unchanged.
	C.snapshot = lambda c: None
	self.pids, self.pcache = None, None # prompt-prefix KV cache: consecutive frames share it

	def frame(self, req):
	import base64
	import copy
	import time
	import torch
	from transformers import DynamicCache
	from beamboot import chip8_fast as F
	from beamboot import chip8_lm as C
	from beamboot.chip8 import Chip8

	mem = base64.b64decode(req["mem"])
	if len(mem) != 4096 or len(req["v"]) != 16 or len(req["keys"]) != 16:
	raise ValueError("bad machine state")
	c = Chip8()
	c.mem[:] = mem
	c.v = [int(x) & 0xFF for x in req["v"]]
	c.i, c.pc, c.dt = int(req["i"]) & 0xFFF, int(req["pc"]) & 0xFFF, int(req["dt"]) & 0xFF
	c.keys = [bool(k) for k in req["keys"]]
	s = F.FastSession.__new__(F.FastSession)
	s.rom, s.ipf, s.c = bytes(mem[0x200:]), 300, c
	t0 = time.perf_counter()
	p = F.checkpoint_prompt(c, 300)
	# Draft (as FastSession._draft) on a copy of the state; keep that copy: it is the state after the
	# frame implied by the drafted tokens.
	after, d, n_rec = copy.deepcopy(c), [], 0
	for kind, ts, _ in C._ref_records(s.rom, 300, 0, lambda fr: (), chip=after):
	d += ts
	if kind in ("instr", "tick"):
	n_rec += 1
	if kind == "tick":
	break
	seq = p + d
	td = time.perf_counter()
	with torch.inference_mode(), self.lock:
	t1 = time.perf_counter() # GPU time is measured after the queue
	# Reuse the KV cache of the longest prompt prefix shared with the previous request (at most
	# len(p) - 1, so the logits for every draft position are recomputed).
	L = 0
	if self.pids is not None:
	n = min(len(self.pids), len(p) - 1)
	while L < n and self.pids[L] == p[L]:
	L += 1
	if L >= 256:
	cache = copy.deepcopy(self.pcache)
	_crop_to(cache, L)
	else:
	cache, L = DynamicCache(config=self.model.config), 0
	out_ = self.model(input_ids=torch.tensor([seq[L:]], device="cuda"), past_key_values=cache,
	use_cache=True, logits_to_keep=len(d) + 1)
	am = out_.logits[0, :len(d)].argmax(-1)
	# first position where the model's own next token differs from the draft (one GPU compare)
	bad = torch.nonzero(am != torch.tensor(d, device="cuda")).flatten()
	rej = int(bad[0]) if len(bad) else None
	_crop_to(cache, len(p))
	self.pids, self.pcache = p, cache
	torch.cuda.synchronize()
	t2 = time.perf_counter()
	sec = t2 - t1
	if rej is None:
	# The model accepted every drafted token, so its tokens are exactly the draft and the state
	# they encode is `after` (token replay == reference: tests/test_chip8_fast.py).
	fin = after
	else:
	# The model disagreed: decode what it did accept, token by token, and report where it stopped.
	replay = F._new_replay(c, 300)
	for t in d[:rej]:
	replay.feed(t)
	fin = replay.c
	out = {"rows": fin.display_rows(), "verified": rej is None, "first_reject": rej,
	"prompt_tokens": len(p), "frame_tokens": len(d), "instructions": max(0, n_rec - 1),
	"seconds": round(sec, 3), "gpu": self.gpu, "draft_seconds": round(td - t0, 3), "cached_prefix": L,
	"state": None if rej is not None else {"mem": base64.b64encode(bytes(fin.mem)).decode(),
	"v": list(fin.v), "i": fin.i, "pc": fin.pc, "dt": fin.dt}}
	out["server_seconds"] = round(time.perf_counter() - t0, 3)
	return out

	@modal.asgi_app()
	def web(self):
	from fastapi import FastAPI, HTTPException, Request
	from fastapi.middleware.cors import CORSMiddleware
	api = FastAPI()
	api.add_middleware(CORSMiddleware, allow_origins=ORIGINS, allow_methods=["POST", "GET"],
	allow_headers=["content-type"])

	@api.get("/health")
	def health():
	return {"ok": True, "gpu": self.gpu, "model_build_seconds": self.build_seconds}

	@api.post("/frame")
	async def frame(request: Request):
	from starlette.concurrency import run_in_threadpool
	try:
	return await run_in_threadpool(self.frame, await request.json())
	except Exception as e: # bad state / unsupported program
	raise HTTPException(status_code=400, detail=str(e)[:200])

	return api