sasikumarM/detraff-env

1import uvicorn2import argparse3from fastapi import FastAPI4import os5 6# OpenEnv imports7try:8    from openenv.core.env_server.http_server import create_app9except ImportError as e:10    raise ImportError(11        "openenv is required. Install dependencies with 'uv sync' or 'pip install openenv-core'"12    ) from e13 14# Internal Project imports15# We use absolute imports to ensure compatibility with Docker and HF Spaces16try:17    from models import DetraffAction, DetraffObservation18    from server.detraff_env_environment import DetraffEnvironment19except ImportError:20    # Fallback for different execution contexts21    from .models import DetraffAction, DetraffObservation22    from .detraff_env_environment import DetraffEnvironment23 24# 1. Create the app instance25app = create_app(26    DetraffEnvironment,27    DetraffAction,28    DetraffObservation,29    env_name="detraff_env",30    max_concurrent_envs=1,31)32 33# 2. Server Available Tasks34@app.get("/tasks")35async def get_tasks():36    """37    Manually serves the task list from openenv.yaml to satisfy the validator.38    """39    yaml_path = os.path.join(os.getcwd(), "openenv.yaml")40    try:41        with open(yaml_path, "r") as f:42            config = yaml.safe_load(f)43        return config.get("tasks", [])44    except Exception:45        # Fallback if YAML is not found46        return [47            {"id": "low_traffic", "name": "low_traffic", "grader": {"type": "reward_threshold", "threshold": 0.8}},48            {"id": "normal_traffic", "name": "normal_traffic", "grader": {"type": "reward_threshold", "threshold": 0.5}},49            {"id": "emergency_peak", "name": "emergency_peak", "grader": {"type": "reward_threshold", "threshold": 0.3}}50        ]51 52# 3. Define the main function (MANDATORY for Validator Step 3)53def main():54    """55    Standard entry point for the environment server.56    The validator calls this function to verify the server can start.57    """58    parser = argparse.ArgumentParser(description="Run the Detraff Env Server")59    parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind to")60    parser.add_argument("--port", type=int, default=8000, help="Port to listen on")61    args = parser.parse_args()62 63    uvicorn.run(app, host=args.host, port=args.port)64 65# 3. Execution block66if __name__ == "__main__":67    main()