Runtime Serialization

We talked a lot about code manipulation, but in doing so we skipped a small but important step: serialization, the ability to encode something in such a way that we can store it, then fetch, decode, and resume its operation later. It was one of John von Neumann's greatest insights that code can, in fact, be stored as data – so entire programs, not just their inputs and outputs, can be serialized and saved into the same machine memory as everything else.

Before that, in what's known as the Harvard architecture, instructions were conceptualized as a distinct kind of information: programming required wiring plugboards of vacuum tubes or feeding punch cards into a reader. Von Neumann's architecture is what gave rise to modern programming; most of the devices we use today adhere to this new paradigm, because only by making code a modifiable medium could we conceive of such things as compilers, installable software, and general-purpose computing. Needless to say, without it we wouldn't have metaprogramming, either.

Now, Python is an interpreted language, so serializing its programs is as simple as treating their source code as strings. Sometimes, however, a program is more than just its code: whatever it does, it also does it on something, and while it may seem self-evident that data can be stored as data, it's actually not as straightforward as it sounds. You see, because the way information is laid out in memory differs from how it's encoded on a disk, some data types are more serializeable than others.

To wit, numbers can be represented in binary form, as a sequence of bits over a fixed number of bytes – one byte for small numbers between 0 and 255, also good enough to represent ASCII characters, and 4 or 8 bytes for ranges that suffice for most purposes. Strings, however, are trickier, because of their variable length: given one byte per character, how many of those bytes do we need to read to reconstruct a particular string? In C, they decided to use a terminator: a special nul-byte, '\0', which denotes the string is over; in Pascal, a length prefix is used instead, with the first 4 bytes reserved for a number that indicates how long this particular string runs. More complex data structures – from lists and dictionaries to images and 3D videos – are all subject to the same constraints and considerations; different formats make different decisions with different trade-offs as to how exactly their objects should be flattened onto a one-dimensional sequence of bits.

Pickling Objects

Conveniently, Python solves most of these challenges with the standard pickle module: a serialization protocol uniquely tailored to Python's data model, which is able to capture and restore most of it without us having to worry about how. All the native types are handled out of the box:

>>> import pickle
>>> pickle.dumps({"x": [1]})
b'\x80\x04\x95\r\x00\x00\x00\x00\x00\x00\x00}\x94\x8c\x01x\x94]\x94K\x01as.'
>>> pickle.loads(_)
{'x': [1]}

And for custom objects, a simple convention allows classes to articulate their state in simpler, native terms – then reconstruct themselves from such on demand:

>>> class A:
...     def __init__(self, x):
...         self.x = x
...     def __repr__(self):
...         return f"A({self.x})"
...     def __getstate__(self):
...         return {"x": self.x}
...     def __setstate__(self, state):
...         self.x = state["x"]

>>> a = A(1)
>>> pickle.dumps(a)
b'\x80\x04\x95\x1f\x00\x00\x00\x00\x00\x00\x00\x8c\x08__main__\x94\x8c\x01A…'
>>> pickle.loads(_)
A(1)

In fact, since most Python objects are just wrappers around a simple dictionary, anyway, even this is usually unnecessary; the default implementation of __getstate__ would just return self.__dict__, and that of __setstate__ would inject it back with self.__dict__.update(state):

>>> class A:
...     def __init__(self, x):
...         self.x = x
...     def __repr__(self):
...         return f"A({self.x})"
...     # no __getstate__ or __setstate__, and yet…

>>> a = A(1)
>>> pickle.dumps(a)
b'\x80\x04\x95\x1f\x00\x00\x00\x00\x00\x00\x00\x8c\x08__main__\x94\x8c\x01A…'
>>> pickle.loads(_)
A(1)

The only catch is that the class has to be importable, so Python can restore the wrapper as well as its state – which is why you might've glimpsed __main__ and A in the serialized string above. Note that it does so without calling the constructor directly; something along the lines of:

cls = __import__(f"{module_name}.{class_name}")
obj = object.__new__(cls)
obj.__dict__.update(state)

Advanced Pickling

Still, __getstate__ and __setstate__ can come in handy on some occasions: as I mentioned, not all data is serializeable out of the box. Files, for instance, or sockets, or any other "I/O objects" that represent a live, temporary handle to some external resource – those things can't really be described by a static sequence of bytes. So if we have:

>>> class A:
...     def __init__(self, path):
...         self.path = path
...         self.file = open(path)

