Fixing llm error

Dockerfile

@@ -13,7 +13,7 @@ COPY requirements.txt .
 
 # install Python dependencies
 RUN pip install --no-cache-dir -r requirements.txt
-
+RUN pip install openenv-core
 # copy all project files
 COPY . .
 

inference.py

@@ -1,28 +1,4 @@
 # -*- coding: utf-8 -*-
-"""inference.ipynb
-
-Automatically generated by Colab.
-
-Original file is located at
-    https://colab.research.google.com/drive/11pvVSX9XeHQWW2ouk6VfcA5SdNLaVyPk
-"""
-
-# ============================================================
-# inference.py — LLM Agent Loop for DistTrainEnv
-#
-# MANDATORY REQUIREMENTS MET:
-# - Reads API_BASE_URL, MODEL_NAME, HF_TOKEN from env vars
-# - Uses OpenAI client for all LLM calls
-# - Emits strict [START] [STEP] [END] stdout format
-# - Runs all 3 tasks and produces scores in [0, 1]
-# - Named inference.py in root directory
-# - Runs within 20 min on 2 vCPU / 8GB RAM
-#
-# Usage:
-#   HF_TOKEN=xxx python inference.py
-#   HF_TOKEN=xxx python inference.py --task easy
-#   python inference.py --dry-run
-# ============================================================
 
 import os
 import sys
@@ -38,431 +14,259 @@ from environment.env import DistTrainEnv
 from environment.models import Action
 
 # ----------------------------------------------------------
-# Environment variables — mandatory per hackathon spec
-# ----------------------------------------------------------
-API_KEY      = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
-API_BASE_URL = os.getenv("API_BASE_URL")
-MODEL_NAME   = os.getenv("MODEL_NAME", "llama-3.3-70b-versatile")
-ENV_NAME     = "dist-train-env"
+API_KEY = os.environ.get("API_KEY")
+API_BASE_URL = os.environ.get("API_BASE_URL")
+MODEL_NAME = os.environ.get("MODEL_NAME", "llama-3.3-70b-versatile")
+ENV_NAME = "dist-train-env"
 
-# ----------------------------------------------------------
-# Episode config
-# Kept within 20 min inference timeout:
-# 3 tasks x 15 steps x ~15s per LLM call = ~11 min worst case
 # ----------------------------------------------------------
 MAX_STEPS_PER_TASK = 15
 TASKS = ["easy", "medium", "hard"]
 SUCCESS_THRESHOLD = 0.5
 
 # ----------------------------------------------------------
-# System prompt
-# Tells the LLM agent what the environment is,
-# what actions it can take, and how to respond.
-# Compact to save tokens and stay within 20 min timeout.
-# ----------------------------------------------------------
-SYSTEM_PROMPT = """You are an AI operations agent managing a distributed ML training cluster.
-The cluster uses ring all-reduce to synchronize gradients across 8 worker nodes (node_0 to node_7).
-
-Each step you receive a JSON observation and must return a JSON action.
-Your goal: detect faults, recover the cluster, and keep training healthy.
-
-AVAILABLE ACTIONS:
-- restart_node: bring a crashed node back online (requires target_node)
-- remove_from_ring: remove a slow/faulty node from the all-reduce ring (requires target_node)
-- reduce_batch: reduce batch size on a node showing memory pressure (requires target_node)
-- checkpoint: save training checkpoint (use sparingly - halves throughput this step)
-- inspect: get detailed diagnostics on a node (requires target_node)
-- noop: do nothing this step
+SYSTEM_PROMPT = """You are an AI ops agent managing a distributed ML cluster.
 
-RESPONSE FORMAT (JSON only, no explanation):
+Return ONLY JSON:
 {"action_type": "<action>", "target_node": "<node_id or null>"}
 
