run-codestral.py

from fastapi import Request
from fastapi.responses import JSONResponse
from llama_cpp.server.settings import ServerSettings, ModelSettings
from llama_cpp.server.app import create_app
import uvicorn
import llama_cpp
import asyncio

def dprint(*args, **kwargs):
    print(*args, **kwargs)

model_path = '/home/zbuntu/projects/Codestral-22B-v0.1-Q5_K_M.gguf'

server_settings = ServerSettings()
server_settings.host = "127.0.0.1"
server_settings.port = 4321

model_setting: dict = {
    "model": model_path,
    "model_alias": "codegemma-7b8bit",
    "n_gpu_layers": -1,
    "rope_scaling_type": llama_cpp.LLAMA_ROPE_SCALING_TYPE_LINEAR,
    "rope_freq_base": 1000000.0,
    "rope_freq_scale": 1,
    "n_ctx": 16000,
    "n_threads":1
}

app = create_app(
    server_settings=server_settings,
    model_settings=[ModelSettings(**model_setting)],
)
 
lock = asyncio.Lock()
queue = asyncio.Queue(maxsize=1)  # Set the queue size to 1

async def wait_for_disconnect(request: Request, queue):
    dprint('starting to wait for disconnect')
    while True:
        if await request.is_disconnected():
            break
        await asyncio.sleep(0.5)
    dprint('client disconnected')
    await queue.get()
    dprint('request removed from queue')


@app.middleware("http")
async def synchronize_requests(request: Request, call_next):
    try:
        dprint('adding request to queue')
        await queue.put(request)  # Add the request to the queue

        dprint('request added to queue')
        async with lock:
            dprint('acquired lock serving request')
            try: 
                # Process the request
                response = await call_next(request)
                dprint('response received')
                asyncio.create_task(wait_for_disconnect(request, queue))
                dprint('wait for disconnect thread created')
                # print(response)
                return response
            finally:
                dprint('request processed')
    except asyncio.QueueFull:
        dprint('exceptoin encountered')
        # If the queue is full, return a 503 error response
        return JSONResponse(status_code=503, content={"error": "Service unavailable, try again later"})


uvicorn.run(
    app,
    host=server_settings.host,
    port=server_settings.port,
    ssl_keyfile=server_settings.ssl_keyfile,
    ssl_certfile=server_settings.ssl_certfile,
)