feat: ExperimentLoader for MLflow-based model/data loading + checkpoint discovery

[adamimos]

Summary

This PR adds a high-level loader to reconstruct models and access run data from MLflow, plus checkpoint discovery on persisters.

Key Features

• ExperimentLoader with from_mlflow(run_id, tracking_uri)
  • load_config() returns saved Hydra config (with fallback search for config.yaml)
  • load_metrics(pattern) returns tidy DataFrame (metric, step, value, timestamp)
  • list_checkpoints() and latest_checkpoint()
  • load_model(step) reconstructs model and loads weights via configured persister
• MLflowRunReader for read-only access to config/params/tags/metrics/artifacts
• Persister checkpoint discovery methods implemented for Local and S3:
  • list_checkpoints(), latest_checkpoint(), checkpoint_exists(), uri_for_step()
• PyTorch support: loader falls back to torch.nn.Module (e.g., transformer_lens) and sets eval() by default

Files / Modules

• simplexity/loaders/experiment_loader.py
• simplexity/logging/mlflow_reader.py, simplexity/logging/run_reader.py
• Persister updates: local_equinox, local_penzai, local_pytorch, s3
• Notebook: notebooks/experiment_loader_demo.ipynb (smoke test)
• README: “Loading From MLflow” section with example usage

Tests

• Local persisters: tests/persistence/test_checkpoint_discovery_local.py
• S3 persister: tests/persistence/test_checkpoint_discovery_s3.py (uses existing mocks)
• MLflow reader: tests/logging/test_mlflow_reader.py (mocked client)
• Loader: tests/loaders/test_experiment_loader.py (checkpoint listing without model instantiation)

Notes

• Config artifact expected as config.yaml; reader searches recursively if not at root
• S3 discovery parses step directories via paginator; URIs include filename when known
• Device handling: if model cfg has device: auto, loader resolves to cuda/mps/cpu

Follow-ups (optional)

• Log checkpoint_step and checkpoint_uri during training for traceability
• Tiny CLI and notebook utilities for convenience
• Artifact resolver for generic URIs if needed later

[claude]

Code Review: ExperimentLoader PR

I've conducted a comprehensive review of this pull request based on the repository's CLAUDE.md guidelines. Overall, this is a well-architected and thoroughly tested feature that adds significant value to the project.

✅ Strengths

Code Quality & Architecture

• Excellent separation of concerns: The ExperimentLoader delegates appropriately to RunReader and ModelPersister abstractions
• Clean protocol-based design: S3Client and S3Paginator protocols provide proper abstraction over boto3
• Proper error handling: Graceful fallbacks (config.yaml search, JAX→PyTorch model instantiation)
• Type safety: Comprehensive type hints and proper use of Protocol classes
• Functional patterns: Follows JAX/functional programming conventions as required

Testing Coverage

• Comprehensive test suite: Well-structured tests for local/S3 persisters and MLflow integration
• Proper mocking: Uses appropriate mocking for external dependencies (MLflow, S3)
• Edge case coverage: Tests checkpoint discovery without model instantiation (good defensive programming)
• CI integration: All tests appear properly integrated with existing test infrastructure

Documentation & Usability

• Clear README section: Well-written documentation with practical examples
• Jupyter notebook demo: Provides hands-on usage examples
• Comprehensive docstrings: Good coverage of public API methods
• Intuitive API design: from_mlflow() factory method and simple method names

🔧 Technical Considerations

Performance & Resource Management

• Efficient S3 pagination: Uses boto3 paginators for large checkpoint listings
• Proper cleanup: MLflowRunReader uses tempfile.TemporaryDirectory with proper cleanup
• Device resolution: Intelligent fallback for PyTorch device selection (cuda → mps → cpu)

Security & Best Practices

• ✅ No hardcoded credentials: Uses AWS configuration files and environment variables
• ✅ Proper input validation: Handles missing configs and checkpoints gracefully
• ✅ No sensitive data logging: Clean error messages without exposing internal details

Extensibility

• Plugin architecture: Easy to extend with new RunReader implementations
• Framework agnostic: Supports both JAX (PredictiveModel) and PyTorch fallbacks
• Configurable persistence: Works with all existing persister implementations

🚨 Minor Suggestions

