Method test_queue_buffer

tests/buffer/queue_test.py:34–97

(self, name, use_priority_queue)

        ]
    )
    async def test_queue_buffer(self, name, use_priority_queue):
        config = ExperienceBufferConfig(
            name=name,
            schema_type="experience",
            storage_type=StorageType.QUEUE.value,
            max_read_timeout=3,
            path=BUFFER_FILE_PATH,
            batch_size=self.train_batch_size,
        )
        config.replay_buffer.enable = use_priority_queue
        config = config.to_storage_config()
        writer = QueueWriter(config)
        reader = QueueReader(config)
        self.assertEqual(await writer.acquire(), 1)
        exps = [
            Experience(
                tokens=torch.tensor([float(j) for j in range(i + 1)]),
                prompt_length=i,
                reward=float(i),
                logprobs=torch.tensor([0.1]),
            )
            for i in range(1, self.put_batch_size + 1)
        ]
        for exp in exps:
            exp.info = {"model_version": 0, "use_count": 0}
        for _ in range(self.total_num // self.put_batch_size):
            await writer.write(exps)
        for _ in range(self.total_num // self.train_batch_size):
            exps = await reader.read()
            self.assertEqual(len(exps), self.train_batch_size)
            print(f"finish read {self.train_batch_size} experience")
        exps = [
            Experience(
                tokens=torch.tensor([float(j) for j in range(i + 1)]),
                reward=float(i),
                logprobs=torch.tensor([0.1]),
                action_mask=torch.tensor([j % 2 for j in range(i + 1)]),
            )
            for i in range(1, self.put_batch_size * 2 + 1)
        ]
        for exp in exps:
            exp.info = {"model_version": 1, "use_count": 0}
        await writer.write(exps)
        exps = await reader.read(batch_size=self.put_batch_size * 2)
        self.assertEqual(len(exps), self.put_batch_size * 2)

        def thread_read(reader, result_queue):
            try:
                batch = asyncio.run(reader.read())
                result_queue.put(batch)
            except StopAsyncIteration as e:
                result_queue.put(e)

        result_queue = queue.Queue()
        t = threading.Thread(target=thread_read, args=(reader, result_queue))
        t.start()
        time.sleep(2)  # make sure the thread is waiting for data
        self.assertEqual(await writer.release(), 0)