IOType: registry-based deserialization for Flow client round-trip

The metadata service stores only the artifact encoding string
("iotype:<type_name>"), not the full serializer_info dict. So when the
Flow client reconstructs an artifact, IOTypeSerializer.deserialize had no
iotype_module/iotype_class to import and KeyError'd.

Populate a registry via IOType.__init_subclass__ keyed by type_name, and
have IOTypeSerializer.deserialize resolve the class from the encoding
suffix first. iotype_module/iotype_class hints remain a secondary lookup
for introspecting artifacts whose owning extension isn't loaded locally.

Fixes round-trip of IOType artifacts through Flow(...).latest_run
.end_task[...].data without touching the metadata service or client
core.py.

Author: Saeid Barati

metaflow/io_types/base.py

@@ -26,6 +26,18 @@
 
 _UNSET = object()
 
+# Registry of concrete IOType subclasses keyed by their ``type_name``. Populated
+# by ``IOType.__init_subclass__``. The datastore encodes each artifact's
+# ``type_name`` in its ``SerializationMetadata.encoding`` (``iotype:<name>``),
+# so ``IOTypeSerializer.deserialize`` can recover the class without the
+# metadata service having to persist the Python module+class path.
+_TYPE_REGISTRY = {}
+
+
+def get_iotype_by_name(type_name):
+    """Return the IOType subclass registered under ``type_name``, or None."""
+    return _TYPE_REGISTRY.get(type_name)
+
 
 class IOType(object, metaclass=ABCMeta):
     """
@@ -45,6 +57,17 @@ class IOType(object, metaclass=ABCMeta):
 
     type_name = None  # e.g. "text", "json", "int64" — set by subclasses.
 
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        # Register concrete subclasses so IOTypeSerializer can recover them
+        # from just the ``type_name`` suffix of a stored artifact's encoding.
+        # Abstract intermediates (``type_name`` is still None) don't register.
+        # Last-write-wins: production code is expected to declare each
+        # ``type_name`` on exactly one class; test-local subclasses that reuse
+        # a name harmlessly overwrite each other.
+        if cls.type_name:
+            _TYPE_REGISTRY[cls.type_name] = cls
+
     def __init__(self, value=_UNSET):
         self._value = value
 

metaflow/plugins/datastores/serializers/iotype_serializer.py

@@ -6,16 +6,20 @@
     STORAGE,
     WIRE,
 )
-from metaflow.io_types.base import IOType
+from metaflow.io_types.base import IOType, get_iotype_by_name
 
 
 class IOTypeSerializer(ArtifactSerializer):
     """
     Bridge between :class:`IOType` and the pluggable serializer framework.
 
     Claims any :class:`IOType` instance on save. On load, reconstructs the
-    original subclass from the ``iotype_module`` / ``iotype_class`` hints
-    that were written into ``serializer_info``.
+    original subclass by looking up its ``type_name`` in the global IOType
+    registry (populated via :meth:`IOType.__init_subclass__`). The
+    ``iotype_module`` / ``iotype_class`` hints stored in ``serializer_info``
+    are kept as a secondary lookup path — useful when a subclass isn't yet
+    registered in the reader process, or when inspecting artifacts produced
+    by extensions whose code isn't installed locally.
 
     ``PRIORITY`` is 50 — ahead of the default (100) so this bridge catches
     :class:`IOType` artifacts before any generic catch-all, and always ahead
@@ -74,13 +78,28 @@ def deserialize(cls, data, metadata=None, context=None, format=STORAGE):
             raise NotImplementedError(
                 "IOTypeSerializer only handles the STORAGE format."
             )
-        info = metadata.serializer_info
-        mod = importlib.import_module(info["iotype_module"])
-        iotype_cls = getattr(mod, info["iotype_class"])
+        info = metadata.serializer_info or {}
+        # Primary path: registry lookup by the type_name encoded in the
+        # artifact's encoding. Works whether or not the metadata service
+        # propagates ``serializer_info`` to the reader.
+        type_name = metadata.encoding[len(cls._ENCODING_PREFIX):]
+        iotype_cls = get_iotype_by_name(type_name)
+        # Fallback: explicit module/class hints in serializer_info. Useful for
+        # inspecting artifacts produced by extensions not loaded locally.
+        if iotype_cls is None and "iotype_module" in info and "iotype_class" in info:
+            mod = importlib.import_module(info["iotype_module"])
+            iotype_cls = getattr(mod, info["iotype_class"])
+        if iotype_cls is None:
+            raise ValueError(
+                "IOTypeSerializer could not resolve a class for encoding %r; "
+                "no IOType subclass is registered under type_name %r and "
+                "serializer_info lacks iotype_module/iotype_class hints."
+                % (metadata.encoding, type_name)
+            )
         # Only allow actual IOType subclasses — metadata is untrusted input.
         if not (isinstance(iotype_cls, type) and issubclass(iotype_cls, IOType)):
             raise ValueError(
-                "IOTypeSerializer metadata references '%s.%s' which is not an "
-                "IOType subclass" % (info["iotype_module"], info["iotype_class"])
+                "IOTypeSerializer resolved %r for encoding %r, which is not "
+                "an IOType subclass" % (iotype_cls, metadata.encoding)
             )
         return iotype_cls.deserialize(data, format=STORAGE, metadata=info)