15 · Event 通知与异步互斥

    async def run():
        entered, release = asyncio.Event(), asyncio.Event()
        sample = Sample(tokens=[1, 2], status=Sample.Status.ABORTED)
        sample._queue_lease = {"task_id": "borrowed"}
        group = [sample]
        buffered = []

        async def fetch(count):
            entered.set()
            await release.wait()
            return [group]

        async def rebuffer(value):
            buffered.append(value)

        async def unexpected(*args):
            raise AssertionError("Paused worker admitted a generation request")

        monkeypatch.setattr(sglang_rollout, "generate_and_rm_group", unexpected)
        worker = _GenerationActor.__new__(_GenerationActor)
        worker.running = True
        worker.data_source = SimpleNamespace(get_samples_async=fetch)
        worker._rebuffer = rebuffer
        task = asyncio.create_task(worker._generate_group(0))
        await entered.wait()
        await worker._pause()
        release.set()
        assert await task is None
        assert buffered == [group]
        assert sample._queue_lease == {"task_id": "borrowed"}
        assert sample.tokens == [1, 2]

    asyncio.run(run())

def _build_http_app(sample_manager, args, update_from_disk_fn=None):
    app = web.Application(client_max_size=64 * 1024 * 1024)
    update_state = {
        "in_progress": False,
        "updating_model_path": None,
        "update_future": None,
    }
    update_lock = asyncio.Lock()

    async def detect(_request: web.Request) -> web.Response:
        return web.json_response({"server_type": "megatron_server"})

    async def healthz(_request: web.Request) -> web.Response:
        return web.json_response({"ok": True})

    async def info(_request: web.Request) -> web.Response:
        return web.json_response({"args": _args_to_dict(args)})

    async def get_loads(_request: web.Request) -> web.Response:
        return web.json_response(await _ray_get(sample_manager.get_loads.remote()))

    async def update_from_disk(request: web.Request) -> web.Response:
        if update_from_disk_fn is None:
            return _json_error("update_from_disk is not available during warmup", 503)

        payload = await _read_json_payload(request)
        model_path = _get_update_model_path(payload)
        if model_path is None:
            return _json_error("missing model_path", 400)

        async with update_lock:
            if getattr(args, "load", None) == model_path:
                return web.json_response({"ok": True, "model_path": model_path, "skipped": True})

            if update_state["in_progress"]:
                if update_state["updating_model_path"] == model_path and update_state["update_future"] is not None:
                    update_future = update_state["update_future"]
                    coalesced = True
                else:
                    updating_model_path = update_state["updating_model_path"]
                    return _json_error(f"update_from_disk is already in progress for {updating_model_path}", 409)
            else:
                update_future = asyncio.get_running_loop().create_future()
                update_state["in_progress"] = True
                update_state["updating_model_path"] = model_path
                update_state["update_future"] = update_future
                coalesced = False

            after_loads = await _ray_get(sample_manager.get_loads.remote())
        except TimeoutError as e:
            error = {"ok": False, "status": 503, "error": str(e)}
        except Exception as e:
            error = {"ok": False, "status": 500, "error": f"update_from_disk failed: {e}"}
        finally:
            if error is None:
                result = {
                    "ok": True,
                    "model_path": model_path,
                    "before_loads": before_loads,
                    "after_loads": after_loads,
                    "update_result": update_result,
                }
                # Reflect the freshly loaded checkpoint in /info. The actors restore
                # their own args after loading, so only the server-side copy needs to be
                # kept in sync here.
                args.load = model_path
                args.ref_load = model_path

            async with update_lock:
                if result is not None:
                    update_future.set_result(result)
                else:
                    update_future.set_result(error)
                update_state["in_progress"] = False
                update_state["updating_model_path"] = None
                update_state["update_future"] = None

        if result is not None:
            return web.json_response(result)
        return _json_error(error["error"], error["status"])