supplymind / scripts /submission_preflight_v2.py
Rishav
Prepare SupplyMind finale submission
9432cbb
Raw
History Blame Contribute Delete
4.96 kB
from __future__ import annotations
import json
import sys
from pathlib import Path
from statistics import mean
from typing import Any
from fastapi.testclient import TestClient
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from supplymind_env.api import app
from supplymind_env_v2.models import V2Observation
from supplymind_env_v2.policies import heuristic_joint_policy, no_op_policy
def check(condition: bool, message: str) -> None:
if not condition:
raise AssertionError(message)
def run_episode(client: TestClient, task_id: str, policy_name: str) -> dict[str, Any]:
reset = client.post("/v2/reset", json={"task_id": task_id})
check(reset.status_code == 200, f"/v2/reset failed for {task_id}: {reset.text}")
observation = reset.json()
total_invalid = 0
rounds = 0
done = False
result_body: dict[str, Any] = {}
while not done:
if policy_name == "no_op":
action = no_op_policy(None).model_dump(mode="json")
else:
action = heuristic_joint_policy(V2Observation.model_validate(observation)).model_dump(mode="json")
step = client.post("/v2/step", json=action)
check(step.status_code == 200, f"/v2/step failed for {task_id}/{policy_name}: {step.text}")
result_body = step.json()
observation = result_body["observation"]
done = bool(result_body["done"])
rounds += 1
total_invalid += len(observation.get("feedback", {}).get("invalid_action_details", []))
check(rounds <= 80, f"episode did not terminate for {task_id}/{policy_name}")
summary = result_body["info"].get("episode_summary", {})
return {
"task_id": task_id,
"policy": policy_name,
"rounds": rounds,
"raw_reward": summary.get("raw_reward"),
"graded_score": summary.get("graded_score"),
"invalid_actions": total_invalid,
}
def validate_v2_api(client: TestClient) -> dict[str, Any]:
health = client.get("/health")
check(health.status_code == 200, "/health must return 200")
rules = client.get("/v2/rules")
check(rules.status_code == 200, "/v2/rules must return 200")
check("action_space" in rules.json(), "/v2/rules should expose action_space")
task_summaries = []
for task_id in ("easy", "medium", "hard"):
reset = client.post("/v2/reset", json={"task_id": task_id})
check(reset.status_code == 200, f"/v2/reset JSON body failed for {task_id}")
body = reset.json()
check(body["task_id"] == task_id, f"/v2/reset returned wrong task for {task_id}: {body['task_id']}")
state = client.get("/v2/state")
check(state.status_code == 200, "/v2/state must return 200")
step = client.post("/v2/step", json={})
check(step.status_code == 200, "/v2/step must accept empty joint action")
step_body = step.json()
check("observation" in step_body and "reward" in step_body and "done" in step_body, "/v2/step shape invalid")
task_summaries.append({"task_id": task_id, "reset_ok": True, "empty_step_ok": True})
for role in ("center", "warehouse"):
reset = client.post(f"/v2/{role}/reset", json={"task_id": "train_easy"})
check(reset.status_code == 200, f"/v2/{role}/reset must return 200")
body = reset.json()
check(body.get("role") == role, f"/v2/{role}/reset should expose role")
check("controlled_action_schema" in body, f"/v2/{role}/reset should expose controlled_action_schema")
return {"tasks": task_summaries, "roles": ["center", "warehouse"]}
def load_eval_summary() -> dict[str, Any] | None:
path = ROOT / "results" / "v2_policy_eval.json"
if not path.exists():
return None
return json.loads(path.read_text(encoding="utf-8")).get("summary")
def validate_repo_files() -> list[str]:
required = ["app.py", "Dockerfile", "openenv.yaml", "inference.py", "requirements.txt"]
missing = [name for name in required if not (ROOT / name).exists()]
check(not missing, f"missing required files: {missing}")
return required
def main() -> None:
client = TestClient(app)
files = validate_repo_files()
api = validate_v2_api(client)
episodes = [
run_episode(client, "train_easy", "no_op"),
run_episode(client, "train_easy", "heuristic"),
]
scores = [float(row["graded_score"]) for row in episodes if row["graded_score"] is not None]
report = {
"status": "ok",
"checked_files": files,
"api": api,
"episode_smoke": episodes,
"mean_smoke_score": round(mean(scores), 4) if scores else None,
"latest_eval_summary": load_eval_summary(),
}
out = ROOT / "results" / "submission_preflight_v2.json"
out.parent.mkdir(exist_ok=True)
out.write_text(json.dumps(report, indent=2), encoding="utf-8")
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()