diff --git a/docs/guide/checkpoints.md b/docs/guide/checkpoints.md index b60c7cfef..6d153a165 100644 --- a/docs/guide/checkpoints.md +++ b/docs/guide/checkpoints.md @@ -202,3 +202,4 @@ dc.read_storage("gs://datachain-demo/dogs-and-cats/", anon=True).map( - **Script path matters:** DataChain links runs by the script's absolute path. Moving the script breaks checkpoint linking. - **Threading/multiprocessing:** Checkpoints are automatically disabled when Python threading or multiprocessing is detected. DataChain's built-in `parallel` setting for Python operations is not affected. - **Unhashable callables:** Built-in functions (`len`, `str`), C extensions, and `Mock` objects produce a different hash on each run, so checkpoints using these as Python operations will always recompute. Use regular `def` functions or lambdas instead. +- **Class UDF state:** Constructor callables and custom objects recompute by default. Override [`state_hash()`](python-engine.md#caching-class-based-operations) only when you can provide a complete stable identity. diff --git a/docs/guide/python-engine.md b/docs/guide/python-engine.md index 561131ae6..738bd59df 100644 --- a/docs/guide/python-engine.md +++ b/docs/guide/python-engine.md @@ -151,6 +151,40 @@ class ImageEncoder(Mapper): del self.model ``` +### Caching class-based operations + +DataChain hashes UDF code, schemas, and constructor arguments. Primitive values and +nested built-in containers are handled automatically. Callables and custom objects +receive a unique identity instead, preventing incorrect cache reuse. Override +`state_hash()` when you can identify such state safely: + +```python +import hashlib +import json +from datachain.lib.udf import Mapper + +class Tokenize(Mapper): + def __init__(self, tokenizer, tokenizer_id: str, max_length: int): + self.tokenizer = tokenizer + self.tokenizer_id = tokenizer_id + self.max_length = max_length + + def state_hash(self) -> str: + state = json.dumps( + {"tokenizer_id": self.tokenizer_id, "max_length": self.max_length}, + sort_keys=True, + ) + return hashlib.sha256(state.encode()).hexdigest() + + def process(self, text: str) -> list[str]: + return self.tokenizer(text)[: self.max_length] +``` + +`state_hash()` must return a SHA-256 hexadecimal string covering all instance state +that affects output. When all constructor arguments are supported primitives, call +`super().state_hash()` to include them in an override. UDF code and schemas are always +included; an incomplete hash can reuse an incorrect cached result. + Use class-based operations sparingly; `.setup()` covers most cases. ## Execution and Scale diff --git a/docs/references/llm.md b/docs/references/llm.md index a52a1087e..ca9537206 100644 --- a/docs/references/llm.md +++ b/docs/references/llm.md @@ -139,7 +139,8 @@ Reliability is layered: Materialized `llm.*` columns are cached and versioned, so re-running a chain reads the stored result instead of re-calling the model; the cache invalidates when any output-affecting input changes (model, prompt, schema, the input column, `type`, -params, ...). +params, ...). Custom objects in parameters receive a per-operation identity, safely +disabling reuse across separately created operations. ## No fused predicate diff --git a/src/datachain/hash_utils.py b/src/datachain/hash_utils.py index af90fcd94..154aea4ce 100644 --- a/src/datachain/hash_utils.py +++ b/src/datachain/hash_utils.py @@ -1,6 +1,7 @@ import hashlib import inspect import logging +import struct import textwrap from collections.abc import Sequence from typing import TypeAlias, TypeVar @@ -16,6 +17,46 @@ ColumnLike: TypeAlias = str | T +def normalize_hash_value(value, *, sort_dicts: bool = False): # noqa: PLR0911 + """Return a complete, deterministic representation for stable cache keys.""" + value_type = type(value) + + def normalize(item): + return normalize_hash_value(item, sort_dicts=sort_dicts) + + if value is None: + return ("none",) + if value_type is float: + return ("float", struct.pack("!d", value)) + if value_type in (bool, int, str, bytes): + return (value_type.__name__, value) + if value_type is dict: + items = tuple((normalize(key), normalize(item)) for key, item in value.items()) + return ("dict", tuple(sorted(items, key=repr)) if sort_dicts else items) + if value_type is set: + return ( + "set", + tuple(sorted(map(normalize, value), key=repr)), + ) + if value_type is frozenset: + return ( + "frozenset", + tuple(sorted(map(normalize, value), key=repr)), + ) + if value_type is list: + return ("list", tuple(map(normalize, value))) + if value_type is tuple: + return ("tuple", tuple(map(normalize, value))) + + raise TypeError(f"value of type {value_type.__name__!r} cannot be hashed safely") + + +def hash_value(value) -> str: + """Hash a value after normalizing unordered containers.""" + normalized = normalize_hash_value(value) + return hashlib.sha256(repr(normalized).encode("utf-8")).hexdigest() + + def _serialize_value(val): # noqa: PLR0911 """Helper to serialize arbitrary values recursively.""" if val is None: diff --git a/src/datachain/lib/arrow.py b/src/datachain/lib/arrow.py index 6cbb18f1a..0d9dac38f 100644 --- a/src/datachain/lib/arrow.py +++ b/src/datachain/lib/arrow.py @@ -1,4 +1,5 @@ import math +import pickle from collections.abc import Sequence from itertools import islice from typing import TYPE_CHECKING, Any @@ -33,6 +34,27 @@ DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY = b"DataChain SignalSchema" +def _parse_options_hash_args(options: ParseOptions) -> dict[str, Any]: + return { + "delimiter": options.delimiter, + "double_quote": options.double_quote, + "escape_char": options.escape_char, + "ignore_empty_lines": options.ignore_empty_lines, + "invalid_row_handler": options.invalid_row_handler, + "newlines_in_values": options.newlines_in_values, + "quote_char": options.quote_char, + } + + +def _csv_format_hash_args(format: CsvFileFormat) -> dict[str, Any]: + scan_options = format.default_fragment_scan_options + return { + "parse_options": _parse_options_hash_args(format.parse_options), + "read_options": pickle.dumps(scan_options.read_options), + "convert_options": pickle.dumps(scan_options.convert_options), + } + + def fix_pyarrow_format(format, parse_options=None): # Re-init invalid row handler: https://issues.apache.org/jira/browse/ARROW-17641 if ( @@ -54,6 +76,7 @@ def __init__( output_schema: type["BaseModel"] | None = None, source: bool = True, nrows: int | None = None, + _generated_output_schema: bool = False, **kwargs, ): """ @@ -75,6 +98,25 @@ def __init__( self.parse_options = kwargs.pop("parse_options", None) self.kwargs = kwargs + @classmethod + def _constructor_hash_args(cls, arguments): + arguments = arguments.copy() + generated_output_schema = arguments.pop("_generated_output_schema") + if generated_output_schema: + # Its stable field shape is already part of the output signal schema. + arguments.pop("output_schema", None) + input_schema = arguments.get("input_schema") + if input_schema is not None: + arguments["input_schema"] = input_schema.serialize().to_pybytes() + kwargs = arguments["kwargs"].copy() + for name, value in kwargs.items(): + if isinstance(value, ParseOptions): + kwargs[name] = _parse_options_hash_args(value) + elif isinstance(value, CsvFileFormat): + kwargs[name] = _csv_format_hash_args(value) + arguments["kwargs"] = kwargs + return arguments + def process(self, file: File): if file._caching_enabled: file.ensure_cached() diff --git a/src/datachain/lib/dc/datachain.py b/src/datachain/lib/dc/datachain.py index 9a259924a..0aa08cddf 100644 --- a/src/datachain/lib/dc/datachain.py +++ b/src/datachain/lib/dc/datachain.py @@ -2537,7 +2537,9 @@ def parse_tabular( except ValueError as e: raise DatasetPrepareError(self.name, e) from e + generated_output_schema = False if isinstance(output, dict): + generated_output_schema = True model_name = model_name or column or "" model = dict_to_data_model(model_name, output) output = model @@ -2563,6 +2565,7 @@ def parse_tabular( model, source, nrows, + _generated_output_schema=generated_output_schema, parse_options=parse_options, **kwargs, ), diff --git a/src/datachain/lib/hf.py b/src/datachain/lib/hf.py index 3f6db377a..8f1f55f6d 100644 --- a/src/datachain/lib/hf.py +++ b/src/datachain/lib/hf.py @@ -46,6 +46,19 @@ ) +def _dataset_hash_args(ds: HFDatasetType) -> Any: + if isinstance(ds, (DatasetDict, IterableDatasetDict)): + return ( + type(ds).__name__, + {name: _dataset_hash_args(split) for name, split in ds.items()}, + ) + if isinstance(ds, (Dataset, IterableDataset)): + fingerprint = getattr(ds, "_fingerprint", None) + if fingerprint is not None: + return (type(ds).__name__, fingerprint) + return ds + + class HFClassLabel(DataModel): string: str integer: int @@ -93,6 +106,17 @@ def __init__( self.args = args self.kwargs = kwargs + @classmethod + def _constructor_hash_args(cls, arguments): + # output_schema is a dynamically-created pydantic class with a random + # name suffix; its stable field shape lands in the UDF hash via the + # output signal schema, so drop it here to keep the constructor hash + # deterministic across runs. + arguments = arguments.copy() + arguments.pop("output_schema", None) + arguments["ds"] = _dataset_hash_args(arguments["ds"]) + return arguments + def setup(self): self.ds_dict = stream_splits(self.ds, *self.args, **self.kwargs) diff --git a/src/datachain/lib/udf.py b/src/datachain/lib/udf.py index 926debc8f..7eb0ecfe4 100644 --- a/src/datachain/lib/udf.py +++ b/src/datachain/lib/udf.py @@ -8,6 +8,7 @@ from functools import partial from graphlib import CycleError, TopologicalSorter from typing import TYPE_CHECKING, Any, TypeVar +from uuid import uuid4 import attrs from fsspec.callbacks import DEFAULT_CALLBACK, Callback @@ -16,7 +17,7 @@ from datachain.asyn import AsyncMapper from datachain.cache import temporary_cache from datachain.dataset import RowDict -from datachain.hash_utils import hash_callable +from datachain.hash_utils import hash_callable, hash_value from datachain.lib.convert.flatten import ( classify_field, flatten, @@ -223,6 +224,20 @@ def prefetch(self) -> int: return self.inner.prefetch +def _hash_constructor_args( + arguments: dict[str, Any], *, warn_on_unsupported: bool = True +) -> str: + try: + return hash_value(arguments) + except (TypeError, RecursionError) as exc: + if warn_on_unsupported: + logger.warning( + "%s; cache reuse across UDF instances is disabled", + exc, + ) + return hashlib.sha256(uuid4().bytes).hexdigest() + + class UDFBase(AbstractUDF): """Base class for stateful user-defined functions. @@ -274,11 +289,49 @@ def process(self, file) -> list[float]: is_input_batched = False is_output_batched = False prefetch: int = 0 + # Set in __new__ from bound constructor args (see hash()); default keeps + # mypy happy and acts as a safe fallback for instances that skip __new__. + _constructor_state_hash: str = "" # Class-level default so subclasses that skip super().__init__() still # have this attribute; _MultiSignalMapper sets params/output per entry # but doesn't call _init(), so _func would otherwise be missing. _func: Callable | None = None + def __new__(cls, *args, **kwargs): + instance = super().__new__(cls) + try: + # bound-method signature so `self` is already stripped; correct even + # when __init__ uses *args instead of a named `self` parameter. + bound = inspect.signature(instance.__init__).bind(*args, **kwargs) + except TypeError: + # Pickle allocates an empty instance and restores its attributes + # afterwards, without calling __init__ or supplying its arguments. + return instance + + bound.apply_defaults() + arguments = cls._constructor_hash_args(dict(bound.arguments)) + instance._constructor_state_hash = _hash_constructor_args( + arguments, + warn_on_unsupported=cls.state_hash is UDFBase.state_hash, + ) + return instance + + @classmethod + def _constructor_hash_args(cls, arguments: dict[str, Any]) -> dict[str, Any]: + """Constructor arguments that determine this UDF instance's identity.""" + return arguments + + def state_hash(self) -> str: + """Return a stable SHA-256 hash for state that affects this UDF's output. + + Override this when constructor arguments contain callables or other opaque + objects that DataChain cannot hash safely. The method is called after + ``__init__`` and must account for all per-instance behavioral state. By + default, DataChain returns its automatic constructor-argument hash. Call + ``super().state_hash()`` to include that hash in an override. + """ + return self._constructor_state_hash + def __init__(self): self.params: SignalSchema | None = None self.output = None @@ -304,6 +357,25 @@ def hash(self, include_body: bool = True) -> str: self.params.hash() if self.params else "", self.output.hash(), ] + # For class-based UDFs, mix in constructor state so two instances that + # differ only in constructor args don't collide. + if self._func is None: + state_hash = self.state_hash() + if not isinstance(state_hash, str) or len(state_hash) != 64: + raise ValueError( + "state_hash() must return a SHA-256 hexadecimal string" + ) + try: + state_hash_bytes = bytes.fromhex(state_hash) + except ValueError as exc: + raise ValueError( + "state_hash() must return a SHA-256 hexadecimal string" + ) from exc + if len(state_hash_bytes) != 32: + raise ValueError( + "state_hash() must return a SHA-256 hexadecimal string" + ) + parts.append(state_hash) return hashlib.sha256( b"".join([bytes.fromhex(part) for part in parts]) @@ -596,6 +668,12 @@ class _MultiSignalMapper(Mapper): irrelevant. Cycles raise ``ValueError`` at construction time. """ + @classmethod + def _constructor_hash_args(cls, arguments): + # hash() is fully overridden below; skip the base state hash so we + # don't warn on the callables in signal_map or store unused bytes. + return {} + def __init__( self, signal_map: "dict[str, Callable | UDFBase]", diff --git a/src/datachain/llm/spec.py b/src/datachain/llm/spec.py index 76ebd6e53..564887420 100644 --- a/src/datachain/llm/spec.py +++ b/src/datachain/llm/spec.py @@ -1,11 +1,12 @@ -import re -import warnings +import logging from collections.abc import Callable, Iterator from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, get_args, get_origin +from uuid import uuid4 from pydantic import BaseModel +from datachain.hash_utils import normalize_hash_value from datachain.lib.udf import BindContext, BoundSpec from datachain.llm import engine from datachain.llm.content import MEDIA_VALUES, Media, build_messages, to_text @@ -14,6 +15,8 @@ if TYPE_CHECKING: from datachain.lib.settings import Settings +logger = logging.getLogger("datachain") + class LLMConfigError(engine.LLMError): """Raised when no model can be resolved for a `datachain.llm` operation.""" @@ -27,30 +30,6 @@ def _element_type(schema: Any) -> tuple[Any, bool]: return schema, False -# A value with no custom repr shows its memory address (````), -# which changes each run and would make the cache key unstable. -_DEFAULT_REPR = re.compile(r" object at 0x[0-9a-fA-F]+>") - - -def _canonical(value: Any) -> Any: - """Order-independent form of a value, so the cache key is stable across - processes (``repr`` of a ``set`` or unsorted ``dict`` is not).""" - if isinstance(value, dict): - items = ((k, _canonical(v)) for k, v in value.items()) - return tuple(sorted(items, key=lambda kv: repr(kv[0]))) - if isinstance(value, (set, frozenset)): - return tuple(sorted((_canonical(v) for v in value), key=repr)) - if isinstance(value, (list, tuple)): - return tuple(_canonical(v) for v in value) - if _DEFAULT_REPR.search(repr(value)): - warnings.warn( - f"llm param {type(value).__name__!r} has no stable repr; it breaks " - "caching (full recompute every run).", - stacklevel=2, - ) - return value - - def _is_secret_key(key: Any) -> bool: if not isinstance(key, str): return False @@ -71,6 +50,18 @@ def _without_secrets(value: Any) -> Any: return value +def _normalize_identity_value( + value: Any, fallback: str, *, remove_secrets: bool = False +) -> Any: + try: + if remove_secrets: + value = _without_secrets(value) + return normalize_hash_value(value, sort_dicts=True) + except (TypeError, RecursionError) as exc: + logger.warning("%s; cache reuse for this LLM operation is disabled", exc) + return ("unsupported", fallback) + + @dataclass class LLMSpec(BoundSpec): """A configured `datachain.llm` operation, used inside `.map()` / `.gen()`. @@ -91,6 +82,12 @@ class LLMSpec(BoundSpec): fallback: str | list[str] | None = None include_usage: bool = False params: dict[str, Any] = field(default_factory=dict) + _identity_fallback: str = field( + default_factory=lambda: uuid4().hex, + init=False, + repr=False, + compare=False, + ) def __post_init__(self) -> None: if self.schema is not None: @@ -166,7 +163,7 @@ def return_annotation(self, to_many: bool = False) -> Any: return tuple[out, Usage] if self.include_usage else out # type: ignore[valid-type] def identity(self, model: str, llm_params: Any = None) -> tuple: - """Cache key baked into the UDF hash; changes iff an output-affecting + """Cache key baked into the UDF hash; changes if an output-affecting input (model, prompt, schema, params, llm_params, ...) changes. The schema is keyed by its JSON schema (fields, types, constraints, name); @@ -183,13 +180,17 @@ def identity(self, model: str, llm_params: Any = None) -> tuple: if self.schema is not None: elem, is_list = _element_type(self.schema) if hasattr(elem, "model_json_schema"): - schema_repr = (_canonical(elem.model_json_schema()), is_list) + schema_repr = ( + _normalize_identity_value( + elem.model_json_schema(), self._identity_fallback + ), + is_list, + ) else: schema_repr = str(self.schema) params = self.params if isinstance(llm_params, dict): params = {**llm_params, **self.params} - params = _without_secrets(params) return ( self.kind, model, @@ -200,7 +201,9 @@ def identity(self, model: str, llm_params: Any = None) -> tuple: self.context_col, self.type, self.include_usage, - _canonical(params), + _normalize_identity_value( + params, self._identity_fallback, remove_secrets=True + ), ) def _resolve_model(self, settings: "Settings") -> str: diff --git a/tests/func/test_udf.py b/tests/func/test_udf.py index 5247d0df8..ad4e2f87c 100644 --- a/tests/func/test_udf.py +++ b/tests/func/test_udf.py @@ -8,6 +8,7 @@ import multiprocess as mp import pytest +from pydantic import BaseModel import datachain as dc from datachain.client.fileslice import FileWrapper @@ -256,6 +257,45 @@ def process(self, size): ] +def test_class_agg_instance_state_produces_distinct_results(test_session): + class Out(BaseModel): + key: int + n: int + + class CountAbove(dc.Aggregator): + def __init__(self, limit: int): + super().__init__() + self.limit = limit + + def process(self, key, value): + yield Out(key=key[0], n=sum(1 for v in value if v > self.limit)) + + src = dc.read_values( + key=[1] * 5 + [2] * 5, + value=[1, 2, 3, 4, 5, 1, 2, 3, 4, 5], + session=test_session, + ) + chain_zero = src.agg( + CountAbove(0), + partition_by="key", + params=["key", "value"], + output={"o": Out}, + ) + chain_three = src.agg( + CountAbove(3), + partition_by="key", + params=["key", "value"], + output={"o": Out}, + ) + + assert chain_zero._query.hash() != chain_three._query.hash() + + rows_zero = sorted(chain_zero.to_list("o.key", "o.n")) + rows_three = sorted(chain_three.to_list("o.key", "o.n")) + assert rows_zero == [(1, 5), (2, 5)] + assert rows_three == [(1, 2), (2, 2)] + + @pytest.mark.parametrize( "cloud_type,version_aware", [("s3", True)], diff --git a/tests/unit/lib/test_arrow.py b/tests/unit/lib/test_arrow.py index 6977db372..2fb57b162 100644 --- a/tests/unit/lib/test_arrow.py +++ b/tests/unit/lib/test_arrow.py @@ -6,6 +6,8 @@ import pyarrow.parquet as pq import pytest from datasets import Dataset +from pyarrow.csv import ParseOptions +from pyarrow.dataset import CsvFileFormat import datachain as dc from datachain.lib.arrow import ( @@ -19,6 +21,48 @@ from datachain.lib.hf import HFClassLabel +def test_arrow_generator_constructor_hash(): + first_schema = dict_to_data_model("", {"value": int}) + second_schema = dict_to_data_model("", {"value": int}) + + def make_generator(output_schema, nrows=None): + input_schema = pa.schema({"value": pa.int64()}) + parse_options = ParseOptions(delimiter=";") + return ArrowGenerator( + input_schema=input_schema, + output_schema=output_schema, + nrows=nrows, + _generated_output_schema=True, + parse_options=parse_options, + format=CsvFileFormat(parse_options=parse_options), + ) + + first = make_generator(first_schema) + second = make_generator(second_schema) + limited = make_generator(second_schema, nrows=1) + + assert first._constructor_state_hash == second._constructor_state_hash + assert first._constructor_state_hash != limited._constructor_state_hash + + +def test_arrow_generator_constructor_hash_with_closure_handler(): + def make_handler(prefix): + def handler(row): + return "skip" if row.text.startswith(prefix) else "error" + + return handler + + def make_generator(prefix): + return ArrowGenerator( + parse_options=ParseOptions(invalid_row_handler=make_handler(prefix)) + ) + + comments = make_generator("#") + metadata = make_generator("!") + + assert comments._constructor_state_hash != metadata._constructor_state_hash + + @pytest.mark.parametrize("cache", [True, False]) def test_arrow_generator(tmp_path, catalog, cache): ids = [12345, 67890, 34, 0xF0123] diff --git a/tests/unit/lib/test_hf.py b/tests/unit/lib/test_hf.py index 51181581e..8516d01a7 100644 --- a/tests/unit/lib/test_hf.py +++ b/tests/unit/lib/test_hf.py @@ -1,3 +1,4 @@ +import pytest from datasets import Array2D, Dataset, DatasetDict, Sequence, Value from datachain.lib.data_model import dict_to_data_model @@ -9,6 +10,24 @@ ) +@pytest.mark.parametrize("as_dict", [False, True]) +def test_hf_generator_constructor_hash(as_dict): + first_ds = Dataset.from_dict({"value": [1]}) + second_ds = Dataset.from_dict({"value": [1]}) + if as_dict: + first_ds = DatasetDict({"train": first_ds}) + second_ds = DatasetDict({"train": second_ds}) + first_schema = dict_to_data_model("", {"value": int}) + second_schema = dict_to_data_model("", {"value": int}) + + first = HFGenerator(first_ds, first_schema) + second = HFGenerator(second_ds, second_schema) + limited = HFGenerator(second_ds, second_schema, limit=1) + + assert first._constructor_state_hash == second._constructor_state_hash + assert first._constructor_state_hash != limited._constructor_state_hash + + def test_hf(): ds = Dataset.from_dict({"pokemon": ["bulbasaur", "squirtle"]}) schema, norm_names = get_output_schema(ds.features) diff --git a/tests/unit/lib/test_llm.py b/tests/unit/lib/test_llm.py index 16f63c029..7586e6e21 100644 --- a/tests/unit/lib/test_llm.py +++ b/tests/unit/lib/test_llm.py @@ -524,13 +524,6 @@ def test_identity_stable_across_param_dict_order(): assert a == b -def test_canonical_orders_sets_and_dicts(): - from datachain.llm.spec import _canonical - - assert _canonical({"a", "b", "c"}) == _canonical({"c", "b", "a"}) - assert _canonical({"x": 1, "y": 2}) == _canonical({"y": 2, "x": 1}) - - def test_param_clobber_is_blocked(fake_llm): bind(llm.complete("t", "real prompt", temperature=0.0), llm="real/model")("hi") call = fake_llm.calls[-1] @@ -732,20 +725,41 @@ def test_callable_llm_params_not_in_identity(): assert spec.identity("m", lambda: {"k": "v"}) == spec.identity("m") -def test_opaque_param_value_warns_about_unstable_cache_key(): +def test_opaque_param_value_disables_cache_reuse(): + class Opaque: + def __repr__(self): + return "Opaque()" + + first = llm.complete("t", client=Opaque()).identity("m") + second = llm.complete("t", client=Opaque()).identity("m") + + assert first != second + + +def test_opaque_param_value_has_stable_identity_within_operation(): class Opaque: pass - with pytest.warns(UserWarning, match="no stable repr"): - llm.complete("t", client=Opaque()).identity("m") + spec = llm.complete("t", client=Opaque()) + + assert spec.identity("m") == spec.identity("m") + + +def test_cyclic_param_value_disables_cache_reuse(): + options = {} + options["self"] = options + first = llm.complete("t", options=options) + second = llm.complete("t", options=options) + + assert first.identity("m") == first.identity("m") + assert first.identity("m") != second.identity("m") -def test_stable_param_values_do_not_warn(): - import warnings +def test_stable_param_values_have_stable_identity(): + first = llm.complete("t", temperature=0.0, opt={"a": 1}).identity("m") + second = llm.complete("t", temperature=0.0, opt={"a": 1}).identity("m") - with warnings.catch_warnings(): - warnings.simplefilter("error") - llm.complete("t", temperature=0.0, opt={"a": 1}).identity("m") + assert first == second def test_secret_params_not_in_identity(): diff --git a/tests/unit/lib/test_udf.py b/tests/unit/lib/test_udf.py index 3a1f09389..04526be9c 100644 --- a/tests/unit/lib/test_udf.py +++ b/tests/unit/lib/test_udf.py @@ -1,4 +1,7 @@ +import os import pickle +import subprocess +import sys import pytest from cloudpickle import dumps, loads @@ -8,6 +11,7 @@ import datachain as dc from datachain import Mapper from datachain.dataset import RowDict +from datachain.hash_utils import hash_value from datachain.lib.file import File from datachain.lib.signal_schema import SignalSchema from datachain.lib.udf import JsonSerializationError, UDFBase, UdfError, UdfRunError @@ -16,6 +20,31 @@ from .test_udf_signature import get_sign +class _OpaqueConstructorValue: + pass + + +class _HashableConstructorValue: + def __hash__(self): + return 1 + + +class _CallableConstructorValue: + def __call__(self, value): + return value + + +def _constructor_function(value): + return value + + +def _make_constructor_closure(captured): + def closure(value): + return value, captured + + return closure + + def test_udf_error(): orig_err = UdfError("test error") for err in (orig_err, loads(dumps(orig_err))): @@ -108,7 +137,7 @@ class MyMapper(Mapper): def process(self, key: str) -> int: return len(key) - sign = get_sign(MyMapper, output="res") + sign = get_sign(MyMapper, params=[], output="res") udf = UDFBase._create(sign, sign.output_schema) assert udf.verbose_name == "MyMapper" @@ -135,6 +164,326 @@ def test_udf_verbose_name_unknown(): assert udf.verbose_name == "" +def test_class_udf_hash_varies_with_instance_state(): + class Limited(Mapper): + def __init__(self, limit: int): + super().__init__() + self.limit = limit + + def process(self, x: int) -> int: + return x + self.limit + + sign_a = get_sign(Limited(0), output="y") + sign_b = get_sign(Limited(3), output="y") + udf_a = Mapper._create(sign_a, sign_a.output_schema) + udf_b = Mapper._create(sign_b, sign_b.output_schema) + assert udf_a.hash() != udf_b.hash() + + +def test_class_udf_hash_is_deterministic_across_instances(): + class Limited(Mapper): + def __init__(self, limit: int): + super().__init__() + self.limit = limit + + def process(self, x: int) -> int: + return x + self.limit + + sign_a = get_sign(Limited(3), output="y") + sign_b = get_sign(Limited(3), output="y") + udf_a = Mapper._create(sign_a, sign_a.output_schema) + udf_b = Mapper._create(sign_b, sign_b.output_schema) + assert udf_a.hash() == udf_b.hash() + + +@pytest.mark.parametrize( + "first_config,second_config", + [ + ({"a": 1, "b": 2}, {"b": 2, "a": 1}), + ( + {"options": {"a": 1, "b": 2}}, + {"options": {"b": 2, "a": 1}}, + ), + ], +) +def test_class_udf_hash_preserves_constructor_dict_order(first_config, second_config): + class Configured(Mapper): + def __init__(self, config): + self.config = config + + def process(self, x: int) -> str: + return ",".join(self.config) + + first = Configured(first_config) + second = Configured(second_config) + sign_a = get_sign(first, output="y") + sign_b = get_sign(second, output="y") + udf_a = Mapper._create(sign_a, sign_a.output_schema) + udf_b = Mapper._create(sign_b, sign_b.output_schema) + + assert udf_a.hash() != udf_b.hash() + + +def test_class_udf_unsupported_constructor_value_disables_cache_reuse(caplog): + class Opaque: + def __repr__(self): + return "Opaque()" + + class Configured(Mapper): + def __init__(self, config: Opaque): + self.config = config + + def process(self, x: int) -> int: + return x + + first = Configured(Opaque()) + second = Configured(Opaque()) + sign_a = get_sign(first, output="y") + sign_b = get_sign(second, output="y") + udf_a = Mapper._create(sign_a, sign_a.output_schema) + udf_b = Mapper._create(sign_b, sign_b.output_schema) + + assert udf_a.hash() == udf_a.hash() + assert udf_a.hash() != udf_b.hash() + assert "cache reuse across UDF instances is disabled" in caplog.text + + +@pytest.mark.parametrize( + "config", + [ + pytest.param(_OpaqueConstructorValue(), id="custom-object"), + pytest.param(_HashableConstructorValue(), id="object-with-hash"), + pytest.param(_constructor_function, id="function"), + pytest.param(lambda value: value, id="lambda"), + pytest.param(_make_constructor_closure("captured"), id="closure"), + pytest.param(_CallableConstructorValue(), id="callable-object"), + pytest.param( + {"options": [{"client": _OpaqueConstructorValue()}]}, + id="nested-custom-object", + ), + ], +) +def test_class_udf_unsupported_constructor_values_do_not_reuse_cache(config): + class Configured(Mapper): + def __init__(self, value): + self.value = value + + def process(self, x: int) -> int: + return x + + first = Configured(config) + second = Configured(config) + sign_a = get_sign(first, output="y") + sign_b = get_sign(second, output="y") + udf_a = Mapper._create(sign_a, sign_a.output_schema) + udf_b = Mapper._create(sign_b, sign_b.output_schema) + + assert udf_a.hash() == udf_a.hash() + assert udf_a.hash() != udf_b.hash() + + +def test_class_udf_cyclic_constructor_value_disables_cache_reuse(): + class Configured(Mapper): + def __init__(self, config): + self.config = config + + def process(self, x: int) -> int: + return x + + config = {} + config["self"] = config + first = Configured(config) + second = Configured(config) + sign_a = get_sign(first, output="y") + sign_b = get_sign(second, output="y") + udf_a = Mapper._create(sign_a, sign_a.output_schema) + udf_b = Mapper._create(sign_b, sign_b.output_schema) + + assert udf_a.hash() == udf_a.hash() + assert udf_a.hash() != udf_b.hash() + + +@pytest.mark.parametrize( + "first_key,second_key,matches", + [ + ("tokenizer-v1", "tokenizer-v1", True), + ("tokenizer-v1", "tokenizer-v2", False), + ], +) +def test_class_udf_state_hash_overrides_opaque_constructor_fallback( + first_key, second_key, matches, caplog +): + class Opaque: + pass + + class Configured(Mapper): + def __init__(self, config: Opaque, cache_key: str): + self.config = config + self.cache_key = cache_key + + def state_hash(self) -> str: + return hash_value(self.cache_key) + + def process(self, x: int) -> int: + return x + + first = Configured(Opaque(), first_key) + second = Configured(Opaque(), second_key) + sign_a = get_sign(first, output="y") + sign_b = get_sign(second, output="y") + udf_a = Mapper._create(sign_a, sign_a.output_schema) + udf_b = Mapper._create(sign_b, sign_b.output_schema) + + assert (udf_a.hash() == udf_b.hash()) is matches + assert "cache reuse across UDF instances is disabled" not in caplog.text + + +def test_class_udf_state_hash_can_extend_automatic_constructor_hash(): + class Configured(Mapper): + def __init__(self, limit: int, extra: str): + self.limit = limit + self.extra = extra + + def state_hash(self) -> str: + return hash_value((super().state_hash(), self.extra)) + + def process(self, x: int) -> int: + return x + self.limit + + first = Configured(3, "shared") + second = Configured(5, "shared") + sign_a = get_sign(first, output="y") + sign_b = get_sign(second, output="y") + udf_a = Mapper._create(sign_a, sign_a.output_schema) + udf_b = Mapper._create(sign_b, sign_b.output_schema) + + assert udf_a.hash() != udf_b.hash() + + +@pytest.mark.parametrize( + "invalid_hash", + ["tokenizer-v1", None, "ab" * 16 + " " * 32], +) +def test_class_udf_state_hash_rejects_invalid_hash(invalid_hash): + class Configured(Mapper): + def state_hash(self) -> str: + return invalid_hash + + def process(self, x: int) -> int: + return x + + udf = Configured() + sign = get_sign(udf, output="y") + udf = Mapper._create(sign, sign.output_schema) + + with pytest.raises( + ValueError, match=r"state_hash\(\) must return a SHA-256 hexadecimal string" + ): + udf.hash() + + +@pytest.mark.parametrize( + "args,kwargs,matches_default", + [ + ((), {}, True), + ((3,), {}, True), + ((), {"limit": 3}, True), + ((4,), {}, False), + ], +) +def test_class_udf_captures_normalized_constructor_arguments( + args, kwargs, matches_default +): + class Limited(Mapper): + def __init__(self, limit: int = 3): + self.limit = limit + + def process(self, x: int) -> int: + return x + self.limit + + baseline = Limited(limit=3) + udf = Limited(*args, **kwargs) + + assert ( + udf._constructor_state_hash == baseline._constructor_state_hash + ) is matches_default + + +def test_class_udf_constructor_hash_survives_cloudpickle_roundtrip(): + class Limited(Mapper): + def __init__(self, limit: int): + self.limit = limit + + def process(self, x: int) -> int: + return x + self.limit + + udf = Limited(3) + restored = loads(dumps(udf)) + + assert restored._constructor_state_hash == udf._constructor_state_hash + + +def test_class_udf_hash_is_deterministic_across_processes(): + code = """ +from datachain import Mapper +from datachain.lib.signal_schema import SignalSchema +from datachain.lib.udf_signature import UdfSignature + +class Limited(Mapper): + def __init__(self, limit=3, labels=frozenset({"a", "b", "c"})): + self.limit = limit + self.labels = labels + + def process(self, x): + return x + self.limit + +udf = Limited() +sign = UdfSignature( + udf, + SignalSchema({"x": int}), + SignalSchema({"y": int}), +) +print(Mapper._create(sign, sign.output_schema).hash()) +""" + + hashes = { + subprocess.check_output( # noqa: S603 + [sys.executable, "-c", code], + env={**os.environ, "PYTHONHASHSEED": seed}, + text=True, + ).strip() + for seed in ("1", "2", "random") + } + + assert len(hashes) == 1 + + +def test_class_udf_hash_without_body_ignores_process_implementation(): + def make_udf(add: bool): + class Stateful(Mapper): + def __init__(self, limit: int = 3): + self.limit = limit + + if add: + + def process(self, x: int) -> int: + return x + self.limit + + else: + + def process(self, x: int) -> int: + return x * self.limit + + sign = get_sign(Stateful(), output="y") # type: ignore[arg-type] + return Mapper._create(sign, sign.output_schema) + + added = make_udf(add=True) + multiplied = make_udf(add=False) + + assert added.hash(include_body=False) == multiplied.hash(include_body=False) + assert added.hash() != multiplied.hash() + + def test_udf_verbose_name_multi_signal_mapper(test_session): chain = dc.read_values(name=["foo.txt"], session=test_session).map( stem=lambda name: name.rsplit(".", 1)[0], diff --git a/tests/unit/test_datachain_hash.py b/tests/unit/test_datachain_hash.py index e9ac86bee..55a3691f0 100644 --- a/tests/unit/test_datachain_hash.py +++ b/tests/unit/test_datachain_hash.py @@ -1,6 +1,6 @@ import pandas as pd import pytest -from pydantic import BaseModel +from pydantic import BaseModel, field_validator import datachain as dc from datachain import func @@ -107,6 +107,31 @@ def test_read_csv_single_file_is_deterministic(test_session, tmp_dir): assert h1 == h2 +def test_read_csv_user_output_model_behavior_changes_hash(test_session, tmp_dir): + def make_model(uppercase): + class Row(BaseModel): + value: str + + @field_validator("value") + @classmethod + def transform(cls, value): + return value.upper() if uppercase else value.lower() + + return Row + + path = tmp_dir / "test.csv" + pd.DataFrame({"value": ["MiXeD"]}).to_csv(path, index=False) + + uppercase = dc.read_csv( + path.as_uri(), output=make_model(True), session=test_session + ) + lowercase = dc.read_csv( + path.as_uri(), output=make_model(False), session=test_session + ) + + assert uppercase._query.hash() != lowercase._query.hash() + + def test_read_csv_multi_file_glob_is_deterministic(test_session, tmp_dir): pd.DataFrame({"a": [1, 2]}).to_csv(tmp_dir / "a.csv", index=False) pd.DataFrame({"a": [3, 4]}).to_csv(tmp_dir / "b.csv", index=False) diff --git a/tests/unit/test_hash_utils.py b/tests/unit/test_hash_utils.py index 8bda8d3d1..45d5c6189 100644 --- a/tests/unit/test_hash_utils.py +++ b/tests/unit/test_hash_utils.py @@ -1,3 +1,5 @@ +import struct + import pytest from sqlalchemy import ( Float, @@ -16,7 +18,84 @@ from sqlalchemy import func as sa_func from datachain import C, func -from datachain.hash_utils import hash_callable, hash_column_elements +from datachain.hash_utils import ( + hash_callable, + hash_column_elements, + hash_value, + normalize_hash_value, +) + + +@pytest.mark.parametrize( + "value,expected", + [ + ( + {"y": 2, "x": 1}, + ( + "dict", + ( + (("str", "y"), ("int", 2)), + (("str", "x"), ("int", 1)), + ), + ), + ), + ( + {"outer": {"y": 2, "x": 1}}, + ( + "dict", + ( + ( + ("str", "outer"), + ( + "dict", + ( + (("str", "y"), ("int", 2)), + (("str", "x"), ("int", 1)), + ), + ), + ), + ), + ), + ), + ( + {"c", "b", "a"}, + ("set", (("str", "a"), ("str", "b"), ("str", "c"))), + ), + (frozenset({2, 1}), ("frozenset", (("int", 1), ("int", 2)))), + ([1, 2], ("list", (("int", 1), ("int", 2)))), + ((1, 2), ("tuple", (("int", 1), ("int", 2)))), + (None, ("none",)), + (True, ("bool", True)), + (1, ("int", 1)), + (1.5, ("float", struct.pack("!d", 1.5))), + ("value", ("str", "value")), + (b"value", ("bytes", b"value")), + ], +) +def test_normalize_hash_value(value, expected): + assert normalize_hash_value(value) == expected + + +def test_normalize_hash_value_rejects_unsupported_value(): + class Opaque: + pass + + with pytest.raises(TypeError, match="cannot be hashed safely"): + normalize_hash_value(Opaque()) + + +@pytest.mark.parametrize( + "first_bits,second_bits", + [ + ("7ff8000000000001", "7ff8000000000002"), + ("7ff8000000000001", "fff8000000000001"), + ], +) +def test_hash_value_distinguishes_float_bit_patterns(first_bits, second_bits): + first = struct.unpack("!d", bytes.fromhex(first_bits))[0] + second = struct.unpack("!d", bytes.fromhex(second_bits))[0] + + assert hash_value(first) != hash_value(second) def double(x): diff --git a/tests/unit/test_query_steps_hash.py b/tests/unit/test_query_steps_hash.py index b5f24d1a5..88242d860 100644 --- a/tests/unit/test_query_steps_hash.py +++ b/tests/unit/test_query_steps_hash.py @@ -322,7 +322,7 @@ def test_subtract_hash(test_session, numbers_dataset, on): DoubleMapper(), ["x"], {"double": int}, - "b58c9679ed454d3f54b4a754585727697a9aea9e4725bd12a842a774b5087963", + "7c901f584e52f41ac22a5fed332da2c3093d7d03f0fc16a780a9ad11eb456379", ), ], ) @@ -362,7 +362,7 @@ def test_udf_mapper_hash( TripleGenerator(), ["x"], {"triple": int}, - "01201327b1926788e6242d2be5383c63b97ec018232ab0844f047cf64ec2dfca", + "208bd4e553088f51983cadac3f645893bc44ffec7ccd625e28d7dfa0ea000170", ), ], )