Skip to content

Nested pydantic models lose their class identity in another process #1910

Description

@dmpetrov

Description

A dataset row containing a plain pydantic.BaseModel reads back, in a process that did not write it, as a different class rebuilt from the stored schema. isinstance fails and passing the value into anything annotated with that type raises a ValidationError that contradicts itself.

# models.py
from pydantic import BaseModel

class Thresholds(BaseModel):
    name: str
    limit: float

class Scenario(BaseModel):
    key: int
    thresholds: Thresholds
# write.py  (process 1)
import datachain as dc
from models import Scenario, Thresholds

dc.read_values(key=[1, 2]).map(
    s=lambda key: Scenario(key=key, thresholds=Thresholds(name="v1", limit=0.5)),
    output={"s": Scenario},
).save("nested_demo")
# read.py  (process 2)
import datachain as dc
from models import Scenario, Thresholds

row = dc.read_dataset("nested_demo").to_list("s")[0][0]
print(type(row.thresholds))                    # <class 'datachain.lib.signal_schema.Thresholds'>
print(isinstance(row.thresholds, Thresholds))  # False
Scenario(key=99, thresholds=row.thresholds)    # ValidationError
ValidationError: 1 validation error for Scenario
thresholds
  Input should be a valid dictionary or instance of Thresholds
  [type=model_type, input_value=Thresholds(name='v1', limit=0.5), input_type=Thresholds]

Write and read in one process and it prints <class 'models.Thresholds'> / True. Subclassing dc.DataModel also works. Plain BaseModel in a fresh process always fails

Root cause

SignalSchema._deserialize_custom_type (src/datachain/lib/signal_schema.py:419-460) returns
ModelStore.get(name, version) if it hits, and otherwise rebuilds the class from the stored field
spec. DataModel registers itself in ModelStore at class definition
(data_model.py:__pydantic_init_subclass__); a plain BaseModel never registers, so a fresh process
always misses and rebuilds. In the writing process the lookup hits, which is why the bug is invisible
locally and only appears once producing and consuming steps are separate jobs, the normal case in
production.

The rebuilt class is then registered under the bare name@version, which is also the mechanism behind
#1872 (same-named models shadowing each other across datasets).

Suggested fix

Synthesizing a class should be the last resort, not the first. In _deserialize_custom_type, before
rebuilding:

  1. ModelStore.get(name, version) as today, but keyed additionally by a structural fingerprint so
    different shapes cannot collide (this also closes Custom model name collisions: same-named generated models shadow each other across datasets #1872);
  2. resolve the already-loaded class: sys.modules.get(module) -> getattr(module, name), accept
    if the field spec matches. The module is already persisted: _get_bases stores the full MRO with
    __module__, so bases[0] is ["Thresholds", "models", null] in existing datasets. No schema
    change, no migration, works on data written by current versions;
  3. otherwise scan the live BaseModel.__subclasses__() closure by name + field spec (covers
    notebooks, __main__, re-exports);
  4. only then rebuild, and warn once when a same-named class with a different shape is loaded in this
    process, instead of silently returning a class that will not match.

Important: resolve only from classes already imported in this process. Do not import a module path
taken from dataset metadata: datasets are shared, and that would be code execution from data. No user
action is required either way, since using the type at all means the user has already imported it.

Two follow-ups in the same area: user validators will start running on read once the real class is
resolved (correct, but a behaviour change worth a release note), and there is currently no
cross-process test for this. A two-process fixture (write in one, read in another) over
BaseModel/DataModel x nested/list/optional would have caught it.

Version Info


Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions