diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml index f61f7a827..0a3581718 100644 --- a/.github/workflows/unit-tests.yml +++ b/.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 }} diff --git a/pyproject.toml b/pyproject.toml index 43875c0f4..b9f4a3881 100644 --- a/pyproject.toml +++ b/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 = [] diff --git a/src/codegen/runner/models/apis.py b/src/codegen/runner/models/apis.py index d33cadd6f..2b7cb918a 100644 --- a/src/codegen/runner/models/apis.py +++ b/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): diff --git a/src/codegen/runner/sandbox/ephemeral_server.py b/src/codegen/runner/sandbox/ephemeral_server.py index 85be575f2..12af608df 100644 --- a/src/codegen/runner/sandbox/ephemeral_server.py +++ b/src/codegen/runner/sandbox/ephemeral_server.py @@ -1,12 +1,15 @@ 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 @@ -14,11 +17,30 @@ 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) diff --git a/src/codegen/runner/sandbox/server.py b/src/codegen/runner/sandbox/server.py index 6d63bc24c..2a97cd853 100644 --- a/src/codegen/runner/sandbox/server.py +++ b/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 diff --git a/uv.lock b/uv.lock index 6f88f0385..ed2e747f6 100644 --- a/uv.lock +++ b/uv.lock @@ -392,6 +392,7 @@ dependencies = [ { name = "tree-sitter-typescript" }, { name = "typing-extensions" }, { name = "unidiff" }, + { name = "uvicorn", extra = ["standard"] }, { name = "watchfiles" }, { name = "wrapt" }, { name = "xmltodict" }, @@ -491,6 +492,7 @@ requires-dist = [ { name = "types-tabulate", marker = "extra == 'types'", specifier = ">=0.9.0.20240106" }, { name = "typing-extensions", specifier = ">=4.12.2" }, { name = "unidiff", specifier = ">=0.7.5" }, + { name = "uvicorn", extras = ["standard"], specifier = ">=0.30.0" }, { name = "watchfiles", specifier = ">=1.0.0,<1.1.0" }, { name = "wrapt", specifier = ">=1.16.0,<2.0.0" }, { name = "xmltodict", specifier = ">=0.13.0,<1.0.0" },