Address review feedback: wire/storage formats + lazy registry

Wire vs storage format
- Added WIRE and STORAGE constants in metaflow.datastore.artifacts.serializer.
- ArtifactSerializer.serialize/deserialize now accept a ``format`` kwarg so a
  single class can own both the storage path (datastore blobs + metadata) and
  the wire path (string for CLI args, protobuf payloads, cross-process IPC).
- PickleSerializer implements STORAGE only; WIRE raises NotImplementedError
  with an explanation (pickle bytes are not safe as a wire payload).
- Serializers that want wire support implement it on the same class; there's
  no need for a second class per format.

Lazy import-hook registry
- New module metaflow/datastore/artifacts/lazy_registry.py with
  SerializerConfig, an importlib.abc.MetaPathFinder interceptor, and public
  register_serializer_for_type / load_serializer_class entry points.
- If the target type's module is already in sys.modules, registration is
  immediate. Otherwise a hook is installed on sys.meta_path and registration
  fires the first time the user's code imports the target module. This
  defers the cost of serializer-module imports (torch, pyarrow, fastavro,
  ...) until those dependencies are actually in play.
- find_spec temporarily removes the interceptor from sys.meta_path during
  its lookup to avoid recursion.

Review nits
- serializer.py: rename SerializationMetadata.type field to obj_type to avoid
  shadowing the type() builtin (dict key inside _info stays "type" for
  backward compatibility with existing datastores).
- serializer.py: SerializerStore skips TYPE=None AND any subclass that is
  still abstract, via inspect.isabstract().
- serializer.py: get_ordered_serializers memoizes the sorted list and
  invalidates on new registration. Drops the O(n²) list.index tiebreaker —
  Python 3.7+ dicts preserve insertion order.
- task_datastore.py: lift SerializerStore / SerializationMetadata imports to
  module top (were on hot paths: __init__ and load_artifacts).
- task_datastore.py: the "no serializer claimed this artifact" branch now
  raises DataException with a message pointing at the PickleSerializer
  fallback invariant, instead of the misleading UnpicklableArtifactException.
- task_datastore.py: "no deserializer claimed this artifact" now looks up
  ``serializer_info["source"]`` and hints at a missing extension.
- SerializedBlob.compress_method: removed. It was documented as "not wired
  into the save path" — shipping an unwired knob invites extension authors
  to rely on it and discover it is a no-op. Will come back with its
  consuming code in a later change.
- serialize() docstring rewritten: the side-effect-free contract is there
  to support retries, caching, and parallel dispatch, and to keep
  serializers testable; I/O belongs in hooks.

Tests
- New test/unit/test_lazy_serializer_registry.py (9 tests covering config
  validation, eager registration for already-imported types, deferred
  registration through the import hook, and the recursion guard).
- test_artifact_serializer.py gains format-dispatch coverage (STORAGE and
  WIRE round-trip on a toy dual-format serializer; PickleSerializer raises
  on WIRE).
- Removed tests for compress_method and adjusted tests that poked at
  _registration_order so they go through the public API only.

Author: Saeid Barati

metaflow/datastore/artifacts/__init__.py

@@ -3,4 +3,12 @@
     SerializationMetadata,
     SerializedBlob,
     SerializerStore,
+    STORAGE,
+    WIRE,
+)
+from .lazy_registry import (
+    SerializerConfig,
+    load_serializer_class,
+    register_serializer_config,
+    register_serializer_for_type,
 )

metaflow/datastore/artifacts/lazy_registry.py

