mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
e9cfba2c17
* feat(lens): isolate trace storage and investigation in a Rust service
* fix(lens): include Rust sources in the image build context
* feat(lens): wire service setup, scoped delivery receipts and lease attempts
* fix(lens): complete service routing and reject stale investigation results
* fix(lens): retry key propagation and validate isolated Compose setup
* fix(lens): seed through isolated ingestion and preserve upstream queue fixes
* chore: sync schema.prisma copies from root
* fix(lens): bind nullable due timestamps as text for Prisma
* chore(ui): remove stale lint suppressions
* fix(lens): address CI failures and review findings
* refactor(lens): remove retired Python worker and run evaluations in Rust
* fix(lens): reuse control connections and satisfy review checks
* test(lens): install and upgrade both Helm charts on Kubernetes
* test(lens): run connection reuse coverage as an integration test
* fix(ui): upgrade Next.js to 16.3.8 security release
* fix(lens): fence stale attempts and preserve reviewed evidence
* Revert "fix(ui): upgrade Next.js to 16.3.8 security release"
This reverts commit 2f79a51b25.
* fix(lens): stop failed investigations and stream history excerpts
* test(lens): cover model tool and result contracts
* test(lens): fix retired routes and reuse installation build artifacts
* test(lens): use portable grep in Helm installation smoke
* test(lens): wait for migrations before forwarding Helm services
* fix(lens): keep failed evidence reads retryable
* fix(lens): preserve sandbox output during process exit
---------
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
105 lines
3.9 KiB
Python
105 lines
3.9 KiB
Python
import asyncio
|
|
import os
|
|
import secrets
|
|
import socket
|
|
from collections.abc import Awaitable, Callable
|
|
from contextlib import suppress
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import uvicorn
|
|
from fastapi import Depends, FastAPI, Header, HTTPException
|
|
|
|
from litellm.proxy.lens.models import Claim, ExecutionContent, ModelRequest, ModelResult, Progress, Result, Sample
|
|
from litellm.proxy.lens.release import PROTOCOL_VERSION
|
|
|
|
|
|
async def run_worker(
|
|
binary: Path,
|
|
claim: Claim,
|
|
sample: Sample,
|
|
read: Callable[[str, str, int], Awaitable[ExecutionContent]],
|
|
model: Callable[[ModelRequest], Awaitable[ModelResult]],
|
|
progress: Callable[[Progress], Awaitable[None]],
|
|
) -> Result:
|
|
token: Final = secrets.token_urlsafe(32)
|
|
release: Final = "lens-evaluation"
|
|
|
|
def auth(authorization: str = Header()) -> None:
|
|
if not secrets.compare_digest(authorization, "Bearer " + token):
|
|
raise HTTPException(401, "Invalid worker credential")
|
|
|
|
app: Final = FastAPI(dependencies=[Depends(auth)])
|
|
completed: Final = asyncio.Future[Result]()
|
|
|
|
@app.post("/lens/worker/claim")
|
|
async def take(protocol_version: int, worker_release: str) -> Claim:
|
|
if protocol_version != PROTOCOL_VERSION or worker_release != release:
|
|
raise HTTPException(409, "Incompatible worker")
|
|
return claim
|
|
|
|
@app.get("/lens/worker/{lens_id}/{job_id}/sample")
|
|
async def sampled(lens_id: str, job_id: str) -> Sample:
|
|
return sample
|
|
|
|
@app.get("/lens/worker/{lens_id}/{job_id}/reviews")
|
|
async def reviews(lens_id: str, job_id: str) -> tuple[()]:
|
|
return ()
|
|
|
|
@app.get("/lens/worker/{lens_id}/{job_id}/content")
|
|
async def content(
|
|
lens_id: str, job_id: str, execution_id: str, cursor: str = "", offset: int = 1
|
|
) -> ExecutionContent:
|
|
return await read(execution_id, cursor, max(0, offset - 1))
|
|
|
|
@app.post("/lens/worker/{lens_id}/{job_id}/model")
|
|
async def infer(lens_id: str, job_id: str, body: ModelRequest) -> ModelResult:
|
|
return await model(body)
|
|
|
|
@app.post("/lens/worker/{lens_id}/{job_id}/progress")
|
|
async def update(lens_id: str, job_id: str, body: Progress) -> bool:
|
|
await progress(body)
|
|
return True
|
|
|
|
@app.post("/lens/worker/{lens_id}/{job_id}/heartbeat")
|
|
async def heartbeat(lens_id: str, job_id: str) -> bool:
|
|
return True
|
|
|
|
@app.post("/lens/worker/{lens_id}/{job_id}/result")
|
|
async def result(lens_id: str, job_id: str, body: Result) -> bool:
|
|
if not completed.done():
|
|
completed.set_result(body)
|
|
return True
|
|
|
|
with socket.socket() as listener:
|
|
listener.bind(("127.0.0.1", 0))
|
|
server: Final = uvicorn.Server(uvicorn.Config(app, log_level="error", access_log=False))
|
|
serving: Final = asyncio.create_task(server.serve(sockets=[listener]))
|
|
try:
|
|
while not server.started:
|
|
if serving.done():
|
|
await serving
|
|
raise RuntimeError("Evaluation gateway failed to start")
|
|
await asyncio.sleep(0.01)
|
|
process: Final = await asyncio.create_subprocess_exec(
|
|
str(binary.resolve()),
|
|
env={
|
|
**os.environ,
|
|
"LITELLM_URL": f"http://127.0.0.1:{listener.getsockname()[1]}",
|
|
"LENS_WORKER_TOKEN": token,
|
|
"LITELLM_RELEASE_TAG": release,
|
|
},
|
|
)
|
|
try:
|
|
exit_code: Final = await process.wait()
|
|
if exit_code != 0 or not completed.done():
|
|
raise RuntimeError(f"Rust worker exited without a result (exit {exit_code})")
|
|
return completed.result()
|
|
finally:
|
|
if process.returncode is None:
|
|
process.kill()
|
|
await process.wait()
|
|
finally:
|
|
server.should_exit = True
|
|
with suppress(asyncio.CancelledError):
|
|
await serving
|