Type Safety Enhancement

Consider adding more specific return types for checkpoint-related methods:

# Current
def list_checkpoints(self) -> list[int]:

# Suggested for consistency
def latest_checkpoint(self) -> int | None:  # ✅ Already implemented correctly

Error Message Clarity

The PyTorch fallback error message is excellent and informative. Consider similar detailed error messages for S3 configuration failures.

Resource Management

The temporary directory cleanup in MLflowRunReader looks correct, but consider adding explicit cleanup in __exit__ method for extra safety.

🏆 Code Quality Score: 9/10

This PR demonstrates excellent engineering practices:

• Follows all CLAUDE.md conventions (120 char lines, type hints, pytest patterns)
• Proper abstraction layers and separation of concerns
• Comprehensive testing with good coverage
• Clean, readable code with minimal complexity
• Excellent documentation and examples

📋 Recommendation

✅ APPROVE - This PR is ready for merge. The implementation is solid, well-tested, and adds significant value to the project. The code quality meets high standards and follows all repository conventions.

The minor suggestions above are truly optional and don't block the merge. Great work on this feature!

🤖 Generated with Claude Code

[claude]

Code Review: PR #76 - ExperimentLoader for MLflow-based model/data loading

✅ Overall Assessment: APPROVED

This is a well-structured, comprehensive feature that adds significant value to the Simplexity ecosystem. The implementation follows the project's conventions and provides a clean interface for MLflow integration.

---

🔍 Code Quality & Best Practices

Strengths:

• Clean Architecture: Well-organized module structure with proper separation of concerns (loaders/, logging/, persistence updates)
• Type Safety: Proper type annotations throughout the codebase, adheres to project standards
• Error Handling: Robust error handling with informative error messages
• Documentation: Clear docstrings following Google style conventions
• Interface Design: Builder pattern usage (from_mlflow()) provides clean initialization

Follows Project Conventions:

• ✅ Snake_case naming for functions and files
• ✅ PascalCase for classes (ExperimentLoader, MLflowRunReader)
• ✅ Proper import organization
• ✅ Line length under 120 characters
• ✅ No inline comments (code is self-documenting)

---

🧪 Test Coverage

Comprehensive Testing:

• Unit Tests: Good coverage across all new components
  • test_experiment_loader.py - Core loader functionality
  • test_mlflow_reader.py - MLflow integration with mocks
  • test_checkpoint_discovery_local.py & test_checkpoint_discovery_s3.py - Persistence layer
• Mock Usage: Proper mocking of external dependencies (MLflow client, S3)
• Edge Cases: Tests cover error conditions and fallback scenarios

Suggestions:

• Consider adding integration tests that exercise the full workflow
• Test device resolution logic (device: auto → cuda/mps/cpu)

---

🔒 Security Considerations

Well-Handled:

• ✅ Credentials: Uses existing config.ini pattern, no hardcoded secrets
• ✅ S3 Access: Leverages existing S3Persister security model
• ✅ MLflow Integration: Uses standard MLflow client with proper authentication
• ✅ Input Validation: Proper validation of run_ids and tracking URIs

Note:

The security model appropriately delegates to existing components rather than introducing new attack vectors.

---

⚡ Performance Considerations

Efficient Design:

• Lazy Loading: Models and configs loaded on-demand
• Caching: Checkpoint discovery results can be cached
• Stream Processing: MLflow artifact downloads handled efficiently
• Pagination: S3 checkpoint discovery uses proper pagination

Minor Concerns:

• Large model loading could benefit from progress indicators
• Consider connection pooling for repeated MLflow API calls

---

🏗️ Architecture & Design

Excellent Patterns:

• Single Responsibility: Each class has clear, focused purpose
• Dependency Injection: Config-based initialization allows flexibility
• Fallback Mechanisms: Graceful degradation (JAX → PyTorch fallback)
• Protocol Usage: Clean interfaces for extensibility

Smart Decisions:

• Device auto-resolution for deployment scenarios
• Recursive config.yaml search for artifact discovery
• Support for both step numbers and "latest" checkpoint selection

---

📋 Specific Observations

ExperimentLoader (simplexity/loaders/experiment_loader.py):

