Add pluggable artifact serializer framework

[saeidbarati157]

Summary

Replace the hardcoded pickle logic in TaskDataStore with a pluggable artifact serializer framework. Existing pickle behavior is preserved exactly — PickleSerializer is a built-in universal fallback so flows that never register a custom serializer behave identically to before.

• ArtifactSerializer ABC with priority-based dispatch, SerializationMetadata namedtuple, SerializedBlob value type.
• SerializerStore metaclass auto-registers subclasses; caches the sorted dispatch list and invalidates on new registration.
• Extensions register serializers via ARTIFACT_SERIALIZERS_DESC (standard Metaflow plugin pattern).
• A lazy import-hook registry lets extensions defer serializer-module imports (e.g. torch, pyarrow) until the user's own code imports the target type.

Design

• Dispatch. Serializers are sorted by PRIORITY (lower first; ties broken by registration order). On save, the first can_serialize(obj) == True wins. On load, can_deserialize(metadata) routes by the encoding field in artifact metadata.
• Wire vs storage format. serialize(obj, format=...) / deserialize(data, format=..., **kw) accept a format kwarg — STORAGE (default) returns (List[SerializedBlob], SerializationMetadata) for the datastore save path; WIRE returns a str for CLI args, protobuf payloads, and cross-process IPC. One class owns both representations. PickleSerializer implements STORAGE only — pickle bytes are not a wire-safe payload — and raises NotImplementedError for WIRE.
• Side-effect-free contract. serialize() must not perform I/O or mutate global state: it may be called multiple times (caching, retries, parallel dispatch) and must stay safely idempotent. Side effects belong in hooks, not serializers.
• Lazy registration. metaflow.datastore.artifacts.lazy_registry.register_serializer_for_type("torch.Tensor", "my_ext.TorchSerializer") stores a declarative config; an importlib.abc.MetaPathFinder watches for torch to be imported and only then loads the serializer class. If the target module is already in sys.modules, registration is immediate.
• Backward compatible. Old artifacts with gzip+pickle-v2 / gzip+pickle-v4 encoding load correctly. Missing encoding defaults to gzip+pickle-v2.

Error paths

• save_artifacts: if no serializer claims an object, raise DataException with a message pointing at the PickleSerializer fallback invariant (instead of the misleading UnpicklableArtifactException for a non-pickle failure).
• load_artifacts: if no deserializer claims an artifact, the error includes serializer_info["source"] when present, hinting at a missing extension.

Test plan

• 29 unit tests for ArtifactSerializer / SerializerStore / SerializedBlob / SerializationMetadata, including wire/storage dispatch and the ABC's abstract-method enforcement.
• 36 unit tests for PickleSerializer (round-trips across scalars, containers, custom classes; WIRE rejection; encoding validation).
• 9 unit tests for the lazy registry (config validation, eager registration, deferred registration via the import hook, recursion guard).
• 6 integration tests for TaskDataStore round-trips, custom serializer priority, and backward compat.
• 245 existing unit tests pass (only pre-existing spin failures remain, unrelated).

[greptile-apps]

Greptile Summary

Replaces the hardcoded pickle logic in TaskDataStore with a pluggable serializer framework. PickleSerializer is preserved as the universal fallback so existing flows are unaffected. The _serializers property now re-queries the live registry on every access, meaning serializers activated by the lazy import hook after TaskDataStore construction are automatically visible.

• New framework (ArtifactSerializer, SerializerStore, SerializationMetadata, SerializedBlob): priority-based dispatch, ABCMeta enforcement, state-machine lifecycle, setup_imports hook with sys.meta_path retry.
• PickleSerializer moved to metaflow/plugins/datastores/serializers/ and wired in via ARTIFACT_SERIALIZERS_DESC; legacy gzip+pickle-v{2,4} encodings are accepted for backward compat on load.
• task_datastore.py: save_artifacts dispatches via the registry, validates blob shape before mutating _info, auto-injects source into serializer_info; load_artifacts routes by encoding with graceful fallback to gzip+pickle-v2 for ancient artifacts.

