server/app.py

from __future__ import annotations

import os

try:
    from openenv.core.env_server.http_server import create_app
except Exception as e:  # pragma: no cover
    raise ImportError(
        "openenv is required for the web interface. Install dependencies with `uv sync`."
    ) from e

from fastapi.responses import HTMLResponse

from models import GPUSchedulerAction, GPUSchedulerObservation
from server.custom_web import CUSTOM_WEB_HTML
from server.gpu_scheduler_openenv_environment import GPUSchedulerOpenEnvEnvironment

os.environ.setdefault("ENABLE_WEB_INTERFACE", "true")

app = create_app(
    GPUSchedulerOpenEnvEnvironment,
    GPUSchedulerAction,
    GPUSchedulerObservation,
    env_name="gpu_scheduler_openenv",
    max_concurrent_envs=2,
)


def _replace_default_web_route() -> None:
    app.router.routes[:] = [
        route
        for route in app.router.routes
        if not (getattr(route, "path", None) == "/web" and "GET" in getattr(route, "methods", set()))
    ]


_replace_default_web_route()


@app.get("/web", response_class=HTMLResponse, include_in_schema=False)
def custom_web() -> HTMLResponse:
    return HTMLResponse(CUSTOM_WEB_HTML)


def main(host: str = "0.0.0.0", port: int = 8000):
    import uvicorn

    uvicorn.run(app, host=host, port=port)


if __name__ == "__main__":
    main()