Repository: LAION-AI/Open-Assistant
File: inference/server/oasst_inference_server/routes/workers.py (Python)

Code before:
            worker_info=worker_info,
            work_queue=work_queue,
            blocking_work_queue=blocking_work_queue
        )

    async def close(self):
        try:
            await self.blocking_work_queue.redis_client.close()
        except Exception:
            logger.warning("Error while closing redis client")

        try:
            await self.websocket.close()
        except Exception:
            logger.warning("Error while closing websocket")



class FutureType(str, Enum):
    WORK_REQUEST = 'WORK_REQUEST'
    WORKER_RESPONSE = 'WORKER_RESPONSE'

AwaitResult = TypeVar('AwaitResult')

class FuturesManager:

    _context: HandleWorkerContext
    _futures: set[asyncio.Future]

    def __init__(self, context: HandleWorkerContext):
        self._context = context
        self._futures = set()


    def ensure_listening_for_work_requests(self):
        requests_in_progress = self._context.num_requests_in_progress
        if requests_in_progress < self._context.config.max_parallel_requests:
            future = self._as_typed_future(FutureType.WORK_REQUEST, self._context.blocking_work_queue.dequeue(timeout=0))
            self._futures.add(future)

    def ensure_listening_to_worker_responses(self):
        future = self._as_typed_future(FutureType.WORKER_RESPONSE, worker_utils.receive_worker_response(websocket=self._context.websocket))
        self._futures.add(future)

    async def wait_for_event_or_timeout(self):
        completed_futures, pending_futures = await asyncio.wait(
            self._futures, timeout=settings.worker_ping_interval, return_when=asyncio.FIRST_COMPLETED
        )
        self._futures = pending_futures

        return [future.result() for future in completed_futures]

Review comment (suggestion):
This is the test that has so far been used to identify a dequeued work request, which has now been replaced by a check for `type_ == FutureType.WORK_REQUEST`

Code after:
            blocking_work_queue=blocking_work_queue,
        )

    async def close(self):
        try:
            await self.blocking_work_queue.redis_client.close()
        except Exception:
            logger.warning("Error while closing redis client")

        try:
            await self.websocket.close()
        except Exception:
            logger.warning("Error while closing websocket")


class FutureType(str, Enum):
    WORK_REQUEST = "WORK_REQUEST"
    WORKER_RESPONSE = "WORKER_RESPONSE"


AwaitResult = TypeVar("AwaitResult")


class FuturesManager:
    _context: HandleWorkerContext
    _futures: set[asyncio.Future]

    def __init__(self, context: HandleWorkerContext):
        self._context = context
        self._futures = set()

    def ensure_listening_for_work_requests(self):
        requests_in_progress = self._context.num_requests_in_progress
        if requests_in_progress < self._context.config.max_parallel_requests:
            future = self._as_typed_future(
                FutureType.WORK_REQUEST, self._context.blocking_work_queue.dequeue(timeout=0)
            )
            self._futures.add(future)

    def ensure_listening_to_worker_responses(self):
        future = self._as_typed_future(
            FutureType.WORKER_RESPONSE, worker_utils.receive_worker_response(websocket=self._context.websocket)
        )
        self._futures.add(future)

    async def wait_for_event_or_timeout(self):
        completed_futures, pending_futures = await asyncio.wait(
            self._futures, timeout=settings.worker_ping_interval, return_when=asyncio.FIRST_COMPLETED
        )
        self._futures = pending_futures