Add IOType ABC, Json, Struct, and the datastore bridge

On top of Netflix#3117, this ships the typed-artifact layer. Scope kept
deliberately small:

- ``IOType`` ABC — the contract extension authors target.
- ``Json`` and ``Struct`` — two concrete types with clear standalone
  value in core: wire format for CLI/IPC, cross-language JSON bytes
  on storage, no pickle code-execution risk. ``Struct`` also walks
  directly-nested ``@dataclass`` fields so ``Outer(inner=Inner(...))``
  round-trips back to its original type (generic containers like
  ``List[Inner]`` come back as raw JSON values — wrap those
  explicitly when you need richer reconstruction).
- ``IOTypeSerializer`` — the bridge that plugs any ``IOType`` instance
  into the ``ArtifactSerializer`` dispatch added by Netflix#3117 so save/load
  through the datastore just works.

What's intentionally *not* in this PR

- Primitive wrappers (Int32/Int64/Float32/Float64/Bool/Text). Standard
  Python numbers and strings flow through ``PickleSerializer``
  unchanged. Wrapping is opt-in, for cases where you want
  constraints/metadata attached.
- ``Tensor``. Pulls in numpy + byte-order/dtype opinions; belongs in an
  extension that can own those choices.
- ``List`` / ``Map`` / ``Enum``. Thin wrappers whose value over plain
  JSON is mostly schema emission — not enough on their own for core.
- Rich schema emission from ``Struct.to_spec()``. Extensions that ship
  primitive wrappers can override to emit fully-typed schemas; core
  just returns ``{"type": "struct"}``.

Contract

``serialize(format=...)`` / ``deserialize(data, format=..., **kw)``
mirror the ``ArtifactSerializer`` signature from Netflix#3117 and use the same
``WIRE`` / ``STORAGE`` constants, so one subclass owns both
representations:

- ``STORAGE`` → ``(List[SerializedBlob], metadata_dict)`` for
  persisting through the datastore.
- ``WIRE`` → ``str`` for CLI args, protobuf payloads, and
  cross-process IPC.

Subclasses implement four hooks (``_wire_serialize``,
``_wire_deserialize``, ``_storage_serialize``,
``_storage_deserialize``). Instantiating without the hooks raises
``TypeError``.

``IOTypeSerializer`` is registered via ``ARTIFACT_SERIALIZERS_DESC``
with ``PRIORITY=50`` — ahead of the default 100 so it catches
``IOType`` instances before a generic catch-all, and always ahead of
the ``PickleSerializer`` fallback (9999). It implements only
``STORAGE``; wire encoding is produced by calling
``IOType.serialize(format=WIRE)`` directly.

Safety

- ``Struct._storage_deserialize`` and ``IOTypeSerializer.deserialize``
  both require the class named in artifact metadata to be an actual
  class (``isinstance(..., type)``) before any further checks. This
  excludes module-level dataclass *instances* (``is_dataclass`` alone
  returns ``True`` for those) and other callables that could be
  invoked with attacker-controlled kwargs.
- Importing the metadata-named module can still run module-level
  side-effect code; the ``Struct`` docstring calls this out so callers
  don't load artifacts from untrusted sources.

Tests

- ``test_base.py`` — abstract instantiation, WIRE/STORAGE dispatch,
  invalid format, equality/hash, spec.
- ``test_json_type.py`` — wire and storage round-trips.
- ``test_struct_type.py`` — dataclass round-trip, dict round-trip,
  directly-nested dataclass round-trip, container-field pass-through,
  rejection of non-dataclass and dataclass-instance metadata.
- ``test_iotype_serializer.py`` — bridge
  ``can_serialize``/``can_deserialize``, round-trip through dataclass
  reconstruction, rejection of non-IOType classes in metadata, WIRE
  not supported on the bridge.

Author: Saeid Barati

metaflow/io_types/__init__.py

@@ -0,0 +1,3 @@
+from .base import IOType
+from .json_type import Json
+from .struct_type import Struct

metaflow/io_types/base.py

