debug server

fixed_server.py

@@ -0,0 +1,318 @@
+import asyncio
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import asynccontextmanager
+import json
+from typing import Annotated
+
+from fastapi import Depends, FastAPI, Header, HTTPException, Request
+from loguru import logger
+from pydantic import BaseModel, Field
+from tqdm import tqdm
+
+from utils.proof_utils import split_proof_header
+from utils.repl_cache import LRUReplCache
+
+from .config import settings
+from .healthcheck import router
+from .kimina_api import router as kimina_router
+from .leanrepl import LeanCrashError, LeanREPL
+
+repls = {}
+semaphore = asyncio.Semaphore(settings.MAX_CONCURRENT_REQUESTS)
+repl_cache = LRUReplCache(max_size=settings.MAX_REPLS)
+
+
+async def _repl_creater():
+    while True:
+        if len(repl_cache.create_queue) > 0:
+            repl_to_create = Counter(repl_cache.create_queue)
+            repl_cache.create_queue = []
+
+            for header, amount in tqdm(repl_to_create.items(), desc="Creating REPLs"):
+                tasks = []
+                creating_repls = []
+                try:
+                    for _ in range(amount):
+                        repl = LeanREPL()
+                        creating_repls.append(repl)
+                        tasks.append(asyncio.to_thread(repl.create_env, header, 600))
+
+                    responses = await asyncio.gather(*tasks)
+                    logger.info(
+                        f"Created {len(responses)} {str([header])[:50]} repls with response {str(responses)[:50]}"
+                    )
+                except LeanCrashError:
+                    for repl in creating_repls:
+                        repl_cache.close_queue.put(repl)
+
+                # put the repls in the cache
+                for repl in creating_repls:
+                    await repl_cache.put(header, repl)
+
+        repl_cache.evict_if_needed()
+        await asyncio.sleep(10)
+
+
+async def _repl_cleaner():
+    while True:
+        await asyncio.sleep(1)
+        while not repl_cache.close_queue.empty():
+            id, repl = repl_cache.close_queue.get()
+            await asyncio.to_thread(repl.close)
+            logger.info(f"Closed {id} repl")
+
+
+async def _stat_printer():
+    update_interval = 15
+    while True:
+        await asyncio.sleep(update_interval)
+        await repl_cache.print_status(update_interval)
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """App lifespan context manager"""
+    app.state.executor = ThreadPoolExecutor(max_workers=5000)
+    asyncio.get_running_loop().set_default_executor(app.state.executor)
+
+    # Repl cache manager tasks
+    relp_cache_tasks = [
+        asyncio.create_task(_repl_cleaner()),
+        asyncio.create_task(_repl_creater()),
+        asyncio.create_task(_stat_printer()),
+    ]
+
+    # Prefill repl_cache, The pre-filled amount should not be greater than settings.MAX_REPLS.
+    # repl_cache.create_queue.extend(["import Mathlib"] * int(settings.MAX_REPLS / 2))
+    # TODO: Make it an initialization parameter.
+    repl_cache.create_queue.extend(
+        ["import Mathlib\nimport Aesop"] * int(settings.MAX_REPLS)
+    )
+
+    try:
+        yield
+    finally:
+        # Cancel cache manager task
+        for task in relp_cache_tasks:
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+
+        # Close thread pools
+        app.state.executor.shutdown(wait=True)
+
+
+app = FastAPI(lifespan=lifespan)
+
+
+# ------ Dependencies ------
+def validate_api_access(request: Request, authorization: str = Header(None)) -> None:
+    api_key = settings.API_KEY
+    if api_key is None:
+        return
+
+    if authorization is None or not authorization.startswith("Bearer "):
+        raise HTTPException(
+            status_code=401, detail="Missing or invalid Authorization header"
+        )
+
+    token = authorization.split("Bearer ")[-1]
+    if token != api_key:
+        raise HTTPException(status_code=403, detail="Invalid API Key")
+
+
+require_access_dep = Annotated[None, Depends(validate_api_access)]
+
+
+# ------ Schemas ------
+class Code(BaseModel):
+    custom_id: str | int
+    proof: str = Field(None)
+    code: str = Field(None)  # To be backward compatibility with autoformalizer client
+
+    def get_proof_content(self) -> str:
+        return self.proof if self.proof is not None else self.code
+
+
+class VerifyRequestBody(BaseModel):
+    codes: list[Code]
+    timeout: int = 300
+    infotree_type: str | None = None
+    disable_cache: bool = False
+
+
+# ------ Endpoint ------
+@app.get("/")
+async def root(require_access_dep: require_access_dep):
+    return {"status": "ok"}
+
+
+@app.post("/verify")
+async def verify(
+    body: VerifyRequestBody,
+    access: require_access_dep,
+):
+    """verify the proof code."""
+    codes = body.codes
+    timeout = body.timeout
+    infotree_type = body.infotree_type
+    disable_cache = body.disable_cache
+
+    tasks = [
+        process_one_code_with_repl_fast(
+            code, timeout, infotree_type, disable_cache=disable_cache
+        )
+        for code in codes
+    ]
+
+    # Await the results of all the tasks concurrently
+    results_data = await asyncio.gather(*tasks)
+
+    results = []
+    for result in results_data:
+        custom_id, error, response = result
+        results.append(
+            {
+                "custom_id": custom_id,
+                "error": error,
+                "response": response,
+            }
+        )
+
+    return {"results": results}
+
+
+async def process_one_code_with_repl_fast(
+    code: Code,
+    timeout: int,
+    infotree_type: str | None,
+    disable_cache: bool = False,
+):
+    # Throttle the incoming request
+    async with semaphore:
+        error_msg = None
+        response = None
+
+        custom_id = code.custom_id
+        proof = code.get_proof_content()
+
+        if proof is None:
+            logger.warning(f"[{custom_id}] No code provided")
+            return custom_id, "No code provided", response
+
+        proof_header, proof_body = split_proof_header(proof)
+
+        log_message = {
+            'custom_id': custom_id,
+            'proof_header': proof_header,
+            'proof_body': proof_body,
+            'timeout': timeout,
+        }
+        logger.debug(
+            f"[{custom_id}] Processing code: {json.dumps(log_message)}"
+        )
+
+        # if we can not found the proof header, create a new repl
+        if len(proof_header.strip()) == 0 or disable_cache:
+            lean_repl = LeanREPL()
+            try:
+                response = await asyncio.to_thread(
+                    lean_repl.one_pass_verify, proof, timeout, infotree_type
+                )
+            except LeanCrashError as e:
+                error_msg = str(e)
+                log_message["error"] = error_msg
+                logger.error(
+                    f"[{custom_id}] Error raised in one_pass_verify with 1-shot repl: {json.dumps(log_message)}"
+                )
+            finally:
+                del lean_repl
+            return custom_id, error_msg, response
+
+        # Get lean repl instance from the lrucache
+        grep_id, repl = await repl_cache.get(proof_header)
+
+        # If we can not get the repl from the lrucache, we will create a new repl
+        if grep_id is None:
+            repl = LeanREPL()
+
+            # And import the proof header
+            try:
+                response = await asyncio.to_thread(
+                    repl.create_env, proof_header, timeout
+                )
+            except LeanCrashError as e:
+                error_msg = str(e)
+                log_message["error"] = error_msg
+                logger.error(
+                    f"[{custom_id}] Error raised in one_pass_verify with 1-shot repl: {json.dumps(log_message)}"
+                )
+                del repl
+                return custom_id, error_msg, response
+
+        try:
+            response = await asyncio.to_thread(
+                repl.extend_env,
+                0,
+                proof_body,
+                timeout,
+                infotree_type,
+            )
+        except LeanCrashError as e:
+            error_msg = str(e)
+            log_message["error"] = error_msg
+            logger.error(
+                f"[{custom_id}] Error raised while extending repl env with proof: {json.dumps(log_message)}"
+            )
+            if grep_id is not None:
+                logger.error(f"[{custom_id}] Removing repl from cache: {grep_id}")
+                await repl_cache.destroy(proof_header, grep_id, repl)
+            else:
+                del repl
+            return custom_id, error_msg, response
+
+        exceeds_limit = False
+        if (
+            settings.REPL_MEMORY_CHECK_INTERVAL is not None
+            and settings.REPL_MEMORY_LIMIT_GB is not None
+            and repl.run_command_total % settings.REPL_MEMORY_CHECK_INTERVAL == 0
+        ):
+            # Check if the REPL exceeds memory limit after execution
+            exceeds_limit = await asyncio.to_thread(
+                repl.exceeds_memory_limit, settings.REPL_MEMORY_LIMIT_GB
+            )
+
+        if exceeds_limit:
+            logger.warning(
+                f"REPL exceeds memory limit after execution, destroying it. last verified proof: {json.dumps(log_message)}"
+            )
+
+            if grep_id is None:
+                del repl
+            else:
+                logger.warning(f"Removing repl from cache: {grep_id}")
+                await repl_cache.destroy(proof_header, grep_id, repl)
+        else:
+            # release back to the cache if memory is within limits
+            if grep_id is None:
+                await repl_cache.put(proof_header, repl)
+            else:
+                await repl_cache.release(proof_header, grep_id, repl)
+
+        return custom_id, error_msg, response
+
+
+@app.post("/one_pass_verify_batch")
+async def one_pass_verify_batch(
+    body: VerifyRequestBody,
+    access: require_access_dep,
+):
+    """Backward compatible endpoint: accepts both 'proof' / 'code' fields."""
+    return await verify(body, access)
+
+
+app.include_router(router)
+app.include_router(kimina_router)

