feat: add optional trained model inference to env space

Add lazy LoRA adapter inference endpoints and GPU model dependencies while preserving the OpenEnv reset/step server routes.

Made-with: Cursor

Dockerfile

@@ -4,7 +4,9 @@ WORKDIR /app
 
 # Install deps first (cached layer — only invalidated when requirements change)
 COPY requirements-server.txt .
-RUN pip install --no-cache-dir -r requirements-server.txt
+RUN pip install --no-cache-dir \
+    torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 && \
+    pip install --no-cache-dir -r requirements-server.txt
 
 # Copy source (excludes venv/, training/, assets/, .git/ via .dockerignore)
 COPY . .
@@ -14,5 +16,5 @@ RUN python -m data.loader || echo "facts.json generation skipped — will use fa
 
 EXPOSE 7860
 
-# 2 workers: handles concurrent sessions without threading issues
-CMD ["uvicorn", "environment.server:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "2"]
+# One worker so the optional trained-model adapter is loaded at most once.
+CMD ["uvicorn", "environment.server:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "1"]

environment/model_inference.py

@@ -0,0 +1,120 @@
+import os
+import threading
+from pathlib import Path
+from typing import Any
+
+import torch
+from pydantic import BaseModel
+
+from environment.actions import EpisodeObservation
+
+
+MODEL_ID = os.getenv("MODEL_ID", "Siddh12334/qwen-1.5b-context-corruption")
+MAX_NEW_TOKENS = int(os.getenv("MODEL_MAX_NEW_TOKENS", "128"))
+_LOCK = threading.Lock()
+_MODEL = None
+_TOKENIZER = None
+
+
+class InferenceRequest(BaseModel):
+    observation: EpisodeObservation
+
+
+class InferenceResponse(BaseModel):
+    text: str
+    loaded_model: str
+
+
+def configure_runtime_dirs():
+    root = Path(os.getenv("MODEL_RUNTIME_DIR", "/tmp/context-corruption-model"))
+    cache = root / "cache"
+    env_dirs = {
+        "HOME": root,
+        "XDG_CACHE_HOME": cache,
+        "HF_HOME": cache / "huggingface",
+        "HF_HUB_CACHE": cache / "huggingface" / "hub",
+        "TRANSFORMERS_CACHE": cache / "huggingface" / "transformers",
+    }
+    for path in env_dirs.values():
+        path.mkdir(parents=True, exist_ok=True)
+    for key, path in env_dirs.items():
+        os.environ.setdefault(key, str(path))
+
+
+def model_status() -> dict[str, Any]:
+    return {
+        "model_id": MODEL_ID,
+        "loaded": _MODEL is not None,
+        "cuda_available": torch.cuda.is_available(),
+        "cuda_device": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
+    }
+
+
+def _load_model():
+    global _MODEL, _TOKENIZER
+    if _MODEL is not None and _TOKENIZER is not None:
+        return _MODEL, _TOKENIZER
+
+    with _LOCK:
+        if _MODEL is not None and _TOKENIZER is not None:
+            return _MODEL, _TOKENIZER
+
+        configure_runtime_dirs()
+        from peft import AutoPeftModelForCausalLM
+        from transformers import AutoTokenizer, BitsAndBytesConfig
+
+        quantization_config = None
+        if torch.cuda.is_available():
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.bfloat16,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_use_double_quant=True,
+            )
+
+        _TOKENIZER = AutoTokenizer.from_pretrained(MODEL_ID)
+        _MODEL = AutoPeftModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            device_map="auto" if torch.cuda.is_available() else "cpu",
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+            quantization_config=quantization_config,
+            low_cpu_mem_usage=True,
+        )
+        _MODEL.eval()
+        return _MODEL, _TOKENIZER
+
+
+def _format_prompt(observation: EpisodeObservation) -> list[dict[str, str]]:
+    docs_text = "\n\n".join(
+        f"[Doc {doc.id}] {doc.title}\n{doc.content}" for doc in observation.documents
+    )
+    system = (
+        "You are an epistemic agent. Answer the question and identify corrupted documents. "
+        'Respond ONLY as JSON: {"answer": "...", "suspicious_docs": [0], "confidence": 0.8}'
+    )
+    user = f"Question: {observation.question}\n\nDocuments:\n{docs_text}"
+    return [
+        {"role": "system", "content": system},
+        {"role": "user", "content": user},
+    ]
+
+
+def run_inference(observation: EpisodeObservation) -> InferenceResponse:
+    model, tokenizer = _load_model()
+    messages = _format_prompt(observation)
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    with torch.inference_mode():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=MAX_NEW_TOKENS,
+            do_sample=False,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+    generated_ids = output_ids[0][inputs["input_ids"].shape[-1]:]
+    text = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
+    return InferenceResponse(text=text, loaded_model=MODEL_ID)

environment/server.py

@@ -7,6 +7,7 @@ load_dotenv()
 
 from environment.actions import ContextCorruptionAction, EpisodeObservation
 from environment.env import ContextCorruptionEnv
+from environment.model_inference import InferenceRequest, model_status, run_inference
 
 _difficulty_env = os.getenv("DIFFICULTY")
 _difficulty = int(_difficulty_env) if _difficulty_env else None
@@ -20,5 +21,15 @@ app = create_app(
     max_concurrent_envs=_max_sessions,
 )
 
+
+@app.get("/model/status")
+def get_model_status():
+    return model_status()
+
+
+@app.post("/model/infer")
+def infer_with_trained_model(request: InferenceRequest):
+    return run_inference(request.observation)
+
 if __name__ == "__main__":
     uvicorn.run("environment.server:app", host="0.0.0.0", port=7860, reload=False)

requirements-server.txt

@@ -6,3 +6,8 @@ datasets>=2.0.0
 faker>=18.0.0
 python-dotenv>=1.0.0
 websockets>=15.0.0
+transformers>=5.5.0
+peft>=0.19.0
+accelerate>=1.0.0
+bitsandbytes>=0.45.0
+safetensors>=0.4.0