main.py

from fastapi import FastAPI, HTTPException
import uvicorn

from inference_types import (
    GenerateRequest,
    GenerateResponse,
    BatchGenerateRequest,
    BatchGenerateResponse,
)

from inference import batch_inference

app = FastAPI()


@app.post("/generate")
async def generate(generate_request: GenerateRequest) -> GenerateResponse:
    try:
        print("Request received:", generate_request)
        batch_request = BatchGenerateRequest(
            batch_messages=[generate_request.messages],
            max_output_length=generate_request.max_output_length,
            temperature=generate_request.temperature,
            top_p=generate_request.top_p,
            top_k=generate_request.top_k,
            repetition_penalty=generate_request.repetition_penalty,
            streaming=generate_request.streaming,
        )
        batch_response: BatchGenerateResponse = batch_inference(batch_request)
        response = GenerateResponse(
            response=batch_response.batch_responses[0][0],
            generation_tokens=batch_response.batch_generation_tokens[0],
            generation_time=batch_response.generation_time,
            tokens_per_second=batch_response.max_throughput_tokens_per_second,
        )
        print("Response:", response)
        return response
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/batch_generate")
async def batch_generate(
    batch_generate_request: BatchGenerateRequest,
) -> BatchGenerateResponse:
    try:
        print("Batch request received:", batch_generate_request)
        batch_response = batch_inference(batch_generate_request)
        print("Batch response:", batch_response)
        return batch_response
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    from mpi4py import MPI

    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    if rank == 0:
        print("Inference server ready!")
        uvicorn.run(app, host="0.0.0.0", port=8000)