kimina_api.py

@@ -0,0 +1,144 @@
+from fastapi import APIRouter, HTTPException, status, Depends
+from loguru import logger
+from typing import List, Optional, Union, Dict, Any
+from pydantic import BaseModel, Field
+
+from .server import require_access_dep, process_one_code_with_repl_fast, Code, VerifyRequestBody
+from .config import settings
+
+router = APIRouter()
+
+
+# ------ Kimina Client Compatible Schemas ------
+class Snippet(BaseModel):
+    id: str
+    code: str = Field(alias="snip")
+
+    class Config:
+        populate_by_name = True
+
+
+class CheckRequest(BaseModel):
+    snippets: List[Snippet]
+    timeout: int = 300
+    debug: bool = False
+    reuse: bool = True
+    infotree: Optional[Dict[str, Any]] = None
+
+
+class Message(BaseModel):
+    severity: str
+    pos: Optional[Dict[str, int]] = None
+    endPos: Optional[Dict[str, int]] = None
+    data: str
+
+
+class Sorry(BaseModel):
+    pos: Optional[Dict[str, int]] = None
+    endPos: Optional[Dict[str, int]] = None
+    goal: str
+
+
+class ReplResponse(BaseModel):
+    id: str
+    error: Optional[str] = None
+    response: Optional[Dict[str, Any]] = None
+
+    def analyze(self):
+        class Status:
+            def __init__(self, status_value):
+                self.value = status_value
+                self.status = status_value
+
+            @property
+            def status(self):
+                return self._status
+
+            @status.setter
+            def status(self, value):
+                self._status = value
+
+        if self.error:
+            return Status("error")
+        if self.response is None:
+            return Status("unknown")
+
+        # Check if there are error messages
+        messages = self.response.get("messages", [])
+        for msg in messages:
+            if msg.get("severity") == "error":
+                return Status("invalid")
+
+        return Status("valid")
+
+
+class CheckResponse(BaseModel):
+    results: List[ReplResponse]
+
+
+# ------ Kimina Client Compatible Endpoint ------
+@router.post("/api/check")
+async def kimina_check(
+    request: CheckRequest,
+    access: None = Depends(require_access_dep)
+):
+    """
+    Kimina client compatible endpoint for code verification.
+    Converts Kimina client format to server format.
+    """
+    try:
+        # Convert Kimina format to server format
+        codes = []
+        for snippet in request.snippets:
+            code_obj = Code(
+                custom_id=snippet.id,
+                code=snippet.code
+            )
+            codes.append(code_obj)
+
+        # Create server request body
+        server_request = VerifyRequestBody(
+            codes=codes,
+            timeout=request.timeout,
+            infotree_type=request.infotree.get("type") if request.infotree else None,
+            disable_cache=not request.reuse
+        )
+
+        # Process using existing verify function infrastructure
+        tasks = [
+            process_one_code_with_repl_fast(
+                code,
+                server_request.timeout,
+                server_request.infotree_type,
+                disable_cache=server_request.disable_cache
+            )
+            for code in server_request.codes
+        ]
+
+        # Await the results of all the tasks concurrently
+        results_data = await asyncio.gather(*tasks)
+
+        # Convert results to Kimina format
+        results = []
+        for result in results_data:
+            custom_id, error, response = result
+            repl_response = ReplResponse(
+                id=str(custom_id),
+                error=error,
+                response=response
+            )
+            results.append(repl_response)
+
+        return CheckResponse(results=results)
+
+    except Exception as e:
+        logger.error(f"Error in kimina_check endpoint: {e}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Internal server error: {str(e)}"
+        )
+
+
+# Import asyncio for the gather function
+import asyncio
+

