workflow inherits BaseModel

tierkreis/tests/controller/test_types.py

@@ -12,10 +12,9 @@
     bytes_from_ptype,
     is_ptype,
     ptype_from_bytes,
-    Workflow,
 )
 from tierkreis.controller.data.graph import GraphData
-from tierkreis.controller.data.models import TKR
+from tierkreis.controller.data.models import TKR, Workflow
 from tests.controller.typed_graphdata import typed_doubler
 
 

tierkreis/tests/workers/graph/stubs.py

@@ -1,8 +1,7 @@
 """Code generated from graph namespace. Please do not edit."""
 
 from typing import NamedTuple
-from tierkreis.controller.data.models import TKR
-from tierkreis.controller.data.types import Workflow
+from tierkreis.controller.data.models import TKR, Workflow
 
 
 class ApplyTwiceInput(NamedTuple):

tierkreis/tierkreis/builder.py

@@ -23,8 +23,9 @@
     TNamedModel,
     dict_from_tmodel,
     init_tmodel,
+    Workflow,
 )
-from tierkreis.controller.data.types import PType, Workflow
+from tierkreis.controller.data.types import PType
 
 
 @dataclass
@@ -151,7 +152,7 @@ def finish_with_outputs(self, outputs: Outputs) -> Workflow[Inputs, Outputs]:
         :type outputs: Outputs
         """
         self.data.output(inputs=dict_from_tmodel(outputs))
-        return Workflow(self.data, self.outputs_type)
+        return Workflow(data=self.data, outputs_type=self.outputs_type)
 
     def embed[A: TModel, B: TModel](
         self, other_fg: Workflow[A, B], inputs: A, outputs_type: type[B]

tierkreis/tierkreis/cli/run_workflow.py

@@ -6,7 +6,8 @@
 
 from tierkreis.controller import run_graph
 from tierkreis.controller.data.graph import GraphData
-from tierkreis.controller.data.types import PType, Workflow
+from tierkreis.controller.data.types import PType
+from tierkreis.controller.data.models import Workflow
 from tierkreis.controller.executor.shell_executor import ShellExecutor
 from tierkreis.controller.executor.uv_executor import UvExecutor
 from tierkreis.controller.storage.filestorage import ControllerFileStorage

tierkreis/tierkreis/controller/data/models.py

@@ -24,8 +24,15 @@
     RestrictedNamedTuple,
     ValueRef,
 )
+from tierkreis.controller.data.graph import GraphData
 from tierkreis.controller.data.types import PType
 
+from pydantic import (
+    BaseModel,
+    model_serializer,
+    SkipValidation,
+)
+
 TKR_PORTMAPPING_FLAG = "__tkr_portmapping__"
 
 
@@ -142,3 +149,36 @@ def init_tmodel[T: TModel](tmodel: type[T], input_fn: Callable[[str], ValueRef])
         return cast("T", model(*args))
     (ref,) = fields.values()
     return tmodel(*ref)
+
+
+class Workflow[Inputs: TModel, Outputs: TModel](BaseModel):
+    data: GraphData
+    outputs_type: SkipValidation[type[Outputs]]
+
+    def __class_getitem__(cls, args):
+        inp, out = args
+        # Pydantic objects to passing in type arguments that are themselves parametrized,
+        # e.g. Workflow[TKR[int], TKR[int]]. The recommendation is to use just `TKR` but
+        # there seems to be caching (i.e. multiple calls with (TKR, TKR) produce the same object)
+        # - hence we use `str` to make sure we get separate instances, and then override
+        # the stored strings with the actual (parametrized) types.
+        if get_origin(inp) == TKR:
+            inp = str(inp)
+        if get_origin(out) == TKR:
+            out = str(out)
+        new_args = (inp, out)
+        annot = super().__class_getitem__(new_args)  # type: ignore # signature says type[Any], but str works ok
+        # This taken from
+        # https://github.com/pydantic/pydantic/blob/812516d71a8696d5e29c5bdab40336d82ccde412/pydantic/_internal/_generics.py#L214-L218
+        pydantic_generic_metadata = getattr(annot, "__pydantic_generic_metadata__")
+        # This assert fails because caching means the `super()` call may return an
+        # instance that we have already mutated
+        # assert pydantic_generic_metadata["args"] == new_args # No, due to caching
+        pydantic_generic_metadata["args"] = args
+        return annot
+
+    @model_serializer
+    def serialize_model(self) -> dict[str, object]:
+        # Just serialize the underlying GraphData, not this frontend/builder type.
+        # We'll reinstate the Workflow wrapper and outputs_type on deserialization.
+        return self.data.model_dump(mode="json")

tierkreis/tierkreis/controller/data/types.py

@@ -7,7 +7,6 @@
 from base64 import b64decode, b64encode
 from collections import defaultdict
 from collections.abc import Mapping, Sequence
-from dataclasses import dataclass
 from inspect import Parameter, _empty, isclass
 from types import NoneType, UnionType
 from typing import (
@@ -22,7 +21,6 @@
     get_args,
     get_origin,
     runtime_checkable,
-    TYPE_CHECKING,
 )
 
 from pydantic import BaseModel, ValidationError
@@ -36,10 +34,6 @@
 )
 from tierkreis.exceptions import TierkreisError
 
-if TYPE_CHECKING:
-    from tierkreis.controller.data.graph import GraphData
-    from tierkreis.controller.data.models import TModel
-
 
 @runtime_checkable
 class NdarraySurrogate(Protocol):
@@ -112,7 +106,6 @@ def from_list(cls, arg: list, /) -> "Self":
     | ListConvertible
     | NdarraySurrogate
     | BaseModel  # Includes GraphData
-    | Workflow  # So, special case: a Workflow is just a GraphData, discard the type info
 )
 type JsonType = Container[ElementaryType]
 logger = logging.getLogger(__name__)
@@ -228,12 +221,10 @@ def is_ptype(annotation: Any) -> TypeIs[type[PType]]:
             annotation,
             (DictConvertible, ListConvertible, NdarraySurrogate, BaseModel, Struct),
         )
-        or (isclass(origin) and issubclass(origin, Workflow))
         or annotation in get_args(ElementaryType.__value__)
     ):
         return True
 
-    origin = get_origin(annotation)
     if origin is not None:
         return is_ptype(origin) and all(is_ptype(x) for x in get_args(annotation))
 
@@ -252,12 +243,12 @@ def ser_from_ptype(ptype: PType, annotation: type[PType] | None) -> JsonType:
     :return: The serialized ptype.
     :rtype: JsonType
     """
