14 · 回到 Slime：串起创建、等待与清理

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

    async def _start(self, rollout_id):
        self.rollout_id = rollout_id
        self.running = True
        if self.producer is None or self.producer.done():
            if self.producer is not None:
                self.producer.result()
            # Keep completed results bounded even when training stops consuming.
            # At most one queued result plus one result per pool slot stays local.
            self.outputs = asyncio.Queue(maxsize=1)
            self.producer = asyncio.create_task(self._produce())

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

    async def _pause(self):
        self.running = False

    async def next(self):
        future = asyncio.run_coroutine_threadsafe(self._next(), get_async_loop().loop)
        return await asyncio.wrap_future(future)

    async def _next(self):
        take = asyncio.create_task(self.outputs.get())
        try:
            await asyncio.wait((take, self.producer), return_when=asyncio.FIRST_COMPLETED)
            if take.done():
                return take.result()
            self.producer.result()  # Propagate custom generation errors.
            if not self.outputs.empty():
                return await take
            return None
        finally:
            take.cancel()
            await asyncio.gather(take, return_exceptions=True)

    async def _produce(self):
        async def generate_slot(group):
            while self.running or group is not None:
                output = await self._generate_group(self.rollout_id, group)
                group = None
                if output is not None:  # Aborted groups have been returned to the durable queue.
                    await self.outputs.put(output)
                del output

        tasks = []
        try:
            for offset in range(0, self.capacity, 64):
                if not self.running:
                    break
                groups = await self.data_source.get_samples_async(min(64, self.capacity - offset))
                tasks.extend(asyncio.create_task(generate_slot(group)) for group in groups)
                # Spread startup callbacks across loop turns to reduce asyncio
                # pressure. Each slot then pulls its own next group immediately.
                await asyncio.sleep(0.001)
            await asyncio.gather(*tasks)
        finally:
            for task in tasks:
                task.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)

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