@@ -0,0 +1,216 @@
+"""
+Lazy serializer registry driven by an import hook.
+
+Extensions ship serializers whose implementation modules may import optional
+heavy dependencies (``torch``, ``pyarrow``, ``fastavro``, ``protobuf``, ...).
+Loading those serializer modules unconditionally at ``metaflow`` import time
+would force every user to pay for dependencies they may not have installed.
+
+This module defers both the serializer class import and its instantiation
+until one of two things happens:
+
+1. The target type's module is already present in :data:`sys.modules` when
+   registration is called — registration then happens immediately.
+2. The target type's module is imported later by the user's code. An
+   ``importlib`` hook watches for that event and performs registration the
+   first time the module is loaded.
+
+The hook is installed on :data:`sys.meta_path` and removes itself from the
+path during its own ``find_spec`` call to avoid recursion.
+"""
+
+import importlib
+import importlib.abc
+import importlib.machinery
+import importlib.util
+import sys
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class SerializerConfig:
+    """
+    Declarative entry recording *which* serializer handles *which* type,
+    without actually importing the serializer class. The class is imported on
+    first use by :func:`load_serializer_class`.
+
+    Parameters
+    ----------
+    canonical_type : str
+        ``"module.ClassName"`` — e.g. ``"builtins.dict"``, ``"torch.Tensor"``.
+    serializer : str
+        Dotted import path to the serializer class, e.g.
+        ``"my_extension.serializers.TorchSerializer"``.
+    priority : int
+        Dispatch priority. Lower is tried first. Matches the existing
+        ``ArtifactSerializer.PRIORITY`` convention.
+    extra_kwargs : dict
+        Optional kwargs passed to the serializer class ``__init__``.
+    """
+
+    canonical_type: str
+    serializer: str
+    priority: int = 100
+    extra_kwargs: dict = field(default_factory=dict)
+
+    def __post_init__(self):
+        if not self.canonical_type:
+            raise ValueError("canonical_type cannot be empty")
+        if not self.serializer or "." not in self.serializer:
+            raise ValueError("serializer must be in 'module.ClassName' format")
+
+    @property
+    def serializer_module(self):
+        return ".".join(self.serializer.split(".")[:-1])
+
+    @property
+    def serializer_class(self):
+        return self.serializer.split(".")[-1]
+
+
+# Module-level registry. Keyed by canonical_type -> SerializerConfig.
+# A separate dict caches instantiated classes so repeated lookups are O(1).
+_registered_configs = {}
+_loaded_serializers = {}
+
+
+def register_serializer_config(config):
+    """Store a config immediately. The serializer class is not imported yet."""
+    _registered_configs[config.canonical_type] = config
+    # Any previously cached class for this type is now stale.
+    _loaded_serializers.pop(config.canonical_type, None)
+
+
+def load_serializer_class(canonical_type):
+    """
+    Resolve and cache the serializer class for ``canonical_type``. Returns
+    ``None`` if no config is registered for that type.
+    """
+    cached = _loaded_serializers.get(canonical_type)
+    if cached is not None:
+        return cached
+    config = _registered_configs.get(canonical_type)
+    if config is None:
+        return None
+    module = importlib.import_module(config.serializer_module)
+    cls = getattr(module, config.serializer_class)
+    _loaded_serializers[canonical_type] = cls
+    return cls
+
+
+def iter_registered_configs():
+    """Iterate all registered configs. Deterministic order (insertion)."""
+    return list(_registered_configs.values())
+
+
+class _WrappedLoader(importlib.abc.Loader):
+    """Delegating loader that fires a callback after ``exec_module``.
+
+    Only ``create_module`` and ``exec_module`` are overridden. Other loader
+    attributes (``get_resource_reader``, ``get_filename``, ``get_data``,
+    ``is_package``, ``get_source``, ...) are forwarded to the wrapped loader
+    via ``__getattr__`` so importers that poke at those interfaces continue
+    to work transparently.
+    """
+
+    def __init__(self, original_loader, interceptor):
+        self._original = original_loader
+        self._interceptor = interceptor
+
+    def create_module(self, spec):
+        return self._original.create_module(spec)
+
+    def exec_module(self, module):
+        self._original.exec_module(module)
+        self._interceptor._on_module_imported(module)
+
+    def __getattr__(self, name):
+        return getattr(self._original, name)
+
+
+class _SerializerImportInterceptor(importlib.abc.MetaPathFinder):
+    """
+    :class:`importlib.abc.MetaPathFinder` that watches for a fixed set of
+    module names and fires :func:`_on_module_imported` once each has been
+    fully executed.
+    """
+
+    def __init__(self):
+        # module_name -> list[SerializerConfig]
+        self._pending = {}
+        self._processed = set()
+
+    def watch(self, module_name, config):
+        self._pending.setdefault(module_name, []).append(config)
+
+    def find_spec(self, fullname, path, target=None):
+        if fullname not in self._pending:
+            return None
+        # Remove ourselves from the path during the lookup below so Python's
+        # normal finders (not us) can resolve the real spec. Reinstall after.
+        was_installed = self in sys.meta_path
+        if was_installed:
+            sys.meta_path.remove(self)
+        try:
+            spec = importlib.util.find_spec(fullname)
+        finally:
+            if was_installed:
+                sys.meta_path.insert(0, self)
+        if spec is None or spec.loader is None:
+            return None
+        spec.loader = _WrappedLoader(spec.loader, self)
+        return spec
+
+    def _on_module_imported(self, module):
+        module_name = module.__name__
+        if module_name in self._processed:
+            return
+        self._processed.add(module_name)
+        for config in self._pending.get(module_name, ()):
+            class_name = config.canonical_type.rsplit(".", 1)[-1]
+            if hasattr(module, class_name):
+                register_serializer_config(config)
+
+
+_interceptor = _SerializerImportInterceptor()
+
+
+def _ensure_interceptor_installed():
+    if _interceptor in sys.meta_path:
+        sys.meta_path.remove(_interceptor)
+    sys.meta_path.insert(0, _interceptor)
+
+
+def register_serializer_for_type(canonical_type, serializer, **kwargs):
+    """
+    Public entry point for extensions.
+
+    If the target type's module is already loaded, the config is stored
+    immediately. Otherwise, an import hook is installed and registration is
+    deferred to the first ``import`` of the target module.
+
+    ``canonical_type`` is ``"module.ClassName"``. ``serializer`` is a dotted
+    path to the serializer class. Additional ``priority`` / ``extra_kwargs``
+    forwarded into :class:`SerializerConfig`.
+    """
+    config = SerializerConfig(
+        canonical_type=canonical_type, serializer=serializer, **kwargs
+    )
+    module_name, class_name = canonical_type.rsplit(".", 1)
+    mod = sys.modules.get(module_name)
+    if mod is not None and hasattr(mod, class_name):
+        register_serializer_config(config)
+        return
+    _ensure_interceptor_installed()
+    _interceptor.watch(module_name, config)
+
+
+def _reset_for_tests():
+    """Clear all module-level state. Intended for unit tests only."""
+    _registered_configs.clear()
+    _loaded_serializers.clear()
+    _interceptor._pending.clear()
+    _interceptor._processed.clear()
+    if _interceptor in sys.meta_path:
+        sys.meta_path.remove(_interceptor)

