Lesson 2: Writing Workers

In the previous example you wrote a worker only using Tierkreis base functionality. One of the main benefits of using Tierkreis is that you can easily transform existing code and gain the benefits of Tierkeis, e.g., checkpointing and repeteability.

In this example we’re going to look at a common problem and use as tasks in a graph. For this we’re writing a custom worker “my_example_worker”. We’re going to use pytket to:

  1. Define a symbolic circuit

  2. Substitute its symbolic parameters at runtime

Compiling and running circuits will be part of Lesson 3.

Prerequisite

If you haven’t done so, set up a Tierkreis project and install the additional dependencies for pytket and sympy. Ruff is an optional dependency for formatting the autogenerated code.

uv init
uv add tierkreis pytket pytket-qiskit sympy ruff
uv run tkr project init

Setting up the worker

Using the the Tierkreis cli is the easiest the way to set up a worker:

uv run tkr init worker -n my_example_worker

Which will set up the packages in a convenient way for you.

The cli creates the folder structure as seen in the first lesson. There are two files in tkr/workers/my_example_worker which you need to care about which are:

  • tkr_my_example_worker_impl/impl.py here you will implement your workers functionality

  • api/api.py is an autogenerated file which contains the task definitions we will use in the graph.

In impl.py you will find the following:

worker = Worker("my_example_worker")

@worker.task()
def your_worker_task(value: int) -> int:
    return value

As you can see, generating a task for Tierkeis simply means adding @worker.task() to a function. We’re going to use the same pattern to expose pytkets functionality to in our new worker.

Defining the tasks

Important

Not all cells in the following section are not necessary to run this code example. These are example tasks that are meant to sit in your own copy of the my_example_workers impl.py. If you’re running this example from the repository, it will contain a copy of the worker already installed.

For the following assume we’re going to use a simple quantum circuits with three symbols \(a,b,c\). Writing this as a worker task means wrapping it with a task function.

# The worker is added here for validity of the example
# Put the task into impl.py
from pytket.circuit import Circuit, fresh_symbol
from tierkreis import Worker

worker = Worker("pytket_example_worker")


@worker.task()
def symbolic_circuit() -> Circuit:
    a = fresh_symbol("a")
    b = fresh_symbol("b")
    c = fresh_symbol("c")
    circ = Circuit(3)
    circ.Rz(a, 0)
    circ.Rz(b, 0)
    circ.Rz(c, 0)
    circ.measure_all()
    return circ

Using plain pytket you could know substitute the symbols like so:

# For explanation only, this will be a task
from sympy import Symbol

circ.symbol_substitution({Symbol("a"): -1, Symbol("b"): 0, Symbol("c"): 1})

writing this as a worker task means wrapping it with a task function that has a circuit input and output.

Important

Use type hints so that Tierkreis can validate the task during construction. Since values are represented by edges, we need a return value, its not sufficient to mutate the state.

# Put the task into impl.py
from sympy import Symbol


@worker.task()
def substitute(circuit: Circuit, a: float, b: float, c: float) -> Circuit:
    circuit.symbol_substitution({Symbol("a"): a, Symbol("b"): b, Symbol("c"): c})
    return circuit

Similarly we could now define other tasks using pytket and @worker.task() e.g.,compilation and simulation.

Generating stubs

To generate the APIs for all workers, you can use the cli with

uv run tkr init stubs

Depending on your development environment it might be necessary to resatart your language server or uv sync --all-extras to pick up the update changes.

Opaque Types

Tierkreis can use any type that is serializable as in and outputs, e.g., the Circuit type from pytket library. To make such types available without bleeding dependencies into graph code, Tiekreis wraps them as OpaqueType with a reference to the original implementation. In this example the circuit inputs of the tasks would be

circuit: TKR[OpaqueType["pytket._tket.circuit.Circuit"]]

Using the tasks

Now you can use the newly declared tasks in a graph similar to how you used the builtin functionality. You have to import the task API from the worker first which you then can use with a task node. First we declare the graph

# Constructing, put into graphs/main.py
from typing import NamedTuple
from tierkreis.builder import Graph
from tierkreis.controller.data.models import TKR


class PytketInputs(NamedTuple):
    a: TKR[float]
    b: TKR[float]
    c: TKR[float]


graph = Graph(PytketInputs, TKR[Circuit])

and then add the tasks:

# Constructing, put into graphs/main.py
from my_example_worker import substitute, symbolic_circuit  # noqa: F811

circuit = graph.task(symbolic_circuit())
substituted = graph.task(
    substitute(circuit, graph.inputs.a, graph.inputs.a, graph.inputs.a)  # type: ignore
)
workflow = graph.finish_with_outputs(substituted)  # type: ignore

Running the graph

As before you know can run the graph, the circuit we have defined already above

# Running, put into graphs/main.py
from uuid import UUID


from tierkreis.controller import run_graph
from tierkreis.executor import ShellExecutor
from tierkreis.storage import FileStorage, read_outputs


storage = FileStorage(workflow_id=UUID(int=12346), name="Pytket example graph")
storage.clean_graph_files()
executor = ShellExecutor(registry_path=None, workflow_dir=storage.workflow_dir)
run_graph(storage, executor, workflow, {"a": -1, "b": 0, "c": 1})
result = read_outputs(workflow, storage)
print(result)

{'bits': [['c', [0]], ['c', [1]], ['c', [2]]], 'commands': [{'args': [['q', [1]], ['c', [1]]], 'op': {'type': 'Measure'}}, {'args': [['q', [2]], ['c', [2]]], 'op': {'type': 'Measure'}}, {'args': [['q', [0]]], 'op': {'params': ['3.0'], 'type': 'Rz'}}, {'args': [['q', [0]]], 'op': {'params': ['3.0'], 'type': 'Rz'}}, {'args': [['q', [0]]], 'op': {'params': ['3.0'], 'type': 'Rz'}}, {'args': [['q', [0]], ['c', [0]]], 'op': {'type': 'Measure'}}], 'created_qubits': [], 'discarded_qubits': [], 'implicit_permutation': [[['q', [0]], ['q', [0]]], [['q', [1]], ['q', [1]]], [['q', [2]], ['q', [2]]]], 'phase': '0.0', 'qubits': [['q', [0]], ['q', [1]], ['q', [2]]]}