We need to account for the file explicitly:

>>> class A:
...     # … same as before
...     def __getstate__(self):
...         state = self.__dict__.copy()
...         del state["file"]
...         return state
...     def __setstate__(self, state):
...         self.__dict__.update(state)
...         self.file = open(self.path)

In other cases, object construction is more complicated still: a singleton class, for example, should always have one instance at most:

>>> class A:
...     _instance = None
...     @classmethod
...     def get(cls):
...         if cls._instance is None:
...             cls._instance = cls()
...         return cls._instance

>>> # use a = A.get() instead of a = A()

Pickling-then-unpickling this instance would invariably create a new object, and hijacking __setstate__ would not help either, because by the time it's called the rogue object will have already been created. So instead, we can use __reduce__: it returns a string (to import a global), or a tuple with a callable and its arguments, to construct the object in a non-standard way. This tuple can also include a state to be injected afterwards, as well as a couple of iterables to aid in the creation of list-like and dict-like objects, whose state is a sequence of items rather than an attributes dictionary. In our case, however, the basic pair will suffice:

>>> class A:
...     # … same as before
...     def __reduce__(self):
...         return self.get, ()

>>> a1 = A.get()
>>> pickle.dumps(a1)
b'\x80\x04\x956\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x94\x8c\x07getattr…'
>>> a2 = pickle.loads(_) # calls A.get()
>>> a1 is a2
True

In fact, an even simpler way to have a singleton is by using a unique, global variable to begin with – so we could have done:

>>> class A:
...     def __reduce__(self):
...         return "a"

>>> a = A() # only ever use this
>>> pickle.dumps(a)
b'\x80\x04\x95\x12\x00\x00\x00\x00\x00\x00\x00\x8c\x08__main__\x94\x8c\x01a…'
>>> a2 = pickle.loads(_)
>>> a is a2
True

Minute details aside, the point I'm trying to illustrate is that Python has a builtin way to take care of serialization for us, and a few handy hooks ready for whenever it might fall short. This means almost any Python program can be flattened into a sequence of bytes that can be written to a file or sent over the wire, only to resume operation at a different time or on a different machine – all with very little effort on our part. Incidentally, this is precisely how the standard multiprocessing library distributes shared functions and values across separate processes – but, as always, there are some interesting caveats to address.

Pickle Limitations

One such is a warning that is often mentioned, about how we should never run (that is, "unpickle") payloads from an untrusted source. The protocol's dynamicity, which makes it so simple and flexible all at once, also means that a hacker can craft an object in such a way that, when deserialized, it will execute some arbitrary code in our program.

Refer back to the first singleton example, where we used __reduce__ to replace the standard object construction with a custom method; see how its pickled data contains builtins and getattr amidst the noise? That's because its reconstruction would look like this:

>>> a = builtins.getattr("A", "get")()

That is, to serialize the construction method A.get, Python has to resort to a more basic mechanism that would let it reconstruct it before it can be used for further reconstruction. If we were to do this instead:

>>> class A:
...     def __reduce__(self):
...         return print, ("hacked!",)
>>> a = A()
>>> pickle.dumps(a)
b'\x80\x04\x95$\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x94\x8c\x05print…
  \x94\x93\x94\x8c\x07hacked!\x94\x85\x94R\x94.'

Then to deserialize our object, Python would run:

>>> a = builtins.print("hacked!")

Which, unsurprisingly, results not in an object, but in:

>>> pickle.loads(_)
hacked!

And that, even if the malicious A definition is not part of our codebase, because it has been reduced to generic enough constructs as to be bootstrappable in any Python environment. An actual hack might look like:

>>> class A:
...     def __reduce__(self):
...         return os.system, ("rm -rf /",)

Which would delegate object construction to the builtin os.system, deleting our filesystem instead.

Another limitation of the Pickle protocol is that try as they might, the core developers couldn't guarantee that absolutely everything could be reconstructed. Dynamic enough objects, like functions defined in a local scope, don't play along:

>>> def f():
...     def g():
...         return 1
...     print(pickle.dumps(g))
>>> f()
AttributeError: can't get local object 'f.<locals>.g'

Because Python simply doesn't provide hooks deep enough to cleave their definition out of an ephemeral scope. Incidentally, that's the reason we can't distribute such objects across multiple processes:

