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

Code before:
            except Exception:
                pass
            raise fastapi.WebSocketException(e.status_code, e.detail)

        logger.info(f"handle_worker: {worker_id=}")
        worker_info = await worker_utils.receive_worker_info(websocket)
        logger.info(f"handle_worker: {worker_info=}")


        work_queue = queueing.work_queue(deps.redis_client, worker_info.config.compat_hash)

        # Make a new Redis client that we can use for the blocking requests without affecting the non-blocking calls
        redis_client = deps.make_redis_client()
        blocking_work_queue = queueing.work_queue(redis_client, worker_info.config.compat_hash)

        return cls(
            websocket=websocket,
            worker_id=worker_id,
            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()


Review comment (suggestion):
Moved to `FuturesManager.ensure_listening_for_worker_response()` with similar changes to return the `FutureType` in addition to the response.

Code after:
                pass
            raise fastapi.WebSocketException(e.status_code, e.detail)

        logger.info(f"handle_worker: {worker_id=}")
        worker_info = await worker_utils.receive_worker_info(websocket)
        logger.info(f"handle_worker: {worker_info=}")

        work_queue = queueing.work_queue(deps.redis_client, worker_info.config.compat_hash)

        # Make a new Redis client that we can use for the blocking requests without affecting the non-blocking calls
        redis_client = deps.make_redis_client()
        blocking_work_queue = queueing.work_queue(redis_client, worker_info.config.compat_hash)

        return cls(
            websocket=websocket,
            worker_id=worker_id,
            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