Add back health check to ephemeral_server (#277)

Author: caroljung-cg

.github/workflows/unit-tests.yml

@@ -22,12 +22,12 @@ jobs:
         uses: ./.github/actions/setup-backend
       - name: Run ATS and Tests
         uses: ./.github/actions/run_ats
-        timeout-minutes: 5
+        timeout-minutes: 15
         with:
           default_tests: "tests/unit"
           codecov_static_token: ${{ secrets.CODECOV_STATIC_TOKEN }}
           codecov_token: ${{ secrets.CODECOV_TOKEN }}
-          collect_args: "--timeout 5"
+          collect_args: "--timeout 15"
           codecov_flags: unit-tests
   codemod-tests:
     # changing the following value will significantly affect github's cost. Be careful and consult with the team before changing it.
@@ -53,7 +53,7 @@ jobs:
         uses: ./.github/actions/setup-oss-repos
       - name: Run ATS and Tests
         uses: ./.github/actions/run_ats
-        timeout-minutes: 10
+        timeout-minutes: 15
         with:
           default_tests: "tests/integration/codemod/test_codemods.py"
           codecov_static_token: ${{ secrets.CODECOV_STATIC_TOKEN }}

pyproject.toml

@@ -60,6 +60,7 @@ dependencies = [
   "starlette<1.0.0,>=0.16.0",
   "tqdm>=4.67.1",
   "tomlkit>=0.13.2",
+  "uvicorn[standard]>=0.30.0",
 ]
 
 license = { text = "Apache-2.0" }
@@ -138,7 +139,7 @@ dev-dependencies = [
   "pytest-benchmark[histogram]>=5.1.0",
   "loguru>=0.7.3",
 ]
-#skeyring-provider = "subprocess"
+
 
 [tool.uv.workspace]
 members = []

src/codegen/runner/models/apis.py

@@ -18,12 +18,10 @@
 
 
 class ServerInfo(BaseModel):
-    repo_id: int = 0
-    container_id: str = ""
+    repo_name: str | None = None
     is_running_codemod: bool = False
     is_shutting_down: bool = False
     warmup_state: WarmupState = WarmupState.PENDING
-    label: str | None = ""
 
 
 class UtilizationMetrics(BaseModel):

src/codegen/runner/sandbox/ephemeral_server.py

@@ -1,24 +1,46 @@
 import logging
 import tempfile
+from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
 
+from codegen.runner.enums.warmup_state import WarmupState
 from codegen.runner.models.apis import (
     RUN_ON_STRING_ENDPOINT,
     GetRunOnStringRequest,
     GetRunOnStringResult,
+    ServerInfo,
 )
 from codegen.runner.sandbox.executor import SandboxExecutor
 from codegen.sdk.codebase.factory.get_session import get_codebase_session
 from codegen.sdk.enums import ProgrammingLanguage
 from codegen.shared.compilation.string_to_code import create_execute_function_from_codeblock
 
 logger = logging.getLogger(__name__)
-app = FastAPI()
+
+server_info: ServerInfo
+
+
+@asynccontextmanager
+async def lifespan(server: FastAPI):
+    global server_info
+    server_info = ServerInfo(warmup_state=WarmupState.COMPLETED)
+    logger.info("Ephemeral server is ready to accept requests")
+    yield
+    logger.info("Shutting down fastapi server")
+
+
+app = FastAPI(lifespan=lifespan)
+
+
+@app.get("/")
+def health() -> ServerInfo:
+    return server_info
 
 
 @app.post(RUN_ON_STRING_ENDPOINT)
 async def run_on_string(request: GetRunOnStringRequest) -> GetRunOnStringResult:
+    server_info.is_running_codemod = True
     logger.info(f"====[ run_on_string ]====\n> Codemod source: {request.codemod_source}\n> Input: {request.files}\n> Language: {request.language}\n")
     language = ProgrammingLanguage(request.language.upper())
     with get_codebase_session(tmpdir=tempfile.mkdtemp(), files=request.files, programming_language=language) as codebase:
@@ -27,3 +49,9 @@ async def run_on_string(request: GetRunOnStringRequest) -> GetRunOnStringResult:
         result = await executor.execute(code_to_exec)
         logger.info(f"Result: {result}")
         return GetRunOnStringResult(result=result)
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)

src/codegen/runner/sandbox/server.py

@@ -7,7 +7,6 @@
 import psutil
 from fastapi import FastAPI
 
-from codegen.runner.constants.envvars import CUSTOMER_REPO_ID
 from codegen.runner.enums.warmup_state import WarmupState
 from codegen.runner.models.apis import (
     BRANCH_ENDPOINT,
@@ -38,10 +37,11 @@ async def lifespan(server: FastAPI):
     global runner
 
     try:
-        server_info = ServerInfo(repo_id=int(os.getenv(CUSTOMER_REPO_ID)), container_id=os.getenv("MODAL_TASK_ID"))
+        repo_config = get_repo_config()
+        server_info = ServerInfo(repo_name=repo_config.full_name)
         logger.info(f"Starting up sandbox fastapi server for repo_id={server_info.repo_id} in container ID={server_info.container_id}")
 
-        runner = SandboxRunner(container_id=server_info.container_id, repo_config=get_repo_config())
+        runner = SandboxRunner(container_id=server_info.container_id, repo_config=repo_config)
         server_info.warmup_state = WarmupState.PENDING
         await runner.warmup()
         server_info.warmup_state = WarmupState.COMPLETED