>>> def distribute():
...     def f(i):
...         print(i)
...     with multiprocessing.Pool() as pool:
...         pool.map(f, range(3))
>>> distribute()
AttributeError: can't get local object 'distribute.<locals>.f'

It's not a problem with the multiprocessing library – it's just that it can't serialize the work it attempts to distribute; or, put more accurately, it can serialize it, but the workers receiving the payload are not able to reconstruct it reliably into their own memory. Moving f to the global scope should resolve this issue, because then they'd be able to find it – at least, if it was defined in a proper file. In the interpreter, this still wouldn't work, for reasons that are hopefully starting to make sense: __main__.f is not something that other processes can find, because __main__ is as ephemeral a namespace as local scopes.

More interesting, though, are objects that simply can't be serialized, no matter their scope. Take generators:

>>> def gen():
...     i = 1
...     while True:
...         yield i
...         i += 1
>>> g = gen()
>>> next(g)
1
>>> pickle.dumps(g)
TypeError: cannot pickle 'generator' object

It would be pretty neat if we could play a generator until a certain point, then take it (in its paused state) and save it for later, so that loading it and playing it again would resume from the point at which it stopped. Unfortunately, that's not something Python can do: the generator's state, as well as the pointer to its next instruction, are part of the program's internal runtime, and as such they can't be flattened into a simple sequence of bytes. Or can they?

Patching Memory

This should be possible, at least in theory: a machine's memory, after all, is also a sequence of bytes, which is why virtual machines can be paused, even mid-process, only to be fully restored and resumed at a later time. So what if some data isn't stored in the exact same format as on disk? As long as we spread the same zeroes and ones that we had before in its memory space, how would the computer know the difference, and why would it care?

In case of a generator, our task is easier still: we don't need to restore the entire memory image, only the parts that pertain to its execution – namely, the local variables and instruction pointer. Both are available via the generator's frame object:

>>> g.gi_frame.f_locals
{'i': 1}
>>> g.gi_frame.f_lasti
10

Where the instruction pointer is simply a bytecode index, keeping track of the position at which the generator has paused. Unsurprisingly, it's a yield statement:

>>> g.gi_frame.f_code.co_code[10]
86
>>> opcode.opname[86]
'YIELD_VALUE'

And note that I'm actually using Python 3.10 here; serializing generators got considerably harder in later versions, due to all the optimizations we've mentioned before: 3.11 reworked frame objects almost beyond recognition, and 3.13 went as far as removing PyFrame_LocalsToFast – a function we're about to lean on – from the public C API altogether (PEP 667). So treat this section as a version-specific exhibit, and run its code on CPython 3.10, like the examples insist; the point it makes about memory being just bytes is timeless, even if the offsets aren't. It's going to be tricky enough even in this version, because this instruction pointer is not writable:

>>> g.gi_frame.f_lasti = 0
AttributeError: attribute 'f_lasti' of 'frame' objects is not writable

And the locals, while they can be changed, also seem to ignore those changes:

>>> g.gi_frame.f_locals["i"] = 10
>>> next(g)
2

Still, remember how I said that memory is just a sequence of bytes, too? That means we can always brute force our way in by accessing and overwriting it directly; we just need to be careful in how we do it, lest we crash the interpreter. It's fragile, and highly version specific, but it can be done; especially in CPython, where the memory address of an object is revealed to us by its ID:

>>> id(g.gi_frame)
4343682448

Before we can write to it, though, we need to give it a shape: and here we'd have to replicate the exact structure of the frame object, per our Python version, so that the values we write land at the right offsets in relation to everything else. This can be done with ctypes, after rummaging through the interpreter's source code to unearth the relevant fields and their types:

>>> import ctypes
>>> class PyObject(ctypes.Structure):
...     _fields_ = [
...         ('ob_refcnt', ctypes.c_ssize_t),
...         ('ob_type', ctypes.c_void_p),
...     ]
>>> class PyFrameObject(ctypes.Structure):
...     _fields_ = [
...         ('ob_base', PyObject),
...         ('ob_size', ctypes.c_ssize_t),
...         ('f_back', ctypes.c_void_p),
...         ('f_code', ctypes.c_void_p),
...         ('f_builtins', ctypes.c_void_p),
...         ('f_globals', ctypes.c_void_p),
...         ('f_locals', ctypes.c_void_p),
...         ('f_valuestack', ctypes.c_void_p),
...         ('f_trace', ctypes.c_void_p),
...         ('f_stackdepth', ctypes.c_int),
...         ('f_trace_lines', ctypes.c_char),
...         ('f_trace_opcodes', ctypes.c_char),
...         ('f_gen', ctypes.c_void_p),
...         ('f_lasti', ctypes.c_int),
...         # we won't write beyond this point, so might as well stop here
...     ]

