update agents and tests

antonkulaga · antonkulaga · commit 3ddd8b6d3bec · 2026-01-20T18:13:21.000+02:00
diff --git a/AGENTS.md b/AGENTS.md
@@ -221,12 +221,23 @@ Datasets are typically uploaded to the `just-dna-seq` organization on Hugging Fa
 - **Delete redundant tests**: If test A (e.g., set equality) fully covers test B (e.g., count check), keep only test A.
 - **Domain constants are OK**: Hardcoding expected enum values or well-known constants from specs is fine; hardcoding row counts or unique counts derived from data inspection is not.
 
+#### Verifying Bug-Catching Claims
+
+When claiming a test "would have caught" a bug, **demonstrate it**:
+
+1. **Isolate the buggy logic** in a test or script
+2. **Run it and show failure** against correct expectations
+3. **Then show the fix passes** the same test
+
+Never claim "tests would have caught this" without running the buggy code against the test.
+
 #### Anti-Patterns to Avoid
 
 - Testing only "happy path" with trivial data
 - Hardcoding expected values that drift from source (use derived ground truth)
 - Mocking data transformations instead of running real pipelines
 - Ignoring edge cases (nulls, empty strings, boundary values, unicode, malformed data)
+- **Claiming tests "would catch bugs" without demonstrating failure on buggy code**
 
 **Meaningless Tests to Avoid** (common AI-generated anti-patterns):
 
diff --git a/src/prepare_annotations/core/dagster_io_managers.py b/src/prepare_annotations/core/dagster_io_managers.py
@@ -124,60 +124,64 @@ class ModuleIOManager(IOManager):
     """
     IO Manager for OakVar module assets.
     
+    This IO manager stores the actual paths returned by assets, so downstream
+    assets can load them correctly without needing to guess/reconstruct paths.
+    
     Handles paths for:
     - Module SQLite databases in data/modules/just_{module}/
     - Converted parquet files in data/output/modules/{module}/
+    
+    Path metadata is stored in a JSON file alongside the Dagster storage.
     """
     
-    # Mapping of asset keys to actual SQLite filenames (as they exist in GitHub repos)
-    SQLITE_FILENAMES: dict[str, str] = {
-        "longevitymap": "longevitymap.sqlite",
-        "lipidmetabolism": "lipid_metabolism.sqlite",
-        "vo2max": "vo2max.sqlite",
-        "superhuman": "superhuman.sqlite",
-        "coronary": "coronary.sqlite",
-    }
+    def __init__(self) -> None:
+        # Path registry persisted to disk
+        self._registry_path = MODULES_OUTPUT_DIR / ".asset_paths.json"
     
-    def _get_asset_path(self, asset_key: str) -> Path:
-        """Resolve the path for a module asset based on its key."""
-        if asset_key.endswith("_sqlite"):
-            module = asset_key.replace("_sqlite", "")
-            sqlite_filename = self.SQLITE_FILENAMES.get(module, f"{module}.sqlite")
-            return MODULES_DIR / f"just_{module}" / sqlite_filename
-        
-        # Handle longevitymap_with_ensembl specifically
-        if asset_key == "longevitymap_with_ensembl":
-            return MODULES_OUTPUT_DIR / "longevitymap" / "longevitymap_ensembl_joined.parquet"
-            
-        # Standard module assets: {module}_{type} (e.g., longevitymap_annotations)
-        parts = asset_key.split("_")
-        if len(parts) >= 2:
-            module = parts[0]
-            type_name = "_".join(parts[1:])
-            return MODULES_OUTPUT_DIR / module / f"{type_name}.parquet"
-            
-        return MODULES_OUTPUT_DIR / asset_key
+    def _load_registry(self) -> dict[str, str]:
+        """Load the asset path registry from disk."""
+        if self._registry_path.exists():
+            import json
+            return json.loads(self._registry_path.read_text())
+        return {}
+    
+    def _save_registry(self, registry: dict[str, str]) -> None:
+        """Save the asset path registry to disk."""
+        import json
+        self._registry_path.parent.mkdir(parents=True, exist_ok=True)
+        self._registry_path.write_text(json.dumps(registry, indent=2))
 
     def handle_output(self, context: OutputContext, obj: Any) -> None:
-        """Log where the asset was stored."""
+        """Store the asset path in registry for later retrieval."""
+        asset_key = context.asset_key.to_user_string()
+        
         if isinstance(obj, Path):
+            # Store the actual path returned by the asset
+            registry = self._load_registry()
+            registry[asset_key] = str(obj)
+            self._save_registry(registry)
             context.log.info(f"Module asset stored at: {obj}")
         else:
             context.log.info(f"Module asset materialized: {type(obj).__name__}")
 
     def load_input(self, context: InputContext) -> Any:
