lab notes

Lab note #051 Cell composition and lazy evaluation

Wil

• 8 min read
Lab note #051 Cell composition and lazy evaluation

Last week, I merged the runtime worker into the kernel worker. Originally, I wanted that to run multiple kernels at once, so users can run languages in a mix of cells, or they can run more than one cell in parallel. It was too hard to trace execution over a message bus between workers. I had to do things like store promises with IDs and resolve them based on the ID after the result came back. But I think this is too big to bite at the moment, so I merged both into the same worker. However, I intend to keep as much runtime-specific code in Javascript, and userland code in Python. The interface between the Python kernel and the Javascript runtime driving the incremental computation is still a bit fuzzy at the moment.

Confused at how Algebraic Effects Fit

So good progress in shaping the notebook, but got stuck a bit last week. I wanted to consolidate the idea of using algebraic effects in a React-like system, but couldn't quite reconcile the two ideas until I made some compromises. In React's hook-based component pattern, you build a component with hooks and return a component tree.

A surprising perspective I've discovered through Witten's writings is that the returned component tree can be viewed as a continuation. To the user, it just feels like you're declaring how a component should be composed of other components. But another perspective that's just as valid is that you're communicating to the React runtime what computations you want it to execute next–a continuation. Because I'm composing cells in notebook, I have a cell tree, but conceptually, it's the same as a component tree.

Hooks are weak-forms of algebraic effects. The basic ones provided by React either keep state (`useState`) or execute side effects (`useEffect`). Unlike algebraic effects that are native their host language, hooks have rules. As far as I can tell, this is for ease of implementation, as well as priming front-end developers for a world of concurrent rendering. From another perspective, it's a way of generating dynamic nodes (what people call self-adjusting computation) in a computational DAG that's more readily understandable (to the average web dev) than what Jane Street's Incremental library does.

Algebraic effects are like resumable exceptions. True implementations can use delimited continuations to resume the computation after a handler has executed on a raised effect.

My mind confused all of these concepts. If hooks are weak-forms of algebraic effects, how do restrictions on hooks change if they were real algebraic effects? If component trees are continuations, should it use the same continuation mechanism as algebraic effects? If so, does it mean we have the same unit of work between algebraic effects and component trees, and are they similar enough to warrant it? If the component tree represents a computational DAG, how do you gather up the results in a tree?

In addition, my naive implementation of algebraic effects didn't support multiple resumption, due to being based on generators, which are one-shot continuations, rather than delimited continuations. I almost pulled out all of the algebraic effect code and thought about just reimplementing straight up React.

However, I commited to a couple design decisions that helped make it more clear to me.

  • Both raised algebraic effects and cell trees returned from cells are descriptions of computation that need to happen next. For now, I'll call them effect description tree
  • The effect description tree are expanded, like how expressions are expanded. Algebraic effects are raised and handled sequentially before the cell's cell tree is returned.
  • Then the expressions are converted into fibers, which are stateful units of work.
  • The fibers are passed to the scheduler which decides what order to evaluate the different fibers that are in the queue within a specific frame rate budget. For now, it will just evaluate in FIFO order.
  • Once the scheduler decides the order, it pulls out however many fibers from the queue and passes it to the rebuilder, which decides which fibers to reduce, like how expressions are reduced, if at all. This is where incrementality happens.
  • If the computational result of a fiber is more effect description trees, they are pushed in the back of the queue.
  • Results of leaf cells can be communicated and reduced back up the tree, looking for an ancestor Then cell, through which the computational DAG can further reduce or further expand.

So unlike build tools, where the entire computational DAG is known ahead of time, this scheme discovers and expands the computational DAG over time, and it does it conditionally based on the state kept in the fiber and the results returned by the effect handlers.

Cell composition

I spent some time doing some dev ex with how a cell tree could be composed. I first started with some JSX, and inspiration from use.GPU's Gather component and its idea of "yeet-reduce". Yeet-reduce is just a way of keeping data-flow unidirectional in a tree-like computational graph, by allowing child components to communicate results back up to some parent, and the parent doing a reduction on those results.

Here's a short example of how it might look to fetch docs from a url and writing them to file, and only when that's finished do we start taking every doc, chunking them, and generating embeddings.

return <Suspense>
        <FetchDocsStage url={url} corpus_path={corpus_path} />
        <Then>
            <DirectoryReader stream={(doc) => 
                <HtmlChunker doc={doc} stream={(chunk) => 
                    <GenerateEmbeddings chunk={chunk} />
                } />
            } />
        </Then>
    </Suspense>

Here, we first have a stage in the pipeline where we fetch documents from a url and store it into a path in FetchDocsStage. When it's finished, the result is a flag that indicates it's finished. That flag gets communicated back up to Suspense, which then passed it to Then. Inside of Then, we continue the pipeline with reading a directory full of docs, chunking each doc, and then generating embeddings from each chunk. In this case, Then acts as an onReady callback for when we're finished fetching all the docs and writing it into some permanent storage.

Here, I imply that nested cells are computed in sequence, and sibling cells can be computed concurrently. However, the user can change this default with Then, as it's understood Then is special and is a continuation of the child cell computations. We can probably add other kinds of syntax sugar in the future that's equivalent to andThen or |> (called the forward pipe operator in Elm)