Now, if I got everything right, casting the address of our frame into this mold will allow us to access its memory in such a way that will actually be writable – but first, let's read it to see that the values align correctly:

>>> f = PyFrameObject.from_address(id(g.gi_frame))
>>> f.f_lasti
5

This seems off, since g.gi_frame.f_lasti was 10, but it's actually good news: internally, Python stores the instruction pointer as an instruction counter, meaning we're currently at the fifth instruction – and since each instruction takes two bytes, its offset in co_code is indeed 5 * 2 = 10. To verify it further, we can check that f_locals points to the same address where the locals are stored:

>>> f.f_locals
4312924480
>>> id(g.gi_frame.f_locals)
4312924480

And it does indeed. And to test that our mold is writable, we can do:

>>> f.f_lasti = 3
>>> g.gi_frame.f_lasti
6 # changed to 2 * 3!

Generator Genealogy

Our serializer can thus be exceedingly simple:

>>> def serialize_generator(g):
...     return {
...         "lasti": g.gi_frame.f_lasti,
...         "locals": g.gi_frame.f_locals.copy(),
...     }

And our deserializer can take these values, recreate the generator, and patch them in:

>>> def deserialize_generator(gen, state):
...     g = gen()
...     f = PyFrameObject.from_address(id(g.gi_frame))
...     f.f_lasti = state["lasti"] // 2
...     f.f_locals = id(state["locals"])
...     return g

Let's give it a go:

>>> state = serialize_generator(g)
>>> g_restored = deserialize_generator(gen, state)
>>> next(g_restored)
UnboundLocalError: local variable 'i' referenced before assignment

Ugh. What happens here is that Python actually converts its local variables to an optimized array for high performance access; us just patching them in bypasses this process, confusing the interpreter as to what was assigned and what wasn't. Incidentally, this is also the reason our writing to gi_frame.f_locals was ignored: the original values were already converted, and Python continued to use them without noticing our change. Luckily, this also gives us a way at a solution: instead of patching the locals in, we can update the originals after all, as long as we also manually trigger said conversion:

>>> def deserialize_generator(gen, state, *args):
...     # … same as before
...     g.gi_frame.f_locals.update(state["locals"])
...     ctypes.pythonapi.PyFrame_LocalsToFast(
...         ctypes.py_object(g.gi_frame),
...         ctypes.c_int(0),
...     )
...     return g

And there you have it:

>>> g_restored = deserialize_generator(gen, state)
>>> next(g_restored)
3 # right where we stopped!

This is cool – but it's still not exactly what I had promised, because state is a dictionary, not something we can write to disk. And of course, we can pickle it with no problem, but to unpickle it – if you check out the deserializer's signature – we need the original gen; that is, we have to recreate the generator before applying our clever patch.

In principle, the generator function can be pickled; the problem is to infer it automatically, because generators don't retain any obvious reference to their creator. What's worse, if said creator expects some arguments (like, if gen would accept an initial i to start from), we'd have to somehow retain those, too; only then can we pickle the whole lot and honestly claim to have serialized everything properly.

So, how do we find something that has no pointer? Another trick is in order: while there is nothing as convenient as g.gi_func, we do have g.gi_code – that is, the generator's code object, which it shares with the function that created it. This code object doesn't have a pointer to this function, either, but at least its unique enough that we can ask the garbage collector for its referrers – all the objects that currently hold a reference to it – and investigate those. As you can see, this list includes the generator objects as well as their creator, but only the latter is an actual function, so we can do:

>>> import gc
>>> from types import FunctionType

>>> def find_generator_function(g):
...     for obj in gc.get_referrers(g.gi_code):
...         if isinstance(obj, FunctionType) and obj.__code__ is g.gi_code:
...             return obj
...     raise RuntimeError(f"no generator function found for {g}")

>>> find_generator_function(g)
<function gen at 0x…>