metaflow/datastore/artifacts/serializer.py

@@ -1,9 +1,16 @@
+import inspect
 from abc import ABCMeta, abstractmethod
 from collections import namedtuple
 
 
+# Serialization formats. STORAGE produces (blobs, metadata) for the datastore;
+# WIRE produces a str for CLI args, protobuf payloads, and cross-process IPC.
+STORAGE = "storage"
+WIRE = "wire"
+
+
 SerializationMetadata = namedtuple(
-    "SerializationMetadata", ["type", "size", "encoding", "serializer_info"]
+    "SerializationMetadata", ["obj_type", "size", "encoding", "serializer_info"]
 )
 
 
@@ -21,20 +28,14 @@ class SerializedBlob(object):
         The blob data (bytes) or a reference key (str).
     is_reference : bool, optional
         If None, auto-detected from value type: str -> reference, bytes -> new data.
-    compress_method : str
-        Compression method for new blobs. Ignored for references. Default "gzip".
-        NOTE: Not yet wired into the save path — ContentAddressedStore currently
-        always applies gzip. This field is forward-looking for when per-blob
-        compression control is needed (e.g., multi-blob IOType support).
     """
 
-    def __init__(self, value, is_reference=None, compress_method="gzip"):
+    def __init__(self, value, is_reference=None):
         if not isinstance(value, (str, bytes)):
             raise TypeError(
                 "SerializedBlob value must be str or bytes, got %s" % type(value).__name__
             )
         self.value = value
-        self.compress_method = compress_method
         if is_reference is None:
             self.is_reference = isinstance(value, str)
         else:
@@ -55,27 +56,57 @@ class SerializerStore(ABCMeta):
     """
 
     _all_serializers = {}
-    _registration_order = []
+    _ordered_cache = None
 
     def __init__(cls, name, bases, namespace):
         super().__init__(name, bases, namespace)