start.sh

@@ -1,12 +1,7 @@
 #!/bin/bash
 set -e
 
-echo "LEAN_SERVER_REPL_PATH: $LEAN_SERVER_REPL_PATH"
-echo "LEAN_SERVER_PROJECT_DIR: $LEAN_SERVER_PROJECT_DIR"
-
-ls -la $LEAN_SERVER_REPL_PATH
-ls -la $LEAN_SERVER_PROJECT_DIR
-
+export PATH="~/.elan/bin:${PATH}"
 which lake
 which elan
 lake --version
@@ -14,6 +9,12 @@ elan --version
 lean --version
 elan toolchain list
 
+echo "LEAN_SERVER_REPL_PATH: $LEAN_SERVER_REPL_PATH"
+echo "LEAN_SERVER_PROJECT_DIR: $LEAN_SERVER_PROJECT_DIR"
+
+ls -la $LEAN_SERVER_REPL_PATH
+ls -la $LEAN_SERVER_PROJECT_DIR
+
 cd /home/user/app/kimina-lean-server
 python3 -m server &
 
@@ -22,6 +23,9 @@ ps aux | grep python
 netstat -tlnp | grep 8888
 ps aux | grep server
 
+free -h
+top -p 15
+
 echo "curl -v http://localhost:8888/"
 curl -v http://localhost:8888/
 echo "curl -v http://localhost:8888/health"
@@ -30,6 +34,8 @@ echo "curl -v http://localhost:8888/api/check"
 curl -v http://localhost:8888/api/check
 echo "curl -v http://localhost:8888/docs"
 curl -v http://localhost:8888/docs
+echo "curl -v http://localhost:8888/verify"
+curl -v http://localhost:8888/verify
 
 cd /home/user/app
 ls -l