feat: Checkpoint pipeline infrastructure (Phase 1)

[JesperDramsch]

Summary

This PR implements Phase 1 of the checkpoint pipeline infrastructure, establishing the core abstractions and pipeline pattern for flexible checkpoint handling in Anemoi.

Changes

• Core abstractions: Added for state management and base class for all pipeline stages
• Pipeline orchestration: Implemented with support for async/sync execution and stage management
• Dynamic component discovery: Created with automatic discovery of checkpoint sources, loaders, and modifiers using module inspection
• Comprehensive error handling: Added exception hierarchy for checkpoint operations
• Utility functions: Implemented helpers for download retry, validation, and metadata extraction

Technical Details

The component catalog uses a hybrid approach for identifying abstract classes:

• ABC inheritance detection
• Abstract method checking
• Name-based convention (Base* prefix)

This ensures true dynamic discovery without hardcoded component lists, making the system easily extensible.

Testing

• Unit tests for all core components
• Integration tests for pipeline execution
• Tests for dynamic component discovery with various abstract class patterns

Related Issues

Closes #493

Next Steps

This is Phase 1 of a multi-phase implementation:

• Phase 2: Three parallel layers (Checkpoint Acquisition Checkpoint Acquisition Layer - Multi-source checkpoint loading (S3, HTTP, local, MLFlow) #458, Loading Orchestration Checkpoint Loading Orchestration (Phase 2) #494, Model Transformation Model Transformation Layer - Post-loading modifications (freezing, transfer learning, adapters) #410)
• Phase 3: System integration and migration (Checkpoint System Integration and Migration (Phase 3) #495)

---

📚 Documentation preview 📚: https://anemoi-training--501.org.readthedocs.build/en/501/

---

📚 Documentation preview 📚: https://anemoi-graphs--501.org.readthedocs.build/en/501/

---

📚 Documentation preview 📚: https://anemoi-models--501.org.readthedocs.build/en/501/

[JesperDramsch]

Phase 1: Checkpoint Pipeline Infrastructure - Progress Update

Summary

This PR implements the foundational infrastructure for the new checkpoint architecture, establishing the core abstractions, pipeline pattern, and utilities that all checkpoint operations will build upon.

Related Issues: #493 (Phase 1), Part of 5-phase checkpoint architecture refactoring
Branch: feature/checkpoint-pipeline-infrastructure
Status: Core complete, configuration integration pending review

🎯 Motivation

The current checkpoint handling in anemoi-training is monolithic and difficult to extend. This PR establishes a clean, three-layer pipeline architecture:

┌─────────────────────────────────────────────────┐
│        Model Transformation Layer               │
│         (Post-loading modifications)            │
├─────────────────────────────────────────────────┤
│         Loading Orchestration Layer             │
│    (Strategies for applying checkpoints)        │
├─────────────────────────────────────────────────┤
│        Checkpoint Acquisition Layer             │
│      (Obtaining checkpoint from sources)        │
├─────────────────────────────────────────────────┤
│         Pipeline Infrastructure  ← THIS PR      │
│         (Core abstractions & context)           │
└─────────────────────────────────────────────────┘

📦 Changes Included

New Files (VERIFIED)

training/src/anemoi/training/checkpoint/
├── __init__.py                    # Package initialization (2.6KB)
├── base.py                        # Core abstractions (16KB)
├── pipeline.py                    # Pipeline orchestrator (26KB)
├── catalog.py                     # Component discovery (22KB)
├── exceptions.py                  # Exception hierarchy (19KB)
├── formats.py                     # Multi-format support (19KB)
└── utils.py                       # Async utilities (17KB)

training/tests/checkpoint/
├── conftest.py                    # Test fixtures (9.7KB)
├── test_base.py                   # Core tests (8.1KB)
├── test_pipeline.py               # Pipeline tests (10KB)
├── test_catalog.py                # Catalog tests (11KB)
├── test_formats.py                # Format tests (26KB)
├── test_utils.py                  # Utility tests (37KB)
└── test_exceptions.py             # Exception tests (32KB)

Modified Files (VERIFIED)

training/pyproject.toml            # Added optional checkpoint dependency (line 67)

Pending Additions (Not Yet in PR)

training/src/anemoi/training/schemas/training.py   # TODO: Add CheckpointPipelineConfig
config/training/checkpoint_pipeline/*.yaml         # TODO: Create config templates

🔑 Key Features

1. Pipeline Pattern

Clean, composable stages for checkpoint processing:

# Example usage (programmatic)
pipeline = CheckpointPipeline(
    stages=[
        CheckpointSource(...),      # Phase 2
        LoadingStrategy(...),       # Phase 2
        ModelModifier(...),         # Phase 2
    ],
    async_execution=True
)

context = await pipeline.execute(initial_context)

2. Component Catalog

Automatic discovery of pipeline components using reflection:

# No manual registration needed!
ComponentCatalog.list_sources()    # Auto-discovers CheckpointSource subclasses
ComponentCatalog.list_loaders()    # Auto-discovers LoadingStrategy subclasses
ComponentCatalog.list_modifiers()  # Auto-discovers ModelModifier subclasses

3. Multi-Format Support

Seamless handling of different checkpoint formats:

• PyTorch Lightning checkpoints
• Standard PyTorch checkpoints
• SafeTensors format (optional)
• Raw state_dict formats

Auto-detection and conversion included.

4. Comprehensive Error Handling

Rich exception hierarchy for debugging:

try:
    context = await pipeline.execute(context)
except CheckpointNotFoundError as e:
    logger.error(f"Checkpoint not found: {e}")
except CheckpointIncompatibleError as e:
    logger.error(f"Incompatible checkpoint: {e}")

5. Async-First Design

Efficient async operations with sync compatibility:

# Async mode (default)
context = await pipeline.execute(context)

# Sync mode (when needed)
pipeline = CheckpointPipeline(stages, async_execution=False)
context = pipeline.execute(context)

🧪 Testing Strategy

Coverage Metrics (VERIFIED)

• Total Tests: 213
• Passing: 211 (99.1%)
• Skipped: 2
• Code Coverage: 80% (798/1003 lines)

Test Execution

# Run all checkpoint tests
pytest training/tests/checkpoint/ -v

# Run with coverage
pytest training/tests/checkpoint/ --cov=anemoi.training.checkpoint --cov-report=html

# Quick verification
pytest training/tests/checkpoint/ -q
# Output: 211 passed, 2 skipped in 46.46s

Test Coverage by Module

• base.py: Core abstractions
• pipeline.py: Pipeline orchestration with Hydra
• catalog.py: Component discovery
• exceptions.py: Error hierarchy
• formats.py: Multi-format detection/conversion
• utils.py: Async downloads, validation, checksums

🔄 Integration Points

With PR #464 (Checkpoint Acquisition)

# PR #464 will implement CheckpointSource subclasses
class S3Source(PipelineStage):  # Uses base.py abstractions
    async def process(self, context: CheckpointContext) -> CheckpointContext:
        # Load from S3, populate context.checkpoint_data
        ...

With PR #494 (Loading Orchestration)

# PR #494 will implement LoadingStrategy subclasses
class TransferLearningLoader(PipelineStage):  # Uses base.py abstractions
    async def process(self, context: CheckpointContext) -> CheckpointContext:
        # Apply checkpoint to model with flexibility
        ...

With PR #442 (Model Modifiers)

# PR #442 will implement ModelModifier as PipelineStage
class FreezingModifier(PipelineStage):  # Uses base.py abstractions
    async def process(self, context: CheckpointContext) -> CheckpointContext:
        # Freeze model layers post-loading
        ...

⚠️ Breaking Changes

None. This PR is pure infrastructure - no existing code is modified or broken.

📋 Configuration Schema

Status: ⚠️ Not yet integrated - needs to be added before merge

Planned addition to training/src/anemoi/training/schemas/training.py:

class CheckpointPipelineConfig(BaseModel):
    """Configuration for checkpoint pipeline"""
    stages: List[Dict[str, Any]] = Field(default_factory=list)
    async_execution: bool = Field(default=True)
    cache_dir: Optional[str] = Field(default=None)
    max_retries: int = Field(default=3)
    timeout: int = Field(default=300)

🔍 Review Checklist

Code Quality ✅

• All code follows project style guidelines (ruff)
• Type hints on all functions (mypy validated)
• Comprehensive docstrings (NumPy style)
• No code duplication

Testing ✅

• Unit tests for all new functionality
• 99.1% test pass rate (211/213)
• 80% overall test coverage achieved
• All critical paths covered

Documentation ⚠️

• Public API documented in code
• User guide needed
• Migration guide needed
• Architecture diagrams included

Configuration ⚠️

• Schema integration pending
• Config templates pending
• Optional dependencies properly marked
• Graceful degradation when optional deps missing

Compatibility ✅

• No breaking changes to existing API
• Backward compatible
• Optional dependencies properly marked
• Graceful degradation

🚀 Deployment Notes

Dependencies

Required: Already in pyproject.toml

• torch >= 2.2.0
• omegaconf
• hydra-core

Optional: Added in this PR (pyproject.toml line 67)

[project.optional-dependencies]
checkpoint = ["safetensors>=0.4.0"]

Migration Path

No migration needed - this is new infrastructure. Existing checkpoint code remains unchanged.

📊 Performance Impact

• Memory: Minimal overhead (~1-2MB for pipeline infrastructure)
• Speed: Async operations enable concurrent checkpoint processing
• CPU: No significant impact, efficient component discovery

🔜 Next Steps

Before Merge (Recommended)

1. Add configuration schema to training/src/anemoi/training/schemas/training.py
2. Create config templates in config/training/checkpoint_pipeline/
3. Add basic user documentation

After Merge

1. Phase 2 PRs can immediately build on this:

   • PR Refactor restart of training #464: Checkpoint Acquisition (in progress)
   • PR Checkpoint Loading Orchestration (Phase 2) #494: Loading Orchestration (ready to start)
   • PR [WIP] 410 refactor of model initialisation ie weight loading model freezing transfer learning #442: Model Transformation (in progress)

2. Documentation: Comprehensive user guides

3. Phase 3: Integration & migration utilities

📝 Honest Assessment

What's Strong

• Core pipeline architecture is solid and tested
• Component discovery eliminates boilerplate
• Multi-format support is comprehensive
• Error handling is thorough
• Test coverage is good (80%, 99.1% pass rate)
• All abstractions are in place for Phase 2

What Needs Work

• Configuration schema not integrated yet
• User documentation is minimal
• Config templates don't exist
• Some edge cases not covered in tests (20% uncovered)

Production Readiness

• For developers building Phase 2: ✅ Ready now
• For end users: ⚠️ Needs config integration and docs

---

Reviewers: Please focus on:

1. API design of CheckpointContext and PipelineStage
2. Component discovery mechanism in catalog.py
3. Exception hierarchy completeness
4. Whether to require config schema integration before merge

[JesperDramsch]

Thank you @anaprietonem and @sahahner for the thorough review! Here's how I'm addressing your feedback:

Already Addressed

• ✅ Safetensors removed (commit 15d1845) - keeping formats simple
• ✅ Added device logging to checkpoint loading utilities (commit 506103d)

Will Address Before Merge

1. Execution model docs - Adding clear documentation for standalone/callback patterns
2. Callback naming - Adding docstring cross-references for clarity

Architectural Decisions

4. ComponentCatalog location - Keeping in training for Phase 1 (tightly coupled to checkpoint base classes). Requesting input from @gareth-j @jjlk for future consideration.
5. Callback unification - These are complementary (save vs load), not duplicates. Will improve naming/docs.
6. Cloud authentication - Documented as Phase 2 scope (PR Checkpoint Acquisition Layer - Multi-source checkpoint loading (S3, HTTP, local, MLFlow) #458). Will follow standard SDK patterns.

Future Work

7. Exception consolidation - Opening issue to audit and consolidate across packages
8. aiohttp optional - Moving to optional-dependencies.remote

Let me know if you'd like any of these addressed differently!

[anaprietonem]

Thanks Jesper for addressing the comments and implementing this. Nice work!! LGTM and can be merged.