17 · 线程池 Future 的结果与退出

        pool = ThreadPoolExecutor(max_workers=NUM_WORKERS)
        inflight: deque = deque()
        try:
            for name, tensor in self._iter_hf_tensors():
                flat = tensor.detach().contiguous().view(torch.uint8).reshape(-1)
                nbytes = int(flat.numel())
                if use_pinned and nbytes <= max_bytes:
                    buf = free_q.get()  # blocks when all buffers are in flight -> backpressures the gather
                    buf[:nbytes].copy_(flat, non_blocking=True)
                    accelerator.current_stream().synchronize()
                    payload, pinned = buf, True
                else:
                    payload, pinned = flat.cpu().numpy(), False
                self.total_bytes += nbytes
                inflight.append(pool.submit(diff_and_compress, name, payload, nbytes, pinned))
                if len(inflight) >= 2 * NUM_WORKERS:
                    collect(inflight.popleft())
            while inflight:
                collect(inflight.popleft())
        finally:
            pool.shutdown()

        with ThreadPoolExecutor(max_workers=len(_probes)) as pool:
            _records["continuous"] = list(pool.map(lambda probe: probe.verify(), _probes))
        start_probes([probe.url for probe in _probes])
    elif rollout_id == 2 and step_id == 0:
        assert _probes
        with ThreadPoolExecutor(max_workers=len(_probes)) as pool:
            _records["next_update"] = list(
                pool.map(
                    lambda probe: probe.verify_flush() if args.flush_cache_interval == 2 else probe.verify(), _probes
                )
            )

    def prefetch(self) -> None:
        with self._lock:
            if self._cpu_tensor is not None or self._future is not None:
                return
            executor = self.prefetcher.executor if self.prefetcher is not None else None
            if executor is not None:
                self._future = executor.submit(self._load_and_prepare)

    def _get_cpu_tensor(self) -> torch.Tensor:
        with self._lock:
            tensor = self._cpu_tensor
            future = self._future
        if tensor is None:
            tensor = future.result() if future is not None else self._load_and_prepare()
            with self._lock:
                self._cpu_tensor = tensor
                self._future = None
        if self.prefetcher is not None:
            self.prefetcher.prefetch_after(self.index)
        return tensor

    def release(self) -> None:
        with self._lock:
            future = self._future
            self._future = None
            self._cpu_tensor = None
        if future is not None and not future.done():
            future.cancel()

    def remote(method):
        return SimpleNamespace(remote=lambda *args: asyncio.run_coroutine_threadsafe(method(*args), loop))

    def ray_wait(refs, num_returns, timeout):
        done, pending = wait(refs, timeout=timeout, return_when=FIRST_COMPLETED)
        return list(done)[:num_returns], list(pending)

    def kill(worker, no_restart=True):
        if worker.actor.producer is not None:
            loop.call_soon_threadsafe(worker.actor.producer.cancel)

    monkeypatch.setattr(ray, "wait", ray_wait)
    monkeypatch.setattr(ray, "get", lambda ref: ref.result())
    monkeypatch.setattr(ray, "cancel", lambda ref: ref.cancel())

        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            futures = {
                executor.submit(read_file, rel_path, full_path, ops): rel_path
                for rel_path, (full_path, ops) in [
                    (r, (self.fs.concat_path(self.path, r), ops)) for r, ops in file_ops.items()
                ]
            }
            for future in tqdm(as_completed(futures), total=len(futures), desc="Pre-loading files"):
                result = future.result()
                if result:
                    rel_path, data, is_full_file = result
                    with self._cache_lock:
                        # Cache by file path - for chunked files, store the full file
                        self._read_cache[rel_path] = data

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        if self._executor is not None:
            self._executor.shutdown(wait=True, cancel_futures=True)
        for source in self.sources:
            source.release()

    def close(self):
        if self._closed:
            return
        self._closed = True
        self._reader.shutdown(wait=True)
        with self._lock:
            leases = list(self._leases.values())
        ray.get(self.controller.release.remote(leases))
        self._stop.set()
        self._heartbeats.join(timeout=5)

def test_slow_worker_control_does_not_block_collection_or_close(scheduler_factory, control):
    scheduler, _, _ = scheduler_factory(3, [1, 1])
    stalled = Future()
    entered = threading.Event()
    original = getattr(scheduler.workers[0], control).remote

    def stall(*args):
        entered.set()
        return stalled

    if control == "pause":
        scheduler.generate(0, prefetch=2)
    getattr(scheduler.workers[0], control).remote = stall
    with ThreadPoolExecutor(2) as callers:
        request = (
            callers.submit(scheduler.generate, 0, prefetch=2)
            if control == "start"
            else callers.submit(scheduler.pause)
        )
        closing = None
        try:
            assert entered.wait(5)
            if control == "start":
                result = request.result(timeout=2)
                assert all(group[0].metadata["worker"] == 1 for group in result.samples)
            closing = callers.submit(scheduler.close)
            closing.result(timeout=2)
            if control == "pause":
                with pytest.raises(RuntimeError, match="scheduler is closed"):
                    request.result(timeout=2)
        finally:
            # Also let the old, blocking implementation exit after test failure.
            if not stalled.done():
                stalled.set_result(None)
            getattr(scheduler.workers[0], control).remote = original
            scheduler.close()