-        if cls.TYPE is not None:
-            if cls.TYPE not in SerializerStore._all_serializers:
-                SerializerStore._registration_order.append(cls.TYPE)
-            SerializerStore._all_serializers[cls.TYPE] = cls
+        # Skip the abstract base and any subclass that didn't implement all
+        # abstract methods — registering a partially-abstract class would
+        # blow up only at dispatch time.
+        if cls.TYPE is None or inspect.isabstract(cls):
+            return
+        SerializerStore._all_serializers[cls.TYPE] = cls
+        SerializerStore._ordered_cache = None
 
     @staticmethod
     def get_ordered_serializers():
         """
         Return serializer classes sorted by (PRIORITY, registration_order).
 
-        This ordering is deterministic for a given set of loaded serializers.
+        Python 3.7+ dicts preserve insertion order, so enumerating
+        ``_all_serializers.values()`` yields registration order. A stable sort
+        on PRIORITY preserves that tiebreaker.
+
+        Serializers registered via the lazy registry are materialized here
+        too: each registered class is imported on demand and folded into the
+        dispatch order. Without this step, a lazy
+        ``register_serializer_for_type`` call would be silently ignored
+        at dispatch time.
         """
-        order = SerializerStore._registration_order
-        return sorted(
-            SerializerStore._all_serializers.values(),
-            key=lambda s: (s.PRIORITY, order.index(s.TYPE)),
-        )
+        # Imported locally to avoid a circular import between this module and
+        # ``lazy_registry`` (which depends on the ArtifactSerializer ABC).
+        from .lazy_registry import iter_registered_configs, load_serializer_class
+
+        lazy_classes = []
+        for cfg in iter_registered_configs():
+            cls = load_serializer_class(cfg.canonical_type)
+            if cls is not None:
+                lazy_classes.append(cls)
+
+        if SerializerStore._ordered_cache is None or lazy_classes:
+            # De-duplicate: lazy classes typically also self-register via the
+            # metaclass, but when loaded outside normal import flow they may
+            # not. ``dict.fromkeys`` preserves first-seen order while dropping
+            # duplicates.
+            combined = list(
+                dict.fromkeys(
+                    list(SerializerStore._all_serializers.values()) + lazy_classes
+                )
+            )
+            SerializerStore._ordered_cache = sorted(
+                combined, key=lambda s: s.PRIORITY
+            )
+        return SerializerStore._ordered_cache
 
 
 class ArtifactSerializer(object, metaclass=SerializerStore):
@@ -135,36 +166,49 @@ def can_deserialize(cls, metadata):
 
     @classmethod
     @abstractmethod
-    def serialize(cls, obj):
+    def serialize(cls, obj, format=STORAGE):
         """
-        Serialize obj to blobs and metadata. Must be side-effect-free.
+        Serialize obj. Must be side-effect-free: this method may be invoked
+        multiple times (caching, retries, parallel dispatch) and must not
+        perform I/O, mutate global state, or register the object elsewhere.
+        Side effects that need to happen at persist time belong in hooks,
+        not in the serializer.
 
         Parameters
         ----------
         obj : Any
             The Python object to serialize.
+        format : str
+            Either ``STORAGE`` (default) or ``WIRE``.
+            - ``STORAGE`` returns a tuple ``(List[SerializedBlob], SerializationMetadata)``
+              for persisting through the datastore.
+            - ``WIRE`` returns a ``str`` representation for CLI args, protobuf
+              payloads, and cross-process IPC. Serializers that cannot provide
+              a wire encoding should raise ``NotImplementedError``.
 
         Returns
         -------
-        tuple
-            (List[SerializedBlob], SerializationMetadata)
+        tuple or str
         """
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
-    def deserialize(cls, blobs, metadata, context):
+    def deserialize(cls, data, metadata=None, context=None, format=STORAGE):
         """
-        Deserialize blobs back to a Python object.
+        Deserialize back to a Python object.
 
         Parameters
         ----------
-        blobs : List[bytes]
-            The raw blob data.
-        metadata : SerializationMetadata
-            Metadata stored alongside the artifact.
-        context : Any
+        data : Union[List[bytes], str]
+            ``List[bytes]`` when ``format=STORAGE``; ``str`` when ``format=WIRE``.
+        metadata : SerializationMetadata, optional
+            Metadata stored alongside the artifact. Required for STORAGE,
+            ignored for WIRE.
+        context : Any, optional
             Optional context for deserialization (e.g., task vs client loading).
+        format : str
+            Either ``STORAGE`` (default) or ``WIRE``.
 
         Returns
         -------