Upload data_loader.py with huggingface_hub

data_loader.py

@@ -43,19 +43,42 @@ def load_data() -> tuple[pd.DataFrame, pd.DataFrame, dict]:
     if "walltime_used" in jobs.columns:
         jobs["walltime_hours"] = jobs["walltime_used"] / 3600.0
 
-    # Timeout classification — interactive vs batch
+    # Behavioral outcome classification
     if "launch_method" in jobs.columns and "last_state" in jobs.columns:
         import numpy as np
 
-        from config import INTERACTIVE_METHODS
+        from config import INTERACTIVE_METHODS, QUICK_EXIT_SECONDS
 
-        is_timeout = jobs["last_state"] == "TIMEOUT"
         is_interactive = jobs["launch_method"].isin(INTERACTIVE_METHODS)
-        jobs["timeout_category"] = np.where(
-            ~is_timeout,
-            jobs["last_state"],
-            np.where(is_interactive, "Interactive Timeout", "Batch Timeout"),
+        wt = jobs.get("walltime_used", pd.Series(dtype="float64"))
+        state = jobs["last_state"]
+
+        # Start with batch classification (maps exit state directly)
+        outcome = state.map(
+            {
+                "COMPLETED": "Completed",
+                "FAILED": "Failed",
+                "TIMEOUT": "Timed Out",
+                "OUT_OF_MEMORY": "Out of Memory",
+            }
+        ).fillna("Cancelled")  # All CANCELLED variants + NODE_FAIL → "Cancelled"
+
+        # Override for interactive jobs
+        is_quick = is_interactive & (wt < QUICK_EXIT_SECONDS)
+        is_failed_interactive = is_interactive & state.isin({"FAILED", "OUT_OF_MEMORY"})
+        is_user_ended = (
+            is_interactive & ~is_quick & ~is_failed_interactive & state.str.startswith("CANCELLED")
         )
+        is_session_expired = (
+            is_interactive & ~is_quick & ~is_failed_interactive & (state == "TIMEOUT")
+        )
+
+        outcome = np.where(is_quick, "Quick Exit", outcome)
+        outcome = np.where(is_failed_interactive, "Failed", outcome)
+        outcome = np.where(is_user_ended, "User Ended", outcome)
+        outcome = np.where(is_session_expired, "Session Expired", outcome)
+
+        jobs["outcome_category"] = outcome
 
     # Queue wait time
     if "submit_time" in jobs.columns and "start_time" in jobs.columns: