api/app.py

import logging
import os
from collections import Counter, defaultdict
import multiprocessing
from datetime import datetime
from concurrent.futures import ProcessPoolExecutor, as_completed
from typing import Dict, List, Tuple
import gc

from fastapi import FastAPI, HTTPException
from fastapi.responses import RedirectResponse

from api.code_execution import untrusted_check

Result = Tuple[str, List[bool]]

def create_app() -> FastAPI:

    level = os.environ.get("LOG_LEVEL", default=logging.INFO)
    logging.basicConfig(level=level)
    logger = logging.getLogger(__name__)

    app = FastAPI()

    @app.get("/")
    def root():
        return RedirectResponse("/docs")

    @app.get("/health", status_code=204)
    def health():
        return

    @app.post("/evaluate/")
    async def evaluate(
        samples: List[dict],
        calibrate: bool = True,
        parallel: int = -1,
        min_time_limit: float = 1,
        max_as_limit: int = 30 * 1024,
        max_data_limit: int = 30 * 1024,
        max_stack_limit: int = 10,
        no_gt: bool = True,
    ) -> dict:
        """
        Evaluate the correctness of the solutions in the given samples data.
        """
        try:
            if parallel < 1:
                n_workers = min(4, max(1, multiprocessing.cpu_count() - 1))

            else:
                n_workers = parallel

            if not no_gt:
                expected_time = get_groundtruth()
            else:
                expected_time = {}

            results = {
                "date": datetime.now().strftime("%Y-%m-%d %H:%M"),
                "eval": {},
            }

            with ProcessPoolExecutor(max_workers=n_workers) as executor:
                futures = []
                completion_id = Counter()
                n_samples = 0
                eval_results = defaultdict(list)  # task_id ->
                remainings = set()

                for i, sample in enumerate(samples):
                    # TODO: investigate why HTTPException detail is not passed to client.

                    required_keys = ["task_id", "res_id", "test", "solution", "entry_point"]
                    if calibrate:
                        required_keys.append("code_prompt")
                    for key in required_keys:
                        if key not in sample:
                            raise HTTPException(status_code=400, detail=f"'{key}' not in sample {i}!")

                    if not isinstance(sample["solution"], str):
                        raise HTTPException(status_code=400, detail="Solution must be a string!")

                    sample["_identifier"] = (
                        sample["task_id"] + f" (line {i+1} )"
                    )

                    task_id = sample["task_id"]

                    solution = sample["solution"]

                    if calibrate:
                        solution = sample["code_prompt"] + "\n    pass\n" + solution
                    remainings.add(sample["_identifier"])
                    args = (
                        completion_id[task_id],
                        sample["res_id"],
                        task_id,
                        solution,
                        sample["test"],
                        sample["entry_point"],
                        max_as_limit,
                        max_data_limit,
                        max_stack_limit,
                        sample["_identifier"],
                        min_time_limit,
                        expected_time.get(task_id) if expected_time.get(task_id) else 20
                    )
                    futures.append(executor.submit(check_correctness, *args))
                    completion_id[task_id] += 1
                    n_samples += 1

                if n_samples != len(remainings):
                    logger.warning(f"Sample count mismatch: {n_samples} samples but {len(remainings)} unique identifiers")

                for future in as_completed(futures):
                    try:
                        result = future.result()
                    except Exception as e:
                        logger.error(f"Worker process failed: {e}")
                        continue
                    remainings.discard(result["_identifier"])
                    eval_results[result["task_id"]].append(result)
                    del future, result
                    gc.collect()

            # sort the results for each problem by completion_id
            for task_id, task_results in eval_results.items():
                task_results.sort(key=lambda x: x["completion_id"])
                results["eval"][task_id] = []
                for res in task_results:
                    stat, details = res["base"]
                    results["eval"][task_id].append(
                        {
                            "res_id": res["res_id"],
                            "task_id": task_id,
                            "solution": res["solution"],
                            "status": stat,
                            "details": details,
                        }
                    )
            return results
        except HTTPException:
            raise
        except Exception as e:
            logger.exception(f"Evaluate endpoint failed: {e}")
            raise HTTPException(status_code=500, detail=str(e))

    return app

def check_correctness(
    completion_id: int,
    res_id: int,
    task_id: str,
    solution: str,
    test: str,
    entry_point: str,
    max_as_limit: float,
    max_data_limit: float,
    max_stack_limit: float,
    identifier=None,
    min_time_limit: float = 0.1,
    gt_time_limit: float = 2.0,
) -> Dict[str, Result]:  
    ret = {
        "completion_id": completion_id,
        "res_id": res_id,
        "task_id": task_id,
        "_identifier": identifier,
        "solution": solution,
    }
    ret["base"] = untrusted_check(
        solution,
        test,
        entry_point,
        max_as_limit,
        max_data_limit,
        max_stack_limit,
        min_time_limit,
        gt_time_limit,
    )
    return ret


def get_groundtruth():
    raise HTTPException(status_code=405, detail="Groundtruth execution is not implemented yet!")