Unified artifact serializer registration via setup_imports()

Replace the two-path registration scheme (ARTIFACT_SERIALIZERS_DESC +
register_serializer_for_type) with a single declarative mechanism.
Extensions ship one serializer file per serializer; heavy imports
live in a setup_imports() classmethod whose ModuleNotFoundError
parks the serializer on an import hook and retries once the awaited
module appears in sys.modules.

Author-facing API
-----------------

    class TorchSerializer(ArtifactSerializer):
        TYPE = "torch"
        PRIORITY = 50

        @classmethod
        def setup_imports(cls, context=None):
            cls.lazy_import("torch")
            cls.lazy_import("pyarrow", alias="pa")

        @classmethod
        def can_serialize(cls, obj):
            return isinstance(obj, cls.torch.Tensor)
        ...

    # mfextinit_myext.py
    ARTIFACT_SERIALIZERS_DESC = [
        ("torch", "my_ext.serializers.torch_serializer.TorchSerializer"),
    ]

Mechanism
---------

- SerializerStore.bootstrap() walks every extension's
  ARTIFACT_SERIALIZERS_DESC directly (the serializer category owns its
  own lifecycle rather than flowing through resolve_plugins), applies
  ENABLED_ARTIFACT_SERIALIZER +/- toggles, and drives each entry
  through a state machine:
      known → importing → class_loaded → importing_deps → active
  with terminal states `broken` / `disabled` and the retry state
  `pending_on_imports`.
- lazy_import(module_path, alias=None) imports the module, stashes it
  on cls, returns it. Rejects reserved names (TYPE, PRIORITY, dispatch
  methods, leading underscore) and double-assignment within one
  setup_imports() call.
- A sys.meta_path interceptor watches the pending modules; when any
  imports, SerializerStore._on_module_imported re-runs the lifecycle
  for every parked entry waiting on that module. Loop guard: the same
  ImportError.name raising twice → broken.
- Dispatch reads _active_serializers only. Exceptions raised by
  can_serialize / can_deserialize on a currently-active serializer are
  caught, counted in the record's dispatch_error_count, and the
  serializer is skipped for that artifact only — preserving the
  PickleSerializer fallback.
- On PRIORITY ties, last-registered wins; a lexicographic tiebreak on
  class_path provides cross-environment reproducibility.

Diagnostic API
--------------

metaflow.datastore.artifacts.list_serializer_status() returns a list
of per-entry dicts with name, class_path, state, awaiting_modules,
last_error, priority, type, import_trigger, and dispatch_error_count.
Primary use case: answering "why isn't my custom serializer active?"
without reading source.

Removed
-------

- register_serializer_for_type() (public API)
- SerializerConfig, register_serializer_config, iter_registered_configs,
  load_serializer_class, plus the backing declarative-config plumbing
  inside lazy_registry.py — nothing in the unified path uses them.
- ARTIFACT_SERIALIZERS = resolve_plugins("artifact_serializer") in
  metaflow/plugins/__init__.py (no readers in the codebase).
- "artifact_serializer" from _plugin_categories in
  metaflow/extension_support/plugins.py.

Testing
-------

- 102 unit tests across test_artifact_serializer, test_pickle_serializer,
  test_serializer_integration, test_serializer_lifecycle, and
  test_serializer_public_api. Coverage: state-machine transitions,
  retry and loop guard, disabled toggle, dispatch exception handling,
  lazy_import reserved-name + collision guards, setup_imports signature
  compatibility (both def setup_imports(cls) and
  def setup_imports(cls, context=None) are supported), subclass
  inheritance, and public API surface assertions.
- SerializerStore._reset_for_tests() clears registry + interceptor state
  and walks MRO to delattr lazy-imported attributes for test isolation.

Author: romain-intel

metaflow/datastore/artifacts/__init__.py

@@ -5,9 +5,4 @@
     SerializedBlob,
     SerializerStore,
 )
-from .lazy_registry import (
-    SerializerConfig,
-    load_serializer_class,
-    register_serializer_config,
-    register_serializer_for_type,
-)
+from .diagnostic import list_serializer_status, SerializerState

metaflow/datastore/artifacts/diagnostic.py

