Example #1
0
def _play_game_against_neural_mcts(
    devices: List[torch.device],
    models: List[torch.jit.ScriptModule],
    context: tube.Context,
    actor_channel: tube.DataChannel,
) -> None:
    nb_devices = len(devices)
    context.start()
    dcm = DataChannelManager([actor_channel])
    while not context.terminated():
        batch = dcm.get_input(max_timeout_s=1)
        if len(batch) == 0:
            continue
        assert len(batch) == 1

        # split in as many part as there are devices
        batches_s = torch.chunk(batch[actor_channel.name]["s"],
                                nb_devices,
                                dim=0)
        futures = []
        reply_eval = {"v": None, "pi": None}
        # multithread
        with ThreadPoolExecutor(max_workers=nb_devices) as executor:
            for device, model, batch_s in zip(devices, models, batches_s):
                futures.append(
                    executor.submit(_forward_pass_on_device, device, model,
                                    batch_s))
            results = [future.result() for future in futures]
            reply_eval["v"] = torch.cat([result["v"] for result in results],
                                        dim=0)
            reply_eval["pi"] = torch.cat([result["pi"] for result in results],
                                         dim=0)
        dcm.set_reply(actor_channel.name, reply_eval)
    dcm.terminate()
Example #2
0
def _play_game_against_mcts(context: tube.Context) -> None:
    context.start()
    while not context.terminated():
        time.sleep(1)