13 · Future、回调与跨线程事件循环

    async def wait(self):
        loop = asyncio.get_running_loop()
        future = loop.create_future()
        self.pending.append(future)
        if not self.scheduled:
            self.scheduled = True
            loop.call_soon(self._release)
        await future

    def _release(self):
        for _ in range(min(64, len(self.pending))):
            future = self.pending.popleft()
            if not future.done():
                future.set_result(None)
        if self.pending:
            # Let socket callbacks and timers run before waking more tasks.
            asyncio.get_running_loop().call_later(0.001, self._release)
        else:
            self.scheduled = False

    async def _rebuffer(self, group):
        if not hasattr(self, "_pending_rebuffers"):
            self._pending_rebuffers = deque()
            self._rebuffer_task = None
        future = asyncio.get_running_loop().create_future()
        self._pending_rebuffers.append((group, future))
        if self._rebuffer_task is None:
            self._rebuffer_task = asyncio.create_task(self._flush_rebuffers())
        await asyncio.shield(future)

    async def _flush_rebuffers(self):
        try:
            await asyncio.sleep(0.01)
            while self._pending_rebuffers:
                batch = [self._pending_rebuffers.popleft() for _ in range(min(64, len(self._pending_rebuffers)))]
                try:
                    await asyncio.to_thread(self.data_source.add_samples, [group for group, _ in batch])
                except Exception as error:
                    for _, future in [*batch, *self._pending_rebuffers]:
                        future.set_exception(error)
                    self._pending_rebuffers.clear()
                    return
                for _, future in batch:
                    future.set_result(None)
        finally:
            self._rebuffer_task = None

class AsyncLoopThread:
    def __init__(self):
        self.loop = asyncio.new_event_loop()
        self._thread = threading.Thread(target=self._start_loop, daemon=True)
        self._thread.start()

    def _start_loop(self):
        asyncio.set_event_loop(self.loop)
        self.loop.run_forever()

    def run(self, coro):
        # Schedule a coroutine onto the loop and block until it's done
        return asyncio.run_coroutine_threadsafe(coro, self.loop).result()

    def _run() -> None:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            runner = web.AppRunner(app, **(runner_kwargs or {}))
            loop.run_until_complete(runner.setup())
            site = web.TCPSite(runner, host, port)
            loop.run_until_complete(site.start())
            actual_port = port
            for sock in site._server.sockets:  # type: ignore[attr-defined]
                actual_port = sock.getsockname()[1]
                break
            box["loop"] = loop
            box["runner"] = runner
            box["port"] = actual_port
            started.set()
            loop.run_forever()
        except BaseException as e:  # pragma: no cover
            err_box.append(e)
            started.set()
            raise

    async def start(self, rollout_id):
        future = asyncio.run_coroutine_threadsafe(self._start(rollout_id), get_async_loop().loop)
        await asyncio.wrap_future(future)

    def stop(self) -> None:
        async def _shutdown() -> None:
            await self.runner.cleanup()

        try:
            fut = asyncio.run_coroutine_threadsafe(_shutdown(), self.loop)
            fut.result(timeout=10)
        except Exception:
            pass
        self.loop.call_soon_threadsafe(self.loop.stop)
        self.thread.join(timeout=5)

    async def close(self):
        future = asyncio.run_coroutine_threadsafe(self._close(), get_async_loop().loop)
        await asyncio.wrap_future(future)

    async def _close(self):
        self.running = False
        if self.producer is not None:
            self.producer.cancel()
            await asyncio.gather(self.producer, return_exceptions=True)
        if getattr(self, "_rebuffer_task", None) is not None:
            await self._rebuffer_task
        # to_thread writes can outlive cancelled awaiters. Drain them before
        # closing their files or killing this process.
        await asyncio.get_running_loop().shutdown_default_executor()
        self.data_source.close()

        tok = self.tokenizer
        s = self.store.setdefault(sid, Session())
        task = asyncio.current_task()
        self.inflight.setdefault(sid, set()).add(task)
        t0 = time.monotonic()
        try:
            translated, tools_schema = self._translate(body)
            prompt_ids = _render_token_ids(translated, tok, tools=tools_schema, add_generation_prompt=True)

            turn = await call_sglang_generate(prompt_ids, s, body, adapter=self, session_id=sid)

    async def shutdown_session(self, sid: str, *, wait_timeout: float = 5.0) -> None:
        """Mark a sid closed and drain its in-flight turn tasks."""
        self.closed.add(sid)
        tasks = [t for t in self.inflight.pop(sid, ()) if not t.done()]
        if not tasks:
            return

        async def _drain() -> None:
            _, pending = await asyncio.wait(tasks, timeout=wait_timeout)
            for task in pending:
                task.cancel()
            if pending:
                await asyncio.gather(*pending, return_exceptions=True)

        loop = tasks[0].get_loop()
        try:
            await asyncio.wrap_future(asyncio.run_coroutine_threadsafe(_drain(), loop))
        except Exception:
            self.logger.exception("[%s] sid=%s shutdown drain failed", self.log_prefix, sid)