But since we're doing this in Python, and I'm wary of introducing pre-processors at this stage:

return Suspense(
    FetchDocsStage(url=url, corpus_path=corpus_path)
    Then(
        DirectoryReader(lambda doc:
          HtmlChunker(lambda chunk:
              GenerateEmbeddings(doc=doc, chunk=chunk)
          ), doc=doc)
        )
    )
)

I was also concerned about how nested components would pass values from one stage to the next. In React, we get around it by using render props. For now, I just allow a cell's children to be a tuple of other cells, or a lambda.

# Different options of passing down stuff into children

# data passed implicitly from parent to child
return DownloadDocs(
    WriteToFile(path = path),
    url = url
)

# using a "render prop"
return DownloadDocs(
    url = url,
    render = lambda doc:
        WriteToFile(data = doc, path = path),
)

# using syntax sugar to chain cells together
return DownloadDocs(
    url = url,
).andThen(
    WriteToFile(data = response, url = url)
)

# ended up picking this--the children can be a lambda
return DownloadDocs(
    lambda response:
        WriteToFile(data = response, url = url),
    url = url,
)

It's not ideal, but it's good enough. Honestly, I could do with Ruby's block syntax here. It'd be so nice.

Haskell's Lazy Evaluation

The other thing I looked at was how lazy evaluation in Haskell works, as well as watching a video on the topic. It struck me that incremental computation is very much related to lazy evaluation, but I never came across any blog posts or papers in literature that was inspired by it. [1]

Haskell's lazy evaluation evaluates pure functions in a non-strict way, which means arguments are not evaluated first. Instead the function is expanded into what's called a weak head normal form, which I simplistically think of as replacing the function with its body definition referenced by all its unevaluated arguments. Keep doing this until we come across a function definition that's evaluated strictly, which is usually a pattern match.

Pattern match represents a divergence in the output of a pure function based on some condition. So in order to figure out which branch the output is going to be, we need to know the value of that condition. To determine which argument expressions need to be evaluated, we find the ones that are needed to find the conditional value of the pattern natch.

Once we have the conditional value of the pattern match, we can choose the selected branch and skip computing any of the other branches. In addition, any previously evaluated pure functions were memoized. Given the same inputs, we always get the same output and can skip computation this way also.

What does lazy evaluation mean with respect to evaluating pure functions that can raise algebraic effects?

Pure functions that raise algebraic effects can be viewed as a pure function that takes the effect handler output as yet another input. It's "Functional Core, Imperative Shell", but with an inverted control. Imperative results don't need to be executed first and then passed down through the functional core. Instead imperative results can be requested from inside the functional core, and resumed once the imperative result is handled.

Asking GPT about current implementations of algebraic effects in Koka, Eff, OCaml, and Unison, I come to understand that when effects are raised, it's essentially a blocking call. Because the compilers don't know if any subsequent part of the function depends on the result of the effect, they wait.

If a user wants to perform concurrent concurrent algebraic effects, they typically do one of two things:

  • raise a single composite effect that composes multiple effects intended to be handled concurrently
  • effects are raised in the order they're encountered, but instead of returning a result, they can return asyncs or threads. These asyncs/threads can be awaited/joined further down the pure function body.

In both instances, algebraic effects are raised in sequence, as they're encountered within a pure function. This is basically strict evaluation.

How can we do lazy evaluation for algebraic effect handling? That would effectively mean that when an algebraic effect is yielded, we don't immediately handle the effect. Instead, we store all raised effect descriptions inside of the cell tree effect description. Then we treat the cell tree as an expression and do a body expansion.

But this means we need cell components that represent conditionals which can drive the reduction of our cell tree evaluation, in the same way that pattern matchine drives lazy evaluation in Haskell. I think I'd need a Match cell in place of Python's in-house match keyword.

return Match(term)(
    Case(pattern)(
        FetchDocs(url=url1)
    ),
    Case(pattern)(
    	FetchData(url=url2)
    )
)

Would web devs building pipelines want to do this? Wouldn't they be confused as to why they can't use the in-house match? I think so.

In addition, React web devs are used to being able to inspect the values of state variables inside the body of their component. With this scheme, the values in the body are effect descriptions, a representation of a future value yet to be computed. This also means some kind of extra step every time to know whether you're dealing with a future value (like an async/thunk) or the current computed value.

def SomeCell(*children, **args):
    # In lazy eval, `data` is like an async/thunk/effect description
    data = yield from useFetchJson("https://api.weather.com")
    
    # hence to print it, you'd have to raise an effect to force it.
    yield from print(data)
    
    return SomeOtherCell(
        data = Eval(data)
    )

I'm not convinced that would be a good idea. That seems like this would render the rest of the Python language unusable, because evaluation then never happens in the body of a cell, but inside of the cell tree.

So unless I somehow have access to the abstract syntax tree of the python cell code, or I'm writing my own language, for now, lazy evaluation of algebraic effects will be out of reach. And for now, it will be up to the user to decide the evaluation order of effects by hand within a cell, and the notebook runtime will decide the evaluation order of the different cells.

Resources:

[1] Maybe I didn't look hard enough in literature, I dunno. But I never saw references to lazy eval when looking at incremental computation.

Sign up for more like this.

Enter your email
Subscribe