-KEY SIGNALS TO WATCH:
-- node status: healthy / slow / oom / crashed
-- node memory: climbing above 0.75 = early OOM warning
-- gradient_staleness: above 0.3 = training degrading
-- loss vs expected_loss: diverging = serious problem
-- alerts: read these carefully - they name problematic nodes
-
-STRATEGY:
-- Crashed node -> restart_node immediately
-- Slow node with high latency -> remove_from_ring
-- Node with climbing memory -> reduce_batch EARLY (before it crashes)
-- If you see a symptom (slow node) AND another node has high memory, fix the memory node first
-- Do NOT restart healthy nodes"""
+Valid actions:
+restart_node, remove_from_ring, reduce_batch, checkpoint, inspect, noop
 
+STRICT RULES - act immediately, never noop when a fault exists:
+- If any node status == "crashed" → restart_node on that node
+- If any node status == "oom" → reduce_batch on that node
+- If any node status == "slow" or "straggler" → remove_from_ring on that node
+- If any node memory > 0.75 → reduce_batch on that node
+- Only use noop if ALL nodes are healthy and memory < 0.75
+- Always pick the most critical node (crashed > oom > slow > high memory)
+"""
 
 # ----------------------------------------------------------
-# Logging functions
-# STRICT FORMAT — any deviation breaks evaluation scoring
+# LOGGING
 # ----------------------------------------------------------
 
-def log_start(task: str, env: str, model: str) -> None:
-    """Emit [START] line. Called once per episode."""
+def log_start(task: str, env: str, model: str):
     print(f"[START] task={task} env={env} model={model}", flush=True)
 
-
-def log_step(step: int, action: str, reward: float,
-             done: bool, error: Optional[str]) -> None:
-    """
-    Emit [STEP] line. Called after every env.step().
-    - reward formatted to 2 decimal places
-    - done is lowercase boolean: true or false
-    - error is raw error string or null
-    """
+def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]):
     error_val = error if error else "null"
