Skip to content

ARTIFACT_STORE_SUMMARY.md

Latest commit

 

History

History
359 lines (259 loc) · 10.8 KB

ARTIFACT_STORE_SUMMARY.md

File metadata and controls

359 lines (259 loc) · 10.8 KB
Raw

Artifact Store Implementation Summary

Overview

Successfully implemented configurable artifact storage for AstroML with support for local filesystem, AWS S3, and Google Cloud Storage (GCS). The system uses fsspec for robust cloud storage handling and integrates seamlessly with MLflow tracking.

What Was Implemented

1. Core Artifact Storage System

File: astroml/storage/artifact_store.py

  • ArtifactStore - Abstract base class defining the storage interface
  • LocalArtifactStore - Local filesystem implementation
  • S3ArtifactStore - AWS S3 implementation
  • GCSArtifactStore - Google Cloud Storage implementation
  • create_artifact_store() - Factory function for creating stores from URIs

Key Features:

  • Unified API across all backends
  • fsspec-based implementation for reliability
  • Support for save, load, exists, delete, list operations
  • Full URI support (file://, s3://, gs://)

2. Configuration System

File: astroml/storage/config.py

  • ArtifactStorageConfig - Main configuration class
  • LocalStorageConfig - Local storage settings
  • S3StorageConfig - S3 settings with credential support
  • GCSStorageConfig - GCS settings with credential support

Features:

  • Pydantic-based validation
  • Environment variable support for credentials
  • URI generation from config
  • Dict serialization/deserialization

3. MLflow Integration

File: astroml/tracking/mlflow_tracker.py (Enhanced)

New Parameters:

  • artifact_uri - URI for artifact storage
  • artifact_store - Pre-configured ArtifactStore instance

New Methods:

  • log_model_artifact() - Returns artifact URI
  • save_artifact() - Save arbitrary artifacts
  • load_artifact() - Load artifacts from store

Backward Compatibility:

  • All existing code continues to work
  • Artifact store is optional
  • MLflow logging unchanged

4. Training Configuration

File: astroml/training/config.py (Enhanced)

  • Added artifact_storage: ArtifactStorageConfig field
  • Integrates with Hydra configuration system
  • Allows per-experiment artifact storage configuration

5. Configuration Files

Created example configurations:

  • configs/artifact_storage/local.yaml - Local storage
  • configs/artifact_storage/s3.yaml - S3 storage
  • configs/artifact_storage/gcs.yaml - GCS storage

6. Dependencies

Updated requirements files:

  • requirements.txt - Added fsspec, s3fs, gcsfs
  • requirements-cpu.txt - Added fsspec, s3fs, gcsfs

7. Documentation

  • ARTIFACT_STORAGE.md - Comprehensive configuration and usage guide
  • ARTIFACT_STORE_INTEGRATION.md - Integration guide with migration path
  • examples/train_with_artifact_store.py - Example training script

8. Tests

File: tests/test_artifact_store.py

Comprehensive test coverage:

  • Local storage tests (save, load, exists, delete, list)
  • S3 storage tests (mocked)
  • GCS storage tests (mocked)
  • Factory function tests
  • Configuration tests

Architecture

┌─────────────────────────────────────────────────────────────┐
│                    Training Script                          │
└────────────────────┬────────────────────────────────────────┘
                     │
                     ▼
┌─────────────────────────────────────────────────────────────┐
│              MLflowTracker (Enhanced)                        │
│  - log_model_artifact()                                     │
│  - save_artifact()                                          │
│  - load_artifact()                                          │
└────────────────────┬────────────────────────────────────────┘
                     │
        ┌────────────┴────────────┐
        │                         │
        ▼                         ▼
┌──────────────────┐    ┌──────────────────────┐
│  MLflow Tracking │    │  ArtifactStore       │
│  (mlruns/)       │    │  (Configurable)      │
└──────────────────┘    └──────────┬───────────┘
                                   │
                    ┌──────────────┼──────────────┐
                    │              │              │
                    ▼              ▼              ▼
            ┌──────────────┐ ┌──────────┐ ┌──────────────┐
            │ Local FS     │ │   S3     │ │     GCS      │
            │ (file://)    │ │(s3://)   │ │  (gs://)     │
            └──────────────┘ └──────────┘ └──────────────┘

Usage Examples

Basic Usage

from astroml.storage import create_artifact_store
from astroml.tracking import MLflowTracker

# Create artifact store
store = create_artifact_store("s3://my-bucket/models")

# Initialize tracker
tracker = MLflowTracker(artifact_store=store)

# Save model
uri = tracker.log_model_artifact(model, checkpoint_path="best.pth")
print(f"Model saved to: {uri}")

With Hydra Configuration

from hydra import compose, initialize_config_dir
from astroml.storage import create_artifact_store

cfg = compose(config_name="config")
artifact_uri = cfg.training.artifact_storage.get_artifact_uri()
store = create_artifact_store(artifact_uri)

Direct Store Usage

from astroml.storage import S3ArtifactStore

store = S3ArtifactStore("my-bucket", "models")

# Save
uri = store.save("local_model.pth", "exp1/model.pth")

# Load
store.load("exp1/model.pth", "downloaded.pth")

# List
artifacts = store.list_artifacts("exp1")

# Delete
store.delete("exp1/model.pth")

Key Features

  1. Multiple Backends

    • Local filesystem (development)
    • AWS S3 (production)
    • Google Cloud Storage (multi-cloud)

  2. Unified Interface

    • Same API regardless of backend
    • Easy to switch backends via configuration

  3. fsspec Integration

    • Robust cloud storage handling
    • Automatic multipart uploads for large files
    • Consistent error handling

  4. Configuration-Driven

    • Define backend via YAML
    • Environment variable support
    • Credential management

  5. MLflow Integration

    • Seamless logging to both MLflow and artifact store
    • Optional - doesn't break existing code
    • Returns artifact URIs for tracking

  6. Backward Compatible

    • All existing code continues to work
    • Artifact store is optional enhancement
    • No breaking changes

File Structure

astroml/
├── storage/
│   ├── __init__.py
│   ├── artifact_store.py      # Core implementations
│   └── config.py              # Configuration classes
├── tracking/
│   └── mlflow_tracker.py      # Enhanced with artifact store
└── training/
    └── config.py              # Enhanced with artifact storage config

configs/
└── artifact_storage/
    ├── local.yaml
    ├── s3.yaml
    └── gcs.yaml

tests/
└── test_artifact_store.py     # Comprehensive tests

examples/
└── train_with_artifact_store.py  # Example training script

Documentation:
├── ARTIFACT_STORAGE.md        # Configuration guide
├── ARTIFACT_STORE_INTEGRATION.md  # Integration guide
└── ARTIFACT_STORE_SUMMARY.md  # This file

Dependencies Added

fsspec>=2024.2.0      # Filesystem abstraction
s3fs>=2024.2.0        # S3 support
gcsfs>=2024.2.0       # GCS support

Testing

Run tests with:

pytest tests/test_artifact_store.py -v

Test coverage includes:

  • Local storage operations
  • S3 operations (mocked)
  • GCS operations (mocked)
  • Factory function
  • Configuration validation

Migration Path

  1. Install dependencies: pip install -r requirements.txt
  2. Update training scripts: Add artifact store initialization
  3. Configure backend: Create artifact_storage config
  4. Test locally: Use local storage first
  5. Deploy to cloud: Switch to S3/GCS in production

Performance Considerations

  • Local Storage: Fastest, no network overhead
  • S3: Good for AWS environments, supports multipart uploads
  • GCS: Good for GCP environments, similar performance to S3

For large models (>1GB):

  • Use multipart uploads (automatic)
  • Compress models before upload
  • Use regional buckets

Security Considerations

  1. Credentials Management

    • Use environment variables for credentials
    • Never commit credentials to version control
    • Use IAM roles in production

  2. Access Control

    • Restrict bucket access via IAM policies
    • Use service accounts for CI/CD
    • Enable bucket versioning for recovery

  3. Encryption

    • S3: Enable server-side encryption
    • GCS: Enable default encryption
    • Consider client-side encryption for sensitive models

Future Enhancements

Potential improvements:

  1. Model registry integration
  2. Artifact versioning and tagging
  3. Automatic cleanup policies
  4. Artifact compression
  5. Parallel uploads for large files
  6. Artifact signing and verification
  7. Cost tracking and optimization
  8. Additional cloud providers (Azure, MinIO)

Troubleshooting

Common Issues

  1. ModuleNotFoundError: fsspec

    • Solution: pip install -r requirements.txt

  2. NoCredentialsError (S3)

    • Solution: Set AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY

  3. DefaultCredentialsError (GCS)

    • Solution: Set GOOGLE_APPLICATION_CREDENTIALS

  4. PermissionError

    • Solution: Verify IAM permissions for credentials

See ARTIFACT_STORAGE.md for detailed troubleshooting.

Conclusion

The artifact storage system provides a flexible, extensible solution for managing model artifacts across different storage backends. It integrates seamlessly with existing MLflow tracking while maintaining full backward compatibility.

The implementation follows best practices:

  • Abstract base class for extensibility
  • Factory pattern for object creation
  • Configuration-driven design
  • Comprehensive error handling
  • Full test coverage
  • Clear documentation

This enables teams to:

  • Develop locally with filesystem storage
  • Deploy to production with S3/GCS
  • Switch backends without code changes
  • Track artifacts across experiments
  • Manage model lifecycle efficiently