Merge branch 'main' into nk/auto_deps_auto_callable_model

Signed-off-by: Nijat Khanbabayev <nijat.khanbabayev@cubistsystematic.com>

Author: Nijat Khanbabayev

ccflow/callable.py

@@ -63,6 +63,7 @@
     "FlowOptionsDeps",
     "FlowOptionsOverride",
     "ModelEvaluationContext",
+    "TransparentModelEvaluationContext",
     "EvaluatorBase",
     "Evaluator",
     "WrapperModel",
@@ -215,6 +216,19 @@ def __deps__(
         Implementations should be decorated with Flow.call.
         """
 
+    def _evaluation_identity_payload(
+        self,
+        context: Any,
+        child_evaluation_key: Callable[[Any, Any], bytes],
+    ) -> Optional[Any]:
+        """Return a narrower identity payload for cache/graph keys when available.
+
+        Returning ``None`` keeps the model on the existing structural key path.
+        This is intentionally narrow and internal: only models whose effective
+        invocation can be described declaratively should override it.
+        """
+        return None
+
 
 CallableModelType = TypeVar("CallableModelType", bound=_CallableModel)
 
@@ -303,7 +317,7 @@ def get_evaluation_context(model: CallableModelType, context: ContextType, as_di
             if as_dict:
                 return dict(model=evaluator, context=evaluation_context)
             else:
-                return ModelEvaluationContext(model=evaluator, context=evaluation_context)
+                return evaluator.make_evaluation_context(evaluation_context)
 
         # The decorator implementation
         def wrapper(model, context=Signature.empty, *, _options: Optional[FlowOptions] = None, **kwargs):
@@ -588,11 +602,11 @@ def load_prices(
         return flow_model(*args, **kwargs)
 
     @staticmethod
-    def transform(*args, **kwargs):
-        """Decorator that turns a top-level function into a serializable with_inputs() transform factory."""
-        from .flow_model import flow_transform
+    def context_transform(*args, **kwargs):
+        """Decorator that turns a top-level function into a serializable with_context() transform factory."""
+        from .flow_model import flow_context_transform
 
-        return flow_transform(*args, **kwargs)
+        return flow_context_transform(*args, **kwargs)
 
 
 # *****************************************************************************
@@ -698,10 +712,47 @@ def __deps__(self, context: ModelEvaluationContext) -> GraphDepList:
     def __exit__(self):
         pass
 
+    def is_transparent(self, context: ModelEvaluationContext) -> bool:
+        """Whether this evaluator does NOT modify the return value for the given context.
+
+        Transparent evaluators may add side effects (logging, caching, timing,
+        dependency ordering) but always return the same value as ``context()``.
+        This allows cache key computation and dependency graph deduplication to
+        skip these layers.
+
+        Override this method to return ``True`` for evaluators that are always
+        transparent, or implement context-dependent logic for evaluators that
+        are only sometimes transparent.
+        """
+        return False
+
+    def make_evaluation_context(self, context: ModelEvaluationContext, **kwargs) -> ModelEvaluationContext:
+        """Create a ModelEvaluationContext wrapping this evaluator around the given context.
+
+        Returns a ``TransparentModelEvaluationContext`` when ``is_transparent(context)``
+        is ``True``, signaling that this layer can be skipped for cache key computation.
+        """
+        if self.is_transparent(context):
+            return TransparentModelEvaluationContext(model=self, context=context, **kwargs)
+        return ModelEvaluationContext(model=self, context=context, **kwargs)
+
+
+class TransparentModelEvaluationContext(ModelEvaluationContext):
+    """A ModelEvaluationContext layer that is safe to skip for cache key computation.
+
+    Created by ``EvaluatorBase.make_evaluation_context()`` when the evaluator's
+    ``is_transparent()`` returns ``True``. Signals that this evaluator layer does
+    not modify the return value and can be ignored when computing cache keys or
+    deduplicating dependency graph nodes.
+    """
+
 
 class Evaluator(EvaluatorBase):
     """A higher-order model that evaluates a function on a CallableModel and a Context."""
 
+    def is_transparent(self, context: ModelEvaluationContext) -> bool:
+        return True
+
     @override
     def __call__(self, context: ModelEvaluationContext) -> ResultType:
         return context()

ccflow/evaluators/common.py

@@ -8,11 +8,18 @@
 from typing import Any, Callable, Dict, List, Optional, Set, Union
 
 import dask.base
-from pydantic import Field, PrivateAttr, field_validator
+from pydantic import Field, PrivateAttr, ValidationError, field_validator
 from typing_extensions import override
 
 from ..base import BaseModel, make_lazy_result
-from ..callable import CallableModel, ContextBase, EvaluatorBase, ModelEvaluationContext, ResultType
+from ..callable import (
+    CallableModel,
+    ContextBase,
+    EvaluatorBase,
+    ModelEvaluationContext,
+    ResultType,
+    TransparentModelEvaluationContext,
+)
 
 __all__ = [
     "cache_key",
@@ -30,6 +37,13 @@
 log = logging.getLogger(__name__)
 
 
+class _EvaluationIdentityFallback(Exception):
+    """Internal signal to stay on the structural evaluation-key path."""
+
+
+_EXPECTED_IDENTITY_FAILURES = (TypeError, ValueError, ValidationError)
+
+
 def combine_evaluators(first: Optional[EvaluatorBase], second: Optional[EvaluatorBase]) -> EvaluatorBase:
     """Helper function to combine evaluators into a new evaluator.
 
@@ -53,16 +67,25 @@ def combine_evaluators(first: Optional[EvaluatorBase], second: Optional[Evaluato
 
 
 class MultiEvaluator(EvaluatorBase):
-    """An evaluator that combines multiple evaluators."""
+    """An evaluator that combines multiple evaluators.
+
+    Each child evaluator is wrapped in a ModelEvaluationContext using its own
+    ``make_evaluation_context()`` method, so transparent children produce
+    ``TransparentModelEvaluationContext`` layers that can be skipped during
+    cache key computation.
+    """
 
     evaluators: List[EvaluatorBase] = Field(
         description="The list of evaluators to combine. The first evaluator in the list will be called first during evaluation."
     )
 
+    def is_transparent(self, context: ModelEvaluationContext) -> bool:
+        return all(e.is_transparent(context) for e in self.evaluators)
+
     @override
     def __call__(self, context: ModelEvaluationContext) -> ResultType:
         for evaluator in self.evaluators:
-            context = ModelEvaluationContext(model=evaluator, context=context, options=context.options)
+            context = evaluator.make_evaluation_context(context, options=context.options)
         return context()
 
 
@@ -71,6 +94,9 @@ class FallbackEvaluator(EvaluatorBase):
 
     evaluators: List[EvaluatorBase] = Field(description="The list of evaluators to try (in order).")
 
+    def is_transparent(self, context: ModelEvaluationContext) -> bool:
+        return all(e.is_transparent(context) for e in self.evaluators)
+
     @override
     def __call__(self, context: ModelEvaluationContext) -> ResultType:
         for evaluator in self.evaluators:
@@ -120,6 +146,9 @@ class LoggingEvaluator(EvaluatorBase):
     log_result: bool = Field(False, description="Whether to log the result of the evaluation")
     format_config: FormatConfig = Field(FormatConfig(), description="Configuration for formatting the result of the evaluation if log_result=True")
 
+    def is_transparent(self, context: ModelEvaluationContext) -> bool:
+        return True
+
     @field_validator("log_level", mode="before")
     @classmethod
     def _validate_log_level(cls, v: Union[int, str]) -> int:
@@ -194,13 +223,105 @@ def _format_result(self, result: ResultType) -> str:
             return f"{msg_str}{pformat(result_dict, **self.format_config.pformat_config)}"
 
 
+def _unwrap_evaluation_context(flow_obj: ModelEvaluationContext) -> tuple[ModelEvaluationContext, str, List[CallableModel]]:
+    fn = flow_obj.fn
+    non_transparent = []
+    while isinstance(flow_obj.context, ModelEvaluationContext):
+        fn = flow_obj.fn if flow_obj.fn != "__call__" else fn
+        if not isinstance(flow_obj, TransparentModelEvaluationContext):
+            non_transparent.append(flow_obj.model)
+        flow_obj = flow_obj.context
+    return flow_obj, fn if fn != "__call__" else flow_obj.fn, non_transparent
+
+
+def _structural_evaluation_key(flow_obj: ModelEvaluationContext) -> bytes:
+    flow_obj, fn, non_transparent = _unwrap_evaluation_context(flow_obj)
+    d = flow_obj.model_dump(mode="python")
+    d["fn"] = fn
+    if non_transparent:
+        d["_evaluators"] = [e.model_dump(mode="python") for e in non_transparent]
+    return dask.base.tokenize(d).encode("utf-8")
+
+
+class _EvaluationKeyBuilder:
+    def __init__(self) -> None:
+        self._memo: Dict[tuple[int, str], bytes] = {}
+        self._active: set[tuple[int, str]] = set()
+
+    def build(self, context: ModelEvaluationContext) -> bytes:
+        try:
+            return self._build(context)
+        except _EvaluationIdentityFallback:
+            return _structural_evaluation_key(context)
+
+    def _build(self, context: ModelEvaluationContext) -> bytes:
+        inner, fn, non_transparent = _unwrap_evaluation_context(context)
+        if fn != "__call__":
+            raise _EvaluationIdentityFallback("Only __call__ evaluations support narrowed identity.")
+        if non_transparent:
+            raise _EvaluationIdentityFallback("Non-transparent evaluator layers stay on the structural key path.")
+        if inner.options.get("validate_result", True) is False:
+            raise _EvaluationIdentityFallback("validate_result=False stays on the structural key path.")
+        return self._key_for_model(inner.model, inner.context)
+
+    def _key_for_model(self, model: CallableModel, context: Any) -> bytes:
+        memo_token = self._memo_token(model, context)
+        cached = self._memo.get(memo_token)
+        if cached is not None:
+            return cached
+        if memo_token in self._active:
+            raise _EvaluationIdentityFallback("Recursive cycle detected while deriving evaluation identity.")
+
+        self._active.add(memo_token)
+        try:
+            payload = self._identity_payload(model, context)
+            if payload is None:
+                raise _EvaluationIdentityFallback("Model did not provide a narrowed identity payload.")
+            key = dask.base.tokenize(("ccflow_evaluation_identity_v1", payload)).encode("utf-8")
+            self._memo[memo_token] = key
+            return key
+        finally:
+            self._active.discard(memo_token)
+
+    def _identity_payload(self, model: CallableModel, context: Any) -> Optional[Any]:
+        try:
+            return model._evaluation_identity_payload(context, self._child_evaluation_key)
+        except _EXPECTED_IDENTITY_FAILURES as exc:
+            raise _EvaluationIdentityFallback(str(exc)) from exc
+
+    def _child_evaluation_key(self, model: CallableModel, context: Any) -> bytes:
+        try:
+            evaluation = model.__call__.get_evaluation_context(model, context)
+        except _EXPECTED_IDENTITY_FAILURES as exc:
+            raise _EvaluationIdentityFallback(str(exc)) from exc
+        return self.build(evaluation)
+
+    def _memo_token(self, model: CallableModel, context: Any) -> tuple[int, str]:
+        if hasattr(context, "model_dump"):
+            context_value = context.model_dump(mode="python")
+        else:
+            context_value = context
+        return (id(model), dask.base.tokenize((type(context), context_value)))
+
+
+def _evaluation_key(flow_obj: ModelEvaluationContext) -> bytes:
+    return _EvaluationKeyBuilder().build(flow_obj)
+
+
 def cache_key(flow_obj: Union[ModelEvaluationContext, ContextBase, CallableModel]) -> bytes:
-    """Returns a key suitable for use in caching.
+    """Returns a structural key suitable for caching and dependency graph deduplication.
+
+    For ``ModelEvaluationContext`` inputs, strips ``TransparentModelEvaluationContext``
+    layers (evaluators that don't modify the return value) so that the key depends
+    only on the underlying model, context, fn, options, and any non-transparent
+    evaluators in the chain.
 
     Args:
         flow_obj: The object to be tokenized to form the cache key.
     """
-    if isinstance(flow_obj, (ModelEvaluationContext, ContextBase, CallableModel)):
+    if isinstance(flow_obj, ModelEvaluationContext):
+        return _structural_evaluation_key(flow_obj)
+    elif isinstance(flow_obj, (ContextBase, CallableModel)):
         return dask.base.tokenize(flow_obj.model_dump(mode="python")).encode("utf-8")
     else:
         raise TypeError(f"object of type {type(flow_obj)} cannot be serialized by this function!")
@@ -213,9 +334,17 @@ class MemoryCacheEvaluator(EvaluatorBase):
     _cache: Dict[bytes, ResultType] = PrivateAttr({})
     _ids: Dict[bytes, ModelEvaluationContext] = PrivateAttr({})
 
+    def is_transparent(self, context: ModelEvaluationContext) -> bool:
+        return True
+
     def key(self, context: ModelEvaluationContext):
-        """Function to convert a ModelEvaluationContext to a key"""
-        return cache_key(context)
+        """Function to convert a ModelEvaluationContext to a cache key.
+
+        Delegates to the shared evaluation-key builder, which narrows generated
+        ``@Flow.model`` identities when safe and otherwise falls back to
+        ``cache_key()`` semantics.
+        """
+        return _evaluation_key(context)
 
     @property
     def cache(self):
@@ -252,7 +381,7 @@ class CallableModelGraph(BaseModel):
 
 
 def _build_dependency_graph(evaluation_context: ModelEvaluationContext, graph: CallableModelGraph, parent_key: Optional[bytes] = None):
-    key = cache_key(evaluation_context)
+    key = _evaluation_key(evaluation_context)
     if parent_key:
         graph.graph[parent_key].add(key)
     if key not in graph.ids:
@@ -275,7 +404,7 @@ def get_dependency_graph(evaluation_context: ModelEvaluationContext) -> Callable
     Args:
         evaluation_context: The model and context to build the graph for.
     """
-    root_key = cache_key(evaluation_context)
+    root_key = _evaluation_key(evaluation_context)
     graph = CallableModelGraph(ids={}, graph={}, root_id=root_key)
     _build_dependency_graph(evaluation_context, graph)
     return graph
@@ -289,6 +418,9 @@ class GraphEvaluator(EvaluatorBase):
 
     _is_evaluating: bool = PrivateAttr(False)
 
+    def is_transparent(self, context: ModelEvaluationContext) -> bool:
+        return True
+
     @override
     def __call__(self, context: ModelEvaluationContext) -> ResultType:
         import graphlib