def _log_timeout_diagnostic(t0: float, instance_id: str) -> None:
    # Dump pending-task names when the wall-clock guard fires. Must not crash.
    try:
        elapsed = time.time() - t0
        pending = [t for t in asyncio.all_tasks() if not t.done()]
        stuck = []
        for t in pending[:5]:  # cap to avoid log spam
            coro = getattr(t, "_coro", None)
            stuck.append(getattr(coro, "__qualname__", repr(coro)))
        logger.warning(
            "[coding_agent_rl] %s: wall_clock_timeout after %.1fs "
            "(guard=%ds); %d tasks pending; sample of stuck: %s",
            instance_id,
            elapsed,
            CONFIG.rollout_guard_sec,
            len(pending),
            stuck,
        )
    except Exception:  # pragma: no cover - diag must never crash
        pass

    async def _loop(self) -> None:
        active_tasks: set[asyncio.Task] = set()
        max_concurrent = self.concurrency
        gid_counter = 0

        while self.running:
            try:
                # Reap done tasks
                if active_tasks:
                    done = {t for t in active_tasks if t.done()}
                    for t in done:
                        try:
                            t.result()  # results already handled in callback
                        except Exception as e:  # noqa: BLE001
                            logger.warning("fully-async task crashed: %r", e)
                    active_tasks -= done
                    if done:
                        # Done callbacks requeue ABORTED groups. Let them run
                        # before asking the data source for replacement work.
                        await asyncio.sleep(0)

                    groups = self.data_buffer.get_samples(1)
                    if not groups:
                        break
                    for group in groups:
                        gid = gid_counter
                        gid_counter += 1
                        task = asyncio.create_task(
                            generate_and_rm_group(
                                self.args,
                                group,
                                sampling_params=self.state.sampling_params.copy(),
                                evaluation=False,
                            )
                        )
                        task.add_done_callback(self._make_done_cb(gid))
                        active_tasks.add(task)

    def _make_done_cb(self, gid: int):
        def _cb(done_task: asyncio.Task) -> None:
            try:
                result = done_task.result()
            except asyncio.CancelledError:
                return
            except Exception as error:
                self.error = error
                return
            if not isinstance(result, list):
                logger.warning(
                    "fully-async: generate_and_rm_group returned %r, expected list[Sample]; dropping",
                    type(result).__name__,
                )
                return
            # Aborted group → requeue, don't ship to training.
            if any(getattr(s, "status", None) == Sample.Status.ABORTED for s in result):
                try:
                    self.data_buffer.add_samples([result])
                except Exception:  # noqa: BLE001
                    logger.exception("fully-async: failed to requeue aborted group")
                return
            self.output_queue.put((gid, result))

        return _cb