Now that we have the function, all that's left is to figure out how to call it – and here, the key insight is that it doesn't actually matter all that much. Whatever arguments we pass in, they're going to become part of the local scope, and that scope will be overwritten, anyway; all we need is for the invocation to work, and that means we need to match the arity – make sure we pass in the correct number of, say, Nones, both positionally and by keyword. Since all arguments can be passed by position except keyword-only ones (and variadic parameters, if there are any, will default to () and {}), this turns out to be rather simple:

>>> import inspect

>>> def placeholder_args(f):
...     s = inspect.signature(f)
...     num_args, kwonly_names = 0, []
...     for param in s.parameters.values():
...         if param.kind == param.KEYWORD_ONLY:
...             kwonly_names.append(param.name)
...         elif param.kind not in (param.VAR_POSITIONAL, param.VAR_KEYWORD):
...             num_args += 1
...     args = [None] * num_args
...     kwargs = dict.fromkeys(kwonly_names)
...     return args, kwargs

So, for example:

>>> def f(a, /, b, c=1, *d, e, f=2, **g):
...     pass

>>> args, kwargs = placeholder_args(f)
>>> args
[None, None, None]
>>> kwargs
{'e': None, 'f': None}
>>> f(*args, **kwargs) # works

Pickling Generators

With that, we're finally ready to pickle a generator:

>>> def pickle_generator(g):
...     return pickle.dumps({
...         "func": find_generator_function(g),
...         "lasti": g.gi_frame.f_lasti,
...         "locals": g.gi_frame.f_locals.copy(),
...     })

>>> def unpickle_generator(payload):
...     state = pickle.loads(payload)
...     args, kwargs = placeholder_args(state["func"])
...     g = state["func"](*args, **kwargs)
...     f = PyFrameObject.from_address(id(g.gi_frame))
...     f.f_lasti = state["lasti"] // 2
...     g.gi_frame.f_locals.update(state["locals"])
...     ctypes.pythonapi.PyFrame_LocalsToFast(
...         ctypes.py_object(g.gi_frame),
...         ctypes.c_int(0),
...     )
...     return g

And there you have it:

>>> def gen(i):
...     while True:
...         yield i
...         i += 1

>>> g = gen(10)
>>> next(g)
10
>>> pickle_generator(g)
b'\x80\x04\x95;\x00\x00\x00\x00\x00\x00\x00}\x94(\x8c\x04func\x94\x8c\x08…'
>>> g_restored = unpickle_generator(_)
>>> next(g_restored)
11

Fast-Forward Generators

This is all pretty rad. Not so much because pickling generators is such an urgent need (although it has its uses), but because even experienced Python developers would swear this is impossible, and yet there it is; we just did it. This attitude – that anything can be done, really, because at the end of the day it's all just a sequence of ones and zeros, so if we rearrange them just right they will mean whatever we want them to mean – this attitude is, in a nutshell, the secret to the virtuosity of some of the best hackers I know.

There is, however, a complementary attitude – and this one is the secret to the lucidity of some of the best system designers I know. It entails an aversion to complexity, a faith that there must be a better way, and the daring to reimagine things so radically as to arrive at completely novel solutions. When both attitudes are mastered, they interplay beautifully: we've already seen it in action when we figured out a way to rewrite a function's bytecode, then ditched our complicated procedure in favor of simply rewiring its global scope. Here, I'd like to call on it again, because upgrading our current method to a newer version of Python would be a nightmare: 3.11 introduced a massive refactoring of how frame objects are represented internally, and 3.13 removed PyFrame_LocalsToFast from the public C API entirely – and while it would still be possible to patch up (everything is, remember?), it would also get pretty ugly, and break again on the next release; in other words, there must be a better way.

What if instead of reconstructing a generator, we'd replay it from the very beginning – but in such a way as to fast-forward to the moment at which it was last paused? This way, we won't have to deal with any murky internals, just keep track of how many yield statements happened, and breeze through them when our generator is restored. Suppose, then, that we have:

>>> def serializeable(func):
...     def call(*args, **kwargs):
...         return serializeable_generator(func, *args, **kwargs)
...     return call

And we use it like so:

>>> g = serializeable(gen)(10)

This would give us a new object, which retains all the relevant information:

>>> class serializeable_generator:
...     def __init__(self, func, *args, **kwargs):
...         self.func = func
...         self.args = args
...         self.kwargs = kwargs
...         self.steps = 0
...         self.g = self.func(*self.args, **self.kwargs)

And emulates an iterator as well:

>>> class serializeable_generator:
...     # … same as before
...     def __iter__(self):
...         return self
...     def __next__(self):
...         self.steps += 1
...         return next(self.g)

