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, self-describing schemas via ``to_spec()``, no pickle
  code-execution risk).
- ``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 ``to_spec()`` — not enough on their own for core.

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``). ``type_name`` + ``to_spec()`` support JSON
schema generation. 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.

Tests

- ``test/unit/io_types/test_base.py`` — abstract instantiation,
  WIRE/STORAGE dispatch, invalid format, equality/hash, spec.
- ``test/unit/io_types/test_json_type.py`` — round-trips.
- ``test/unit/io_types/test_struct_type.py`` — dataclass round-trip,
  dict round-trip, ``to_spec()`` with dataclass fields.
- ``test/unit/io_types/test_iotype_serializer.py`` — bridge
  ``can_serialize``/``can_deserialize``, round-trip through
  dataclass reconstruction, rejection of non-IOType classes in
  metadata (security), 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,119 @@
+import dataclasses
+import json
+import typing
+
+from ..datastore.artifacts.serializer import SerializedBlob
+from .base import IOType
+
+
+def _iotype_for_annotation(annotation):
+    """
+    Resolve a Python type annotation to an IOType class. Only IOType
+    subclasses are recognized; plain Python types return ``None`` and the
+    caller falls back to a string representation in :meth:`Struct.to_spec`.
+    """
+    if isinstance(annotation, type) and issubclass(annotation, IOType):
+        return annotation
+    return None
+
+
+class Struct(IOType):
+    """
+    Structured type mapping to Python @dataclass.
+
+    Wire: JSON string. Storage: JSON UTF-8 bytes.
+
+    Wraps a @dataclass instance. Fields are inferred from dataclass annotations
+    with implicit scalar mapping (str->Text, int->Int64, float->Float64, bool->Bool).
+
+    Parameters
+    ----------
+    value : dataclass instance or dict, optional
+        The wrapped value. Dataclass instances are serialized via dataclasses.asdict.
+        Plain dicts are serialized directly as JSON.
+    dataclass_type : type, optional
+        The @dataclass class for type descriptor use (no value).
+    """
+
+    type_name = "struct"
+
+    def __init__(self, value=None, dataclass_type=None):
+        if value is not None 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:
+            import importlib
+
+            mod = importlib.import_module(dc_module)
+            dc_type = getattr(mod, dc_class)
+            # Security: only allow actual dataclasses, not arbitrary classes
+            if not dataclasses.is_dataclass(dc_type):
+                raise ValueError(
+                    "Struct metadata references '%s.%s' which is not a dataclass"
+                    % (dc_module, dc_class)
+                )
+            return cls(dc_type(**data), dataclass_type=dc_type)
+        # Fallback: return as plain dict wrapped in Struct
+        return cls(data)
+
+    def to_spec(self):
+        spec = {"type": self.type_name}
+        if self._dataclass_type is not None and dataclasses.is_dataclass(
+            self._dataclass_type
+        ):
+            # Use typing.get_type_hints() to resolve string annotations
+            # (handles `from __future__ import annotations`)
+            try:
+                hints = typing.get_type_hints(self._dataclass_type)
+            except Exception:
+                hints = {}
+            fields = []
+            for f in dataclasses.fields(self._dataclass_type):
+                annotation = hints.get(f.name, f.type)
+                field_iotype = _iotype_for_annotation(annotation)
+                if field_iotype is not None:
+                    field_spec = field_iotype().to_spec()
+                else:
+                    field_spec = {"type": str(annotation)}
+                fields.append({"name": f.name, **field_spec})
+            spec["fields"] = fields
+        return spec

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,81 @@
+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))
+        return (
+            blobs,
+            SerializationMetadata(
+                obj_type=obj.type_name,
+                size=size,
+                encoding=cls._ENCODING_PREFIX + obj.type_name,
+                serializer_info={
+                    "iotype_module": obj.__class__.__module__,
+                    "iotype_class": obj.__class__.__name__,
+                    **meta_dict,
+                },
+            ),
+        )
+
+    @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)