@@ -0,0 +1,55 @@
+"""Per-entry diagnostic records for the artifact-serializer lifecycle."""
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import List, Optional
+
+
+class SerializerState(str, Enum):
+    KNOWN = "known"
+    IMPORTING = "importing"
+    CLASS_LOADED = "class_loaded"
+    IMPORTING_DEPS = "importing_deps"
+    ACTIVE = "active"
+    PENDING_ON_IMPORTS = "pending_on_imports"
+    BROKEN = "broken"
+    DISABLED = "disabled"
+
+
+@dataclass
+class SerializerRecord:
+    name: str
+    class_path: str
+    state: SerializerState = SerializerState.KNOWN
+    awaiting_modules: List[str] = field(default_factory=list)
+    last_error: Optional[str] = None
+    priority: Optional[int] = None
+    type: Optional[str] = None
+    import_trigger: Optional[str] = None
+    dispatch_error_count: int = 0
+
+    def as_dict(self):
+        return {
+            "name": self.name,
+            "class_path": self.class_path,
+            "state": self.state.value,
+            "awaiting_modules": list(self.awaiting_modules),
+            "last_error": self.last_error,
+            "priority": self.priority,
+            "type": self.type,
+            "import_trigger": self.import_trigger,
+            "dispatch_error_count": self.dispatch_error_count,
+        }
+
+
+def list_serializer_status():
+    """Return a list of per-serializer diagnostic records as dicts.
+
+    One entry per tuple in ``ARTIFACT_SERIALIZERS_DESC`` (post-toggle),
+    including entries in ``pending_on_imports``, ``broken``, and
+    ``disabled`` states. Used for debugging "why isn't my custom
+    serializer active?".
+    """
+    from .serializer import SerializerStore
+
+    return [rec.as_dict() for rec in SerializerStore._records.values()]

metaflow/datastore/artifacts/lazy_registry.py

@@ -1,22 +1,24 @@
 """
-Lazy serializer registry driven by an import hook.
+Import-hook plumbing that the serializer registry uses to retry a serializer's
+``setup_imports`` after one of its required modules becomes importable.
 
 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.
+Loading those modules unconditionally at ``metaflow`` import time would force
+every user to pay for dependencies they may not have installed. When
+``SerializerStore.bootstrap_entries`` encounters such a missing module, it
+parks the entry in ``pending_on_imports`` state and installs a watch here.
+The first time the awaited module is imported by the user's code, this
+interceptor fires ``SerializerStore._on_module_imported`` so the registry can
+retry activation.
+
+The interceptor is installed on :data:`sys.meta_path` and removes itself from
+the path during its own ``find_spec`` call to avoid recursion.
+
+This module has no public API — extensions declare serializers through
+``ARTIFACT_SERIALIZERS_DESC`` in their ``mfextinit_*`` file and interact with
+the registry via the state-machine public surface in
+:mod:`metaflow.datastore.artifacts.serializer`.
 """
 
 import importlib
@@ -25,84 +27,6 @@
 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``.
@@ -131,21 +55,26 @@ def __getattr__(self, 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.
+    :class:`importlib.abc.MetaPathFinder` that watches a set of module names
+    and notifies :class:`SerializerStore` once each has finished executing.
     """
 
     def __init__(self):
-        # module_name -> list[SerializerConfig]
-        self._pending = {}
+        # Module names to watch on behalf of SerializerStore records parked
+        # via _park_on_import_error. Firing calls
+        # SerializerStore._on_module_imported.
+        self._watched = set()
+        # Modules we have already notified about, to avoid firing twice if
+        # the same module gets imported through multiple paths.
         self._processed = set()
 
-    def watch(self, module_name, config):
-        self._pending.setdefault(module_name, []).append(config)
+    def watch(self, module_name):
+        """Watch ``module_name``. When it finishes executing,
+        :meth:`SerializerStore._on_module_imported` is called."""
+        self._watched.add(module_name)
 
     def find_spec(self, fullname, path, target=None):
-        if fullname not in self._pending:
+        if fullname not in self._watched:
             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.
@@ -167,10 +96,16 @@ def _on_module_imported(self, module):
         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)
+        if module_name not in self._watched:
+            return
+        try:
+            from .serializer import SerializerStore
+
+            SerializerStore._on_module_imported(module_name, module)
+        except Exception:
+            # A broken callback must not break the host's import. The record
+            # itself will be marked BROKEN via _retry_bootstrap.
+            pass
 
 
 _interceptor = _SerializerImportInterceptor()
@@ -182,35 +117,9 @@ def _ensure_interceptor_installed():
     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._watched.clear()
     _interceptor._processed.clear()
     if _interceptor in sys.meta_path:
         sys.meta_path.remove(_interceptor)