This way, serialization is as simple as:

>>> class serializeable_generator:
...     # … same as before
...     def serialize(self):
...         return pickle.dumps({
...             "func": self.func,
...             "args": self.args,
...             "kwargs": self.kwargs,
...             "steps": self.steps,
...         })

And deserialization fast-forwards through as many steps as were advanced:

>>> class serializeable_generator:
...     # … same as before
...     def fast_forward(self, steps):
...         self.steps = steps
...         for _ in range(self.steps):
...             next(self.g)

>>> def deserialize(payload):
...     state = pickle.loads(payload)
...     g = serializeable_generator(
...         state["func"],
...         *state["args"],
...         **state["kwargs"],
...     )
...     g.fast_forward(state["steps"])
...     return g

Once again, read it and weep:

>>> next(g)
10
>>> g.serialize()
b'\x80\x04\x95?\x00\x00\x00\x00\x00\x00\x00}\x94(\x8c\x04func\x94\x8c\x08…'
>>> g_restored = deserialize(_)
>>> next(g_restored)
11

Snapshots

Needless to say, this is much simpler and much more robust than what we had before – but it's also not exactly the same, is it? If our generator would take a long time to compute certain values, or generate certain side-effects as it runs along, all of this would recur every time we replay it, fast-forward or no. Take this contrived example:

>>> def gen():
...     x = random.randint(1, 100)
...     print(f"selected {x}")
...     yield # pause
...     if x % 2 == 0:
...         print("calculating...")
...         time.sleep(10) # emulate heavy calculation
...         print(f"calculated {x + 1}")
...     else:
...         print("sending request...")
...         response = requests.post(URL, data=str(x)).json()
...         print(f"received {response['result']}")
...     yield # pause

It has all the problems described above, and then some: it might take 10 seconds, it might send a request (I/O is always a side-effect we wouldn't want to repeat), and on top of everything, it's indeterministic – it might choose one branch or another every time we re-run it. Fortunately, all of this can be mitigated with one simple convention: that we put any heavy, non-idempotent, or non-deterministic behavior in separate functions, and decorate those with @snapshot.

The snapshot decorator, in turn, will keep track of the index of each function call in a given run and cache its results. This way, when we replay a generator, we won't have to recalculate any of its problematic values: they should be available instantly, predictably, and with no side-effects from that cache, without us having to re-invoke any functions that had already run.

To make sure each cache is matched to the particular generator in whose context it is considered, let's use the standard contextvars module, which lets us manage values per synchronous flow, be it a thread or a task. Why not a plain module attribute? Because generators nest and interleave: the moment one snapshotted generator drives another, or two of them alternate on the same thread, a single global "current generator" would be overwritten mid-flight – a ContextVar, set on entry and reset on exit, keeps the bookkeeping correct per flow. Before the serializeable generator delegates to the original one, then, it should do:

>>> from contextvars import ContextVar
>>> current_generator = ContextVar("current_generator", default=None)

>>> class serializeable_generator:
...     # … same as before
...     def __next__(self):
...         self.steps += 1
...         return self._step()
...     def _step(self):
...         token = current_generator.set(self)
...         try:
...             return next(self.g)
...         finally:
...             current_generator.reset(token)

That is, set itself as the "current generator" for this particular step, then roll back to the previous one when it's over; this way, nested calls will replace and restore each other seamlessly. Now, if we add a couple caches:

>>> class serializeable_generator:
...     # … same as before
...     def __init__(self, func, *args, **kwargs):
...         # … same as before
...         self.rcache = []
...         self.wcache = []

We can write the snapshot decorator to read values from rcache if it isn't empty, or compute and write them to wcache if it is:

>>> def snapshot(func):
...     def wrapper(*args, **kwargs):
...         g = current_generator.get()
...         if g and g.rcache:
...             return g.rcache.pop(0)
...         value = func(*args, **kwargs)
...         if g:
...             g.wcache.append(value)
...         return value
...     return wrapper

All that's left is to include the full cache (that's wcache) in the serialized payload, and restore both when we fast-forward, so wcache keeps accumulating, and rcache can be consumed by calls for which a snapshot already exists:

>>> class serializeable_generator:
...     # … same as before
...     def serialize(self):
...         return pickle.dumps({
...             # … same as before
...             "cache": self.wcache,
...         })
...     def fast_forward(self, cache, steps):
...         self.rcache = cache.copy()
...         self.wcache = cache.copy()
...         self.steps = steps
...         for _ in range(self.steps):
...             self._step()

