better file uploading

agent/core/session_uploader.py

@@ -2,20 +2,34 @@
 """
 Standalone script for uploading session trajectories to HuggingFace.
 This runs as a separate process to avoid blocking the main agent.
+Uses individual file uploads to avoid race conditions.
 """
 
 import json
 import os
 import sys
+from datetime import datetime
 from pathlib import Path
 
 
-def upload_session_to_dataset(session_file: str, repo_id: str, max_retries: int = 3):
-    """Upload a single session file to HuggingFace dataset"""
+def upload_session_as_file(
+    session_file: str, repo_id: str, max_retries: int = 3
+) -> bool:
+    """
+    Upload a single session as an individual JSONL file (no race conditions)
+
+    Args:
+        session_file: Path to local session JSON file
+        repo_id: HuggingFace dataset repo ID
+        max_retries: Number of retry attempts
+
+    Returns:
+        True if successful, False otherwise
+    """
     try:
-        from datasets import Dataset, load_dataset
+        from huggingface_hub import HfApi
     except ImportError:
-        print("Error: datasets library not available", file=sys.stderr)
+        print("Error: huggingface_hub library not available", file=sys.stderr)
         return False
 
     try:
@@ -28,8 +42,17 @@ def upload_session_to_dataset(session_file: str, repo_id: str, max_retries: int
         if upload_status == "success":
             return True
 
-        # Prepare row for upload
-        row = {
+        hf_token = os.getenv("HF_TOKEN")
+        if not hf_token:
+            # Update status to failed
+            data["upload_status"] = "failed"
+            with open(session_file, "w") as f:
+                json.dump(data, f, indent=2)
+            return False
+
+        # Prepare JSONL content (single line)
+        # Store messages and events as JSON strings to avoid schema conflicts
+        session_row = {
             "session_id": data["session_id"],
             "session_start_time": data["session_start_time"],
             "session_end_time": data["session_end_time"],
@@ -38,49 +61,78 @@ def upload_session_to_dataset(session_file: str, repo_id: str, max_retries: int
             "events": json.dumps(data["events"]),
         }
 
-        hf_token = os.getenv("HF_TOKEN")
-        if not hf_token:
-            # Update status to failed
-            data["upload_status"] = "failed"
-            with open(session_file, "w") as f:
-                json.dump(data, f, indent=2)
-            return False
+        # Create temporary JSONL file
+        import tempfile
 
-        # Try to load existing dataset and append
-        for attempt in range(max_retries):
-            try:
+        with tempfile.NamedTemporaryFile(
+            mode="w", suffix=".jsonl", delete=False
+        ) as tmp:
+            json.dump(session_row, tmp)  # Single line JSON
+            tmp_path = tmp.name
+
+        try:
+            # Generate unique path in repo: sessions/YYYY-MM-DD/session_id.jsonl
+            session_id = data["session_id"]
+            date_str = datetime.fromisoformat(data["session_start_time"]).strftime(
+                "%Y-%m-%d"
+            )
+            repo_path = f"sessions/{date_str}/{session_id}.jsonl"
+
+            # Upload with retries
+            api = HfApi()
+            for attempt in range(max_retries):
                 try:
-                    existing_dataset = load_dataset(repo_id, split="train")
-                    new_dataset = Dataset.from_dict(
-                        {k: list(existing_dataset[k]) + [v] for k, v in row.items()}
+                    # Try to create repo if it doesn't exist (idempotent)
+                    try:
+                        api.create_repo(
+                            repo_id=repo_id,
+                            repo_type="dataset",
+                            private=True,
+                            token=hf_token,
+                            exist_ok=True,  # Don't fail if already exists
+                        )
+
+                    except Exception:
+                        # Repo might already exist, continue
+                        pass
+
+                    # Upload the session file
+                    api.upload_file(
+                        path_or_fileobj=tmp_path,
+                        path_in_repo=repo_path,
+                        repo_id=repo_id,
+                        repo_type="dataset",
+                        token=hf_token,
+                        commit_message=f"Add session {session_id}",
                     )
-                except Exception:
-                    # Dataset doesn't exist yet, create new one
-                    new_dataset = Dataset.from_dict({k: [v] for k, v in row.items()})
 
-                # Push to hub
-                new_dataset.push_to_hub(repo_id, private=True, token=hf_token)
-
-                # Update status to success
-                data["upload_status"] = "success"
-                data["upload_url"] = f"https://huggingface.co/datasets/{repo_id}"
-                with open(session_file, "w") as f:
-                    json.dump(data, f, indent=2)
+                    # Update local status to success
+                    data["upload_status"] = "success"
+                    data["upload_url"] = f"https://huggingface.co/datasets/{repo_id}"
+                    with open(session_file, "w") as f:
+                        json.dump(data, f, indent=2)
 
-                return True
+                    return True
 
+                except Exception:
+                    if attempt < max_retries - 1:
+                        import time
+
+                        wait_time = 2**attempt
+                        time.sleep(wait_time)
+                    else:
+                        # Final attempt failed
+                        data["upload_status"] = "failed"
+                        with open(session_file, "w") as f:
+                            json.dump(data, f, indent=2)
+                        return False
+
+        finally:
+            # Clean up temp file
+            try:
+                os.unlink(tmp_path)
             except Exception:
-                if attempt < max_retries - 1:
-                    import time
-
-                    wait_time = 2**attempt
-                    time.sleep(wait_time)
-                else:
-                    # Final attempt failed
-                    data["upload_status"] = "failed"
-                    with open(session_file, "w") as f:
-                        json.dump(data, f, indent=2)
-                    return False
+                pass
 
     except Exception as e:
         print(f"Error uploading session: {e}", file=sys.stderr)
@@ -104,7 +156,7 @@ def retry_failed_uploads(directory: str, repo_id: str):
 
             # Only retry pending or failed uploads
             if upload_status in ["pending", "failed"]:
-                upload_session_to_dataset(str(filepath), repo_id)
+                upload_session_as_file(str(filepath), repo_id)
 
         except Exception:
             pass
@@ -124,7 +176,7 @@ if __name__ == "__main__":
             sys.exit(1)
         session_file = sys.argv[2]
         repo_id = sys.argv[3]
-        success = upload_session_to_dataset(session_file, repo_id)
+        success = upload_session_as_file(session_file, repo_id)
         sys.exit(0 if success else 1)
 
     elif command == "retry":