• Clean builder pattern with from_mlflow()
• Proper error messages for missing dependencies
• Good separation between model loading and metadata access

MLflowRunReader (simplexity/logging/mlflow_reader.py):

• Read-only interface design is appropriate
• Efficient pandas DataFrame creation for metrics
• Robust artifact discovery with fallback search

Persistence Layer Updates:

• Consistent API across all persisters (list_checkpoints(), latest_checkpoint())
• S3 implementation properly handles pagination
• Local persisters use filesystem efficiently

---

🚀 Value Addition

This feature significantly enhances the MLflow integration story:

• Researcher Productivity: Easy model reconstruction from experiments
• Deployment Support: Streamlined path from research to production
• Multi-Framework: Supports both JAX and PyTorch models
• Documentation: Excellent README examples and Jupyter notebook

---

💡 Minor Suggestions for Future Iterations

1. CLI Integration: Consider adding simple CLI commands for common operations
2. Batch Operations: Support for loading multiple experiments/checkpoints
3. Metadata Caching: Cache experiment metadata to reduce API calls
4. Progress Indicators: For large model/artifact downloads

---

🎯 Final Verdict

APPROVED ✅

This PR demonstrates excellent software engineering practices, comprehensive testing, and thoughtful design. It adds substantial value while maintaining the project's quality standards. The implementation is production-ready and well-documented.

Recommendation: Merge after CI passes.

[claude]

Code Review for PR #76: ExperimentLoader for MLflow-based model/data loading + checkpoint discovery

Overall Assessment

This is a well-structured and comprehensive feature addition that follows the project's conventions. The implementation provides a clean abstraction for loading models and data from MLflow experiments.

✅ Code Quality & Best Practices

Strengths:

• Clean Architecture: Good separation of concerns with ExperimentLoader, MLflowRunReader, and persistence layer updates
• Type Safety: Proper use of type hints throughout, including generic types and union types
• Error Handling: Comprehensive exception handling with informative error messages
• Documentation: Clear docstrings and inline comments where needed
• Follows Project Conventions: Adheres to CLAUDE.md guidelines for naming, imports, and structure

Minor Suggestions:

• Line lengths are within the 120-character limit ✓
• Import ordering follows project standards ✓
• Function/method naming uses proper snake_case ✓

⚠️ Potential Issues

Device Resolution Logic:

def _resolve_device(self, device: str) -> str:
    # Consider caching this result to avoid repeated torch imports
    # Current implementation re-imports torch on every call

Exception Handling in _instantiate_persister_only:

• The method silently returns None on any exception, which might hide configuration issues
• Consider logging warnings or being more specific about which exceptions to catch

S3 Checkpoint Discovery:

• The pagination logic looks correct but consider adding bounds checking for very large buckets
• Missing validation that the bucket exists before listing

🔒 Security Considerations

✅ Good Security Practices:

• No hardcoded credentials or secrets
• Proper use of temporary directories with cleanup
• S3 credentials handled through standard AWS config patterns
• Input validation for run IDs and tracking URIs

Recommendations:

• Consider adding validation for S3 bucket names to prevent injection attacks
• The config_path parameter should validate file paths to prevent directory traversal

🚀 Performance Considerations

Efficient Design:

• Config caching with _cached_config prevents redundant MLflow calls ✓
• Lazy loading of models and persisters ✓
• Uses pagination for S3 listing to handle large datasets ✓

Potential Optimizations:

• Consider connection pooling for MLflow client in high-throughput scenarios
• The metrics loading could benefit from date/time range filtering for large experiments

🧪 Test Coverage

Strong Test Coverage:

• Good use of mocking for external dependencies (MLflow, S3)
• Tests cover core functionality without requiring actual MLflow/S3 setup
• Proper isolation with _StubReader for unit testing

Areas for Enhancement:

• Could add integration tests with real MLflow (using test tracking URI)
• Error path testing could be expanded (network failures, missing artifacts)
• Test the PyTorch fallback path more thoroughly

📝 Specific Code Comments

ExperimentLoader

The dual JAX/PyTorch model loading is well-implemented:

# Good fallback pattern - tries JAX first, then PyTorch
try:
    model = typed_instantiate(model_config, PredictiveModel)
except Exception:
    # Clear fallback to PyTorch with good error messages

