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

Code before:

async def initiate_work_for_message(
    *,
    websocket: fastapi.WebSocket,
    work_queue: queueing.RedisQueue,
    message_id: str,
    worker_id: str,
    worker_config: inference.WorkerConfig,
) -> inference.WorkRequest:
    async with deps.manual_create_session() as session:
        cr = chat_repository.ChatRepository(session=session)

        message = await cr.start_work(
            message_id=message_id,
            worker_id=worker_id,
            worker_config=worker_config,
        )
        work_request = await worker_utils.build_work_request(session, message.id)

    logger.info(f"Created {work_request=} with {len(work_request.thread.messages)=}")
    try:
        await worker_utils.send_worker_request(websocket, work_request)
    except Exception as e:
        logger.exception(f"Error while sending work request to worker: {str(e)}")
        async with deps.manual_create_session() as session:
            await cr.reset_work(message_id)
        await work_queue.enqueue(message_id, enforce_max_size=False)
        raise

    return work_request


async def handle_token_response(
    response: inference.TokenResponse,
    work_request_map: WorkRequestContainerMap,
):
    work_response_container = get_work_request_container(work_request_map, response.request_id)
    message_queue = queueing.message_queue(
        deps.redis_client,
        message_id=work_response_container.message_id,
    )
    await message_queue.enqueue(response.json())
    work_response_container.num_responses += 1


async def handle_plugin_intermediate_response(
    response: inference.PluginIntermediateResponse,
    work_request_map: WorkRequestContainerMap,
):
    work_response_container = get_work_request_container(work_request_map, response.request_id)
    message_queue = queueing.message_queue(

Review comment (suggestion):
In order to avoid logical changes as much as possible, I kept this piece of code here as it was, but that now means that it accesses the internal structure of the `FuturesManager`.

I'm considering whether it might be nicer to not throw all the futures into a single set, but instead keep one future in a specific variable for listening to worker responses, and a set for listening to the potentially multiple futures with which we listen for work requests. The futures manager could then have only a single method that ensures that there's at least one pending future of each type, and only that would be called from the outside at the end of each loop

Code after:
                await session.update(worker_response.metrics)
            case "general_error":
                worker_response = cast(inference.GeneralErrorResponse, worker_response)
                await handle_general_error_response(
                    response=worker_response,
                )
                await session.update(worker_response.metrics)
            case "safe_prompt":
                logger.info("Received safe prompt response")
                worker_response = cast(inference.SafePromptResponse, worker_response)
                await handle_safe_prompt_response(
                    response=worker_response,
                    work_request_map=context.work_request_map,
                )
            case "plugin_intermediate":
                worker_response = cast(inference.PluginIntermediateResponse, worker_response)
                await handle_plugin_intermediate_response(
                    work_request_map=context.work_request_map,
                    response=worker_response,
                )
            case _:
                raise RuntimeError(f"Unknown response type: {worker_response.response_type}")
    finally:
        if len(futures._futures) == 0:
            futures.ensure_listening_for_work_requests()
        futures.ensure_listening_to_worker_responses()


async def _reset_or_abort_open_work_requests(context: HandleWorkerContext):
    logger.info(f"Handling {context.num_requests_in_progress} work requests outstanding")
    for container in context.work_request_map.values():
        try:
            message_id = container.message_id
            if container.num_responses == 0:
                logger.warning(f"Marking {message_id=} as pending since no work was done.")
                async with deps.manual_chat_repository() as cr:
                    await cr.reset_work(message_id)
                await context.work_queue.enqueue(message_id, enforce_max_size=False)
            else:
                logger.warning(f"Aborting {message_id=}")
                await abort_message(message_id=message_id, error="Aborted due to worker error.")
        except Exception as e:
            logger.exception(f"Error while trying to reset work for {message_id=}: {str(e)}")


@router.websocket("/work")
async def handle_worker(
    websocket: fastapi.WebSocket,
    api_key: str = worker_utils.api_key_header,
    protocol_version: str = worker_utils.protocol_version_header,
):