@@ -0,0 +1,138 @@
+"""Typed-artifact contract for Metaflow.
+
+This module defines the minimal :class:`IOType` abstract base class. OSS
+Metaflow ships the contract; concrete types (scalars, tensors, enums,
+dataclass-backed structs, etc.) live in extensions — they embody
+deployment-specific opinions about encoding, byte order, and dataclass
+inference that do not belong in core.
+
+:class:`IOType` mirrors the ``format`` argument introduced on
+:class:`metaflow.datastore.artifacts.serializer.ArtifactSerializer` so a
+single subclass can own both representations:
+
+- ``STORAGE`` — blob-based, persisted through the datastore.
+- ``WIRE`` — string-based, for CLI args, protobuf payloads, and
+  cross-process IPC.
+
+Subclasses implement four hooks (``_wire_serialize``, ``_wire_deserialize``,
+``_storage_serialize``, ``_storage_deserialize``); callers use the public
+``serialize(format=...)`` / ``deserialize(data, format=...)`` methods.
+"""
+
+from abc import ABCMeta, abstractmethod
+
+from metaflow.datastore.artifacts.serializer import STORAGE, WIRE
+
+
+_UNSET = object()
+
+
+class IOType(object, metaclass=ABCMeta):
+    """
+    Base class for typed Metaflow artifacts.
+
+    An :class:`IOType` instance plays two roles:
+
+    - **Descriptor** (no value): ``Int64`` in a spec describes an int64
+      field.
+    - **Wrapper** (with value): ``Int64(42)`` wraps a value for typed
+      serialization.
+
+    Subclasses implement four internal operations, dispatched by the
+    ``format`` argument of the public :meth:`serialize` / :meth:`deserialize`
+    methods.
+    """
+
+    type_name = None  # e.g. "text", "json", "int64" — set by subclasses.
+
+    def __init__(self, value=_UNSET):
+        self._value = value
+
+    @property
+    def value(self):
+        """The wrapped Python value, or ``_UNSET`` if this is a pure descriptor."""
+        return self._value
+
+    # -- Public API --------------------------------------------------------
+
+    def serialize(self, format=STORAGE):
+        """
+        Serialize the wrapped value. Must be side-effect-free.
+
+        Parameters
+        ----------
+        format : str
+            ``STORAGE`` (default) returns ``(List[SerializedBlob], dict)``.
+            ``WIRE`` returns a ``str``.
+        """
+        if format == WIRE:
+            return self._wire_serialize()
+        if format == STORAGE:
+            return self._storage_serialize()
+        raise ValueError("format must be %r or %r, got %r" % (STORAGE, WIRE, format))
+
+    @classmethod
+    def deserialize(cls, data, format=STORAGE, **kwargs):
+        """
+        Reconstruct an :class:`IOType` from serialized data.
+
+        Parameters
+        ----------
+        data : Union[str, List[bytes]]
+            ``str`` when ``format=WIRE``; ``List[bytes]`` when ``format=STORAGE``.
+        format : str
+            ``STORAGE`` (default) or ``WIRE``.
+        **kwargs
+            Forwarded to the underlying ``_storage_deserialize`` hook
+            (e.g. metadata the datastore produced at save time).
+        """
+        if format == WIRE:
+            return cls._wire_deserialize(data)
+        if format == STORAGE:
+            return cls._storage_deserialize(data, **kwargs)
+        raise ValueError("format must be %r or %r, got %r" % (STORAGE, WIRE, format))
+
+    # -- Subclass hooks ----------------------------------------------------
+
+    @abstractmethod
+    def _wire_serialize(self):
+        """Value -> string (for CLI args, protobuf, external APIs)."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def _wire_deserialize(cls, s):
+        """String -> :class:`IOType` instance."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _storage_serialize(self):
+        """Value -> ``(List[SerializedBlob], metadata_dict)``. Side-effect-free."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def _storage_deserialize(cls, blobs, **kwargs):
+        """``(List[bytes], metadata)`` -> :class:`IOType` instance."""
+        raise NotImplementedError
+
+    # -- Spec generation ---------------------------------------------------
+
+    def to_spec(self):
+        """JSON type spec. Works with or without a wrapped value."""
+        return {"type": self.type_name}
+
+    # -- Dunder ------------------------------------------------------------
+
+    def __repr__(self):
+        if self._value is _UNSET:
+            return "%s()" % self.__class__.__name__
+        return "%s(%r)" % (self.__class__.__name__, self._value)
+
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return NotImplemented
+        return self._value == other._value
+
+    def __hash__(self):
+        return hash((type(self), self._value))

metaflow/io_types/json_type.py

@@ -0,0 +1,27 @@
+import json
+
+from ..datastore.artifacts.serializer import SerializedBlob
+from .base import IOType
+
+
+class Json(IOType):
+    """JSON type (dict or list). Wire: JSON string. Storage: UTF-8 JSON bytes."""
+
+    type_name = "json"
+
+    def _wire_serialize(self):
+        return json.dumps(self._value, separators=(",", ":"), sort_keys=True)
+
+    @classmethod
+    def _wire_deserialize(cls, s):
+        return cls(json.loads(s))
+
+    def _storage_serialize(self):
+        blob = json.dumps(self._value, separators=(",", ":"), sort_keys=True).encode(
+            "utf-8"
+        )
+        return [SerializedBlob(blob)], {}
+
+    @classmethod
+    def _storage_deserialize(cls, blobs, **kwargs):
+        return cls(json.loads(blobs[0].decode("utf-8")))

metaflow/io_types/struct_type.py

@@ -0,0 +1,127 @@
+import dataclasses
+import importlib
+import json
+import typing
+
+from ..datastore.artifacts.serializer import SerializedBlob
+from .base import IOType, _UNSET
+
+
+def _reconstruct(dc_type, data):
+    """
+    Rebuild a dataclass instance from JSON-decoded ``data``, recursing into
+    fields whose annotation is itself a dataclass. Containerized annotations
+    (``List[Foo]``, ``Dict[str, Foo]``, ``Optional[Foo]``, ...) are left as
+    raw JSON-decoded values; callers that need rich container reconstruction
+    should wrap the field explicitly (e.g. in a ``List`` IOType shipped by
+    an extension).
+    """
+    try:
+        hints = typing.get_type_hints(dc_type)
+    except Exception:
+        hints = {}
+    kwargs = {}
+    for f in dataclasses.fields(dc_type):
+        raw = data.get(f.name)
+        annotation = hints.get(f.name, f.type)
+        if (
+            isinstance(annotation, type)
+            and dataclasses.is_dataclass(annotation)
+            and isinstance(raw, dict)
+        ):
+            kwargs[f.name] = _reconstruct(annotation, raw)
+        else:
+            kwargs[f.name] = raw
+    return dc_type(**kwargs)
+
+
+class Struct(IOType):
+    """
+    Structured type mapping to a Python ``@dataclass``.
+
+    Wire: JSON string. Storage: JSON UTF-8 bytes.
+
+    Wraps a ``@dataclass`` instance. On save, ``dataclasses.asdict`` flattens
+    the whole tree to plain dicts; on load, fields typed as dataclasses are
+    recursively rebuilt into their original types. Generic container
+    annotations (``List[Foo]``, ``Dict[str, Foo]``, ``Optional[Foo]``) are
+    not walked — those fields come back as raw JSON-decoded values. Wrap
+    those explicitly (e.g. via ``List[Struct]`` support shipped by an
+    extension) when you need typed containers.
+
+    .. warning::
+       ``Struct._storage_deserialize`` imports the dataclass module named in
+       the artifact metadata. Metadata written by this class is safe, but
+       metadata supplied from an untrusted source can trigger arbitrary
+       imports (and any import-time side effects those modules carry).
+       Only load artifacts from sources you trust.
+
+    Parameters
+    ----------
+    value : dataclass instance or dict, optional
+        The wrapped value. Dataclass instances are serialized via
+        ``dataclasses.asdict``; plain dicts are serialized directly.
+    dataclass_type : type, optional
+        The ``@dataclass`` class, for type-descriptor use (no value).
+    """
+
+    type_name = "struct"
+
+    def __init__(self, value=_UNSET, dataclass_type=None):
+        if value is not _UNSET and dataclasses.is_dataclass(value):
+            self._dataclass_type = type(value)
+        elif dataclass_type is not None:
+            self._dataclass_type = dataclass_type
+        else:
+            self._dataclass_type = None
+        super().__init__(value)
+
+    def _to_dict(self):
+        """Convert value to dict, handling both dataclass and plain dict."""
+        if dataclasses.is_dataclass(self._value):
+            return dataclasses.asdict(self._value)
+        if isinstance(self._value, dict):
+            return self._value
+        raise TypeError(
+            "Struct value must be a dataclass instance or dict, got %s"
+            % type(self._value).__name__
+        )
+
+    def _wire_serialize(self):
+        return json.dumps(self._to_dict(), separators=(",", ":"), sort_keys=True)
+
+    @classmethod
+    def _wire_deserialize(cls, s):
+        return cls(json.loads(s))
+
+    def _storage_serialize(self):
+        blob = json.dumps(
+            self._to_dict(), separators=(",", ":"), sort_keys=True
+        ).encode("utf-8")
+        meta = {}
+        if self._dataclass_type is not None:
+            meta["dataclass_module"] = self._dataclass_type.__module__
+            meta["dataclass_class"] = self._dataclass_type.__name__
+        return [SerializedBlob(blob)], meta
+
+    @classmethod
+    def _storage_deserialize(cls, blobs, **kwargs):
+        data = json.loads(blobs[0].decode("utf-8"))
+        metadata = kwargs.get("metadata", {})
+        dc_module = metadata.get("dataclass_module")
+        dc_class = metadata.get("dataclass_class")
+        if dc_module and dc_class:
+            mod = importlib.import_module(dc_module)
+            dc_type = getattr(mod, dc_class)
+            # Guard against crafted metadata — require a class that's
+            # actually a dataclass. ``dataclasses.is_dataclass`` alone
+            # returns True for dataclass *instances*; the ``isinstance(..., type)``
+            # check excludes that (and anything else callable).
+            if not (isinstance(dc_type, type) and dataclasses.is_dataclass(dc_type)):
+                raise ValueError(
+                    "Struct metadata references '%s.%s' which is not a dataclass"
+                    % (dc_module, dc_class)
+                )
+            return cls(_reconstruct(dc_type, data), dataclass_type=dc_type)
+        # Fallback: return as plain dict wrapped in Struct
+        return cls(data)

metaflow/plugins/__init__.py

@@ -190,6 +190,7 @@
 # Add artifact serializers here. Ordering is by PRIORITY (lower = tried first).
 # PickleSerializer is the universal fallback (PRIORITY=9999).
 ARTIFACT_SERIALIZERS_DESC = [
+    ("iotype", ".datastores.serializers.iotype_serializer.IOTypeSerializer"),
     ("pickle", ".datastores.serializers.pickle_serializer.PickleSerializer"),
 ]
 

metaflow/plugins/datastores/serializers/iotype_serializer.py

@@ -0,0 +1,86 @@
+import importlib
+
+from metaflow.datastore.artifacts.serializer import (
+    ArtifactSerializer,
+    SerializationMetadata,
+    STORAGE,
+    WIRE,
+)
+from metaflow.io_types.base import IOType
+
+
+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``.
+
+    ``PRIORITY`` is 50 — ahead of the default (100) so this bridge catches
+    :class:`IOType` artifacts before any generic catch-all, and always ahead
+    of the :class:`PickleSerializer` fallback (9999).
+
+    Only the ``STORAGE`` format is implemented on this bridge; ``WIRE`` is
+    handled by callers that talk to :class:`IOType` directly (CLI parsing,
+    protobuf payload construction), not through the datastore.
+    """
+
+    TYPE = "iotype"
+    PRIORITY = 50
+
+    _ENCODING_PREFIX = "iotype:"
+
+    @classmethod
+    def can_serialize(cls, obj):
+        return isinstance(obj, IOType)
+
+    @classmethod
+    def can_deserialize(cls, metadata):
+        return metadata.encoding.startswith(cls._ENCODING_PREFIX)
+
+    @classmethod
+    def serialize(cls, obj, format=STORAGE):
+        if format == WIRE:
+            raise NotImplementedError(
+                "IOTypeSerializer only handles the STORAGE format; wire "
+                "encoding is produced by calling IOType.serialize(format=WIRE) "
+                "directly."
+            )
+        blobs, meta_dict = obj.serialize(format=STORAGE)
+        size = sum(len(b.value) for b in blobs if isinstance(b.value, bytes))
+        # Subclass metadata goes first so the routing keys below always win.
+        # An IOType subclass whose ``_storage_serialize`` happens to return
+        # ``iotype_module`` or ``iotype_class`` in its own meta dict must not
+        # be able to overwrite the routing info the deserialize path needs.
+        serializer_info = {
+            **meta_dict,
+            "iotype_module": obj.__class__.__module__,
+            "iotype_class": obj.__class__.__name__,
+        }
+        return (
+            blobs,
+            SerializationMetadata(
+                obj_type=obj.type_name,
+                size=size,
+                encoding=cls._ENCODING_PREFIX + obj.type_name,
+                serializer_info=serializer_info,
+            ),
+        )
+
+    @classmethod
+    def deserialize(cls, data, metadata=None, context=None, format=STORAGE):
+        if format == WIRE:
+            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"])
+        # 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"])
+            )
+        return iotype_cls.deserialize(data, format=STORAGE, metadata=info)