MLflowRunReader

The recursive config.yaml search is a nice UX touch:

# Fallback search for config.yaml is user-friendly
for path in self.list_artifacts():
    if Path(path).name == "config.yaml":
        # ... download and return

Persistence Layer

Good consistent interface additions across all persisters with proper error handling.

🎯 Summary

This is a high-quality implementation that adds significant value to the project. The code follows best practices, handles edge cases well, and provides good user experience.

Recommendation: APPROVE ✅

The implementation is production-ready with only minor optimization opportunities. The test coverage is adequate for the initial release, and the security considerations are properly addressed.

Great work on the comprehensive documentation and example notebook! 🚀

🤖 Generated with Claude Code

[claude]

Code Review for PR #76: ExperimentLoader for MLflow-based model/data loading + checkpoint discovery

Summary

This PR adds comprehensive MLflow integration with a high-level ExperimentLoader API and checkpoint discovery functionality. The implementation follows good architectural patterns and provides a clean interface for loading models and data from MLflow runs.

✅ Code Quality & Best Practices

Strengths:

• Modular Design: Clear separation of concerns between ExperimentLoader, MLflowRunReader, and persister extensions
• Error Handling: Proper error handling with informative messages (e.g., checkpoint discovery failures)
• Type Safety: Good use of type hints throughout the codebase
• Documentation: Comprehensive docstrings following Google style
• Testing: Good test coverage including mocked MLflow scenarios

Following CLAUDE.md Guidelines:

• ✅ Uses type annotations consistently
• ✅ Follows functional programming patterns with JAX/Equinox compatibility
• ✅ Proper naming conventions (snake_case functions, PascalCase classes)
• ✅ Google-style docstrings implemented
• ✅ Line length appears to stay within 120 characters

🔍 Potential Issues & Improvements

1. Error Handling Enhancement

# In checkpoint discovery, consider more specific exception types
except Exception as e:
    # Could be more specific about S3 vs network vs auth errors

2. Device Resolution Logic
The device auto-resolution in ExperimentLoader.load_model() is helpful but could benefit from:

• Explicit logging when device is auto-resolved
• Fallback handling if CUDA/MPS detection fails

3. Memory Considerations
For large models, consider:

• Lazy loading options
• Memory usage warnings for large checkpoints
• Streaming capabilities for S3-stored large models

4. Configuration Validation
The recursive config.yaml search is clever, but could benefit from:

• Maximum recursion depth limit
• Clear error messages when config structure is unexpected

🚀 Performance Considerations

Positive:

• Efficient S3 paginator usage for checkpoint discovery
• Proper use of MLflow client connection management
• Lazy evaluation patterns where appropriate

Suggestions:

• Consider caching checkpoint lists for repeated calls
• Batch operations where possible for multiple checkpoint operations

🔒 Security Assessment

Good practices observed:

• No hardcoded credentials
• Proper use of environment variables and config files
• S3 credential handling through standard AWS patterns

No security concerns identified - the implementation properly delegates authentication to external systems (MLflow, AWS).

🧪 Test Coverage

Excellent test coverage:

• Unit tests for core functionality
• Mocked external dependencies (MLflow, S3)
• Edge case handling (missing checkpoints, config discovery)
• Integration scenarios in notebook demo

Suggestion: Consider adding integration tests with real (test) MLflow instances for CI/CD pipelines.

📚 Documentation & Usability

Strengths:

• Clear README section with practical examples
• Comprehensive notebook demo
• Good inline documentation
• Helpful error messages

Minor suggestion: Consider adding troubleshooting section to README for common issues (auth failures, missing dependencies, etc.).

🎯 Overall Assessment

This is a high-quality PR that:

• Solves a real need for MLflow integration
• Follows project conventions and best practices
• Provides comprehensive functionality with good error handling
• Includes thorough testing and documentation

Recommendation: ✅ APPROVE

The implementation is production-ready and follows the established patterns in the codebase. The modular design will make it easy to extend and maintain.

📋 Minor Suggestions for Future Iterations

• Consider adding async/await support for large model loading
• Batch checkpoint operations for performance
• Add integration with experiment tracking beyond MLflow
• Consider adding model comparison utilities

Great work on this comprehensive feature addition! 🚀