>>> def deserialize_generator(payload):
...     # … same as before
...     g.fast_forward(state["cache"], state["steps"])
...     return g

It's a bit mind-bending, so let's recap: a generator sets itself as the "current generator" before each step, then runs it – so if this step calls any @snapshot-decorated function, one of two things will happen. If this function is called for the first time, as is always the case on the generator's initial run, rcache will be empty, so the function will execute fully and stash its result in wcache. If we serialize our state at this point, it will include this result – so when we deserialize and replay the generator, it will be ready for us (in rcache and wcache both). Which brings us to the second scenario: when we fast-forward, one result per each function called should be available in rcache, so this time around it won't be recalculated, simply popped out of the cache. Thus, all the steps for which values were already computed will complete instantly, and we'll fast forward to the point where we last paused without any disturbance. As the execution continues, more values will be added to wcache, so if we serialize our state again, everything necessary to replay it up to our new position will have accumulated just the same.

To demonstrate, let's extract our problematic behaviors:

>>> @snapshot
... def random_number():
...     return random.randint(1, 100)

>>> @snapshot
... def heavy_calculation(x):
...     print("calculating...")
...     time.sleep(10)
...     return x + 1

>>> @snapshot
... def http_request(x):
...     print("sending request...")
...     return requests.post(URL, data=str(x)).json()["result"]

And change the original generator to make use their snapshot-producing versions:

>>> def gen():
...     x = random_number()
...     print(f"selected {x}")
...     yield
...     if x % 2 == 0:
...         result = heavy_calculation()
...         print(f"calculated {result}")
...     else:
...         result = http_request(x)
...         print(f"received {result}")
...     yield

Now, if we serialize and replay it, nothing will repeat; whatever random number was selected will be retained when the run is resumed:

>>> g = serializeable(gen)()
>>> next(g)
selected 42 # an even number, so…
>>> g_restored = deserialize_generator(g.serialize())
>>> next(g_restored)
calculating... # and 10 seconds later:
calculated 43

Or:

>>> g = serializeable(gen)()
>>> next(g)
selected 75 # an odd number, so…
>>> g_restored = deserialize_generator(g.serialize())
>>> next(g_restored)
sending request... # suppose the server increments its input by one
received 76

What's more, if our generator would have further yields, serializing, deserializing, and replaying it again would not repeat either the heavy calculation or the request: both the random number and the subsequent result would be cached and reused:

>>> g_restored_again = deserialize_generator(g.serialize())
>>> next(g_restored_again)
selected 75
received 76 # no request sent
...

Durable Execution

I realize this exercise might seem a bit academic; both our silly generator, and the overall purpose of serializing one such, call for some additional context – like, why is any of this even useful?

Turns out, there are quite a few startups whose core value proposition is being able to do exactly that – and the associated buzzword they use is "durable execution". You see, generators are the backbone on which coroutines are built, and asynchronous programming is all the rage in modern Python programming. It works great for small, quick tasks that can be densely interlaced while waiting for I/O, like a web app serving files or database records with some minimal processing, which is what most web apps do. However, in certain scenarios one needs to support long-running tasks: do something, then wait for an email, or a slack message, or a user's manual review. In such situations, the coroutine can await for hours, even days – and if the server happens to crash, update, or otherwise reboot during this time, its state would be lost. Hence durable execution: such use cases require some mechanism to serialize the runtime and persist it, so it can be restored if necessary.

Incidentally, this is exactly what we've been doing; the only difference, really, is that a coroutine's await statement gives back a value. To be fair, so does a generator's yield – if we'd use g.send(value) instead of next(g), the only change we'd have to implement is to swap steps for a list of such values to inject on each step, rather than a simple counter. Of course, we'd also have to rewrite our serializeable_generator to adhere to the interface expected by an event loop that's running tasks, rather than its current schtick of pretending to be an iterator – but fleshing out such a system is quite technical and quite boring, so it's left as an exercise to the reader (this chapter's examples include a small sketch of such an event loop, restoring its tasks from a durable store, if you'd like a head start). Our core insights, of how to fast-forward a piece-wise flow from its beginning and snapshot any turbulent steps so as to restore their results without recomputation – those were the fun part, anyway, and they remain pretty much the same.