Confidence Score: 5/5

The change is safe to merge. Backward compatibility is preserved: PickleSerializer remains the universal fallback, legacy encodings are accepted on load, and the old gzip+pickle artifact format is handled transparently.

All previously flagged critical paths are properly addressed. The two remaining comments are narrow robustness gaps in edge-case retry scenarios that do not affect the primary dispatch path or any existing artifact.

The lazy_import retry idempotency logic and the _call_setup_imports helper in serializer.py are the most likely places to surface unexpected behavior for extension authors writing multi-dependency serializers.

Important Files Changed

Filename	Overview
metaflow/datastore/artifacts/serializer.py	Core serializer framework: SerializerStore metaclass (now correctly using ABCMeta), dispatch logic, lazy-import helpers, and retry state machine. Two minor robustness gaps: lazy_import idempotency check uses module.name which can differ from the import path, and _call_setup_imports gives a confusing error when setup_imports lacks @classmethod.
metaflow/datastore/task_datastore.py	Save/load artifact paths refactored to dispatch through the pluggable serializer registry. _serializers is now a property that re-queries the live registry on each access, correctly handling lazy-registered serializers for long-lived datastore instances. Empty/multi-blob validation and source auto-injection are well-guarded.
metaflow/plugins/datastores/serializers/pickle_serializer.py	Straightforward PickleSerializer using protocol 4. Legacy gzip+pickle-v{2,4} encodings are accepted for backward compat in can_deserialize. WIRE format correctly raises NotImplementedError.
metaflow/datastore/artifacts/lazy_registry.py	sys.meta_path interceptor that fires SerializerStore._on_module_imported when a watched module finishes loading. _WrappedLoader correctly restores spec.loader on both success and exception paths.
metaflow/plugins/init.py	Adds ARTIFACT_SERIALIZERS_DESC declaring PickleSerializer as the core fallback and triggers SerializerStore.bootstrap() so the dispatch pool is populated before any TaskDataStore is constructed.
metaflow/datastore/artifacts/diagnostic.py	New SerializerRecord dataclass and SerializerState enum for per-serializer lifecycle tracking. Clean, standalone implementation.
metaflow/datastore/exceptions.py	UnpicklableArtifactException now accepts optional artifact_name so PickleSerializer can raise it without knowing the artifact name, re-raised with the name attached by the datastore.
test/unit/test_serializer_integration.py	Integration tests covering round-trips, backward compat, priority dispatch, blob-count validation, exception propagation, and post-init dynamic registration.
test/unit/test_artifact_serializer.py	Unit tests for ABC enforcement, lazy_import mechanics, priority ordering, and the double-assignment guard. Tests remain correct after the idempotency refinement.

_{Reviews (15): Last reviewed commit: "Address saikonen review: pickle owns Unp..." | Re-trigger Greptile}

[codecov]

Codecov Report

❌ Patch coverage is 80.69217% with 106 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (master@352f643). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
metaflow/datastore/artifacts/serializer.py	78.11%	55 Missing and 17 partials ⚠️
metaflow/datastore/artifacts/lazy_registry.py	75.71%	10 Missing and 7 partials ⚠️
metaflow/datastore/task_datastore.py	79.76%	10 Missing and 7 partials ⚠️

Additional details and impacted files

@@            Coverage Diff            @@
##             master    #3117   +/-   ##
=========================================
  Coverage          ?   28.00%           
=========================================
  Files             ?      381           
  Lines             ?    52347           
  Branches          ?     9238           
=========================================
  Hits              ?    14661           
  Misses            ?    36750           
  Partials          ?      936           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[romain-intel]

Minor comments.

[saeidbarati157]

Thanks Romain! Addressed in fed2c78:

