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

Code before:

    @property
    def num_requests_in_progress(self) -> int:
        return len(self.work_request_map)

    @classmethod
    async def build(cls, websocket: fastapi.WebSocket, api_key: str, protocol_version: str):
        try:
            worker_utils.get_protocol_version(protocol_version)
            api_key = worker_utils.get_api_key(api_key)
            worker_id = await worker_utils.get_worker_id(api_key=api_key, protocol_version=protocol_version)
        except fastapi.HTTPException as e:
            logger.warning(f"handle_worker: {e.status_code=} {e.detail=}")
            if e.status_code == fastapi.status.HTTP_426_UPGRADE_REQUIRED:
                await worker_utils.send_worker_request(websocket=websocket, request=inference.UpgradeProtocolRequest())
            elif e.status_code == fastapi.status.HTTP_401_UNAUTHORIZED:
                await worker_utils.send_worker_request(websocket=websocket, request=inference.WrongApiKeyRequest())
            try:
                await websocket.close(code=e.status_code, reason=e.detail)
            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:

Review comment (suggestion):
Moved to `HandleWorkerContext.build()` without changes

Code after:
    @property
    def num_requests_in_progress(self) -> int:
        return len(self.work_request_map)

    @classmethod
    async def build(cls, websocket: fastapi.WebSocket, api_key: str, protocol_version: str):
        try:
            worker_utils.get_protocol_version(protocol_version)
            api_key = worker_utils.get_api_key(api_key)
            worker_id = await worker_utils.get_worker_id(api_key=api_key, protocol_version=protocol_version)
        except fastapi.HTTPException as e:
            logger.warning(f"handle_worker: {e.status_code=} {e.detail=}")
            if e.status_code == fastapi.status.HTTP_426_UPGRADE_REQUIRED:
                await worker_utils.send_worker_request(websocket=websocket, request=inference.UpgradeProtocolRequest())
            elif e.status_code == fastapi.status.HTTP_401_UNAUTHORIZED:
                await worker_utils.send_worker_request(websocket=websocket, request=inference.WrongApiKeyRequest())
            try:
                await websocket.close(code=e.status_code, reason=e.detail)
            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")