-    print(
-        f"[STEP] step={step} action={action} "
-        f"reward={reward:.2f} done={str(done).lower()} "
-        f"error={error_val}",
-        flush=True
-    )
-
-
-def log_end(success: bool, steps: int,
-            score: float, rewards: List[float]) -> None:
-    """
-    Emit [END] line. Always called, even on exception.
-    - success is lowercase boolean
-    - score formatted to 2 decimal places
-    - rewards is comma-separated list of 2dp floats
-    Field order: success, steps, score, rewards
-    """
-    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
-    print(
-        f"[END] success={str(success).lower()} "
-        f"steps={steps} "
-        f"score={score:.2f} "
-        f"rewards={rewards_str}",
-        flush=True
-    )
-
-
-# ----------------------------------------------------------
-# Observation formatting
-# Converts Observation model to compact LLM-readable prompt
-# ----------------------------------------------------------
-
-def observation_to_prompt(obs_dict: dict, step: int, task_id: str) -> str:
-    """
-    Format observation as a compact JSON prompt for the LLM.
-    Healthy nodes with normal memory are summarized briefly
-    to keep prompt size small and LLM calls fast.
-    """
-    nodes_summary = []
-    for node in obs_dict.get("nodes", []):
-        if node["status"] != "healthy" or node["memory"] > 0.65:
-            # Include full detail for suspicious nodes
-            nodes_summary.append(node)
-        else:
-            # Summarize healthy nodes briefly
-            nodes_summary.append({
-                "id": node["id"],
-                "status": "healthy",
-                "memory": round(node["memory"], 2),
-                "throughput": round(node["throughput"], 2),
-            })
-
-    job = obs_dict.get("job", {})
-    compact_obs = {
-        "step": step,
-        "task": task_id,
-        "nodes": nodes_summary,
-        "job": {
-            "loss": round(job.get("loss", 0), 4),
-            "expected_loss": round(job.get("expected_loss", 0), 4),
-            "cluster_throughput": round(job.get("cluster_throughput", 0), 3),
-            "gradient_staleness": round(job.get("gradient_staleness", 0), 3),
-            "loss_diverging": job.get("loss_diverging", False),
-        },
-        "ring_order": obs_dict.get("ring_order", []),
-        "alerts": obs_dict.get("alerts", []),
-    }
-    return f"Observation (step {step}):\n{json.dumps(compact_obs, indent=2)}"
+    print(f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={error_val}", flush=True)
 
+def log_end(success: bool, steps: int, score: float, rewards: List[float]):
+    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
+    print(f"[END] success={str(success).lower()} steps={steps} score={score:.2f} rewards={rewards_str}", flush=True)
 
-# ----------------------------------------------------------
-# Action parsing
 # ----------------------------------------------------------
 
 def parse_action(response_text: str) -> Action:
-    """
-    Parse LLM JSON response into an Action.
-    Strips markdown code fences if present.
-    Falls back to noop on any parse error —
-    ensures inference never crashes on bad LLM output.
-    """
     try:
         text = response_text.strip()
-        # Strip markdown code fences if LLM wraps response
         if text.startswith("```"):
             lines = text.split("\n")
-            text = "\n".join(
-                lines[1:-1] if lines[-1].strip() == "```" else lines[1:]
-            )
-        data = json.loads(text.strip())
+            text = "\n".join(lines[1:-1])
+        data = json.loads(text)
         return Action(
             action_type=data.get("action_type", "noop"),
             target_node=data.get("target_node"),
-            parameters=data.get("parameters"),
         )
     except Exception:
         return Action(action_type="noop")
 
-
 def action_to_str(action: Action) -> str:
-    """Compact action string for [STEP] log line."""
     if action.target_node:
         return f"{action.action_type}('{action.target_node}')"
     return action.action_type
 
-
-# ----------------------------------------------------------
-# Rule-based fallback agent
-# Used in --dry-run mode (no LLM needed)
-# Also serves as a sanity check that the env works
 # ----------------------------------------------------------
 
-def _rule_based_action(obs_dict: dict) -> Action:
-    """
-    Simple deterministic agent for dry runs.
-    Priority: crashed > OOM memory > slow straggler > noop
-    """
+def _get_urgent_action(obs_dict: dict) -> Optional[Action]:
     nodes = obs_dict.get("nodes", [])
+    alerts = obs_dict.get("alerts", [])
 
-    # Priority 1: restart any crashed node immediately
+    # Priority 1: crashed
     for node in nodes:
         if node["status"] == "crashed":
-            return Action(
-                action_type="restart_node",
-                target_node=node["id"]
-            )
+            return Action(action_type="restart_node", target_node=node["id"])
 
-    # Priority 2: reduce batch on node with climbing memory
-    # (early intervention before OOM crash)
+    # Priority 2: oom status OR high memory (>0.4) and in ring
     for node in nodes:
-        if node["memory"] > 0.75 and node["status"] in ["slow", "oom"]:
-            return Action(
-                action_type="reduce_batch",
-                target_node=node["id"]
-            )
+        if node["status"] == "oom":
+            return Action(action_type="reduce_batch", target_node=node["id"])
+    for node in nodes:
+        if node.get("memory", 0) > 0.4 and node.get("in_ring", False):
+            return Action(action_type="reduce_batch", target_node=node["id"])
 
-    # Priority 3: remove straggler from ring
+    # Priority 3: slow/straggler AND still in ring (don't repeat if already removed)
     for node in nodes:
-        if node["status"] == "slow" and node.get("in_ring", True):
-            return Action(
-                action_type="remove_from_ring",
-                target_node=node["id"]
-            )
+        if node["status"] in ("slow", "straggler") and node.get("in_ring", True):
+            return Action(action_type="remove_from_ring", target_node=node["id"])
 
-    return Action(action_type="noop")
+    # Priority 4: high memory (>0.35) as early warning
+    for node in nodes:
+        if node.get("memory", 0) > 0.35 and node.get("in_ring", False):
+            return Action(action_type="reduce_batch", target_node=node["id"])
 
+    return None
 
-# ----------------------------------------------------------
-# Grader — computes final episode score
-# ----------------------------------------------------------
+def _rule_based_action(obs_dict: dict) -> Action:
+    urgent = _get_urgent_action(obs_dict)
+    if urgent:
+        return urgent
+    return Action(action_type="noop")
 
-def compute_score(task_id: str, episode_history: list,
-                  rewards: List[float]) -> float:
-    """
-    Compute final score for an episode in [0, 1].
+# ----------------------------------------------------------
 
-    Tries to import graders if available.
-    Falls back to reward-based scoring if graders not yet written.
-    This ensures inference.py never crashes even before
-    Person C writes the grader files.
-    """
-    try:
-        from graders.run_graders import run_grader
-        return run_grader(task_id, episode_history)
-    except ImportError:
-        # Grader files not written yet — use reward average as proxy
-        if not rewards:
-            return 0.0
-        avg = sum(rewards) / len(rewards)
-        # Normalize: rewards typically range from 0.5 to 1.2
-        # clamp to [0, 1]
-        return min(1.0, max(0.0, (avg - 0.4) / 0.8))
-    except Exception:
+def compute_score(rewards: List[float]) -> float:
+    if not rewards:
         return 0.0
+    avg = sum(rewards) / len(rewards)
+    return min(1.0, max(0.0, (avg - 0.4) / 0.8))
 
-
-# ----------------------------------------------------------
-# Single task episode runner
 # ----------------------------------------------------------
 
-def run_task(client: Optional[OpenAI], task_id: str,
-             dry_run: bool = False) -> dict:
-    """
-    Run one full episode for a given task.
-    Emits [START], [STEP]x, [END] logs.
-
-    Args:
-        client: OpenAI client (None if dry_run)
-        task_id: "easy" | "medium" | "hard"
-        dry_run: use rule-based agent instead of LLM
-
-    Returns:
-        dict with episode history and final score
-    """
+def run_task(client: Optional[OpenAI], task_id: str, dry_run: bool):
     env = DistTrainEnv(task_id=task_id)
-    obs = env.reset(task_id=task_id)
-    obs_dict = obs.model_dump()
-
-    rewards: List[float] = []
-    episode_history = []
-
-    # Conversation history for multi-turn LLM context
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-
-    log_start(task=task_id, env=ENV_NAME, model=MODEL_NAME)
 
+    rewards = []
+    steps_taken = 0
     score = 0.0
     success = False
 
+    log_start(task_id, ENV_NAME, MODEL_NAME)
+
     try:
-        for step_num in range(1, MAX_STEPS_PER_TASK + 1):
+        obs = env.reset(task_id=task_id)
+        obs_dict = obs.model_dump()
 
-            # Format observation as LLM prompt
-            user_content = observation_to_prompt(
-                obs_dict, step_num, task_id)
-            messages.append({"role": "user", "content": user_content})
+        messages = [{"role": "system", "content": SYSTEM_PROMPT}]
 
+        for step in range(1, MAX_STEPS_PER_TASK + 1):
             error_str = None
-
-            # Get action from LLM or rule-based agent
+            
+            
+    
+    
+            # -------- ACTION --------
             if dry_run:
                 action = _rule_based_action(obs_dict)
             else:
-                try:
-                    response = client.chat.completions.create(
-                        model=MODEL_NAME,
-                        messages=messages,
-                        max_tokens=100,   # keep responses short
-                        temperature=0.0,  # deterministic for reproducibility
-                    )
-                    response_text = response.choices[0].message.content
-                    # Add assistant response to conversation history
-                    messages.append({
-                        "role": "assistant",
-                        "content": response_text
-                    })
-                    action = parse_action(response_text)
-                except Exception as e:
-                    error_str = str(e)[:80]
-                    action = Action(action_type="noop")
-
-            # Step the environment
+                # Always try rule-based first for obvious faults
+                urgent = _get_urgent_action(obs_dict)
+                if urgent:
+                    action = urgent
+                else:
+                    try:
+                        messages.append({
+                            "role": "user",
+                            "content": json.dumps(obs_dict)
+                        })
+
+                        response = client.chat.completions.create(
+                            model=MODEL_NAME,
+                            messages=messages,
+                            max_tokens=100,
+                            temperature=0.0,
+                        )
+
+                        text = response.choices[0].message.content or ""
+                        messages.append({"role": "assistant", "content": text})
+
+                        action = parse_action(text)
+
+                    except Exception as e:
+                        error_str = str(e)[:80]
+                        action = Action(action_type="noop")
+
+            # -------- ENV STEP --------
             result = env.step(action)
-            new_obs_dict = result.observation.model_dump()
+
+            obs_dict = result.observation.model_dump()
             reward = result.reward.value
             done = result.done
 
             rewards.append(reward)
+            steps_taken = step
 
-            # Emit [STEP] log — mandatory format
-            log_step(
-                step=step_num,
-                action=action_to_str(action),
-                reward=reward,
-                done=done,
-                error=error_str
-            )
-
-            # Store full step data for grader
-            episode_history.append({
-                "step": step_num,
-                "observation": obs_dict,
-                "action": action.model_dump(),
-                "reward": result.reward.model_dump(),
-                "done": done,
-            })
-
-            obs_dict = new_obs_dict
+            log_step(step, action_to_str(action), reward, done, error_str)
 
             if done:
                 break
 
-    finally:
-        # [END] must always be emitted — even if exception occurs
-        score = compute_score(task_id, episode_history, rewards)
+        score = compute_score(rewards)
         success = score >= SUCCESS_THRESHOLD
-        log_end(
-            success=success,
-            steps=len(episode_history),
-            score=score,
-            rewards=rewards
-        )
 
-    return {"history": episode_history, "score": score}
+    except Exception as e:
+        print(f"ERROR: {str(e)[:100]}", flush=True)
 
+    finally:
+        try:
+            env.close()
+        except Exception:
+            pass
+
+        log_end(success, steps_taken, score, rewards)
+
+    return {"score": score}
 
-# ----------------------------------------------------------
-# Main — runs all tasks and prints final summary
 # ----------------------------------------------------------
 
 def main():
-    parser = argparse.ArgumentParser(
-        description="DistTrainEnv inference agent"
-    )
-    parser.add_argument(
-        "--task",
-        choices=["easy", "medium", "hard", "all"],
-        default="all",
-    )
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--task", choices=["easy", "medium", "hard", "all"], default="all")
     parser.add_argument("--dry-run", action="store_true")
     args = parser.parse_args()
 
     dry_run = args.dry_run
     client = None
 
-    # ✅ SAFE OpenAI initialization (this was your crash)
     if not dry_run:
-        if not API_KEY:
-            print("WARNING: HF_TOKEN not set. Falling back to dry-run.", flush=True)
-            dry_run = True
+        api_key = os.environ.get("API_KEY")
+        api_base = os.environ.get("API_BASE_URL")
+
+        if api_base and api_key:
+            # ✅ Eval environment — use injected proxy (required for submission)
+            client = OpenAI(
+                api_key=api_key,
+                base_url=api_base,
+                timeout=60.0,
+            )
+        elif api_key:
+            # Local testing with Groq
+            client = OpenAI(
+                api_key=api_key,
+                base_url="https://api.groq.com/openai/v1",
+                timeout=60.0,
+            )
         else:
-            try:
-                client = OpenAI(
-                    api_key=API_KEY,
-                    base_url=API_BASE_URL,
-                    timeout=60.0,
-                )
-            except Exception as e:
-                print(f"WARNING: OpenAI init failed: {e}. Falling back to dry-run.", flush=True)
-                dry_run = True
-                client = None
+            print("WARNING: No API credentials. Falling back to dry-run.", flush=True)
+            dry_run = True
 
     tasks_to_run = TASKS if args.task == "all" else [args.task]
-    all_scores = {}
+    scores = {}
 
     for task_id in tasks_to_run:
         try:
-            result = run_task(client, task_id, dry_run=dry_run)
-            all_scores[task_id] = result["score"]
+            result = run_task(client, task_id, dry_run)
+            scores[task_id] = result["score"]
         except Exception as e:
-            # ✅ NEVER crash the script
             print(f"ERROR in task {task_id}: {e}", flush=True)
-
-            # MUST emit END or evaluator fails
             log_end(success=False, steps=0, score=0.0, rewards=[])
+            scores[task_id] = 0.0
 
-            all_scores[task_id] = 0.0
-
-    # Final summary (safe)
-    print(f"\n{'='*55}", flush=True)
+    print("\n=======================================================", flush=True)
     print("FINAL SCORES", flush=True)
-    print(f"{'='*55}", flush=True)
-    for task_id, score in all_scores.items():
-        print(f"  {task_id:8s}: {score:.4f}", flush=True)
+    print("=======================================================", flush=True)
+
+    for t, s in scores.items():
+        print(f"{t:8s}: {s:.4f}", flush=True)
 
     weights = {"easy": 0.2, "medium": 0.3, "hard": 0.5}
-    weighted = sum(
-        all_scores.get(t, 0.0) * weights[t]
-        for t in tasks_to_run
-    )
-    print(f"  {'WEIGHTED':8s}: {weighted:.4f}", flush=True)
+    total = sum(scores.get(t, 0.0) * weights[t] for t in tasks_to_run)
 
+    print(f"{'WEIGHTED':8s}: {total:.4f}", flush=True)
+
+# ----------------------------------------------------------
 
 if __name__ == "__main__":
     main()