-        """Load asset by returning its expected path."""
+        """Load asset by looking up its stored path."""
         asset_key = context.upstream_output.asset_key.to_user_string() if context.upstream_output else "unknown"
-        path = self._get_asset_path(asset_key)
         
-        if not path.exists():
-            raise FileNotFoundError(
-                f"Module asset not found at {path}. "
-                f"Materialize the {asset_key} asset first."
-            )
-            
-        context.log.info(f"Loading module data from: {path}")
-        return path
+        # Look up the actual path from registry
+        registry = self._load_registry()
+        if asset_key in registry:
+            path = Path(registry[asset_key])
+            if path.exists():
+                context.log.info(f"Loading module data from registry: {path}")
+                return path
+        
+        # Fallback: raise clear error
+        raise FileNotFoundError(
+            f"Module asset '{asset_key}' not found in registry. "
+            f"Materialize the {asset_key} asset first. "
+            f"Registry path: {self._registry_path}"
+        )
 
 
 @io_manager
diff --git a/src/prepare_annotations/huggingface/uploader.py b/src/prepare_annotations/huggingface/uploader.py
@@ -6,10 +6,12 @@
 """
 
 from pathlib import Path
-from typing import Optional, Dict, List
+from typing import Optional, List
 from eliot import start_action
 from huggingface_hub import HfApi, hf_hub_download, list_repo_files, CommitOperationAdd
 from huggingface_hub.utils import RepositoryNotFoundError, HfHubHTTPError
+from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
+import httpx
 
 from prepare_annotations.core.models import SingleUploadResult, BatchUploadResult
 from prepare_annotations.huggingface.dataset_cards import (
@@ -18,6 +20,10 @@
     save_dataset_card
 )
 
+# Retry configuration for HuggingFace uploads
+HF_UPLOAD_MAX_RETRIES = 3
+HF_UPLOAD_TIMEOUT_SECONDS = 300  # 5 minutes
+
 
 def upload_to_hf_if_changed(
     parquet_file: Path,
@@ -336,20 +342,38 @@ def upload_files_batch(
             if commit_message is None:
                 commit_message = f"Upload {num_uploading} parquet file{'s' if num_uploading > 1 else ''}"
             
-            try:
-                commit_info = api.create_commit(
+            # Retry logic for transient network errors
+            @retry(
+                stop=stop_after_attempt(HF_UPLOAD_MAX_RETRIES),
+                wait=wait_exponential(multiplier=1, min=4, max=60),
+                retry=retry_if_exception_type((httpx.TimeoutException, httpx.NetworkError)),
+                reraise=True,
+            )
+            def create_commit_with_retry():
+                return api.create_commit(
                     repo_id=repo_id,
                     repo_type=repo_type,
                     operations=operations,
                     commit_message=commit_message,
                 )
+            
+            try:
+                commit_info = create_commit_with_retry()
                 
                 action.log(
                     message_type="success",
                     step="commit_created",
                     commit_url=commit_info.commit_url,
                     num_files=num_uploading
                 )
+            except (httpx.TimeoutException, httpx.NetworkError) as e:
+                action.log(
+                    message_type="error",
+                    step="commit_failed_after_retries",
+                    error=str(e),
+                    max_retries=HF_UPLOAD_MAX_RETRIES
+                )
+                raise
             except Exception as e:
                 action.log(
                     message_type="error",
diff --git a/tests/test_module_io_manager.py b/tests/test_module_io_manager.py
@@ -0,0 +1,159 @@
+"""
+Tests for ModuleIOManager path registry.
+
+The IO manager now properly stores the paths returned by assets in a registry,
+instead of trying to guess/reconstruct paths from asset keys.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+
+from prepare_annotations.core.dagster_io_managers import ModuleIOManager
+
+
+@pytest.fixture
+def io_manager(tmp_path: Path) -> ModuleIOManager:
+    """Create an IO manager with a temp registry path."""
+    manager = ModuleIOManager()
+    manager._registry_path = tmp_path / ".asset_paths.json"
+    return manager
+
+
+class TestModuleIOManagerRegistry:
+    """Tests for the path registry functionality."""
+
+    def test_handle_output_stores_path_in_registry(
+        self,
+        io_manager: ModuleIOManager,
+        tmp_path: Path,
+    ) -> None:
+        """handle_output should store the returned Path in the registry."""
+        # Create a mock output context
+        context = MagicMock()
+        context.asset_key.to_user_string.return_value = "coronary_with_ensembl"
+        
+        # The asset returns this path
+        asset_output_path = tmp_path / "coronary" / "coronary_ensembl_joined.parquet"
+        asset_output_path.parent.mkdir(parents=True, exist_ok=True)
+        asset_output_path.touch()
+        
+        # Handle the output
+        io_manager.handle_output(context, asset_output_path)
+        
+        # Registry should contain the path
+        registry = io_manager._load_registry()
+        assert "coronary_with_ensembl" in registry
+        assert registry["coronary_with_ensembl"] == str(asset_output_path)
+
+    def test_load_input_retrieves_path_from_registry(
+        self,
+        io_manager: ModuleIOManager,
+        tmp_path: Path,
+    ) -> None:
+        """load_input should retrieve the path from registry."""
+        # Create a file
+        asset_path = tmp_path / "longevitymap" / "annotations.parquet"
+        asset_path.parent.mkdir(parents=True, exist_ok=True)
+        asset_path.touch()
+        
+        # Manually store in registry (simulating prior handle_output)
+        registry = {"longevitymap_annotations": str(asset_path)}
+        io_manager._save_registry(registry)
+        
+        # Create mock input context
+        context = MagicMock()
+        context.upstream_output.asset_key.to_user_string.return_value = "longevitymap_annotations"
+        
+        # Load should return the path
+        result = io_manager.load_input(context)
+        assert result == asset_path
+
+    def test_load_input_raises_if_not_in_registry(
+        self,
+        io_manager: ModuleIOManager,
+    ) -> None:
+        """load_input should raise FileNotFoundError if asset not in registry."""
+        context = MagicMock()
+        context.upstream_output.asset_key.to_user_string.return_value = "unknown_asset"
+        
+        with pytest.raises(FileNotFoundError, match="not found in registry"):
+            io_manager.load_input(context)
+
+    def test_registry_persists_across_instances(
+        self,
+        tmp_path: Path,
+    ) -> None:
+        """Registry should persist to disk and be readable by new instances."""
+        registry_path = tmp_path / ".asset_paths.json"
+        
+        # First instance stores a path
+        manager1 = ModuleIOManager()
+        manager1._registry_path = registry_path
+        
+        asset_path = tmp_path / "vo2max" / "weights.parquet"
+        asset_path.parent.mkdir(parents=True, exist_ok=True)
+        asset_path.touch()
+        
+        context = MagicMock()
+        context.asset_key.to_user_string.return_value = "vo2max_weights"
+        manager1.handle_output(context, asset_path)
+        
+        # New instance should be able to load
+        manager2 = ModuleIOManager()
+        manager2._registry_path = registry_path
+        
+        context2 = MagicMock()
+        context2.upstream_output.asset_key.to_user_string.return_value = "vo2max_weights"
+        
+        result = manager2.load_input(context2)
+        assert result == asset_path
+
+    def test_multiple_assets_stored_in_registry(
+        self,
+        io_manager: ModuleIOManager,
+        tmp_path: Path,
+    ) -> None:
+        """Multiple assets should be stored in the same registry."""
+        # Store multiple assets
+        assets = {
+            "coronary_annotations": tmp_path / "coronary" / "annotations.parquet",
+            "coronary_studies": tmp_path / "coronary" / "studies.parquet",
+            "coronary_with_ensembl": tmp_path / "coronary" / "coronary_ensembl_joined.parquet",
+        }
+        
+        for asset_key, path in assets.items():
+            path.parent.mkdir(parents=True, exist_ok=True)
+            path.touch()
+            
+            context = MagicMock()
+            context.asset_key.to_user_string.return_value = asset_key
+            io_manager.handle_output(context, path)
+        
+        # All should be in registry
+        registry = io_manager._load_registry()
+        assert len(registry) == 3
+        for asset_key, path in assets.items():
+            assert registry[asset_key] == str(path)
+
+
+class TestModuleIOManagerNonPathOutputs:
+    """Tests for non-Path outputs (e.g., dicts from upload assets)."""
+
+    def test_handle_output_ignores_non_path_objects(
+        self,
+        io_manager: ModuleIOManager,
+    ) -> None:
+        """Non-Path outputs should not be stored in registry."""
+        context = MagicMock()
+        context.asset_key.to_user_string.return_value = "coronary_hf_upload"
+        
+        # Upload assets return dicts, not Paths
+        upload_result = {"repo_id": "just-dna-seq/annotators", "num_uploaded": 3}
+        io_manager.handle_output(context, upload_result)
+        
+        # Should not be in registry
+        registry = io_manager._load_registry()
+        assert "coronary_hf_upload" not in registry