+
+    # ALAN from tierkreis.controller.data.models import Workflow
     if sr := get_serializer(annotation):
         return sr.serializer(ptype)
 
     match ptype:
-        case Workflow():
-            return ser_from_ptype(ptype.data, annotation)
         case bytes() | bytearray() | memoryview():
             return bytes(ptype)
         case bool() | int() | float() | complex() | str() | NoneType() | TypeVar():
@@ -318,6 +309,7 @@ def coerce_from_annotation[T: PType](ser: Any, annotation: type[T] | None) -> T:
     :rtype: T
     """
     from tierkreis.controller.data.graph import GraphData
+    from tierkreis.controller.data.models import Workflow
 
     if annotation is None:
         return ser
@@ -369,16 +361,22 @@ def coerce_from_annotation[T: PType](ser: Any, annotation: type[T] | None) -> T:
     if issubclass(origin, NdarraySurrogate):
         return pickle.loads(ser)
 
+    if issubclass(origin, Workflow):  # Also a BaseModel so do this first
+        # Serialized as GraphData - reinstate erased output type
+        # This taken from
+        # https://github.com/pydantic/pydantic/blob/812516d71a8696d5e29c5bdab40336d82ccde412/pydantic/_internal/_generics.py#L214-L218
+        pydantic_generic_metadata = getattr(annotation, "__pydantic_generic_metadata__")
+        _inputs, outputs = pydantic_generic_metadata.get("args")
+        return Workflow(
+            data=coerce_from_annotation(ser, GraphData), outputs_type=outputs
+        )  # type: ignore
+
     if issubclass(origin, BaseModel):
         if not issubclass(annotation, origin):
             msg = "Invalid subclass relation encountered."
             raise TypeError(msg)
         return annotation(**ser)
 
-    if issubclass(origin, Workflow):
-        _inputs, outputs = get_args(annotation)
-        return annotation(coerce_from_annotation(ser, GraphData), outputs)  # type: ignore
-
     if issubclass(origin, Struct):
         d = {
             k: coerce_from_annotation(ser[k], v)
@@ -467,9 +465,3 @@ def has_default(t: Parameter) -> bool:
     :rtype: bool
     """
     return not (isclass(t.default) and issubclass(t.default, _empty))
-
-
-@dataclass(frozen=True)
-class Workflow[Inputs: TModel, Outputs: TModel]:
-    data: "GraphData"
-    outputs_type: type[Outputs]

tierkreis/tierkreis/namespace.py

@@ -134,8 +134,8 @@ def stubs(self) -> str:
 
 from typing import Literal, NamedTuple, Sequence, TypeVar, Generic, Protocol, Union
 from types import NoneType
-from tierkreis.controller.data.models import TKR, OpaqueType
-from tierkreis.controller.data.types import PType, Struct, Workflow
+from tierkreis.controller.data.models import TKR, OpaqueType, Workflow
+from tierkreis.controller.data.types import PType, Struct
 
 {models_str}