diff --git a/docs/source/core_concepts.md b/docs/source/core_concepts.md index 1d7ce7ec6..c7c298cd2 100644 --- a/docs/source/core_concepts.md +++ b/docs/source/core_concepts.md @@ -30,7 +30,7 @@ From the data dependencies, the runtime environment can infer which tasks it nee ### Types In Tierkreis, it is possible to associate types with edges (data) at construction. -Type information can be used by the `GraphBuilder` to infer additional information about the graph to prevent runtime errors. +Type information can be used while building the `Graph` to detect and prevent runtime errors. ## Execution model diff --git a/docs/source/examples/data/typed_eval.py b/docs/source/examples/data/typed_eval.py index 7285ac4bb..6712495bf 100644 --- a/docs/source/examples/data/typed_eval.py +++ b/docs/source/examples/data/typed_eval.py @@ -1,6 +1,6 @@ from typing import NamedTuple -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.builtins import iadd, itimes from tierkreis.controller.data.core import EmptyModel from tierkreis.controller.data.models import TKR @@ -17,26 +17,23 @@ class DoublerOutput(NamedTuple): def typed_doubler(): - g = GraphBuilder(TKR[int], TKR[int]) + g = Graph(TKR[int], TKR[int]) out = g.task(itimes(a=g.const(2), b=g.inputs)) - g.outputs(out) - return g + return g.finish_with_outputs(out) def typed_doubler_plus_multi(): - g = GraphBuilder(DoublerInput, DoublerOutput) + g = Graph(DoublerInput, DoublerOutput) mul = g.task(itimes(a=g.inputs.x, b=g.const(2))) out = g.task(iadd(a=mul, b=g.inputs.intercept)) - g.outputs(DoublerOutput(a=g.inputs.x, value=out)) - return g + return g.finish_with_outputs(DoublerOutput(a=g.inputs.x, value=out)) def typed_doubler_plus(): - g = GraphBuilder(DoublerInput, TKR[int]) + g = Graph(DoublerInput, TKR[int]) mul = g.task(itimes(a=g.inputs.x, b=g.const(2))) out = g.task(iadd(a=mul, b=g.inputs.intercept)) - g.outputs(out) - return g + return g.finish_with_outputs(out) class TypedEvalOutputs(NamedTuple): @@ -44,7 +41,6 @@ class TypedEvalOutputs(NamedTuple): def typed_eval(): - g = GraphBuilder(EmptyModel, TypedEvalOutputs) + g = Graph(EmptyModel, TypedEvalOutputs) e = g.eval(typed_doubler_plus(), DoublerInput(x=g.const(6), intercept=g.const(0))) - g.outputs(TypedEvalOutputs(typed_eval_output=e)) - return g + return g.finish_with_outputs(TypedEvalOutputs(typed_eval_output=e)) diff --git a/docs/source/examples/errors_and_debugging.ipynb b/docs/source/examples/errors_and_debugging.ipynb index 3d6dd0341..bca425953 100644 --- a/docs/source/examples/errors_and_debugging.ipynb +++ b/docs/source/examples/errors_and_debugging.ipynb @@ -27,18 +27,17 @@ "metadata": {}, "outputs": [], "source": [ - "from tierkreis.builder import GraphBuilder\n", + "from tierkreis.builder import Graph, Workflow\n", "from tierkreis.controller.data.core import EmptyModel\n", "from tierkreis.controller.data.models import TKR\n", "\n", "from error_worker import fail\n", "\n", "\n", - "def error_graph() -> GraphBuilder:\n", - " g = GraphBuilder(EmptyModel, TKR[str])\n", + "def error_graph() -> Workflow:\n", + " g = Graph(EmptyModel, TKR[str])\n", " output = g.task(fail())\n", - " g.outputs(output)\n", - " return g" + " return g.finish_with_outputs(output)" ] }, { diff --git a/docs/source/examples/hamiltonian.ipynb b/docs/source/examples/hamiltonian.ipynb index 177d13f5a..b41a66fdb 100644 --- a/docs/source/examples/hamiltonian.ipynb +++ b/docs/source/examples/hamiltonian.ipynb @@ -69,7 +69,7 @@ "metadata": {}, "outputs": [], "source": [ - "from tierkreis.builder import GraphBuilder\n", + "from tierkreis.builder import Graph\n", "from tierkreis.controller.data.models import TKR\n", "from typing import NamedTuple, Literal\n", "\n", @@ -84,7 +84,7 @@ " ansatz: TKR[Literal[\"pytket._tket.circuit.Circuit\"]]\n", "\n", "\n", - "simulation_graph = GraphBuilder(SymbolicExecutionInputs, TKR[float])\n", + "simulation_graph = Graph(SymbolicExecutionInputs, TKR[float])\n", "substituted_circuit = simulation_graph.task(\n", " substitute(\n", " simulation_graph.inputs.ansatz,\n", @@ -125,7 +125,7 @@ "\n", "\n", "def exp_val():\n", - " g = GraphBuilder(SubmitInputs, TKR[float])\n", + " g = Graph(SubmitInputs, TKR[float])\n", "\n", " circuit = g.inputs.circuit\n", " pauli_string = g.inputs.pauli_string\n", @@ -137,8 +137,7 @@ "\n", " backend_result = g.task(submit_single(compiled_circuit, n_shots))\n", " av = g.task(expectation(backend_result))\n", - " g.outputs(av)\n", - " return g" + " return g.finish_with_outputs(av)" ] }, { @@ -195,14 +194,13 @@ "\n", "\n", "def compute_terms():\n", - " g = GraphBuilder(ComputeTermsInputs, TKR[float])\n", + " g = Graph(ComputeTermsInputs, TKR[float])\n", "\n", " res_0, res_1 = g.task(untuple(g.inputs.value))\n", " prod = g.task(times(res_0, res_1))\n", " sum = g.task(add(g.inputs.accum, prod))\n", "\n", - " g.outputs(sum)\n", - " return g" + " return g.finish_with_outputs(sum)" ] }, { @@ -243,7 +241,7 @@ "outputs": [], "source": [ "computed = simulation_graph.eval(fold_graph(compute_terms()), fold_inputs)\n", - "simulation_graph.outputs(computed)" + "workflow = simulation_graph.finish_with_outputs(computed)" ] }, { @@ -340,11 +338,11 @@ "run_graph(\n", " storage,\n", " multi_executor,\n", - " simulation_graph,\n", + " workflow,\n", " inputs,\n", " polling_interval_seconds=0.2,\n", ")\n", - "output = read_outputs(simulation_graph, storage)\n", + "output = read_outputs(workflow, storage)\n", "print(output)" ] } diff --git a/docs/source/examples/hello_world.py b/docs/source/examples/hello_world.py index 0056ce49d..4e796e386 100644 --- a/docs/source/examples/hello_world.py +++ b/docs/source/examples/hello_world.py @@ -1,10 +1,10 @@ -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.controller.data.models import TKR from hello_world_worker import greet -graph = GraphBuilder(inputs_type=TKR[str], outputs_type=TKR[str]) +graph = Graph(inputs_type=TKR[str], outputs_type=TKR[str]) hello = graph.const("Hello ") output = graph.task(greet(greeting=hello, subject=graph.inputs)) -graph.outputs(output) +graph.finish_with_outputs(output) diff --git a/docs/source/examples/hello_world_graph.ipynb b/docs/source/examples/hello_world_graph.ipynb index 0281e4640..2d2dcc824 100644 --- a/docs/source/examples/hello_world_graph.ipynb +++ b/docs/source/examples/hello_world_graph.ipynb @@ -29,7 +29,7 @@ "metadata": {}, "outputs": [], "source": [ - "from tierkreis.builder import GraphBuilder\n", + "from tierkreis.builder import Graph\n", "from tierkreis.controller.data.models import TKR" ] }, @@ -50,7 +50,7 @@ "metadata": {}, "outputs": [], "source": [ - "graph = GraphBuilder(inputs_type=TKR[str], outputs_type=TKR[str])" + "graph = Graph(inputs_type=TKR[str], outputs_type=TKR[str])" ] }, { @@ -116,7 +116,7 @@ "metadata": {}, "outputs": [], "source": [ - "graph.outputs(output)" + "workflow = graph.finish_with_outputs(output)" ] }, { @@ -148,7 +148,7 @@ "\n", "def main(input_value: str) -> None:\n", " run_workflow(\n", - " graph.data,\n", + " workflow,\n", " {\"value\": input_value},\n", " name=\"hello_world\",\n", " run_id=100, # Assign a fixed uuid for our workflow.\n", diff --git a/docs/source/examples/hpc.ipynb b/docs/source/examples/hpc.ipynb index ff7a20e69..87138b41c 100644 --- a/docs/source/examples/hpc.ipynb +++ b/docs/source/examples/hpc.ipynb @@ -62,7 +62,7 @@ "metadata": {}, "outputs": [], "source": [ - "from tierkreis.builder import GraphBuilder\n", + "from tierkreis.builder import Graph, Workflow\n", "from substitution_worker import substitute\n", "from tierkreis.pytket_worker import (\n", " add_measure_all,\n", @@ -72,9 +72,9 @@ "from tierkreis.aer_worker import submit_single\n", "\n", "\n", - "def symbolic_execution() -> GraphBuilder:\n", + "def symbolic_execution() -> Workflow:\n", " \"\"\"A graph that substitutes 3 parameters into a circuit and gets an expectation value.\"\"\"\n", - " g = GraphBuilder(SymbolicCircuitsInputs, SymbolicCircuitsOutputs)\n", + " g = Graph(SymbolicCircuitsInputs, SymbolicCircuitsOutputs)\n", " a = g.inputs.a\n", " b = g.inputs.b\n", " c = g.inputs.c\n", @@ -88,8 +88,7 @@ " backend_result = g.task(submit_single(circuit=compiled_circuit, n_shots=n_shots))\n", " av = g.task(expectation(backend_result=backend_result))\n", "\n", - " g.outputs(SymbolicCircuitsOutputs(expectation=av))\n", - " return g" + " return g.finish_with_outputs(SymbolicCircuitsOutputs(expectation=av))" ] }, { diff --git a/docs/source/examples/multiple_outputs.ipynb b/docs/source/examples/multiple_outputs.ipynb index 7f00abae4..cdb1512de 100644 --- a/docs/source/examples/multiple_outputs.ipynb +++ b/docs/source/examples/multiple_outputs.ipynb @@ -56,17 +56,17 @@ "metadata": {}, "outputs": [], "source": [ - "from tierkreis.builder import GraphBuilder\n", + "from tierkreis.builder import Graph\n", "from tierkreis.controller.data.models import TKR\n", "from multiple_outputs_worker import new_opaque_point, OpaquePoint # noqa: F811\n", "\n", - "g = GraphBuilder(TKR[float], TKR[OpaquePoint])\n", + "g = Graph(TKR[float], TKR[OpaquePoint])\n", "diag = g.task(new_opaque_point(g.inputs, g.inputs, g.inputs))\n", "\n", "# Uncommenting the following line will give an error.\n", "# diag.x ## Cannot access attribute \"x\" for class \"TKR[OpaquePoint]\"\n", "\n", - "g.outputs(diag)" + "workflow = g.finish_with_outputs(diag)" ] }, { @@ -92,8 +92,8 @@ "\n", "storage = FileStorage(UUID(int=400), do_cleanup=True)\n", "executor = UvExecutor(Path().parent / \"example_workers\", storage.logs_path)\n", - "run_graph(storage, executor, g, 4)\n", - "outputs = read_outputs(g, storage)\n", + "run_graph(storage, executor, workflow, 4)\n", + "outputs = read_outputs(workflow, storage)\n", "assert outputs == {\"x\": 4, \"y\": 4, \"z\": 4}" ] }, @@ -144,15 +144,15 @@ "metadata": {}, "outputs": [], "source": [ - "from tierkreis.builder import GraphBuilder\n", + "from tierkreis.builder import Graph\n", "from tierkreis.controller.data.models import TKR\n", "from tierkreis.builtins.stubs import add\n", "from multiple_outputs_worker import new_point # noqa: F811\n", "\n", - "g = GraphBuilder(TKR[float], TKR[float])\n", + "g = Graph(TKR[float], TKR[float])\n", "diag = g.task(new_point(g.inputs, g.inputs, g.inputs))\n", "sum = g.task(add(diag.x, g.task(add(diag.y, diag.z))))\n", - "g.outputs(sum)" + "workflow = g.finish_with_outputs(sum)" ] }, { @@ -178,8 +178,8 @@ "\n", "storage = FileStorage(UUID(int=400), do_cleanup=True)\n", "executor = UvExecutor(Path().parent / \"example_workers\", storage.logs_path)\n", - "run_graph(storage, executor, g, 4)\n", - "outputs = read_outputs(g, storage)\n", + "run_graph(storage, executor, workflow, 4)\n", + "outputs = read_outputs(workflow, storage)\n", "assert outputs == 12" ] } diff --git a/docs/source/examples/parallelism.ipynb b/docs/source/examples/parallelism.ipynb index 066a7c403..e3c23853f 100644 --- a/docs/source/examples/parallelism.ipynb +++ b/docs/source/examples/parallelism.ipynb @@ -35,7 +35,7 @@ "outputs": [], "source": [ "from typing import Literal, NamedTuple\n", - "from tierkreis.builder import GraphBuilder\n", + "from tierkreis.builder import Graph\n", "from tierkreis.controller.data.models import TKR, OpaqueType\n", "from tierkreis.builtins import untuple\n", "from tierkreis.aer_worker import (\n", @@ -58,7 +58,7 @@ "\n", "\n", "def aer_simulate_single():\n", - " g = GraphBuilder(SimulateJobInputsSingle, TKR[BackendResult])\n", + " g = Graph(SimulateJobInputsSingle, TKR[BackendResult])\n", " circuit_shots = g.task(untuple(g.inputs.circuit_shots))\n", "\n", " compiled_circuit = g.task(\n", @@ -68,12 +68,11 @@ " )\n", " )\n", " res = g.task(aer_run(compiled_circuit, circuit_shots.b))\n", - " g.outputs(res)\n", - " return g\n", + " return g.finish_with_outputs(res)\n", "\n", "\n", "def qulacs_simulate_single():\n", - " g = GraphBuilder(SimulateJobInputsSingle, TKR[BackendResult])\n", + " g = Graph(SimulateJobInputsSingle, TKR[BackendResult])\n", " circuit_shots = g.task(untuple(g.inputs.circuit_shots))\n", "\n", " compiled_circuit = g.task(\n", @@ -83,8 +82,7 @@ " )\n", " )\n", " res = g.task(qulacs_run(compiled_circuit, circuit_shots.b))\n", - " g.outputs(res)\n", - " return g" + " return g.finish_with_outputs(res)" ] }, { @@ -108,7 +106,7 @@ "\n", "\n", "def compile_simulate_single():\n", - " g = GraphBuilder(SimulateJobInputsSingle, TKR[BackendResult])\n", + " g = Graph(SimulateJobInputsSingle, TKR[BackendResult])\n", "\n", " aer_res = g.eval(aer_simulate_single(), g.inputs)\n", " qulacs_res = g.eval(qulacs_simulate_single(), g.inputs)\n", @@ -116,8 +114,7 @@ " g.task(str_eq(g.inputs.simulator_name, g.const(\"aer\"))), aer_res, qulacs_res\n", " )\n", "\n", - " g.outputs(res)\n", - " return g" + " return g.finish_with_outputs(res)" ] }, { @@ -142,7 +139,7 @@ " compilation_optimisation_level: TKR[int]\n", "\n", "\n", - "g = GraphBuilder(SimulateJobInputs, TKR[list[BackendResult]])" + "g = Graph(SimulateJobInputs, TKR[list[BackendResult]])" ] }, { @@ -207,7 +204,7 @@ "source": [ "res = g.map(compile_simulate_single(), job_inputs)\n", "\n", - "g.outputs(res)" + "workflow = g.finish_with_outputs(res)" ] }, { @@ -268,7 +265,7 @@ "inputs[\"simulator_name\"] = \"aer\"\n", "print(\"Simulating using aer...\")\n", "start = time.time()\n", - "run_graph(storage, executor, g, inputs, polling_interval_seconds=0.1)\n", + "run_graph(storage, executor, workflow, inputs, polling_interval_seconds=0.1)\n", "print(f\"time taken: {time.time() - start}\")" ] }, @@ -292,7 +289,7 @@ "print(\"Simulating using qulacs...\")\n", "storage.clean_graph_files()\n", "start = time.time()\n", - "run_graph(storage, executor, g, inputs, polling_interval_seconds=0.1)\n", + "run_graph(storage, executor, workflow, inputs, polling_interval_seconds=0.1)\n", "print(f\"time taken: {time.time() - start}\")" ] }, diff --git a/docs/source/examples/qsci.ipynb b/docs/source/examples/qsci.ipynb index 7a2f1b25f..48b132a33 100644 --- a/docs/source/examples/qsci.ipynb +++ b/docs/source/examples/qsci.ipynb @@ -343,16 +343,16 @@ "outputs": [], "source": [ "from tierkreis.aer_worker import submit_single\n", - "from tierkreis.builder import GraphBuilder\n", + "from tierkreis.builder import Graph, Workflow\n", "from tierkreis.controller.data.models import TKR, OpaqueType\n", "from tierkreis.quantinuum_worker import compile_circuit_quantinuum\n", "\n", "\n", - "def _compile_and_run() -> GraphBuilder[\n", + "def _compile_and_run() -> Workflow[\n", " TKR[OpaqueType[\"pytket._tket.circuit.Circuit\"]], # noqa: F821\n", " TKR[OpaqueType[\"pytket.backends.backendresult.BackendResult\"]], # noqa: F821\n", "]:\n", - " g = GraphBuilder(\n", + " g = Graph(\n", " TKR[OpaqueType[\"pytket._tket.circuit.Circuit\"]],\n", " TKR[OpaqueType[\"pytket.backends.backendresult.BackendResult\"]],\n", " )\n", @@ -361,8 +361,7 @@ " compiled_circuit = g.task(compile_circuit_quantinuum(g.inputs))\n", " backend_result = g.task(submit_single(compiled_circuit, n_shots))\n", "\n", - " g.outputs(backend_result)\n", - " return g" + " return g.finish_with_outputs(backend_result)" ] }, { @@ -423,7 +422,7 @@ " state_prep,\n", ")\n", "\n", - "qsci_graph = GraphBuilder(QSCIInputs, QSCIOutputs)\n", + "qsci_graph = Graph(QSCIInputs, QSCIOutputs)\n", "# Separate tasks 'make_h_init'+'state_pre' and 'make_h_hsim' run in parallel\n", "ham_init = qsci_graph.task(\n", " make_ham(\n", @@ -474,7 +473,7 @@ " ),\n", ")\n", "\n", - "qsci_graph.outputs(QSCIOutputs(energy))" + "workflow = qsci_graph.finish_with_outputs(QSCIOutputs(energy))" ] }, { @@ -522,7 +521,7 @@ "run_graph(\n", " storage,\n", " multi_executor,\n", - " qsci_graph,\n", + " workflow,\n", " {\n", " k: json.dumps(v).encode()\n", " for k, v in (\n", @@ -556,7 +555,7 @@ " },\n", " polling_interval_seconds=0.01,\n", ")\n", - "output = read_outputs(qsci_graph, storage)" + "output = read_outputs(workflow, storage)" ] } ], diff --git a/docs/source/examples/scipy.ipynb b/docs/source/examples/scipy.ipynb index 334ff5cb6..1f2b46543 100644 --- a/docs/source/examples/scipy.ipynb +++ b/docs/source/examples/scipy.ipynb @@ -61,7 +61,7 @@ "metadata": {}, "outputs": [], "source": [ - "from tierkreis.builder import GraphBuilder\n", + "from tierkreis.builder import Graph\n", "from tierkreis.controller.data.core import EmptyModel\n", "\n", "from scipy_worker import (\n", @@ -72,7 +72,7 @@ " transpose,\n", ")\n", "\n", - "sample_graph = GraphBuilder(EmptyModel, ScipyOutputs)\n", + "sample_graph = Graph(EmptyModel, ScipyOutputs)\n", "onedim = sample_graph.task(linspace(sample_graph.const(0), sample_graph.const(10)))\n", "\n", "pointed = sample_graph.task(add_point(onedim, sample_graph.const(0)))\n", @@ -80,7 +80,7 @@ "\n", "twodim = sample_graph.task(reshape(onedim, sample_graph.const([5, 10])))\n", "a = sample_graph.task(transpose(twodim))\n", - "sample_graph.outputs(ScipyOutputs(a, scalar))" + "workflow = sample_graph.finish_with_outputs(ScipyOutputs(a, scalar))" ] }, { @@ -163,9 +163,9 @@ "\n", "storage = FileStorage(UUID(int=207), do_cleanup=True, name=\"scipy_graph\")\n", "executor = UvExecutor(Path().parent / \"example_workers\", storage.logs_path)\n", - "run_graph(storage, executor, sample_graph, {})\n", + "run_graph(storage, executor, workflow, {})\n", "\n", - "outputs = read_outputs(sample_graph, storage)" + "outputs = read_outputs(workflow, storage)" ] }, { @@ -211,9 +211,9 @@ "source": [ "os.environ[\"SER_METHOD\"] = \"tolist\"\n", "storage.clean_graph_files()\n", - "run_graph(storage, executor, sample_graph, {})\n", + "run_graph(storage, executor, workflow, {})\n", "\n", - "outputs = read_outputs(sample_graph, storage)" + "outputs = read_outputs(workflow, storage)" ] } ], diff --git a/docs/source/examples/signing_graph.ipynb b/docs/source/examples/signing_graph.ipynb index d0dfa2ea4..dc6771624 100644 --- a/docs/source/examples/signing_graph.ipynb +++ b/docs/source/examples/signing_graph.ipynb @@ -123,7 +123,7 @@ "metadata": {}, "outputs": [], "source": [ - "from tierkreis.builder import GraphBuilder\n", + "from tierkreis.builder import Graph\n", "from tierkreis.models import TKR, EmptyModel\n", "\n", "from auth_worker import sign, verify\n", @@ -131,7 +131,7 @@ "\n", "\n", "def signing_graph():\n", - " g = GraphBuilder(EmptyModel, TKR[bool])\n", + " g = Graph(EmptyModel, TKR[bool])\n", " message = g.const(\"dummymessage\")\n", " passphrase = g.const(b\"dummypassphrase\")\n", "\n", @@ -141,9 +141,7 @@ "\n", " signing_result = g.task(sign(private_key, passphrase, message)).hex_signature\n", " verification_result = g.task(verify(public_key, signing_result, message))\n", - " g.outputs(verification_result)\n", - "\n", - " return g" + " return g.finish_with_outputs(verification_result)" ] }, { @@ -218,8 +216,8 @@ "from tierkreis.storage import read_outputs\n", "from tierkreis import run_graph\n", "\n", - "run_graph(storage, executor, signing_graph().get_data(), {})\n", - "is_verified = read_outputs(signing_graph().get_data(), storage)\n", + "run_graph(storage, executor, signing_graph(), {})\n", + "is_verified = read_outputs(signing_graph(), storage)\n", "print(is_verified)" ] }, @@ -250,18 +248,17 @@ "\n", "\n", "def stdinout_graph():\n", - " g = GraphBuilder(EmptyModel, TKR[bytes])\n", + " g = Graph(EmptyModel, TKR[bytes])\n", " message = g.const(b\"dummymessage\")\n", " output = g.task(script(\"tee\", message))\n", "\n", - " g.outputs(output)\n", - " return g\n", + " return g.finish_with_outputs(output)\n", "\n", "\n", "storage.clean_graph_files()\n", "stdinout = StdInOut(registry_path, storage.workflow_dir)\n", - "run_graph(storage, stdinout, stdinout_graph().get_data(), {})\n", - "out = read_outputs(stdinout_graph().get_data(), storage)\n", + "run_graph(storage, stdinout, stdinout_graph(), {})\n", + "out = read_outputs(stdinout_graph(), storage)\n", "print(out)" ] } diff --git a/docs/source/examples/storage_and_executors.ipynb b/docs/source/examples/storage_and_executors.ipynb index 64c698a9c..049985e9c 100644 --- a/docs/source/examples/storage_and_executors.ipynb +++ b/docs/source/examples/storage_and_executors.ipynb @@ -57,7 +57,7 @@ "metadata": {}, "outputs": [], "source": [ - "from tierkreis.builder import GraphBuilder\n", + "from tierkreis.builder import Graph\n", "from tierkreis.controller.data.models import TKR\n", "from tierkreis.quantinuum_worker import (\n", " compile_using_info,\n", @@ -65,7 +65,7 @@ " run_circuit,\n", ")\n", "\n", - "g = GraphBuilder(TKR[Circuit], TKR[BackendResult])\n", + "g = Graph(TKR[Circuit], TKR[BackendResult])\n", "info = g.task(get_backend_info(device_name=g.const(\"H2-1\")))\n", "compiled_circuit = g.task(compile_using_info(g.inputs, info))\n", "results = g.task(\n", @@ -75,7 +75,7 @@ " device_name=g.const(\"H2-1SC\"),\n", " ),\n", ")\n", - "g.outputs(results)" + "workflow = g.finish_with_outputs(results)" ] }, { @@ -171,7 +171,7 @@ "from tierkreis import run_graph\n", "\n", "circuit = circuit_from_qasm(Path().parent / \"data\" / \"ghz_state_n23.qasm\")\n", - "run_graph(storage, executor, g, circuit)" + "run_graph(storage, executor, workflow, circuit)" ] }, { diff --git a/docs/source/examples/types_and_defaults.ipynb b/docs/source/examples/types_and_defaults.ipynb index 746244b67..7f49adb00 100644 --- a/docs/source/examples/types_and_defaults.ipynb +++ b/docs/source/examples/types_and_defaults.ipynb @@ -68,11 +68,11 @@ "outputs": [], "source": [ "from tierkreis.aer_worker import get_compiled_circuit, submit_single\n", - "from tierkreis.builder import GraphBuilder\n", + "from tierkreis.builder import Graph, Workflow\n", "\n", "\n", - "def compile_run_single() -> GraphBuilder[IBMInput, IBMOutput]:\n", - " g = GraphBuilder(IBMInput, IBMOutput)\n", + "def compile_run_single() -> Workflow[IBMInput, IBMOutput]:\n", + " g = Graph(IBMInput, IBMOutput)\n", "\n", " compiled_circuit = g.task(\n", " get_compiled_circuit(\n", @@ -81,8 +81,7 @@ " ),\n", " )\n", " res = g.task(submit_single(compiled_circuit, g.inputs.n_shots))\n", - " g.outputs(IBMOutput(res, g.inputs.n_shots))\n", - " return g" + " return g.finish_with_outputs(IBMOutput(res, g.inputs.n_shots))" ] }, { @@ -163,8 +162,8 @@ " optimisation_level: TKR[int] | None = None # Optional input\n", "\n", "\n", - "def compile_optional() -> GraphBuilder[IBMInputOptional, IBMOutput]:\n", - " g = GraphBuilder(IBMInputOptional, IBMOutput)\n", + "def compile_optional() -> Workflow[IBMInputOptional, IBMOutput]:\n", + " g = Graph(IBMInputOptional, IBMOutput)\n", "\n", " compiled_circuit = g.task(\n", " get_compiled_circuit(\n", @@ -173,8 +172,7 @@ " ),\n", " )\n", " res = g.task(submit_single(compiled_circuit, g.inputs.n_shots))\n", - " g.outputs(IBMOutput(res, g.inputs.n_shots)) # type: ignore\n", - " return g" + " return g.finish_with_outputs(IBMOutput(res, g.inputs.n_shots)) # type: ignore" ] }, { diff --git a/docs/source/executors/index.md b/docs/source/executors/index.md index eaa765456..1fc1ed26e 100644 --- a/docs/source/executors/index.md +++ b/docs/source/executors/index.md @@ -75,7 +75,7 @@ Now as graph definition in `main.py`: ```py from first_worker import some_task -graph = GraphBuilder(TKR[NoneType], TKR[NoneType]) +graph = Graph(TKR[NoneType], TKR[NoneType]) graph.task(some_task) ``` @@ -120,7 +120,7 @@ If different executors are necessary they can be combined using the [`MultipleEx ```py def multiple_graph(): # Both tasks are the same, we just use different names to test the task executor - g = GraphBuilder(TKR[str], TKR[str]) + g = Graph(TKR[str], TKR[str]) first_call = g.data.func( "shell_worker.meet", {"greeting": g.inputs.value_ref()}, @@ -131,8 +131,7 @@ def multiple_graph(): {"greeting": out.value_ref()}, ) output: TKR[str] = TKR(*second_call("value")) - g.outputs(output) - return g + return g.finish_with_outputs(output) def main(): @@ -160,7 +159,7 @@ Alternativerly if you need control on task level, e.g., to execute the same work ```py def task_graph(): # Both tasks are the same, we just use different names to test the task executor - g = GraphBuilder(TKR[str], TKR[str]) + g = Graph(TKR[str], TKR[str]) first_call = g.data.func( "shell_worker.meet", {"greeting": g.inputs.value_ref()}, @@ -171,8 +170,8 @@ def task_graph(): {"greeting": out.value_ref()}, ) output: TKR[str] = TKR(*second_call("value")) - g.outputs(output) - return g + return g.finish_with_outputs(output) + def main(): g = task_graph() storage = ControllerFileStorage(UUID(int=305), name="Task") diff --git a/docs/source/executors/shell.md b/docs/source/executors/shell.md index fd4f9d95a..6a6ce4423 100644 --- a/docs/source/executors/shell.md +++ b/docs/source/executors/shell.md @@ -67,11 +67,11 @@ Since we don't have stubs for the script (which we could generate by providing a ```{code} ipython3 from tierkreis.models import EmptyModel, TKR -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from ..tutorial.auth_stubs import sign -def signing_graph()-> GraphBuilder[EmptyModel, TKR[str]]: - g = GraphBuilder(EmptyModel, TKR[str]) +def signing_graph()-> Graph[EmptyModel, TKR[str]]: + g = Graph(EmptyModel, TKR[str]) # Define the fix inputs to the graph message = g.const("dummymessage") passphrase = g.const(b"dummypassphrase") @@ -89,8 +89,7 @@ def signing_graph()-> GraphBuilder[EmptyModel, TKR[str]]: # Finally we use the auth worker to sign the message signing_result = g.task(sign(private_key, passphrase, message)).hex_signature - g.outputs(signing_result) - return g + return g.finish_with_outputs(signing_result) ``` Running the graph follows all the usual steps. diff --git a/docs/source/graphs/builtins.md b/docs/source/graphs/builtins.md index 56c1e5947..df78aed8f 100644 --- a/docs/source/graphs/builtins.md +++ b/docs/source/graphs/builtins.md @@ -14,15 +14,15 @@ pip install tierkreis ## Constructing the graph -First we instantiate a `GraphBuilder` object. +First we instantiate a `Graph` object. The arguments to the constructor describe the inputs and outputs of the graph respectively. The following graph has no inputs and outputs a single integer. ```{code-cell} ipython3 -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.models import EmptyModel, TKR -g = GraphBuilder(EmptyModel, TKR[int]) +g = Graph(EmptyModel, TKR[int]) ``` In general a graph can have a multiple inputs and multiple outputs @@ -33,7 +33,7 @@ In order to keep a clear separation between the types used in the Tierkreis grap (The `TKR[A]` wrapper type indicates that an edge in the graph contains a value of type `A`.) ``` -We can add constants to a graph using `GraphBuilder.const`. +We can add constants to a graph using `Graph.const`. ```{code-cell} ipython3 one = g.const(1) @@ -43,7 +43,7 @@ two = g.const(2) The constants will be added into the data structure defining the graph. In particular if the graph is serialized then these constants will be hard-coded into that serialization. -We can add tasks using `GraphBuilder.task`: +We can add tasks using `Graph.task`: ```{code-cell} ipython3 from tierkreis.builtins import iadd @@ -54,15 +54,15 @@ three = g.task(iadd(g.const(1), g.const(2))) In this example we import the type stubs provided by the Tierkreis library for the built-in functions. This allows us to use the [pyright](https://github.com/microsoft/pyright) static analysis tool to check that the input and outputs types of the tasks are what we expect them to be. -To finish, we specify the outputs of the graph. +To finish, we convert the graph into a runnable `Workflow` by specifying the outputs: ```{code-cell} ipython3 -g.outputs(three) +workflow = g.finish_with_outputs(three) ``` ## Running the graph -To run a general Tierkreis graph we need to set up:- +To run a general Tierkreis graph (`Workflow`) we need to set up:- - a way to store and share inputs and outputs (the 'storage' interface) - a way to run tasks (the 'executor' interface) @@ -103,6 +103,6 @@ With the storage and executor specified we can now run a graph using `run_graph` from tierkreis import run_graph from tierkreis.storage import read_outputs -run_graph(storage, executor, g.get_data(), {}) -print(read_outputs(g, storage)) +run_graph(storage, executor, workflow, {}) +print(read_outputs(workflow, storage)) ``` diff --git a/docs/source/graphs/eval.md b/docs/source/graphs/eval.md index b78d83bff..6594f8603 100644 --- a/docs/source/graphs/eval.md +++ b/docs/source/graphs/eval.md @@ -14,13 +14,13 @@ pip install tierkreis ## Graph -We can run graphs from within other graphs by using `GraphBuilder.eval`. +We can run graphs from within other graphs by using `Graph.eval`. Recall the `fib_step` graph that we wrote in the [previous tutorial](./loop.md) ```{code-cell} ipython3 from typing import NamedTuple -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.builtins import iadd from tierkreis.models import TKR @@ -30,9 +30,9 @@ class FibData(NamedTuple): b: TKR[int] -fib_step = GraphBuilder(FibData, FibData) -sum = fib_step.task(iadd(fib_step.inputs.a, fib_step.inputs.b)) -fib_step.outputs(FibData(fib_step.inputs.b, sum)) +builder = Graph(FibData, FibData) +sum = builder.task(iadd(builder.inputs.a, builder.inputs.b)) +fib_step = builder.finish_with_outputs(FibData(builder.inputs.b, sum)) ``` We create a graph `fib4` that calls `fib_step` three times. @@ -41,10 +41,10 @@ The graph will have no inputs and gives a single integer as output: ```{code-cell} ipython3 from tierkreis.models import EmptyModel -fib4 = GraphBuilder(EmptyModel, TKR[int]) +fib4 = Graph(EmptyModel, TKR[int]) ``` -The `GraphBuilder.eval` method takes a `GraphBuilder` object as its first argument +The `Graph.eval` method takes a `Workflow` object as its first argument and the appropriately typed input data as the second object. ```{code-cell} ipython3 @@ -56,7 +56,7 @@ We can iterate manually as follows: ```{code-cell} ipython3 third = fib4.eval(fib_step, second) fourth = fib4.eval(fib_step, third) -fib4.outputs(fourth.b) +workflow = fib4.finish_with_outputs(fourth.b) ``` In the [next tutorial](./loop.md) we will see how to iterate programmatically. @@ -77,6 +77,6 @@ storage = FileStorage(UUID(int=99), name="Nested graphs using Eval") executor = ShellExecutor(Path("."), workflow_dir=storage.workflow_dir) storage.clean_graph_files() -run_graph(storage, executor, fib4.get_data(), {}) -print(read_outputs(fib4, storage)) +run_graph(storage, executor, workflow, {}) +print(read_outputs(workflow, storage)) ``` diff --git a/docs/source/graphs/inputs.md b/docs/source/graphs/inputs.md index 8261193b4..1a8f2d26a 100644 --- a/docs/source/graphs/inputs.md +++ b/docs/source/graphs/inputs.md @@ -19,14 +19,14 @@ pip install tierkreis ### Elementary types Like Python functions, Tierkreis graphs can have input and output arguments. -We use the constructor of `GraphBuilder` to indicate that our function takes a single integer to a single integer: +We use the constructor of `Graph` to indicate that our function takes a single integer to a single integer: ```{code-cell} ipython3 -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.models import TKR # f(x) = 2x + 1 -f = GraphBuilder(TKR[int], TKR[int]) +f = Graph(TKR[int], TKR[int]) ``` The implementation of this graph can be done entirely using Tierkreis built-in functions: @@ -36,7 +36,7 @@ from tierkreis.builtins import iadd, itimes double = f.task(itimes(f.const(2), f.inputs)) f_out = f.task(iadd(double, f.const(1))) -f.outputs(f_out) +f = f.finish_with_outputs(f_out) ``` ### Nested types within a single output @@ -65,8 +65,8 @@ The contents of `A` will not in general be accessible to the graph builder code. ```{code-cell} ipython3 from tierkreis.models import EmptyModel -init_data = GraphBuilder(EmptyModel, TKR[FibDataStruct]) -init_data.outputs(init_data.const(FibDataStruct(a=0, b=1))) +init_data = Graph(EmptyModel, TKR[FibDataStruct]) +init_workflow = init_data.finish_with_outputs(init_data.const(FibDataStruct(a=0, b=1))) ``` ## Multiple inputs and multiple outputs @@ -85,13 +85,13 @@ To use this in the signature of a graph, we pass it directly in. This way Tierkreis will interpret the different attributes of the `NamedTuple` as different inputs/outputs. ```{code-cell} ipython3 -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.builtins import iadd from tierkreis.models import TKR -fib_step = GraphBuilder(FibData, FibData) +fib_step = Graph(FibData, FibData) sum = fib_step.task(iadd(fib_step.inputs.a, fib_step.inputs.b)) -fib_step.outputs(FibData(fib_step.inputs.b, sum)) +fib_step = fib_step.finish_with_outputs(FibData(fib_step.inputs.b, sum)) ``` Note that we are now able to access the contents of `FibData` in the graph builder. @@ -108,7 +108,7 @@ class FibData(NamedTuple): a: int b: int -fib_step_2 = GraphBuilder(TKR[FibData], TKR[FibData]) +fib_step_2 = Graph(TKR[FibData], TKR[FibData]) ``` However we would then not be able to access attributes of `FibData` in the graph builder code. @@ -127,7 +127,7 @@ If some data is only used in workers and can be passed between them without the ## Combinations of single and multiple inputs We can combine the various types of inputs and outputs in the natural way. -For instance the following are all valid ways to construct a `GraphBuilder` object: +For instance the following are all valid ways to construct a `Graph` object: ```{code-cell} ipython3 class MultiPortInputData(NamedTuple): @@ -138,10 +138,10 @@ class MultiPortOutputData(NamedTuple): a: TKR[str] b: TKR[list[int]] -g = GraphBuilder(TKR[int], TKR[str]) -g = GraphBuilder(MultiPortInputData, MultiPortOutputData) -g = GraphBuilder(TKR[str], MultiPortOutputData) -g = GraphBuilder(MultiPortInputData, TKR[str]) +g = Graph(TKR[int], TKR[str]) +g = Graph(MultiPortInputData, MultiPortOutputData) +g = Graph(TKR[str], MultiPortOutputData) +g = Graph(MultiPortInputData, TKR[str]) ``` # Execution @@ -161,14 +161,14 @@ storage = FileStorage(UUID(int=99), name="Graph inputs and outputs") executor = ShellExecutor(Path("."), workflow_dir=storage.workflow_dir) storage.clean_graph_files() -run_graph(storage, executor, f.get_data(), 10) +run_graph(storage, executor, f, 10) print(read_outputs(f, storage)) storage.clean_graph_files() -run_graph(storage, executor, init_data.get_data(), {}) -print(read_outputs(init_data, storage)) +run_graph(storage, executor, init_workflow, {}) +print(read_outputs(init_workflow, storage)) storage.clean_graph_files() -run_graph(storage, executor, fib_step.get_data(), {'a': 0, 'b': 1}) +run_graph(storage, executor, fib_step, {'a': 0, 'b': 1}) print(read_outputs(fib_step, storage)) ``` diff --git a/docs/source/graphs/loop.md b/docs/source/graphs/loop.md index ea2b95ac0..e78ebb4de 100644 --- a/docs/source/graphs/loop.md +++ b/docs/source/graphs/loop.md @@ -6,8 +6,8 @@ kernelspec: # Iteration using Loop -One way to perform iteration in Tierkreis is to use `GraphBuilder.loop`. -The first argument to `GraphBuilder.loop` is a graph that constitutes the loop body. +One way to perform iteration in Tierkreis is to use `Graph.loop`. +The first argument to `Graph.loop` is a graph that constitutes the loop body. The second is the initial input to the loop. ## Graphs @@ -19,7 +19,7 @@ The signature of the loop body is: ```{code-cell} ipython3 from typing import NamedTuple -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.models import TKR class LoopBodyInput(NamedTuple): @@ -33,7 +33,7 @@ class LoopBodyOutput(NamedTuple): value: TKR[int] should_continue: TKR[bool] -loop_body = GraphBuilder(LoopBodyInput, LoopBodyOutput) +loop_body = Graph(LoopBodyInput, LoopBodyOutput) ``` The loop body input type consists of all the variables that we want the loop body to have access to. @@ -52,18 +52,18 @@ from tierkreis.builtins import iadd, igt, rand_int i = loop_body.task(iadd(loop_body.const(1), loop_body.inputs.i)) value = loop_body.task(iadd(loop_body.inputs.step, loop_body.inputs.value)) pred = loop_body.task(igt(loop_body.inputs.bound, i)) -loop_body.outputs(LoopBodyOutput(i=i, value=value, should_continue=pred)) +body1 = loop_body.finish_with_outputs(LoopBodyOutput(i=i, value=value, should_continue=pred)) ``` -The main graph constructs the initial values for the loop and uses `GraphBuilder.loop` to run the loop. +The main graph constructs the initial values for the loop and uses `Graph.loop` to run the loop. ```{code-cell} ipython3 from tierkreis.models import EmptyModel -f = GraphBuilder(EmptyModel, TKR[int]) +f = Graph(EmptyModel, TKR[int]) init = LoopBodyInput(f.const(0), f.const(0), f.const(2), f.const(10)) -loop_output = f.loop(loop_body, init) -f.outputs(loop_output.value) +loop_output = f.loop(body1, init) +f.finish_with_outputs(loop_output.value) ``` ### Tracking Outputs @@ -76,24 +76,24 @@ We simply provide an additional name to the loop. ```{code-cell} ipython3 from tierkreis.models import EmptyModel -f = GraphBuilder(EmptyModel, TKR[int]) +f = Graph(EmptyModel, TKR[int]) init = LoopBodyInput(f.const(0), f.const(0), f.const(2), f.const(10)) -loop_output = f.loop(loop_body, init, "my_loop") -f.outputs(loop_output.value) +loop_output = f.loop(body1, init, "my_loop") +workflow = f.finish_with_outputs(loop_output.value) ``` We can refer to this name after execution. ### Repeat until success -In addition to bounded iteration, the `GraphBuilder.loop` method can also define a 'repeat until success' loop. +In addition to bounded iteration, the `Graph.loop` method can also define a 'repeat until success' loop. First we need to create the graph that constitutes the body of the loop. As a toy example, we choose a random number between 1 and 10 inclusive and continue if it is less than 10. ```{code-cell} ipython3 from typing import NamedTuple -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.builtins import iadd, igt, rand_int from tierkreis.models import TKR @@ -107,11 +107,11 @@ class LoopBodyOutput(NamedTuple): should_continue: TKR[bool] -body = GraphBuilder(LoopBodyInput, LoopBodyOutput) +body = Graph(LoopBodyInput, LoopBodyOutput) i = body.task(iadd(body.const(1), body.inputs.i)) a = body.task(rand_int(body.const(0), body.const(10))) pred = body.task(igt(body.const(10), a)) -body.outputs(LoopBodyOutput(i=i, should_continue=pred)) +body = body.finish_with_outputs(LoopBodyOutput(i=i, should_continue=pred)) ``` The main graph runs the loop and tells us the iteration on which we found success. @@ -119,9 +119,9 @@ The main graph runs the loop and tells us the iteration on which we found succes ```{code-cell} ipython3 from tierkreis.models import EmptyModel -g = GraphBuilder(EmptyModel, TKR[int]) +g = Graph(EmptyModel, TKR[int]) loop_output = g.loop(body, LoopBodyInput(g.const(0))) -g.outputs(loop_output.i) +rus_workflow = g.finish_with_outputs(loop_output.i) ``` # Execution @@ -140,11 +140,11 @@ storage = FileStorage(UUID(int=99), name="Nested graphs using Eval") executor = ShellExecutor(Path("."), workflow_dir=storage.workflow_dir) storage.clean_graph_files() -run_graph(storage, executor, f.get_data(), {}) -print(read_outputs(f, storage)) -print(read_loop_trace(f, storage, "my_loop")) +run_graph(storage, executor, workflow, {}) +print(read_outputs(workflow, storage)) +print(read_loop_trace(workflow, storage, "my_loop")) storage.clean_graph_files() -run_graph(storage, executor, g.get_data(), {}) -print(read_outputs(g, storage)) +run_graph(storage, executor, rus_workflow, {}) +print(read_outputs(rus_workflow, storage)) ``` diff --git a/docs/source/graphs/map.md b/docs/source/graphs/map.md index f870845df..1fda07689 100644 --- a/docs/source/graphs/map.md +++ b/docs/source/graphs/map.md @@ -33,24 +33,23 @@ The stub files will provide us with type hints in the graph building process lat We can import this stub file to help create our graph. The graph builder manipulates references to values, not the values themselves. -(The one exception to this rule is when we add a constant value to a graph using `GraphBuilder.const`. Then the actual value is added to the graph definition and `GraphBuilder.const` returns a reference to this value.) +(The one exception to this rule is when we add a constant value to a graph using `Graph.const`. Then the actual value is added to the graph definition and `Graph.const` returns a reference to this value.) The references are type checked using the `TKR` type. I.e. a reference to an `int` has the type `TKR[int]`. ```{code-cell} from typing import NamedTuple from tierkreis.models import EmptyModel, TKR -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.builtins import mean from auth_worker import encrypt, EncryptionResult def map_body(): - g = GraphBuilder(TKR[str], EncryptionResult) + g = Graph(TKR[str], EncryptionResult) result = g.task(encrypt(plaintext=g.inputs, work_factor=g.const(2**14))) - g.outputs(result) - return g + return g.finish_with_outputs(result) class GraphOutputs(NamedTuple): @@ -59,7 +58,7 @@ class GraphOutputs(NamedTuple): def graph(): - g = GraphBuilder(EmptyModel, GraphOutputs) + g = Graph(EmptyModel, GraphOutputs) plaintexts = g.const([f"plaintext+{n}" for n in range(20)]) results = g.map(map_body(), plaintexts) @@ -69,8 +68,7 @@ def graph(): av = g.task(mean(values=times)) out = GraphOutputs(ciphertexts=ciphertexts, average_time_taken=av) - g.outputs(out) - return g + return g.finish_with_outputs(out) ``` ## Running the graph @@ -97,7 +95,7 @@ executor = UvExecutor( registry_path=Path("../examples/example_workers"), logs_path=storage.logs_path ) start = time.time() -run_graph(storage, executor, graph().data, {}) +run_graph(storage, executor, graph(), {}) total_time = time.time() - start outputs = read_outputs(graph(), storage) diff --git a/docs/source/logging_and_errors.md b/docs/source/logging_and_errors.md index 877c347e5..e5e7f251c 100644 --- a/docs/source/logging_and_errors.md +++ b/docs/source/logging_and_errors.md @@ -29,17 +29,17 @@ The goal of such errors is to catch errors before potentially running an expensi For example providing an incorrect type to the output ```python -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.controller.data.models import TKR -g = GraphBuilder(TKR[int], TKR[list[int]]) -g.outputs(g.inputs) +g = Graph(TKR[int], TKR[list[int]]) +g.finish_with_outputs(g.inputs) ``` -will provide an error message as an integer cant be converted to a list: +will provide an error message as an integer can't be converted to a list: ``` -Argument of type "TKR[int]" cannot be assigned to parameter "outputs" of type "TKR[list[int]]" in function "outputs" +Argument of type "TKR[int]" cannot be assigned to parameter "outputs" of type "TKR[list[int]]" in function "finish_with_outputs"   "TKR[int]" is not assignable to "TKR[list[int]]"     Type parameter "T@TKR" is covariant, but "int" is not a subtype of "list[int]"       "int" is not assignable to "list[int]" diff --git a/docs/source/visualization.md b/docs/source/visualization.md index b8e101fbe..84cf9f174 100644 --- a/docs/source/visualization.md +++ b/docs/source/visualization.md @@ -48,10 +48,10 @@ If you delete the optional 'extra' output on line 62 then the browser will live Programmatically this is available through: ```py -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis_visualization.visualize_graph import visualize_graph -graph = GraphBuilder() +graph = Graph() visualize_graph(graph) ``` diff --git a/docs/source/worker/hello_world.md b/docs/source/worker/hello_world.md index 57671b94e..e605fc263 100644 --- a/docs/source/worker/hello_world.md +++ b/docs/source/worker/hello_world.md @@ -71,16 +71,16 @@ Since this worker uses the Tierkreis Python library, we can automatically genera ## Graph creation -Now can we import the `greet` function from the stubs file and use it in `GraphBuilder.task`. +Now can we import the `greet` function from the stubs file and use it in `Graph.task`. ```{code-cell} ipython3 -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.models import TKR from hello_stubs import greet -g = GraphBuilder(TKR[str], TKR[str]) +g = Graph(TKR[str], TKR[str]) output = g.task(greet(greeting=g.const("hello "), subject=g.inputs)) -g.outputs(output) +workflow = g.finish_with_outputs(output) ``` ## Execution @@ -103,6 +103,6 @@ storage = FileStorage(UUID(int=99), "hello_world_tutorial", do_cleanup=True) executor = UvExecutor( registry_path=Path("../examples/example_workers"), logs_path=storage.logs_path ) -run_graph(storage, executor, g.data, "world!") -read_outputs(g, storage) +run_graph(storage, executor, workflow, "world!") +read_outputs(workflow, storage) ``` diff --git a/docs/source/worker/native_workers/aer_worker.md b/docs/source/worker/native_workers/aer_worker.md index b72ea8684..8af404117 100644 --- a/docs/source/worker/native_workers/aer_worker.md +++ b/docs/source/worker/native_workers/aer_worker.md @@ -29,7 +29,7 @@ The full api is available in the {py:mod}`API Docs `. The Tierkreis Python package provides a few prepackaged graphs to make it easier to compile and run circuits with Aer. `tierkreis.graphs.simulate.compile_simulate.compile_simulate` is intended for the common use case of compiling a list of circuits in parallel and then running them in parallel. -It can be included within a custom graph using `GraphBuilder.eval` or run as a standalone graph. +It can be included within a custom graph using `Graph.eval` or run as a standalone graph. An example use is in `docs/source/examples/parallelism.ipynb` in the [Tierkreis repo](https://github.com/Quantinuum/tierkreis), which looks like: diff --git a/docs/source/worker/native_workers/ibmq_worker.md b/docs/source/worker/native_workers/ibmq_worker.md index 58ad3e8f9..20a5aa691 100644 --- a/docs/source/worker/native_workers/ibmq_worker.md +++ b/docs/source/worker/native_workers/ibmq_worker.md @@ -43,7 +43,7 @@ class IBMQInput(NamedTuple): def compile_run_single(): - g = GraphBuilder( + g = Graph( IBMQInput, TKR[OpaqueType["pytket.backends.backendresult.BackendResult"]] ) @@ -55,8 +55,7 @@ def compile_run_single(): ) ) res = g.task(run_circuit(compiled_circuit, g.inputs.n_shots. g.inputs.backend)) - g.outputs(res) - return g + return g.finish_with_outputs(res) circuit = ...your circuit here... diff --git a/docs/source/worker/native_workers/nexus_worker.md b/docs/source/worker/native_workers/nexus_worker.md index 8a6eef35c..046cd997e 100644 --- a/docs/source/worker/native_workers/nexus_worker.md +++ b/docs/source/worker/native_workers/nexus_worker.md @@ -38,7 +38,7 @@ The full api is available in the {py:mod}`API Docs `. The Tierkreis Python package provides a couple of prepackaged graphs to make it easier to interact with the Nexus API. `tierkreis.graphs.nexus.submit_poll.nexus_submit_and_poll` is intended to automate the whole process of (parallelised) circuit upload, submission, status polling and result retrieval. -It can be included within a custom graph using `GraphBuilder.eval` or run as a standalone graph. +It can be included within a custom graph using `Graph.eval` or run as a standalone graph. The function `nexus_submit_and_poll` takes an optional argument to specify the minimum delay between successive polls in the status polling loop. The default is to poll every `30` seconds. diff --git a/docs/source/worker/native_workers/pytket_worker.md b/docs/source/worker/native_workers/pytket_worker.md index f64ffe04b..34188c281 100644 --- a/docs/source/worker/native_workers/pytket_worker.md +++ b/docs/source/worker/native_workers/pytket_worker.md @@ -66,7 +66,7 @@ class IBMInput(NamedTuple): def compile_run_single(): - g = GraphBuilder( + g = Graph( IBMInput, TKR[OpaqueType["pytket.backends.backendresult.BackendResult"]] ) @@ -78,8 +78,7 @@ def compile_run_single(): ) ) res = g.task(submit_single(compiled_circuit, g.inputs.n_shots)) - g.outputs(res) - return g + return g.finish_with_outputs(res) circuit = ...your circuit here... diff --git a/docs/source/worker/native_workers/quantinuum_worker.md b/docs/source/worker/native_workers/quantinuum_worker.md index d2487f922..83c712219 100644 --- a/docs/source/worker/native_workers/quantinuum_worker.md +++ b/docs/source/worker/native_workers/quantinuum_worker.md @@ -44,7 +44,7 @@ class QuantinuumInput(NamedTuple): def compile_run_single(): - g = GraphBuilder( + g = Graph( QuantinuumInput, TKR[OpaqueType["pytket.backends.backendresult.BackendResult"]] ) @@ -56,8 +56,7 @@ def compile_run_single(): ) ) res = g.task(run_circuit(compiled_circuit, g.inputs.n_shots, g.inputs.backend)) - g.outputs(res) - return g + return g.finish_with_outputs(res) circuit = ...your circuit here... diff --git a/docs/source/worker/native_workers/qulacs_worker.md b/docs/source/worker/native_workers/qulacs_worker.md index e924f0656..1f23b0be6 100644 --- a/docs/source/worker/native_workers/qulacs_worker.md +++ b/docs/source/worker/native_workers/qulacs_worker.md @@ -28,7 +28,7 @@ The full api is available in the {py:mod}`API Docs `. The Tierkreis Python package provides a few prepackaged graphs to make it easier to compile and run circuits with Qulacs. `tierkreis.graphs.simulate.compile_simulate.compile_simulate` is intended for the common use case of compiling a list of circuits in parallel and then running them in parallel. -It can be included within a custom graph using `GraphBuilder.eval` or run as a standalone graph. +It can be included within a custom graph using `Graph.eval` or run as a standalone graph. An example use is in `docs/source/examples/parallelism.ipynb` in the [Tierkreis repo](https://github.com/Quantinuum/tierkreis), which looks like: diff --git a/tierkreis/tests/controller/defaults_graphs.py b/tierkreis/tests/controller/defaults_graphs.py index a13b237a6..50996a17b 100644 --- a/tierkreis/tests/controller/defaults_graphs.py +++ b/tierkreis/tests/controller/defaults_graphs.py @@ -1,6 +1,6 @@ from typing import NamedTuple -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph, Workflow from tierkreis.builtins import tkr_range from tierkreis.controller.data.models import TKR @@ -18,42 +18,39 @@ class OuterOutputs(NamedTuple): extra_output: TKR[int] | None = None -def omit_input() -> GraphBuilder[Inputs, TKR[list[int]]]: - g = GraphBuilder(Inputs, TKR[list[int]]) +def omit_input() -> Workflow[Inputs, TKR[list[int]]]: + g = Graph(Inputs, TKR[list[int]]) range_1 = g.task(tkr_range(g.inputs.start, g.inputs.stop)) - g.outputs(range_1) - return g + return g.finish_with_outputs(range_1) -def passthru() -> GraphBuilder[Inputs, TKR[list[int]]]: - g = GraphBuilder(Inputs, TKR[list[int]]) +def passthru() -> Workflow[Inputs, TKR[list[int]]]: + g = Graph(Inputs, TKR[list[int]]) range_1 = g.task(tkr_range(g.inputs.start, g.inputs.stop, g.inputs.step)) - g.outputs(range_1) - return g + return g.finish_with_outputs(range_1) -def defaults_omit() -> GraphBuilder[Inputs, OuterOutputs]: - g = GraphBuilder(Inputs, OuterOutputs) +def defaults_omit() -> Workflow[Inputs, OuterOutputs]: + g = Graph(Inputs, OuterOutputs) range_1 = g.eval(omit_input(), Inputs(g.inputs.start, g.inputs.stop)) range_2 = g.eval(passthru(), Inputs(g.inputs.start, g.inputs.stop)) range_3 = g.eval(passthru(), Inputs(g.inputs.start, g.inputs.stop, g.const(2))) - g.outputs(OuterOutputs(range_1, range_2, range_3)) - return g + return g.finish_with_outputs(OuterOutputs(range_1, range_2, range_3)) -def defaults_passthru() -> GraphBuilder[Inputs, OuterOutputs]: - g = GraphBuilder(Inputs, OuterOutputs) +def defaults_passthru() -> Workflow[Inputs, OuterOutputs]: + g = Graph(Inputs, OuterOutputs) range_1 = g.eval(omit_input(), Inputs(g.inputs.start, g.inputs.stop)) range_2 = g.eval(passthru(), Inputs(g.inputs.start, g.inputs.stop)) range_3 = g.eval(passthru(), Inputs(g.inputs.start, g.inputs.stop, g.const(2))) - g.outputs(OuterOutputs(range_1, range_2, range_3, g.inputs.step)) - return g + return g.finish_with_outputs(OuterOutputs(range_1, range_2, range_3, g.inputs.step)) -def defaults_not_none() -> GraphBuilder[Inputs, OuterOutputs]: - g = GraphBuilder(Inputs, OuterOutputs) +def defaults_not_none() -> Workflow[Inputs, OuterOutputs]: + g = Graph(Inputs, OuterOutputs) range_1 = g.eval(omit_input(), Inputs(g.inputs.start, g.inputs.stop)) range_2 = g.eval(passthru(), Inputs(g.inputs.start, g.inputs.stop)) range_3 = g.eval(passthru(), Inputs(g.inputs.start, g.inputs.stop, g.const(2))) - g.outputs(OuterOutputs(range_1, range_2, range_3, g.inputs.start)) - return g + return g.finish_with_outputs( + OuterOutputs(range_1, range_2, range_3, g.inputs.start) + ) diff --git a/tierkreis/tests/controller/loop_graphdata.py b/tierkreis/tests/controller/loop_graphdata.py index 426a87ea1..5bead49cc 100644 --- a/tierkreis/tests/controller/loop_graphdata.py +++ b/tierkreis/tests/controller/loop_graphdata.py @@ -1,7 +1,7 @@ from typing import NamedTuple import tierkreis.builtins.stubs as tkr_builtins -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph, Workflow from tierkreis.controller.data.core import EmptyModel from tierkreis.controller.data.graph import GraphData from tierkreis.models import TKR @@ -71,8 +71,8 @@ class MultipleAccOut(NamedTuple): acc3: TKR[int] -def _loop_body_multiple_acc() -> GraphBuilder[MultipleAcc, MultipleAccOut]: - g = GraphBuilder(MultipleAcc, MultipleAccOut) +def _loop_body_multiple_acc() -> Workflow[MultipleAcc, MultipleAccOut]: + g = Graph(MultipleAcc, MultipleAccOut) acc = g.inputs.acc1 acc2 = g.inputs.acc2 @@ -89,7 +89,7 @@ def _loop_body_multiple_acc() -> GraphBuilder[MultipleAcc, MultipleAccOut]: new_acc2 = g.task(tkr_builtins.iadd(a=acc2, b=two)) new_acc3 = g.task(tkr_builtins.iadd(a=acc3, b=three)) - g.outputs( + return g.finish_with_outputs( MultipleAccOut( should_continue=should_continue, acc1=new_acc, @@ -98,8 +98,6 @@ def _loop_body_multiple_acc() -> GraphBuilder[MultipleAcc, MultipleAccOut]: ), ) - return g - class LoopMultipleAccOut(NamedTuple): acc1: TKR[int] @@ -107,8 +105,8 @@ class LoopMultipleAccOut(NamedTuple): acc3: TKR[int] -def loop_multiple_acc() -> GraphBuilder[EmptyModel, LoopMultipleAccOut]: - g = GraphBuilder(EmptyModel, LoopMultipleAccOut) +def loop_multiple_acc() -> Workflow[EmptyModel, LoopMultipleAccOut]: + g = Graph(EmptyModel, LoopMultipleAccOut) acc1 = g.const(0) acc2 = g.const(0) @@ -117,9 +115,9 @@ def loop_multiple_acc() -> GraphBuilder[EmptyModel, LoopMultipleAccOut]: body = _loop_body_multiple_acc() loop = g.loop(body, MultipleAcc(acc1, acc2, acc3), "my_loop") - g.outputs(LoopMultipleAccOut(acc1=loop.acc1, acc2=loop.acc2, acc3=loop.acc3)) - - return g + return g.finish_with_outputs( + LoopMultipleAccOut(acc1=loop.acc1, acc2=loop.acc2, acc3=loop.acc3) + ) class Scoping(NamedTuple): @@ -132,25 +130,25 @@ class ScopingOut(NamedTuple): current: TKR[int] -def _loop_body_scoping() -> GraphBuilder[Scoping, ScopingOut]: - g = GraphBuilder(Scoping, ScopingOut) +def _loop_body_scoping() -> Workflow[Scoping, ScopingOut]: + g = Graph(Scoping, ScopingOut) one = g.const(1) next_val = g.task(tkr_builtins.iadd(g.inputs.current, one)) should_continue = g.task(tkr_builtins.neq(g.inputs.end, g.inputs.current)) - g.outputs(ScopingOut(should_continue=should_continue, current=next_val)) - - return g + return g.finish_with_outputs( + ScopingOut(should_continue=should_continue, current=next_val) + ) class LoopScopingOut(NamedTuple): result: TKR[int] -def loop_scoping() -> GraphBuilder[EmptyModel, LoopScopingOut]: - g = GraphBuilder(EmptyModel, LoopScopingOut) +def loop_scoping() -> Workflow[EmptyModel, LoopScopingOut]: + g = Graph(EmptyModel, LoopScopingOut) start = g.const(0) end = g.const(10) @@ -158,6 +156,4 @@ def loop_scoping() -> GraphBuilder[EmptyModel, LoopScopingOut]: body = _loop_body_scoping() loop = g.loop(body, Scoping(start, end), "scoped_loop") - g.outputs(LoopScopingOut(result=loop.current)) - - return g + return g.finish_with_outputs(LoopScopingOut(result=loop.current)) diff --git a/tierkreis/tests/controller/test_error.py b/tierkreis/tests/controller/test_error.py index ad5ea9114..e0b351fb6 100644 --- a/tierkreis/tests/controller/test_error.py +++ b/tierkreis/tests/controller/test_error.py @@ -4,7 +4,7 @@ import pytest from tests.workers.failing_worker.stubs import exit_code_1, fail, wont_fail -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph, Workflow from tierkreis.controller import run_graph from tierkreis.controller.data.core import EmptyModel from tierkreis.controller.data.location import Loc @@ -16,28 +16,24 @@ WORKER_PATH = Path(__file__).parent.parent / "workers" -def will_fail_graph() -> GraphBuilder[EmptyModel, TKR[int]]: - graph = GraphBuilder(EmptyModel, TKR[int]) - graph.outputs(graph.task(fail())) - return graph +def will_fail_graph() -> Workflow[EmptyModel, TKR[int]]: + graph = Graph(EmptyModel, TKR[int]) + return graph.finish_with_outputs(graph.task(fail())) -def wont_fail_graph() -> GraphBuilder[EmptyModel, TKR[int]]: - graph = GraphBuilder(EmptyModel, TKR[int]) - graph.outputs(graph.task(wont_fail())) - return graph +def wont_fail_graph() -> Workflow[EmptyModel, TKR[int]]: + graph = Graph(EmptyModel, TKR[int]) + return graph.finish_with_outputs(graph.task(wont_fail())) -def fail_in_eval() -> GraphBuilder[EmptyModel, TKR[int]]: - graph = GraphBuilder(EmptyModel, TKR[int]) - graph.outputs(graph.eval(will_fail_graph(), EmptyModel())) - return graph +def fail_in_eval() -> Workflow[EmptyModel, TKR[int]]: + graph = Graph(EmptyModel, TKR[int]) + return graph.finish_with_outputs(graph.eval(will_fail_graph(), EmptyModel())) -def non_zero_exit_code() -> GraphBuilder[EmptyModel, TKR[int]]: - graph = GraphBuilder(EmptyModel, TKR[int]) - graph.outputs(graph.task(exit_code_1())) - return graph +def non_zero_exit_code() -> Workflow[EmptyModel, TKR[int]]: + graph = Graph(EmptyModel, TKR[int]) + return graph.finish_with_outputs(graph.task(exit_code_1())) def test_raise_error() -> None: @@ -46,7 +42,7 @@ def test_raise_error() -> None: executor = UvExecutor(WORKER_PATH, logs_path=storage.logs_path) storage.clean_graph_files() with pytest.raises(TierkreisError): - run_graph(storage, executor, g.get_data(), {}, n_iterations=1000) + run_graph(storage, executor, g.data, {}, n_iterations=1000) assert storage.node_has_error(Loc("-.N0")) @@ -55,7 +51,7 @@ def test_raises_no_error() -> None: storage = ControllerFileStorage(UUID(int=43), name="wont_fail") executor = UvExecutor(WORKER_PATH, logs_path=storage.logs_path) storage.clean_graph_files() - run_graph(storage, executor, g.get_data(), {}, n_iterations=100) + run_graph(storage, executor, g.data, {}, n_iterations=100) assert not storage.node_has_error(Loc("-.N0")) @@ -65,7 +61,7 @@ def test_nested_error() -> None: executor = UvExecutor(WORKER_PATH, logs_path=storage.logs_path) storage.clean_graph_files() with pytest.raises(TierkreisError): - run_graph(storage, executor, g.get_data(), {}, n_iterations=1000) + run_graph(storage, executor, g.data, {}, n_iterations=1000) assert (storage.logs_path.parent / "-/_error").exists() @@ -75,5 +71,5 @@ def test_non_zero_exit_code() -> None: executor = UvExecutor(WORKER_PATH, logs_path=storage.logs_path) storage.clean_graph_files() with pytest.raises(TierkreisError): - run_graph(storage, executor, g.get_data(), {}, n_iterations=1000) + run_graph(storage, executor, g.data, {}, n_iterations=1000) assert (storage.logs_path.parent / "-/_error").exists() diff --git a/tierkreis/tests/controller/test_read_loop_trace.py b/tierkreis/tests/controller/test_read_loop_trace.py index c38bb7584..bf8abc60b 100644 --- a/tierkreis/tests/controller/test_read_loop_trace.py +++ b/tierkreis/tests/controller/test_read_loop_trace.py @@ -25,7 +25,7 @@ 9, ), ( - loop_multiple_acc().get_data(), + loop_multiple_acc().data, return_value, "multi_acc", 9, diff --git a/tierkreis/tests/controller/test_restart.py b/tierkreis/tests/controller/test_restart.py index 49d3bc40b..7dfa47363 100644 --- a/tierkreis/tests/controller/test_restart.py +++ b/tierkreis/tests/controller/test_restart.py @@ -2,7 +2,7 @@ from uuid import UUID from tierkreis import run_graph -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.builtins import iadd, itimes from tierkreis.controller.data.location import Loc from tierkreis.controller.data.models import TKR @@ -15,12 +15,12 @@ def test_restart() -> None: storage.clean_graph_files() executor = UvExecutor(Path(__file__).parent.parent / "workers", storage.logs_path) - g = GraphBuilder(TKR[int], TKR[int]) + g = Graph(TKR[int], TKR[int]) plus_one = g.task(iadd(g.const(1), g.inputs)) left = g.task(iadd(g.const(1), plus_one)) right = g.task(itimes(g.const(2), plus_one)) out = g.task(iadd(left, right)) - g.outputs(out) + g = g.finish_with_outputs(out) run_graph(storage, executor, g, {"value": 0}) @@ -32,7 +32,7 @@ def test_restart() -> None: Loc().N(left.node_index), Loc().N(right.node_index), Loc().N(out.node_index), - Loc().N(g.get_data().output_idx()), + Loc().N(g.data.output_idx()), ] assert sorted(invalidated) == sorted(expected) diff --git a/tierkreis/tests/controller/test_resume.py b/tierkreis/tests/controller/test_resume.py index af04fe824..1be3ad1f6 100644 --- a/tierkreis/tests/controller/test_resume.py +++ b/tierkreis/tests/controller/test_resume.py @@ -38,7 +38,7 @@ typed_map_simple, embed_graph, ) -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Workflow from tierkreis.controller import run_graph from tierkreis.controller.data.graph import GraphData from tierkreis.controller.data.location import Loc @@ -52,7 +52,7 @@ param_data: list[ tuple[ - GraphData | GraphBuilder, + GraphData | Workflow, dict[str, PType] | PType, str, dict[str, PType] | PType, @@ -135,9 +135,7 @@ ), (embed_graph(), {"s1": "1", "s2": "4", "final": 2}, "embed_graph", 1), ] -params: list[ - tuple[GraphData | GraphBuilder, Any, str, int, dict[str, PType] | PType] -] = [ +params: list[tuple[GraphData | Workflow, Any, str, int, dict[str, PType] | PType]] = [ (graph, output, name, i + 1, inputs) for i, (graph, output, name, inputs) in enumerate(param_data) ] @@ -213,7 +211,7 @@ def test_resume( # noqa: PLR0913 with_worker_param_data: list[ - tuple[GraphData | GraphBuilder, Any, str, dict[str, PType] | PType] + tuple[GraphData | Workflow, Any, str, dict[str, PType] | PType] ] = [ (eval_body_is_from_worker(), 21, "eval_body_is_from_worker", {"value": 10}), (eval_graph_of_graph(), 31, "eval_graph_of_graph", {"value": 3}), @@ -226,7 +224,7 @@ def test_resume( # noqa: PLR0913 ] with_worker_params: list[ - tuple[GraphData | GraphBuilder, Any, str, int, dict[str, PType] | PType] + tuple[GraphData | Workflow, Any, str, int, dict[str, PType] | PType] ] = [ (graph, output, name, i + 1, inputs) for i, (graph, output, name, inputs) in enumerate(with_worker_param_data) diff --git a/tierkreis/tests/controller/test_types.py b/tierkreis/tests/controller/test_types.py index a0c3f851a..9c5bb0318 100644 --- a/tierkreis/tests/controller/test_types.py +++ b/tierkreis/tests/controller/test_types.py @@ -13,6 +13,9 @@ is_ptype, ptype_from_bytes, ) +from tierkreis.controller.data.graph import GraphData +from tierkreis.controller.data.models import TKR, Workflow +from tests.controller.typed_graphdata import typed_doubler class UntupledModel[U, V](BaseModel): @@ -65,6 +68,8 @@ def from_list(cls, args: list) -> "DummyListConvertible": type_list.append(tuple[int, str]) type_list.append(tuple[int | str]) type_list.append(Mapping[str, dict[str, bytes]]) +type_list.append(GraphData) +type_list.append(Workflow[TKR[int], TKR[bool]]) fail_list: Sequence[type] = [] fail_list.append(UUID) @@ -106,6 +111,7 @@ def test_bytes_roundtrip(ptype: PType) -> None: annotated_ptypes: Sequence[tuple[PType, type]] = [ (ptype, type(ptype)) for ptype in ptypes ] + [ # Not possible to deserialise without annotations + (typed_doubler(), Workflow[TKR[int], TKR[int]]), ( [DummyListConvertible(a=1), DummyListConvertible(a=2)], list[DummyListConvertible], diff --git a/tierkreis/tests/controller/typed_graphdata.py b/tierkreis/tests/controller/typed_graphdata.py index 669fadad3..cf86f2848 100644 --- a/tierkreis/tests/controller/typed_graphdata.py +++ b/tierkreis/tests/controller/typed_graphdata.py @@ -1,7 +1,12 @@ from typing import NamedTuple -from tests.workers.graph.stubs import doubler_plus_graph, graph_of_graph, apply_twice -from tierkreis.builder import GraphBuilder, TypedGraphRef +from tests.workers.graph.stubs import ( + doubler_plus_graph, + graph_of_graph, + apply_twice, + ApplyTwiceInput, +) +from tierkreis.builder import Graph, TypedGraphRef, Workflow from tierkreis.builtins import ( conjugate, eq, @@ -15,7 +20,7 @@ tkr_str, ) from tierkreis.controller.data.core import EmptyModel -from tierkreis.controller.data.models import TKR, OpaqueType +from tierkreis.controller.data.models import TKR class DoublerInput(NamedTuple): @@ -28,38 +33,34 @@ class DoublerOutput(NamedTuple): value: TKR[int] -def typed_doubler() -> GraphBuilder[TKR[int], TKR[int]]: - g = GraphBuilder(TKR[int], TKR[int]) +def typed_doubler() -> Workflow[TKR[int], TKR[int]]: + g = Graph(TKR[int], TKR[int]) out = g.task(itimes(a=g.const(2), b=g.inputs)) - g.outputs(out) - return g + return g.finish_with_outputs(out) -def typed_doubler_plus_multi() -> GraphBuilder[DoublerInput, DoublerOutput]: - g = GraphBuilder(DoublerInput, DoublerOutput) +def typed_doubler_plus_multi() -> Workflow[DoublerInput, DoublerOutput]: + g = Graph(DoublerInput, DoublerOutput) mul = g.task(itimes(a=g.inputs.x, b=g.const(2))) out = g.task(iadd(a=mul, b=g.inputs.intercept)) - g.outputs(DoublerOutput(a=g.inputs.x, value=out)) - return g + return g.finish_with_outputs(DoublerOutput(a=g.inputs.x, value=out)) -def typed_doubler_plus() -> GraphBuilder[DoublerInput, TKR[int]]: - g = GraphBuilder(DoublerInput, TKR[int]) +def typed_doubler_plus() -> Workflow[DoublerInput, TKR[int]]: + g = Graph(DoublerInput, TKR[int]) mul = g.task(itimes(a=g.inputs.x, b=g.const(2))) out = g.task(iadd(a=mul, b=g.inputs.intercept)) - g.outputs(out) - return g + return g.finish_with_outputs(out) class TypedEvalOutputs(NamedTuple): typed_eval_output: TKR[int] -def typed_eval() -> GraphBuilder[EmptyModel, TypedEvalOutputs]: - g = GraphBuilder(EmptyModel, TypedEvalOutputs) +def typed_eval() -> Workflow[EmptyModel, TypedEvalOutputs]: + g = Graph(EmptyModel, TypedEvalOutputs) e = g.eval(typed_doubler_plus(), DoublerInput(x=g.const(6), intercept=g.const(0))) - g.outputs(TypedEvalOutputs(typed_eval_output=e)) - return g + return g.finish_with_outputs(TypedEvalOutputs(typed_eval_output=e)) class LoopBodyInput(NamedTuple): @@ -71,61 +72,54 @@ class LoopBodyOutput(NamedTuple): should_continue: TKR[bool] -def loop_body() -> GraphBuilder[LoopBodyInput, LoopBodyOutput]: - g = GraphBuilder(LoopBodyInput, LoopBodyOutput) +def loop_body() -> Workflow[LoopBodyInput, LoopBodyOutput]: + g = Graph(LoopBodyInput, LoopBodyOutput) a_plus = g.task(iadd(a=g.inputs.loop_acc, b=g.const(1))) pred = g.task(igt(a=g.const(10), b=a_plus)) - g.outputs(LoopBodyOutput(loop_acc=a_plus, should_continue=pred)) - return g + return g.finish_with_outputs(LoopBodyOutput(loop_acc=a_plus, should_continue=pred)) -def typed_loop() -> GraphBuilder[EmptyModel, TKR[int]]: - g = GraphBuilder(EmptyModel, TKR[int]) +def typed_loop() -> Workflow[EmptyModel, TKR[int]]: + g = Graph(EmptyModel, TKR[int]) loop = g.loop(loop_body(), LoopBodyInput(loop_acc=g.const(6))) - g.outputs(loop.loop_acc) - return g + return g.finish_with_outputs(loop.loop_acc) -def typed_map_simple() -> GraphBuilder[TKR[list[int]], TKR[list[int]]]: - g = GraphBuilder(TKR[list[int]], TKR[list[int]]) +def typed_map_simple() -> Workflow[TKR[list[int]], TKR[list[int]]]: + g = Graph(TKR[list[int]], TKR[list[int]]) m = g.map(typed_doubler(), g.inputs) - g.outputs(m) - return g + return g.finish_with_outputs(m) -def typed_map() -> GraphBuilder[TKR[list[int]], TKR[list[int]]]: - g = GraphBuilder(TKR[list[int]], TKR[list[int]]) +def typed_map() -> Workflow[TKR[list[int]], TKR[list[int]]]: + g = Graph(TKR[list[int]], TKR[list[int]]) ins = g.map(lambda n: DoublerInput(x=n, intercept=g.const(6)), g.inputs) m = g.map(typed_doubler_plus(), ins) - g.outputs(m) - return g + return g.finish_with_outputs(m) -def typed_destructuring() -> GraphBuilder[TKR[list[int]], TKR[list[int]]]: - g = GraphBuilder(TKR[list[int]], TKR[list[int]]) +def typed_destructuring() -> Workflow[TKR[list[int]], TKR[list[int]]]: + g = Graph(TKR[list[int]], TKR[list[int]]) ins = g.map(lambda n: DoublerInput(x=n, intercept=g.const(6)), g.inputs) m = g.map(typed_doubler_plus_multi(), ins) mout = g.map(lambda x: x.value, m) - g.outputs(mout) - return g + return g.finish_with_outputs(mout) -def tuple_untuple() -> GraphBuilder[EmptyModel, TKR[int]]: - g = GraphBuilder(EmptyModel, TKR[int]) +def tuple_untuple() -> Workflow[EmptyModel, TKR[int]]: + g = Graph(EmptyModel, TKR[int]) t = g.task(tkr_tuple(g.const(1), g.const(2))) ut = g.task(untuple(t)) - g.outputs(g.task(iadd(ut.a, ut.b))) - return g + return g.finish_with_outputs(g.task(iadd(ut.a, ut.b))) -def factorial() -> GraphBuilder[TKR[int], TKR[int]]: - g = GraphBuilder(TKR[int], TKR[int]) +def factorial() -> Workflow[TKR[int], TKR[int]]: + g = Graph(TKR[int], TKR[int]) pred = g.task(igt(g.inputs, g.const(1))) n_minus_one = g.task(iadd(g.const(-1), g.inputs)) rec = g.eval(g.ref(), n_minus_one) out = g.ifelse(pred, g.task(itimes(g.inputs, rec)), g.const(1)) - g.outputs(out) - return g + return g.finish_with_outputs(out) class GCDInput(NamedTuple): @@ -133,71 +127,52 @@ class GCDInput(NamedTuple): b: TKR[int] -def gcd() -> GraphBuilder[GCDInput, TKR[int]]: - g = GraphBuilder(GCDInput, TKR[int]) +def gcd() -> Workflow[GCDInput, TKR[int]]: + g = Graph(GCDInput, TKR[int]) pred = g.task(igt(g.inputs.b, g.const(0))) a_mod_b = g.task(mod(g.inputs.a, g.inputs.b)) rec = g.eval(g.ref(), GCDInput(a=g.inputs.b, b=a_mod_b)) - g.outputs(g.ifelse(pred, rec, g.inputs.a)) - return g + return g.finish_with_outputs(g.ifelse(pred, rec, g.inputs.a)) -def tkr_conj() -> GraphBuilder[TKR[complex], TKR[complex]]: - g = GraphBuilder(TKR[complex], TKR[complex]) +def tkr_conj() -> Workflow[TKR[complex], TKR[complex]]: + g = Graph(TKR[complex], TKR[complex]) z = g.task(conjugate(g.inputs)) - g.outputs(z) - return g + return g.finish_with_outputs(z) -def tkr_list_conj() -> GraphBuilder[TKR[list[complex]], TKR[list[complex]]]: - g = GraphBuilder(TKR[list[complex]], TKR[list[complex]]) +def tkr_list_conj() -> Workflow[TKR[list[complex]], TKR[list[complex]]]: + g = Graph(TKR[list[complex]], TKR[list[complex]]) zs = g.map(tkr_conj(), g.inputs) - g.outputs(zs) - return g - + return g.finish_with_outputs(zs) -def eval_body_is_from_worker() -> GraphBuilder[TKR[int], TKR[int]]: - g = GraphBuilder(TKR[int], TKR[int]) - graph = g.task(doubler_plus_graph()) - graph_ref = TypedGraphRef(graph.value_ref(), TKR[int], TKR[int]) - out = g.eval(graph_ref, g.inputs) - g.outputs(out) - return g - -class ApplyTwiceInput(NamedTuple): - # Note we mangle this like the stub generator would, although the generator - # never sees the graph's inputs as they are hidden in an untyped GraphData. - graph: TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]] # noqa: F821 - value: TKR[int] +def eval_body_is_from_worker() -> Workflow[TKR[int], TKR[int]]: + g = Graph(TKR[int], TKR[int]) + graph = TypedGraphRef(g.task(doubler_plus_graph()), TKR[int]) + out = g.eval(graph, g.inputs) + return g.finish_with_outputs(out) -def eval_from_worker_with_graph_from_worker() -> GraphBuilder[TKR[int], TKR[int]]: - g = GraphBuilder(TKR[int], TKR[int]) +def eval_from_worker_with_graph_from_worker() -> Workflow[TKR[int], TKR[int]]: + g = Graph(TKR[int], TKR[int]) graph = g.task(doubler_plus_graph()) - # This is ok, but we can't pass the graph_ref into ApplyTwiceInput - # graph_ref = TypedGraphRef(graph.value_ref(), TKR[int], TKR[int]) inputs = ApplyTwiceInput(graph=graph, value=g.inputs) - ap2 = g.task(apply_twice()) - ap2_ref = TypedGraphRef(ap2.value_ref(), ApplyTwiceInput, TKR[int]) - out = g.eval(ap2_ref, inputs) - g.outputs(out) - return g + ap2 = TypedGraphRef(g.task(apply_twice()), TKR[int]) + out = g.eval(ap2, inputs) + return g.finish_with_outputs(out) -def eval_graph_of_graph() -> GraphBuilder[TKR[int], TKR[int]]: - g = GraphBuilder(TKR[int], TKR[int]) +def eval_graph_of_graph() -> Workflow[TKR[int], TKR[int]]: + g = Graph(TKR[int], TKR[int]) graph = g.task(doubler_plus_graph()) - # This is ok, but we can't pass the graph_ref into exponentiate_graph: - # graph_ref = TypedGraphRef(graph.value_ref(), TKR[int], TKR[int]) - eg = g.task(graph_of_graph(graph, g.const(3))) - exp_graph = TypedGraphRef(eg.value_ref(), TKR[int], TKR[int]) - out = g.eval(exp_graph, g.inputs) - g.outputs(out) - return g + e1 = g.task(graph_of_graph(graph, g.const(3))) + eg = TypedGraphRef(e1, TKR[int]) + out = g.eval(eg, g.inputs) + return g.finish_with_outputs(out) def embed_graph(): @@ -211,7 +186,7 @@ class OuterOutput(NamedTuple): final: TKR[int] def inner(): - g = GraphBuilder(TKR[int], InnerOutput) + g = Graph(TKR[int], InnerOutput) s = g.task(tkr_str(g.inputs)) div2 = g.task(idivide(a=g.inputs, b=g.const(2))) times3plus1 = g.task( @@ -219,12 +194,10 @@ def inner(): ) even = g.task(eq(g.task(mod(a=g.inputs, b=g.const(2))), g.const(0))) n = g.ifelse(even, div2, times3plus1) - g.outputs(InnerOutput(log=s, nxt=n)) - return g + return g.finish_with_outputs(InnerOutput(log=s, nxt=n)) - g = GraphBuilder(TKR[int], OuterOutput) + g = Graph(TKR[int], OuterOutput) inner_g = inner() - first = g.embed(inner_g, g.inputs) - second = g.embed(inner_g, first.nxt) - g.outputs(OuterOutput(first.log, second.log, second.nxt)) - return g + first = g.embed(inner_g, g.inputs, InnerOutput) + second = g.embed(inner_g, first.nxt, InnerOutput) + return g.finish_with_outputs(OuterOutput(first.log, second.log, second.nxt)) diff --git a/tierkreis/tests/executor/test_executors.py b/tierkreis/tests/executor/test_executors.py index 85d00b004..a60796b7a 100644 --- a/tierkreis/tests/executor/test_executors.py +++ b/tierkreis/tests/executor/test_executors.py @@ -2,7 +2,7 @@ from uuid import UUID from tests.workers.hello_world_worker.stubs import greet -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Workflow, Graph from tierkreis.builtins import neg from tierkreis.consts import PACKAGE_PATH from tierkreis.controller import run_graph @@ -22,15 +22,14 @@ def shell_graph(): - g = GraphBuilder(TKR[str], TKR[str]) + g = Graph(TKR[str], TKR[str]) result = g.data.func( # escape hatch into untyped builder "shell_worker.greet", {"greeting": g.inputs.value_ref()}, ) output: TKR[str] = TKR(*result("value")) # unsafe cast - g.outputs(output) - return g + return g.finish_with_outputs(output) def test_shell_executor(): @@ -78,10 +77,9 @@ def test_suppress_env(): assert "main.sh" in exec_data.launch_command -def builtin_graph() -> GraphBuilder[TKR[bool], TKR[bool]]: - g = GraphBuilder(TKR[bool], TKR[bool]) - g.outputs(g.task(neg(g.inputs))) - return g +def builtin_graph() -> Workflow[TKR[bool], TKR[bool]]: + g = Graph(TKR[bool], TKR[bool]) + return g.finish_with_outputs(g.task(neg(g.inputs))) def test_builtin_executor(): @@ -105,13 +103,11 @@ def test_builtin_executor(): assert exec_data.executor == "builtin" -def hello_graph() -> GraphBuilder[TKR[str], TKR[str]]: - g = GraphBuilder(TKR[str], TKR[str]) +def hello_graph() -> Workflow[TKR[str], TKR[str]]: + g = Graph(TKR[str], TKR[str]) hello = g.const("hello ") output = g.task(greet(greeting=hello, subject=g.inputs)) - g.outputs(output) - - return g + return g.finish_with_outputs(output) def test_uv_executor(): @@ -138,15 +134,14 @@ def test_uv_executor(): def stdinout_graph(): - g = GraphBuilder(TKR[str], TKR[str]) + g = Graph(TKR[str], TKR[str]) result = g.data.func( # escape hatch into untyped builder "stdinout_worker.greet", {"greeting": g.inputs.value_ref()}, ) output: TKR[str] = TKR(*result("value")) # unsafe cast - g.outputs(output) - return g + return g.finish_with_outputs(output) def test_stdinout_executor(): @@ -172,7 +167,7 @@ def test_stdinout_executor(): def task_graph(): # Both tasks are the same, we just use different names to test the task executor - g = GraphBuilder(TKR[str], TKR[str]) + g = Graph(TKR[str], TKR[str]) first_call = g.data.func( "shell_worker.meet", {"greeting": g.inputs.value_ref()}, @@ -183,8 +178,7 @@ def task_graph(): {"greeting": out.value_ref()}, ) output: TKR[str] = TKR(*second_call("value")) - g.outputs(output) - return g + return g.finish_with_outputs(output) def test_task_executor(): @@ -225,7 +219,7 @@ def test_task_executor(): def multiple_graph(): # Both tasks are the same, we just use different names to test the task executor - g = GraphBuilder(TKR[str], TKR[str]) + g = Graph(TKR[str], TKR[str]) first_call = g.data.func( "shell_worker.meet", {"greeting": g.inputs.value_ref()}, @@ -236,8 +230,7 @@ def multiple_graph(): {"greeting": out.value_ref()}, ) output: TKR[str] = TKR(*second_call("value")) - g.outputs(output) - return g + return g.finish_with_outputs(output) def test_multiple_executor(): diff --git a/tierkreis/tests/executor/test_hpc_executor.py b/tierkreis/tests/executor/test_hpc_executor.py index 5f47f70ba..727f91dff 100644 --- a/tierkreis/tests/executor/test_hpc_executor.py +++ b/tierkreis/tests/executor/test_hpc_executor.py @@ -4,7 +4,7 @@ import pytest from tests.executor.stubs import mpi_rank_info -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.controller import run_graph from tierkreis.controller.data.graph import GraphData from tierkreis.controller.data.models import TKR @@ -19,9 +19,9 @@ def mpi_graph() -> GraphData: - builder = GraphBuilder(outputs_type=TKR[str | None]) + builder = Graph(outputs_type=TKR[str | None]) mpi_result = builder.task(mpi_rank_info()) - builder.outputs(mpi_result) + builder.finish_with_outputs(mpi_result) return builder.data diff --git a/tierkreis/tests/workers/graph/main.py b/tierkreis/tests/workers/graph/main.py index 65154cfe1..362ad7cea 100644 --- a/tierkreis/tests/workers/graph/main.py +++ b/tierkreis/tests/workers/graph/main.py @@ -2,57 +2,55 @@ from typing import NamedTuple from tierkreis import Worker -from tierkreis.builder import GraphBuilder, TypedGraphRef +from tierkreis.builder import Graph, Workflow, TypedGraphRef from tierkreis.builtins import iadd, itimes -from tierkreis.controller.data.graph import GraphData from tierkreis.models import TKR worker = Worker("graph") @worker.task() -def doubler_plus_graph() -> GraphData: - g = GraphBuilder(TKR[int], TKR[int]) +def doubler_plus_graph() -> Workflow[TKR[int], TKR[int]]: + g = Graph(TKR[int], TKR[int]) double = g.task(itimes(g.inputs, g.const(2))) out = g.task(iadd(double, g.const(1))) - g.outputs(out) - return g.get_data() + return g.finish_with_outputs(out) @worker.task() # The input graph here is expected to be int->int, but we have no way to express that in the type system. -# (GraphBuilder doesn't work as it's not accepted by the stub generator) -def graph_of_graph(f: GraphData, n: int) -> GraphData: +# (Graph doesn't work as it's not accepted by the stub generator) +def graph_of_graph( + f: Workflow[TKR[int], TKR[int]], n: int +) -> Workflow[TKR[int], TKR[int]]: """Builds a new graph: lambda x: f^n(x) I.e. the graph applies the first argument `f` to the graph's input `n` times. The graph contains the argument graph `f` as a constant.""" - g = GraphBuilder(TKR[int], TKR[int]) + g = Graph(TKR[int], TKR[int]) v = g.inputs - ref = TypedGraphRef(g.const(f).value_ref(), TKR[int], TKR[int]) + ref: TypedGraphRef[TKR[int], TKR[int]] = g._graph_const(f) for _ in range(n): v = g.eval(ref, v) - g.outputs(v) - return g.get_data() + return g.finish_with_outputs(v) + + +class ApplyTwiceInput(NamedTuple): + graph: TKR[Workflow[TKR[int], TKR[int]]] + value: TKR[int] @worker.task() -def apply_twice() -> GraphData: +def apply_twice() -> Workflow[ApplyTwiceInput, TKR[int]]: """Returns a graph for lambda f,x: f(f(x)). That is, `f` and `x` are inputs to the graph, not the worker function building it. """ - - class ApplyTwiceInput(NamedTuple): - graph: TKR[GraphData] - value: TKR[int] - - g = GraphBuilder(ApplyTwiceInput, TKR[int]) - f = TypedGraphRef(g.inputs.graph.value_ref(), TKR[int], TKR[int]) + g = Graph(ApplyTwiceInput, TKR[int]) + f = TypedGraphRef(g.inputs.graph, TKR[int]) run_once = g.eval(f, g.inputs.value) run_twice = g.eval(f, run_once) - g.outputs(run_twice) - return g.get_data() + return g.finish_with_outputs(run_twice) if __name__ == "__main__": diff --git a/tierkreis/tests/workers/graph/stubs.py b/tierkreis/tests/workers/graph/stubs.py index 9217da5c0..a134207d2 100644 --- a/tierkreis/tests/workers/graph/stubs.py +++ b/tierkreis/tests/workers/graph/stubs.py @@ -1,13 +1,18 @@ """Code generated from graph namespace. Please do not edit.""" from typing import NamedTuple -from tierkreis.controller.data.models import TKR, OpaqueType +from tierkreis.controller.data.models import TKR, Workflow + + +class ApplyTwiceInput(NamedTuple): + graph: TKR[Workflow[TKR[int], TKR[int]]] # noqa: F821 # fmt: skip + value: TKR[int] # noqa: F821 # fmt: skip class doubler_plus_graph(NamedTuple): @staticmethod - def out() -> type[TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]]]: # noqa: F821 # fmt: skip - return TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]] # noqa: F821 # fmt: skip + def out() -> type[TKR[Workflow[TKR[int], TKR[int]]]]: # noqa: F821 # fmt: skip + return TKR[Workflow[TKR[int], TKR[int]]] # noqa: F821 # fmt: skip @property def namespace(self) -> str: @@ -15,12 +20,12 @@ def namespace(self) -> str: class graph_of_graph(NamedTuple): - f: TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]] # noqa: F821 # fmt: skip + f: TKR[Workflow[TKR[int], TKR[int]]] # noqa: F821 # fmt: skip n: TKR[int] # noqa: F821 # fmt: skip @staticmethod - def out() -> type[TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]]]: # noqa: F821 # fmt: skip - return TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]] # noqa: F821 # fmt: skip + def out() -> type[TKR[Workflow[TKR[int], TKR[int]]]]: # noqa: F821 # fmt: skip + return TKR[Workflow[TKR[int], TKR[int]]] # noqa: F821 # fmt: skip @property def namespace(self) -> str: @@ -29,8 +34,8 @@ def namespace(self) -> str: class apply_twice(NamedTuple): @staticmethod - def out() -> type[TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]]]: # noqa: F821 # fmt: skip - return TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]] # noqa: F821 # fmt: skip + def out() -> type[TKR[Workflow[ApplyTwiceInput, TKR[int]]]]: # noqa: F821 # fmt: skip + return TKR[Workflow[ApplyTwiceInput, TKR[int]]] # noqa: F821 # fmt: skip @property def namespace(self) -> str: diff --git a/tierkreis/tierkreis/builder.py b/tierkreis/tierkreis/builder.py index 7d3c7398d..c42c912d8 100644 --- a/tierkreis/tierkreis/builder.py +++ b/tierkreis/tierkreis/builder.py @@ -23,6 +23,7 @@ TNamedModel, dict_from_tmodel, init_tmodel, + Workflow, ) from tierkreis.controller.data.types import PType @@ -66,11 +67,9 @@ class TypedGraphRef[Ins: TModel, Outs: TModel]: :attr graph_ref: The graph reference. :attr outputs_type: The output type of the graph. - :attr inputs_type: The input type of the graph. """ - graph_ref: ValueRef - inputs_type: type[Ins] + graph_ref: TKR[Workflow[Ins, Outs]] outputs_type: type[Outs] @@ -115,7 +114,7 @@ def namespace(self) -> str: return exec_script(input=script_input) -class GraphBuilder[Inputs: TModel, Outputs: TModel]: +class Graph[Inputs: TModel, Outputs: TModel]: """Class to construct typed workflow graphs. :attr data: The underlying graph data. @@ -124,7 +123,8 @@ class GraphBuilder[Inputs: TModel, Outputs: TModel]: :attr outputs_type: The output type of the graph. """ - outputs_type: type + inputs_type: type[Inputs] + outputs_type: type[Outputs] inputs: Inputs def __init__( @@ -137,32 +137,28 @@ def __init__( self.outputs_type = outputs_type self.inputs = init_tmodel(self.inputs_type, self.data.input) - def get_data(self) -> GraphData: - """Return the underlying graph from the builder. - - :return: The graph. - :rtype: GraphData - """ - return self.data - def ref(self) -> TypedGraphRef[Inputs, Outputs]: """Return a reference of the typed graph. :return: The ref of the typed graph. :rtype: TypedGraphRef[Inputs, Outputs] """ - return TypedGraphRef((-1, "body"), self.inputs_type, self.outputs_type) + return TypedGraphRef(TKR(-1, "body"), self.outputs_type) - def outputs(self, outputs: Outputs) -> None: + def finish_with_outputs(self, outputs: Outputs) -> Workflow[Inputs, Outputs]: """Set output nodes of a graph. :param outputs: The output nodes. :type outputs: Outputs """ self.data.output(inputs=dict_from_tmodel(outputs)) + return Workflow(self.data, self.outputs_type) - def embed[A: TModel, B: TModel](self, other: "GraphBuilder[A, B]", inputs: A) -> B: - if other.data.graph_output_idx is None: + def embed[A: TModel, B: TModel]( + self, other_fg: Workflow[A, B], inputs: A, outputs_type: type[B] + ) -> B: + other = other_fg.data + if other.graph_output_idx is None: raise ValueError("Can only embed graphs with an output node defined.") ins: Mapping[str, ValueRef] = dict_from_tmodel(inputs) @@ -174,7 +170,7 @@ def reindex(vr: ValueRef) -> ValueRef: if idx in port_map: return port_map[idx][port] if idx not in node_map: - node = other.data.nodes[idx] + node = other.nodes[idx] if node.type == "input": port_map[idx] = {node.name: ins[node.name]} return port_map[idx][port] @@ -187,8 +183,8 @@ def reindex(vr: ValueRef) -> ValueRef: node_map[idx] = func("dummy_port")[0] return (node_map[idx], port) - outputs = other.data.nodes[other.data.graph_output_idx].inputs - return init_tmodel(other.outputs_type, lambda p: reindex(outputs[p])) + outputs = other.nodes[other.graph_output_idx].inputs + return init_tmodel(other_fg.outputs_type, lambda p: reindex(outputs[p])) def const[T: PType](self, value: T) -> TKR[T]: """Add a constant node to the graph. @@ -196,7 +192,10 @@ def const[T: PType](self, value: T) -> TKR[T]: :return: The constant value. :rtype: TKR[T] """ - idx, port = self.data.const(value) + # Workflow exists at build-time only; erase the types at runtime: + idx, port = self.data.const( + value.data if isinstance(value, Workflow) else value + ) return TKR[T](idx, port) def ifelse[A: PType, B: PType]( @@ -255,14 +254,12 @@ def eifelse[A: PType, B: PType]( def _graph_const[A: TModel, B: TModel]( self, - graph: GraphBuilder[A, B], + graph: Workflow[A, B], ) -> TypedGraphRef[A, B]: # TODO @philipp-seitz: Turn this into a public method? - idx, port = self.data.const(graph.data.model_dump()) - return TypedGraphRef[A, B]( - graph_ref=(idx, port), - outputs_type=graph.outputs_type, - inputs_type=graph.inputs_type, + return TypedGraphRef( + self.const(graph), + graph.outputs_type, ) def task[Out: TModel](self, func: Function[Out]) -> Out: @@ -281,30 +278,30 @@ def task[Out: TModel](self, func: Function[Out]) -> Out: def eval[A: TModel, B: TModel]( self, - body: GraphBuilder[A, B] | TypedGraphRef[A, B], + body: Workflow[A, B] | TypedGraphRef[A, B], eval_inputs: A, ) -> B: """Add a evaluation node to the graph. This will evaluate a nested graph with the given inputs. + :param A the input type of the graph. + :param B the output type of the graph. :param body: The graph to evaluate. - :type body: TypedGraphRef[A, B] | GraphBuilder[A, B], - where A are the input type and B the output type of the graph. :param eval_inputs: The inputs to the graph. - :type eval_inputs: A :return: The outputs of the evaluation. - :rtype: B """ - if isinstance(body, GraphBuilder): + if isinstance(body, Workflow): body = self._graph_const(body) - idx, _ = self.data.eval(body.graph_ref, dict_from_tmodel(eval_inputs))("dummy") + idx, _ = self.data.eval( + body.graph_ref.value_ref(), dict_from_tmodel(eval_inputs) + )("dummy") return init_tmodel(body.outputs_type, lambda p: (idx, p)) def loop[A: TModel, B: LoopOutput]( self, - body: TypedGraphRef[A, B] | GraphBuilder[A, B], + body: TypedGraphRef[A, B] | Workflow[A, B], loop_inputs: A, name: str | None = None, ) -> B: @@ -314,22 +311,18 @@ def loop[A: TModel, B: LoopOutput]( To trace intermediate values, use the name attribute in conjunction with read_loop_trace. + :param A the input type of the graph. + :param B the output type of the graph. :param body: The graph to loop. - :type body: TypedGraphRef[A, B] | GraphBuilder[A, B], - where A are the input type and B the output type of the graph. :param loop_inputs: The inputs to the loop graph. - :type loop_inputs: A :param name: An optional name for the loop. - :type name: str | None :return: The outputs of the loop. - :rtype: B """ - if isinstance(body, GraphBuilder): + if isinstance(body, Workflow): body = self._graph_const(body) - graph = body.graph_ref idx, _ = self.data.loop( - graph, + body.graph_ref.value_ref(), dict_from_tmodel(loop_inputs), "should_continue", name, @@ -371,32 +364,28 @@ def _map_graph_full[A: TModel, B: TModel]( body: TypedGraphRef[A, B], ) -> TList[B]: ins = dict_from_tmodel(map_inputs._value) # noqa: SLF001 - idx, _ = self.data.map(body.graph_ref, ins)("x") + idx, _ = self.data.map(body.graph_ref.value_ref(), ins)("x") return TList(init_tmodel(body.outputs_type, lambda s: (idx, s + "-*"))) @overload def map[A: PType, B: TNamedModel]( self, - body: ( - Callable[[TKR[A]], B] | TypedGraphRef[TKR[A], B] | GraphBuilder[TKR[A], B] - ), + body: (Callable[[TKR[A]], B] | TypedGraphRef[TKR[A], B] | Workflow[TKR[A], B]), map_inputs: TKR[list[A]], ) -> TList[B]: ... @overload def map[A: TNamedModel, B: PType]( self, - body: ( - Callable[[A], TKR[B]] | TypedGraphRef[A, TKR[B]] | GraphBuilder[A, TKR[B]] - ), + body: (Callable[[A], TKR[B]] | TypedGraphRef[A, TKR[B]] | Workflow[A, TKR[B]]), map_inputs: TList[A], ) -> TKR[list[B]]: ... @overload def map[A: TNamedModel, B: TNamedModel]( self, - body: TypedGraphRef[A, B] | GraphBuilder[A, B], + body: TypedGraphRef[A, B] | Workflow[A, B], map_inputs: TList[A], ) -> TList[B]: ... @@ -406,26 +395,23 @@ def map[A: PType, B: PType]( body: ( Callable[[TKR[A]], TKR[B]] | TypedGraphRef[TKR[A], TKR[B]] - | GraphBuilder[TKR[A], TKR[B]] + | Workflow[TKR[A], TKR[B]] ), map_inputs: TKR[list[A]], ) -> TKR[list[B]]: ... def map( self, - body: TypedGraphRef | Callable | GraphBuilder, + body: TypedGraphRef | Callable | Workflow, map_inputs: TKR | TList, ) -> Any: """Add a map node to the graph. :param body: The graph to map over. - :type body: TypedGraphRef | Callable | GraphBuilder :param map_inputs: The values to map over. - :type map_inputs: TKR | TList :return: The outputs of the map. - :rtype: Any """ - if isinstance(body, GraphBuilder): + if isinstance(body, Workflow): body = self._graph_const(body) if isinstance(body, Callable): diff --git a/tierkreis/tierkreis/cli/run.py b/tierkreis/tierkreis/cli/run.py index 3047d8e28..e54e20ed9 100644 --- a/tierkreis/tierkreis/cli/run.py +++ b/tierkreis/tierkreis/cli/run.py @@ -10,7 +10,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph from tierkreis.cli.run_workflow import run_workflow from tierkreis.controller.data.graph import GraphData from tierkreis.controller.data.types import PType, ptype_from_bytes @@ -56,13 +56,13 @@ def load_graph(graph_input: str) -> GraphData: graph_name: Any = getattr(module, function_name) if isinstance(graph_name, GraphData): graph_data = graph_name - if isinstance(graph_name, GraphBuilder): + if isinstance(graph_name, Graph): graph_data = graph_name.data if isinstance(graph_name, Callable): graph_object = graph_name() if isinstance(graph_object, GraphData): graph_data = graph_object - if isinstance(graph_object, GraphBuilder): + if isinstance(graph_object, Graph): graph_data = graph_object.data if graph_data is None: logger.error("Could not load object %s as GraphData", graph_data) diff --git a/tierkreis/tierkreis/cli/templates.py b/tierkreis/tierkreis/cli/templates.py index fd4913797..d5098b1a6 100644 --- a/tierkreis/tierkreis/cli/templates.py +++ b/tierkreis/tierkreis/cli/templates.py @@ -151,7 +151,7 @@ def default_graph(worker_name: str) -> str: from pathlib import Path from uuid import UUID -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Workflow, Graph from tierkreis.controller import run_graph from tierkreis.controller.data.models import TKR, OpaqueType from tierkreis.executor import ShellExecutor, UvExecutor @@ -167,11 +167,10 @@ class GraphOutputs(NamedTuple): value: TKR[int] -def workflow() -> GraphBuilder[GraphInputs, GraphOutputs]: - g = GraphBuilder(GraphInputs, GraphOutputs) +def workflow() -> Workflow[GraphInputs, GraphOutputs]: + g = Graph(GraphInputs, GraphOutputs) out = g.task(your_worker_task(g.inputs.value)) - g.outputs(GraphOutputs(value=out)) - return g + return g.finish_with_outputs(GraphOutputs(value=out)) def main() -> None: graph = workflow() @@ -301,10 +300,10 @@ def your_worker_task(value: int) -> int: You can use them as a task in the graph: ```python -def your_graph() -> GraphBuilder[TKR[int], TKR[int]]: - g = GraphBuilder(TKR[int], TKR[int]) +def your_graph() -> Workflow[TKR[int], TKR[int]]: + g = Graph(TKR[int], TKR[int]) out = g.task(your_worker_task(g.inputs)) - g.outputs(out) + g.finish_with_outputs(out) return g ``` If you used the `tkr init project` example, you will see a working graph code example in `tkr/graphs/main.py`. diff --git a/tierkreis/tierkreis/codegen.py b/tierkreis/tierkreis/codegen.py index 4ff70dad0..45a7f03ab 100644 --- a/tierkreis/tierkreis/codegen.py +++ b/tierkreis/tierkreis/codegen.py @@ -106,7 +106,20 @@ def format_model(model: Model) -> str: outs.sort() outs_str = "\n ".join(outs) - bases = ["NamedTuple"] if is_portmapping else ["Struct", "Protocol"] + def is_tmodel() -> bool: + has_ptypes = False + has_tkrs = False + for decl in model.decls: + if decl.t.is_ptype: + has_ptypes = True + else: + has_tkrs = True + assert not (has_ptypes and has_tkrs), ( + "Model decls should be all PTypes or all TKRs" + ) + return has_tkrs + + bases = ["NamedTuple"] if is_portmapping or is_tmodel() else ["Struct", "Protocol"] bases_str = ", ".join(bases) generic_type_str = format_generic_type(model.t, include_bound=True, is_tkr=False) diff --git a/tierkreis/tierkreis/controller/__init__.py b/tierkreis/tierkreis/controller/__init__.py index a989f7607..60abfba7b 100644 --- a/tierkreis/tierkreis/controller/__init__.py +++ b/tierkreis/tierkreis/controller/__init__.py @@ -9,7 +9,7 @@ from time import sleep from typing import TYPE_CHECKING -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Workflow from tierkreis.controller.data.graph import Eval, GraphData from tierkreis.controller.data.location import Loc from tierkreis.controller.data.models import TModel @@ -31,7 +31,7 @@ def run_graph[A: TModel, B: TModel]( storage: ControllerStorage, executor: ControllerExecutor, - g: GraphData | GraphBuilder[A, B], + g: GraphData | Workflow[A, B], graph_inputs: dict[str, PType] | PType, n_iterations: int = 10000, polling_interval_seconds: float = 0.01, @@ -48,7 +48,7 @@ def run_graph[A: TModel, B: TModel]( :param executor: The executor backend for the controller. :type executor: ControllerExecutor :param g: The graph to run. - :type g: GraphData | GraphBuilder[A, B] + :type g: GraphData | Workflow[A, B] :param graph_inputs: The inputs to the graph. If a single PType is provided, it will be provided as the input "value". :type graph_inputs: dict[str, PType] | PType @@ -61,8 +61,8 @@ def run_graph[A: TModel, B: TModel]( :type enable_logging: bool, optional :raises TierkreisError: If the graph encounters errors during execution. """ - if isinstance(g, GraphBuilder): - g = g.get_data() + if isinstance(g, Workflow): + g = g.data if not isinstance(graph_inputs, dict): graph_inputs = {"value": graph_inputs} diff --git a/tierkreis/tierkreis/controller/data/models.py b/tierkreis/tierkreis/controller/data/models.py index df117053b..163d0a9d8 100644 --- a/tierkreis/tierkreis/controller/data/models.py +++ b/tierkreis/tierkreis/controller/data/models.py @@ -8,6 +8,7 @@ Literal, Protocol, Union, + Self, cast, dataclass_transform, get_args, @@ -24,7 +25,14 @@ RestrictedNamedTuple, ValueRef, ) -from tierkreis.controller.data.types import PType +from tierkreis.controller.data.graph import GraphData +from tierkreis.controller.data.types import ( + PType, + ModelConvertible, + coerce_from_annotation, +) + +from pydantic import BaseModel TKR_PORTMAPPING_FLAG = "__tkr_portmapping__" @@ -142,3 +150,19 @@ def init_tmodel[T: TModel](tmodel: type[T], input_fn: Callable[[str], ValueRef]) return cast("T", model(*args)) (ref,) = fields.values() return tmodel(*ref) + + +@dataclass(frozen=True) +class Workflow[Inputs: TModel, Outputs: TModel](ModelConvertible): + data: GraphData + outputs_type: type[Outputs] + + def to_model(self) -> BaseModel: + return self.data + + @classmethod + def from_model(cls, annot: type[Self], arg: BaseModel, /) -> Self: + from tierkreis.controller.data.graph import GraphData + + _inputs, outputs = get_args(annot) + return cls(coerce_from_annotation(arg, GraphData), outputs) diff --git a/tierkreis/tierkreis/controller/data/types.py b/tierkreis/tierkreis/controller/data/types.py index 33ff544f1..52b23d812 100644 --- a/tierkreis/tierkreis/controller/data/types.py +++ b/tierkreis/tierkreis/controller/data/types.py @@ -72,6 +72,22 @@ def from_dict(cls, arg: dict, /) -> "Self": ... +@runtime_checkable +class ModelConvertible(Protocol): + """A protocol for types that can be converted to and from pydantic BaseModels.""" + + def to_model(self) -> BaseModel: + """Convert self to a BaseModel.""" + ... + + @classmethod + def from_model(cls, annot: type[Self], arg: BaseModel, /) -> "Self": + """Construct self from a BaseModel. + + :param annot: The annotation of the type to convert to, i.e. so it can contain type arguments for generic ModelConvertible's.""" + ... + + @runtime_checkable class ListConvertible(Protocol): """A protocol for types that can be converted to and from lists.""" @@ -106,6 +122,7 @@ def from_list(cls, arg: list, /) -> "Self": | ListConvertible | NdarraySurrogate | BaseModel + | ModelConvertible ) type JsonType = Container[ElementaryType] logger = logging.getLogger(__name__) @@ -200,7 +217,8 @@ def is_ptype(annotation: Any) -> TypeIs[type[PType]]: :return: The according TypeIs if the annotation is a PType, otherwise False. :rtype: TypeIs[type[PType]] """ - if get_origin(annotation) is Annotated: + origin = get_origin(annotation) + if origin is Annotated: return is_ptype(get_args(annotation)[0]) if _is_generic(annotation): @@ -218,12 +236,22 @@ def is_ptype(annotation: Any) -> TypeIs[type[PType]]: isclass(annotation) and issubclass( annotation, - (DictConvertible, ListConvertible, NdarraySurrogate, BaseModel, Struct), + ( + DictConvertible, + ModelConvertible, + ListConvertible, + NdarraySurrogate, + BaseModel, + Struct, + ), + ) + or annotation in get_args(ElementaryType.__value__) + or ( # should we check args are ptypes too? + isclass(origin) and issubclass(origin, ModelConvertible) ) - ) 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)) @@ -264,6 +292,8 @@ def ser_from_ptype(ptype: PType, annotation: type[PType] | None) -> JsonType: return {k: ser_from_ptype(p, arg) for k, p in ptype.items()} case DictConvertible(): return ser_from_ptype(ptype.to_dict(), None) + case ModelConvertible(): + return ser_from_ptype(ptype.to_model(), None) case ListConvertible(): return ser_from_ptype(ptype.to_list(), None) case BaseModel(): @@ -305,13 +335,15 @@ def coerce_from_annotation[T: PType](ser: Any, annotation: type[T] | None) -> T: :return: The coerced value. :rtype: T """ + if annotation is None: return ser if ds := get_deserializer(annotation): return ds.deserializer(ser) - if get_origin(annotation) is Annotated: + origin = get_origin(annotation) + if origin is Annotated: return coerce_from_annotation(ser, get_args(annotation)[0]) if _is_union(annotation): @@ -323,7 +355,6 @@ def coerce_from_annotation[T: PType](ser: Any, annotation: type[T] | None) -> T: msg = f"Could not deserialise {ser} as {annotation}" raise TierkreisError(msg) - origin = get_origin(annotation) if origin is None: origin = annotation @@ -341,11 +372,28 @@ def coerce_from_annotation[T: PType](ser: Any, annotation: type[T] | None) -> T: return ser if issubclass(origin, DictConvertible): + # This works only when get_origin(annotation)=None so `origin is annotation`. + # Otherwise, annotation would be a GenericAlias instantiation of origin, + # not a subclass. if not issubclass(annotation, origin): msg = "Invalid subclass relation encountered." raise TypeError(msg) return annotation.from_dict(ser) + if issubclass(origin, ModelConvertible): + if isclass(annotation): + assert issubclass(annotation, ModelConvertible) + assert annotation is origin # because get_origin(annotation) == None + else: + # not a subclass of ModelConvertible as not a class + from typing import _GenericAlias + + assert isinstance(annotation, _GenericAlias) + # The Annotated that we cast to here is an intersection type + return annotation.from_model( + cast(Annotated[type[ModelConvertible], T], annotation), ser + ) + if issubclass(origin, ListConvertible): if not issubclass(annotation, origin): msg = "Invalid subclass relation encountered." diff --git a/tierkreis/tierkreis/graphs/fold.py b/tierkreis/tierkreis/graphs/fold.py index de28d7038..154d91522 100644 --- a/tierkreis/tierkreis/graphs/fold.py +++ b/tierkreis/tierkreis/graphs/fold.py @@ -2,15 +2,14 @@ from typing import NamedTuple, TypeVar -from tierkreis.builder import GraphBuilder, TypedGraphRef +from tierkreis.builder import Graph, Workflow, TypedGraphRef from tierkreis.builtins import head, igt, tkr_len -from tierkreis.controller.data.graph import GraphData from tierkreis.controller.data.models import TKR from tierkreis.controller.data.types import PType class _FoldGraphOuterInputs[A: PType, B: PType](NamedTuple): - func: TKR[GraphData] + func: TKR[Workflow["FoldFunctionInput[A, B]", TKR[B]]] accum: TKR[B] values: TKR[list[A]] @@ -21,16 +20,11 @@ class _FoldGraphOuterOutputs[A: PType, B: PType](NamedTuple): should_continue: TKR[bool] -class _InnerFuncInput[A: PType, B: PType](NamedTuple): - accum: TKR[B] - value: TKR[A] - - -def _fold_graph_outer[A: PType, B: PType]() -> GraphBuilder[ +def _fold_graph_outer[A: PType, B: PType]() -> Workflow[ _FoldGraphOuterInputs[A, B], _FoldGraphOuterOutputs[A, B], ]: - g = GraphBuilder(_FoldGraphOuterInputs[A, B], _FoldGraphOuterOutputs[A, B]) + g = Graph(_FoldGraphOuterInputs[A, B], _FoldGraphOuterOutputs[A, B]) func = g.inputs.func accum = g.inputs.accum @@ -44,17 +38,13 @@ def _fold_graph_outer[A: PType, B: PType]() -> GraphBuilder[ headed = g.task(head(values)) # Apply the function if we were able to pop off a value. - tgd = TypedGraphRef[_InnerFuncInput, TKR[B]]( - func.value_ref(), - _InnerFuncInput, - TKR[B], - ) - applied_next = g.eval(tgd, _InnerFuncInput(accum, headed.head)) - + tgd = TypedGraphRef[FoldFunctionInput[A, B], TKR[B]](func, TKR[B]) + applied_next = g.eval(tgd, FoldFunctionInput(accum, headed.head)) next_accum = g.ifelse(non_empty, applied_next, accum) next_values = g.ifelse(non_empty, headed.rest, values) - g.outputs(_FoldGraphOuterOutputs(next_accum, next_values, non_empty)) - return g + return g.finish_with_outputs( + _FoldGraphOuterOutputs(next_accum, next_values, non_empty) + ) A_co = TypeVar("A_co", bound=PType, covariant=True) @@ -86,26 +76,25 @@ class FoldFunctionInput[A: PType, B: PType](NamedTuple): def fold_graph[A_co: PType, B_co: PType]( - func: GraphBuilder[FoldFunctionInput[A_co, B_co], TKR[B_co]], -) -> GraphBuilder[FoldGraphInputs[A_co, B_co], TKR[B_co]]: + func: Workflow[FoldFunctionInput[A_co, B_co], TKR[B_co]], +) -> Workflow[FoldGraphInputs[A_co, B_co], TKR[B_co]]: """Construct a fold graph. fold : {func: (b -> a -> b)} -> {initial: b} -> {values: list[a]} -> {value: b} fold : { A x B -> B } -> { list[A] x B -> B } :param func: The function to fold over. - :type func: GraphBuilder[FoldFunctionInput[A_co, B_co], TKR[B_co]] + :type func: Graph[FoldFunctionInput[A_co, B_co], TKR[B_co]] :return: A graph implementing the fold function. - :rtype: GraphBuilder[FoldGraphInputs[A_co, B_co], TKR[B_co]] + :rtype: Graph[FoldGraphInputs[A_co, B_co], TKR[B_co]] """ - g = GraphBuilder(FoldGraphInputs[A_co, B_co], TKR[B_co]) + g = Graph(FoldGraphInputs[A_co, B_co], TKR[B_co]) foldfunc = g._graph_const(func) # noqa: SLF001 # TODO @mwpb: include the computation inside the fold ins = _FoldGraphOuterInputs( - TKR(*foldfunc.graph_ref), + foldfunc.graph_ref, g.inputs.initial, g.inputs.values, ) loop = g.loop(_fold_graph_outer(), ins) - g.outputs(loop.accum) - return g + return g.finish_with_outputs(loop.accum) diff --git a/tierkreis/tierkreis/graphs/nexus/submit_poll.py b/tierkreis/tierkreis/graphs/nexus/submit_poll.py index bb17a70d6..98b85d08f 100644 --- a/tierkreis/tierkreis/graphs/nexus/submit_poll.py +++ b/tierkreis/tierkreis/graphs/nexus/submit_poll.py @@ -3,7 +3,7 @@ # ruff: noqa: F821 from typing import NamedTuple -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph, Workflow from tierkreis.builtins import tkr_sleep from tierkreis.controller.data.models import TKR, OpaqueType from tierkreis.nexus_worker import ( @@ -53,22 +53,21 @@ class _LoopOutputs(NamedTuple): should_continue: TKR[bool] -def upload_circuit_graph() -> GraphBuilder[UploadCircuitInputs, TKR[ExecuteJobRef]]: +def upload_circuit_graph() -> Workflow[UploadCircuitInputs, TKR[ExecuteJobRef]]: """Construct a graph to upload a circuit to nexus. :return: A uploading graph. - :rtype: GraphBuilder[UploadCircuitInputs, TKR[ExecuteJobRef]] + :rtype: Graph[UploadCircuitInputs, TKR[ExecuteJobRef]] """ - g = GraphBuilder(UploadCircuitInputs, TKR[ExecuteJobRef]) + g = Graph(UploadCircuitInputs, TKR[ExecuteJobRef]) programme = g.task(upload_circuit(g.inputs.project_name, g.inputs.circuit)) - g.outputs(programme) # type: ignore[arg-type] - return g + return g.finish_with_outputs(programme) # type: ignore[arg-type] def _polling_loop_body( polling_interval: float, -) -> GraphBuilder[TKR[ExecuteJobRef], _LoopOutputs]: - g = GraphBuilder(TKR[ExecuteJobRef], _LoopOutputs) +) -> Workflow[TKR[ExecuteJobRef], _LoopOutputs]: + g = Graph(TKR[ExecuteJobRef], _LoopOutputs) pred = g.task(is_running(g.inputs)) wait = g.ifelse( @@ -78,21 +77,20 @@ def _polling_loop_body( ) results = g.ifelse(pred, g.const([]), g.task(get_results(g.inputs))) - g.outputs(_LoopOutputs(results=results, should_continue=wait)) - return g + return g.finish_with_outputs(_LoopOutputs(results=results, should_continue=wait)) def nexus_submit_and_poll( polling_interval: float = 30.0, -) -> GraphBuilder[JobInputs, TKR[list[BackendResult]]]: +) -> Workflow[JobInputs, TKR[list[BackendResult]]]: """Construct a graph submitting and polling a nexus job. :param polling_interval: The polling interval in seconds, defaults to 30.0 :type polling_interval: float, optional :return: A graph performing submission and polling. - :rtype: GraphBuilder[JobInputs, TKR[list[BackendResult]]] + :rtype: Graph[JobInputs, TKR[list[BackendResult]]] """ - g = GraphBuilder(JobInputs, TKR[list[BackendResult]]) + g = Graph(JobInputs, TKR[list[BackendResult]]) upload_inputs = g.map( lambda x: UploadCircuitInputs(g.inputs.project_name, x), g.inputs.circuits, @@ -110,5 +108,4 @@ def nexus_submit_and_poll( ) res = g.loop(_polling_loop_body(polling_interval), ref) - g.outputs(res.results) - return g + return g.finish_with_outputs(res.results) diff --git a/tierkreis/tierkreis/graphs/simulate/compile_simulate.py b/tierkreis/tierkreis/graphs/simulate/compile_simulate.py index d6b33470c..9b572fc08 100644 --- a/tierkreis/tierkreis/graphs/simulate/compile_simulate.py +++ b/tierkreis/tierkreis/graphs/simulate/compile_simulate.py @@ -9,7 +9,7 @@ from tierkreis.aer_worker import ( run_circuit as aer_run, ) -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph, Workflow from tierkreis.builtins import str_eq, tkr_zip, untuple from tierkreis.controller.data.models import TKR, OpaqueType from tierkreis.qulacs_worker import ( @@ -53,15 +53,15 @@ class SimulateJobInputsSingle(NamedTuple): compilation_optimisation_level: TKR[int] -def aer_simulate_single() -> GraphBuilder[SimulateJobInputsSingle, TKR[BackendResult]]: +def aer_simulate_single() -> Workflow[SimulateJobInputsSingle, TKR[BackendResult]]: """Construct a graph to simulate a single circuit using qiskit aer. This ignores the simulator_name field. :return: The graph for the simulation. - :rtype: GraphBuilder[SimulateJobInputsSingle, TKR[BackendResult]] + :rtype: Graph[SimulateJobInputsSingle, TKR[BackendResult]] """ - g = GraphBuilder(SimulateJobInputsSingle, TKR[BackendResult]) + g = Graph(SimulateJobInputsSingle, TKR[BackendResult]) circuit_shots = g.task(untuple(g.inputs.circuit_shots)) compiled_circuit = g.task( @@ -71,11 +71,10 @@ def aer_simulate_single() -> GraphBuilder[SimulateJobInputsSingle, TKR[BackendRe ), ) res = g.task(aer_run(compiled_circuit, circuit_shots.b)) - g.outputs(res) - return g + return g.finish_with_outputs(res) -def qulacs_simulate_single() -> GraphBuilder[ +def qulacs_simulate_single() -> Workflow[ SimulateJobInputsSingle, TKR[BackendResult], ]: @@ -84,9 +83,9 @@ def qulacs_simulate_single() -> GraphBuilder[ This ignores the simulator_name field. :return: The graph for the simulation. - :rtype: GraphBuilder[SimulateJobInputsSingle, TKR[BackendResult]] + :rtype: Graph[SimulateJobInputsSingle, TKR[BackendResult]] """ - g = GraphBuilder(SimulateJobInputsSingle, TKR[BackendResult]) + g = Graph(SimulateJobInputsSingle, TKR[BackendResult]) circuit_shots = g.task(untuple(g.inputs.circuit_shots)) compiled_circuit = g.task( @@ -96,20 +95,19 @@ def qulacs_simulate_single() -> GraphBuilder[ ), ) res = g.task(qulacs_run(compiled_circuit, circuit_shots.b)) - g.outputs(res) - return g + return g.finish_with_outputs(res) -def compile_simulate_single() -> GraphBuilder[ +def compile_simulate_single() -> Workflow[ SimulateJobInputsSingle, TKR[BackendResult], ]: """CConstruct a graph to simulate a single job on either aer or qulacs. :return: The graph for the simulation. - :rtype: GraphBuilder[ SimulateJobInputsSingle, TKR[BackendResult], ] + :rtype: Graph[ SimulateJobInputsSingle, TKR[BackendResult], ] """ - g = GraphBuilder(SimulateJobInputsSingle, TKR[BackendResult]) + g = Graph(SimulateJobInputsSingle, TKR[BackendResult]) aer_res = g.eval(aer_simulate_single(), g.inputs) qulacs_res = g.eval(qulacs_simulate_single(), g.inputs) @@ -119,17 +117,15 @@ def compile_simulate_single() -> GraphBuilder[ qulacs_res, ) - g.outputs(res) - return g + return g.finish_with_outputs(res) -def compile_simulate() -> GraphBuilder[SimulateJobInputs, TKR[list[BackendResult]]]: +def compile_simulate() -> Workflow[SimulateJobInputs, TKR[list[BackendResult]]]: """Construct a graph to simulate multiple jobs on either aer or qulacs. :return: The graph for the simulation. - :rtype: GraphBuilder[SimulateJobInputs, TKR[list[BackendResult]]] """ - g = GraphBuilder(SimulateJobInputs, TKR[list[BackendResult]]) + g = Graph(SimulateJobInputs, TKR[list[BackendResult]]) circuits_shots = g.task(tkr_zip(g.inputs.circuits, g.inputs.n_shots)) @@ -143,5 +139,4 @@ def compile_simulate() -> GraphBuilder[SimulateJobInputs, TKR[list[BackendResult ) res = g.map(compile_simulate_single(), inputs) - g.outputs(res) - return g + return g.finish_with_outputs(res) diff --git a/tierkreis/tierkreis/idl/models.py b/tierkreis/tierkreis/idl/models.py index 2b05a2e7c..1c3282747 100644 --- a/tierkreis/tierkreis/idl/models.py +++ b/tierkreis/tierkreis/idl/models.py @@ -6,7 +6,7 @@ from typing import Annotated, Self, get_args, get_origin from tierkreis.controller.data.core import RestrictedNamedTuple -from tierkreis.controller.data.types import _is_generic +from tierkreis.controller.data.types import _is_generic, is_ptype type ElementaryType = ( type[int | float | bytes | str | bool | NoneType | Mapping | Sequence] @@ -28,6 +28,8 @@ class GenericType: origin: ElementaryType args: "Sequence[GenericType | str]" + is_ptype: bool = True # Assume yes if not constructed from a `type` + # This might be a dangerous default but seems right for the parser. @classmethod def from_type(cls, t: type) -> "Self": @@ -47,7 +49,7 @@ def from_type(cls, t: type) -> "Self": subargs = [] [subargs.append(str(x)) for x in args if _is_generic(x)] [subargs.append(cls.from_type(x)) for x in args if not _is_generic(x)] - return cls(origin, subargs) + return cls(origin, subargs, is_ptype(t)) @classmethod def _included_structs(cls, t: "GenericType") -> "set[GenericType]": diff --git a/tierkreis/tierkreis/namespace.py b/tierkreis/tierkreis/namespace.py index 6a4041d94..290b12f24 100644 --- a/tierkreis/tierkreis/namespace.py +++ b/tierkreis/tierkreis/namespace.py @@ -134,7 +134,7 @@ 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.models import TKR, OpaqueType, Workflow from tierkreis.controller.data.types import PType, Struct {models_str} diff --git a/tierkreis/tierkreis/storage.py b/tierkreis/tierkreis/storage.py index 68c367ed3..065cfba64 100644 --- a/tierkreis/tierkreis/storage.py +++ b/tierkreis/tierkreis/storage.py @@ -1,6 +1,6 @@ """Implementation to access node storage data.""" -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Workflow from tierkreis.controller.data.graph import GraphData from tierkreis.controller.data.location import Loc from tierkreis.controller.data.models import TModel @@ -38,7 +38,7 @@ def _read_output( def read_outputs[A: TModel, B: TModel]( - graph: GraphData | GraphBuilder[A, B], + graph: GraphData | Workflow[A, B], storage: ControllerStorage, ) -> dict[str, PType] | PType: """Read the outputs of a workflow graph. @@ -46,7 +46,7 @@ def read_outputs[A: TModel, B: TModel]( The bytes are parsed into Python types if possible. :param graph: The graph to read. - :type graph: GraphData | GraphBuilder + :type graph: GraphData | Workflow :param storage: The storage of the workflow run. :type storage: ControllerStorage :return: The output values. If the graph has a single output port named "value" it @@ -55,9 +55,9 @@ def read_outputs[A: TModel, B: TModel]( :rtype: dict[str, PType] | PType """ output_annotation = None - if isinstance(graph, GraphBuilder): + if isinstance(graph, Workflow): output_annotation = graph.outputs_type - graph = graph.get_data() + graph = graph.data out_ports = list(graph.nodes[graph.output_idx()].inputs.keys()) @@ -77,7 +77,7 @@ def read_outputs[A: TModel, B: TModel]( def read_loop_trace( - graph: GraphData | GraphBuilder, + graph: GraphData | Workflow, storage: ControllerStorage, node_name: str, output_name: str | None = None, @@ -87,7 +87,7 @@ def read_loop_trace( This is useful to track intermediate values in a loop. :param graph: The graph to read. - :type graph: GraphData | GraphBuilder + :type graph: GraphData | Workflow :param storage: The storage of the workflow run. :type storage: ControllerStorage :param node_name: The name of the loop node. @@ -101,8 +101,8 @@ def read_loop_trace( of values for the specified output port is returned. :rtype: list[PType | dict[str, list[PType]]] """ - if isinstance(graph, GraphBuilder): - graph = graph.get_data() + if isinstance(graph, Workflow): + graph = graph.data loc = storage.loc_from_node_name(node_name) if loc is None: msg = f"Loop name {node_name} not found in debug data." diff --git a/tierkreis_visualization/tierkreis_visualization/visualize_graph.py b/tierkreis_visualization/tierkreis_visualization/visualize_graph.py index 145839ea5..59bda9273 100644 --- a/tierkreis_visualization/tierkreis_visualization/visualize_graph.py +++ b/tierkreis_visualization/tierkreis_visualization/visualize_graph.py @@ -2,19 +2,20 @@ import uvicorn -from tierkreis.builder import GraphBuilder +from tierkreis.builder import Graph, Workflow from tierkreis.controller.data.graph import GraphData from tierkreis_visualization.app import app_from_graph_data -def visualize_graph(graph_data: GraphData | GraphBuilder) -> None: +def visualize_graph(graph_data: GraphData | Workflow | Graph) -> None: """Visualize a computation graph in a web browser. :param graph_data: The computation graph to visualize. - :type graph_data: GraphData | GraphBuilder + :type graph_data: GraphData | Graph """ - if isinstance(graph_data, GraphBuilder): - graph_data = graph_data.get_data() + # Accessing the .data of a Graph isn't generally good practice but makes sense here. + if isinstance(graph_data, (Workflow, Graph)): + graph_data = graph_data.data app = app_from_graph_data(graph_data) config = uvicorn.Config(app, host="0.0.0.0", port=8000, loop="asyncio") server = uvicorn.Server(config)