diff --git a/tierkreis/tests/controller/test_types.py b/tierkreis/tests/controller/test_types.py index c7e151df0..9c5bb0318 100644 --- a/tierkreis/tests/controller/test_types.py +++ b/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 diff --git a/tierkreis/tests/workers/graph/stubs.py b/tierkreis/tests/workers/graph/stubs.py index b87db639a..a134207d2 100644 --- a/tierkreis/tests/workers/graph/stubs.py +++ b/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): diff --git a/tierkreis/tierkreis/builder.py b/tierkreis/tierkreis/builder.py index 56ee05bb3..8982d2132 100644 --- a/tierkreis/tierkreis/builder.py +++ b/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] diff --git a/tierkreis/tierkreis/cli/run_workflow.py b/tierkreis/tierkreis/cli/run_workflow.py index d402043c3..5e91b7b12 100644 --- a/tierkreis/tierkreis/cli/run_workflow.py +++ 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 diff --git a/tierkreis/tierkreis/controller/data/models.py b/tierkreis/tierkreis/controller/data/models.py index 854bba9b3..4fb26da32 100644 --- a/tierkreis/tierkreis/controller/data/models.py +++ b/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") diff --git a/tierkreis/tierkreis/controller/data/types.py b/tierkreis/tierkreis/controller/data/types.py index ec68ed18d..09036aac2 100644 --- a/tierkreis/tierkreis/controller/data/types.py +++ b/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] diff --git a/tierkreis/tierkreis/namespace.py b/tierkreis/tierkreis/namespace.py index eec25bae7..290b12f24 100644 --- a/tierkreis/tierkreis/namespace.py +++ b/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}