• serializer.py method declarations — added full type annotations to the ABC method signatures.
• context param — stripped from deserialize across the ABC, PickleSerializer, the TaskDataStore caller, and tests. Can re-add if/when we have a concrete use case.
• STORAGE/WIRE as an enum — promoted to SerializationFormat(str, Enum). Subclassing str keeps existing format == 'storage' comparisons working so no downstream migration is needed.
• _serializers_override comment — reworded to describe the test override; moved the property-rationale block next to the property definition itself.
• Multi-blob error — appended "at this time. If you have a need for multi blob serializers, please reach out to the Metaflow team."
• Deserializer-not-found error — now prints the serializer_info dict alongside encoding.

All 77 serializer/pickle/lazy/integration unit tests pass.

[greptile-apps]

Protocol-2 fallback removed — behavioral regression for some objects

The old save_artifacts caught SystemError and OverflowError from pickle.dumps(protocol=4) and automatically retried with protocol=2. That fallback is gone here, so any object that triggers one of those errors under protocol=4 (e.g., certain C-extension types or deeply nested structures) will now raise a raw SystemError/OverflowError instead of being serialized successfully. Artifacts that Metaflow previously saved without issue may now fail to save.

@classmethod
def serialize(cls, obj, format=SerializationFormat.STORAGE):
    if format == SerializationFormat.WIRE:
        raise NotImplementedError(
            "PickleSerializer does not support the WIRE format; pickle "
            "produces opaque binary bytes that are not safe to pass as "
            "CLI args or inline IPC payloads."
        )
    try:
        blob = pickle.dumps(obj, protocol=4)
        encoding = "pickle-v4"
    except (SystemError, OverflowError):
        blob = pickle.dumps(obj, protocol=2)
        encoding = "pickle-v2"
    return (
        [SerializedBlob(blob, is_reference=False)],
        SerializationMetadata(
            obj_type=str(type(obj)),
            size=len(blob),
            encoding=encoding,
            serializer_info={},
        ),
    )

[saikonen]

possibly relevant still.

[saikonen]

actually strike that, I'm forgetting the context for keeping the protocol-2 encoding here but it seems to have been tied to python <=3.6 so dropping the fallback should be fine at this point.

[jefferynidd]

would this serializer store be accessible from other metaflow plugins ? i'm working on updating the checkpoint decorator and it'd be super useful if we could re-use this to serialize artifacts for checkpointing purposes !

we'd have to be able to (1) access the serializations for an artifact someone wants to checkpoint , and (2) save the artifacts in a different location from the usual metaflow artifacts

[saeidbarati157]

Yes, that's exactly what the decoupling enables. The serializer produces (blobs, metadata) independently of where you persist them:

from metaflow.datastore.artifacts.serializer import SerializerStore

# (1) serialize
for s in SerializerStore.get_ordered_serializers():
    if s.can_serialize(obj):
        blobs, metadata = s.serialize(obj)
        break

# (2) deserialize from your own storage
for s in SerializerStore.get_ordered_serializers():
    if s.can_deserialize(metadata):
        return s.deserialize([blob], metadata)

You hand blobs to whatever storage you want (your own S3 path, a separate CAS, etc.).

[jefferynidd]

nice ! this is great , thanks for the detailed example !

[saikonen]

couldn't a non-pickle serializer also raise a TypeError ? One option could be to treat any raised MetaflowException as pass-through, and have the pickle-serializer raise the UnpicklableArtifactException instead?

[romain-intel]

I think the idea was to try to return the same exception as before. I think what we could do is specially catch UnpicklableArtifactException (and yes, raise it directly in the pickle serializer) and for everything else, fall through to the DataException which would maintain backwards compatibility with exception types. Would that work for you?

[saeidbarati157]

Moved pickle.dumps's TypeError handling into PickleSerializer itself. It now raises UnpicklableArtifactException (with artifact_name optional). TaskDataStore.save_artifacts catches it specifically to re-attach the artifact name, lets any other MetaflowException pass through unchanged, and only wraps non-MetaflowException failures in DataException for context. Added three tests covering the three branches.

[saikonen]

possibly relevant still.