amakelov

vor 2 Jahrendiscuss

Thanks! And thanks for sharing the pointer - I think I've seen `mltrace` at some point in the past. The tool has some similarities, but seems different from `mandala` on a philosophical level - `mltrace` seems more opinionated and domain-specific to ML. `mandala`'s goal is to make persistence logic a more generic part of Python semantics, and there's much more emphasis on composing `@op`s freely in whatever complex ways you want. Python functions are great and everyone knows what they do - let's give them more superpowers!

If this is something you're interested in, I can probably give a more detailed comparison if I find the time.

amakelov

vor 2 Jahrendiscuss

Thanks Rachit! Great running into you after all these years!

Being aware of types is certainly a must in a more performance-critical implementation; this project is not at this stage though, opting for simplicity and genericity instead. I've found this best for maintenance until the core stabilizes; plus, it's not a major pain point in my ML projects yet.

Regarding incremental computation: the main idea is simple - if you call a function on some inputs, and this function was called on the same (modulo hashing) inputs in the past and used some set of dependencies that currently have the same code as before (or compatible code, if manually marked by the user), the past call's outputs are reused (key question here: will there be at most one such past call? yes, if functions invoke their dependencies deterministically).

You can probably add some tools to automatically delete old versions and everything that depends on them, but this is definitely a decision the user must make (e.g., you might want to be able to time-travel back to look at old results). I'm happy to answer more nuanced questions about the incremental computation implementation in `mandala` if you have any!

amakelov

vor 2 Jahrendiscuss

Forgot to mention: yes, the dependency tracking is transitive, i.e. if your @op calls a @track-decorated function, which in turn calls another @track-decorated function, then both dependencies will show up, etc.

amakelov

vor 2 Jahrendiscuss

Thanks for sharing! This is a great project. It is quite close to the memoization part of `mandala` and I'll add it to the related work in the README. I think the similarities are:

- using `joblib` to hash arbitrary objects (which is a good fit for ML which inlcudes a lot of numpy arrays, which joblib is optimized for)

- how composition of decorated functions is emphasized - I think that's very important

- wrapping outputs of memoized functions in special objects: this encourages composition, and also makes it possible to run pipelines "lazily" by retracing memoized calls without actually loading large objects in memory

- versioning: in a past version of `mandala`, I used the solution you provided (which is now subsumed by the versioning system, but it still quite helpful)

The differences: - w.r.t. memoization, in `mandala` you can represent data structures in a way transparent to the system. E.g., you can have a memoized function return a list of things, and each thing will have an independent storage address and be usable by downstream memoized functions. Most importantly, this is tracked by the provenance system (though I'm not sure - maybe this is also possible in `provenance`?)

- one big finding (for me) while doing this project is that memoization on its own is not sufficient to manage a complex project; you need some declarative way to understand what has been computed. This is what all the `ComputationFrame` stuff is about.

- finally, the versioning system: as you mention in the `provenance` docs, it's almost impossible to figure out what a Python function depends on, but `mandala` bites this bullet in a restricted sense; you can read about it here: https://amakelov.github.io/blog/deps/

Re:Unison - yes definitely; it is mentioned in the related work on github! A major difference is that Unison hashes the AST of functions; `mandala` is not that smart (currently) and hashes the source code.

amakelov

vor 2 Jahrendiscuss

Great question! The versioning system does something essentially equivalent to what you describe. It currently works as follows:

- When a call to an `@op` is executed, it keeps a stack of @track-decorated functions that are called (you can add some magic - not implemented currently - via import-time decoration to automatically track all functions in your project; I've opted against this by default to make the system more explicit).

- The "version" of the call is a hash of the collection of hashes of the code of the `@op` itself and the code of the dependencies that were accessed

- Note that this is much more reliable than static analysis, and much more performant/foolproof than using `sys.settrace`; see this blog post for discussion: https://amakelov.github.io/blog/deps/

- When a new call is made, it is first checked against the saved calls. The inputs are hashed, and if the system finds a previous call on these hashes where all the dependencies had the same (or compatible) code with the current codebase, this call is re-used. *Assuming deterministic dependencies*, there can be at most 1 such call, so everything is well-defined. I don't think this is an unrealistic assumption, though you should keep it in mind - it is pretty fundamental to how this system works. Otherwise, disambiguating versions based on the state of the code alone is impossible.

- When a dependency changes, you're alerted which `@op`s' versions are affected, and given a choice to either ignore this change or not (in the latter case, only calls that actually used this dependency will be recomputed).

The versioning system is mostly a separate module (not in a way that it can be imported independently of the rest, but it should be pretty doable). I'd love to hear more about your use case. It may not be too difficult to export just versioning as its own thing - though from what you describe, it should also have some component of memoization? As in, you seem to be interested in using this information to invalidate/keep things in the cache?

amakelov

vor 2 Jahrendiscuss

Thanks! Yes I think a caching solution like this is great for notebooks, because it makes it very cheap to re-run the whole thing (as long as you reasonably organized long-running computations into `@op`s), overcoming the notorious state problem.

Graphviz is indeed a lifesaver; and you can similarly think of `mandala` as a "build system for Python objects" (with all the cool and uncool things that come with that; serializing arbitrary Python objects with strong guarantees is hard https://amakelov.github.io/mandala/tutorials/gotchas/).

I've no experience with rust, but I'd be curious to hear about the difficulties that came up. I'd expect to see some overlap with Python!

amakelov

vor 2 Jahrendiscuss

Oh also totally missed the Borges mention the first time - I'm a big fan of his stories!

amakelov

vor 2 Jahrendiscuss

Thanks! Indeed, despite the fact that the main goal is to track ML experiments, the approach taken in `mandala` has a lot in common with e.g. time-travel debugging (https://en.wikipedia.org/wiki/Time_travel_debugging).

In reality, there are many similarities between experiment tracking, debugging, and high-level computation graphs - they're all different ways of getting a handle on what a program did when it was ran.

amakelov

vor 2 Jahrendiscuss

This blog port gives an overview of the core dependency tracking logic: https://amakelov.github.io/blog/deps/

amakelov

vor 2 Jahrendiscuss

Ah, yes, the notorious state problem in notebooks. In your project, do you find the dependencies statically or dynamically?

amakelov

vor 2 Jahrendiscuss

Great question - personally, I mostly use it from notebooks, and I think it's a great fit for that. Bundling experiment tracking with incremental computation makes a lot of sense in a notebook (or any other interactive) environment, because it solves the problem of state: if all your computations are end-to-end memoized, re-running the entire notebook is cheap (I routinely do this "retracing" to just get to some results I want to look at).

That being said, nothing prevents you from running this in a script too, and there are benefits of doing this as well. If your script mostly composes `@op` calls (possibly with some light control flow logic as needed), you get resumability "for free" on the level of individual `@op` calls after a crash. However, the workflow you're describing may run into some features that aren't implemented (yet) if your runs write to different storages. `mandala` makes it easy to get a "holistic" picture of a single `Storage` and compare things there. Comparing across storages will be more awkward. But it shouldn't be too hard to write a function that merges storages (and it's a very natural thing to do, as they're basically big tables).

amakelov

vor 2 Jahrendiscuss

Thanks! Indeed, the ultimate (but very ambitious from the point of view of coordination and infrastructure) vision would be to build the "planetary computer" where everyone can contribute and every computation is transparently reproducible. While researching `mandala`, I ran into some ideas along those lines from the folks at Protocol Labs, e.g. https://www.youtube.com/watch?v=PBo1lEZ_Iik, though I'm not aware of the current status.

amakelov

vor 2 Jahrendiscuss

In order,

1. Yes, you can choose to create a persistent storage by passing `db_path` to `Storage()`. The current implementation is just an SQLite file. To run on many machines, you don't really need to be able to re-import from a dataframe (dumping to a dataframe is meant to be an exit point from `mandala` so that you can do downstream analyses in a format more familiar than `ComputationFrame`) - `ComputationFrame`s can be merged via the union (`|`) operator, see here https://amakelov.github.io/mandala/blog/01_cf/#tidy-tools-me... for an example. Storages don't support merging yet, but it's certainly possible!

2. Already answered in 1.

3. Nope, but I'd be happy to (though I feel like `mandala` took memoization in a substantially different direction). Are you in a position to make an introduction?

4. This project is currently not optimized for performance, though I've used it in projects spanning millions of memoized calls. The typical use case is to decorate functions that take a long time to compute, so the overhead of memoization amortizes. A very quick benchmark on my laptop shows ~6ms per call for in-memory storage, ~9ms for a persistent storage, with a simple arithmetic function that otherwise takes ~0 time.

5. Great question - currently, the dependency tracer is restricted to user-chosen functions to avoid tracking function calls an imported library makes. You could use a bit of magic (import-time automatic decoration) to track all functions in a file or a directory (not implemented right now). The reasoning is that, for a typical multi-month ML project, you usually have a single conda environment so you want to ignore library changes. Similarly, system-level (e.g. environment variables) are also not tracked. I think a very useful feature would be to at least record the versions of each imported library, so that storages can be ported between environments with some guarantees (or warnings).

6. - If an `@op` call was memoized, the underlying Python function call succeeded, so in this sense it can't be "broken"; it's however possible that there was a bug. In this case, you can delete the affected calls and all values that depend on them (if you keep these values, you're left with "zombie" values that don't have a proper computational history). The `ComputationFrame` supports declarative deletion - you build a ComputationFrame that captures the calls you want to delete, and call `.delete_calls()` - though there's still no example of this in the tutorial :) Alternatively, you can change the affected function and mark this as a new version. Then you should be able to delete all calls using the previous version (though, not supported at this moment).

- How the cache is invalidated is detailed here: https://github.com/amakelov/mandala?tab=readme-ov-file#how-i...

amakelov

vor 3 Jahrendiscuss

This is neat and self-contained! But as someone running experiments with a high degree of interactivity, I often have an orthogonal requirement: add more computations to the same cell without recomputing previous computations done in the cell (or in other cells).

For a concrete example, often in an ML project you want to study how several quantities vary across several parameters. A straightforward workflow for this is: write some nested loops, collect results in python dictionaries, finally put everything together in a dataframe and compare (by plotting or otherwise).

However, after looking at the results, maybe you spot some trend and wonder if it will continue if you tweak one of the parameters by using a new value for it; of course, you also want to look at the previous values and bring everything together in the same plot(s). You now have a problem: either re-run the cell (thus losing previous work, which is annoying even if you have to wait 1 minute - you know it's a wasted minute!), or write the new computation in a new cell, possibly with a lot of redundancy (which over time makes the notebook hard to navigate and keep consistent).

So, this and other considerations eventually convinced me that the function is more natural than the cell as an interface/boundary at which caching should be implemented, at least for my use cases (coming from ML research). I wrote a framework based on this idea, with lots of other features (some quite experimental/unusual) to turn this into a feasible experiment management tool - check it out at https://github.com/amakelov/mandala

P.S.: I notice you use `pickle` for the hashing - `joblib.dump` is faster with objects containing numpy arrays, which covers a lot of useful ML things

amakelov

vor 3 Jahrendiscuss

I see two concerns here:

- inputs/outputs being high volume: the inputs/outputs that are large are often also things that don't change over the course of a project (e.g. a dataset or a model). So you don't really need to cache the object itself, just a (typically short string) immutable reference to it. As long as the object can be looked up at runtime, everything's fine;

- detecting changes in data: content hashing is the general way in which you can tell if a result changed; using `joblib.dump` and then hashing the resulting string provides a good starting implementation, though certainly there are some corner cases to be aware of.

Both of these approaches are available/used in mandala (https://github.com/amakelov/mandala; disclosure: I'm the author), which uses content hashing to tell when data (or even code/code dependencies) have changed, and gives you a generic caching decorator for functions which can then look up large objects by reference; this is the way I used it for e.g. my mechanistic interpretability work, which is often of the form one big model + lots of analyses producing tiny artifacts based on it.

amakelov

vor 3 Jahrendiscuss

Author here - thank you for the compliment! Always happy to answer any questions (the docs are admittedly quite sparse) and hear about what people like/dislike about the library!

amakelov

vor 3 Jahrendiscuss

The dependency tracking and the graph/SQL thing are largely independent pieces.

The graph/SQL thing is used to query the memoization tables of the memoized functions by "joining them along a given computational graph". This roughly means that you can point to a computation, and ask: give me a table of the values of such-and-such variables in this computation across all analogous computations in the storage. Here, "analogous" could mean, for example, running the same composition of functions, but with different input parameters. The motivation is to be able to easily ask a broad class of natural queries of your storage (e.g., how given outputs depend on given inputs across all experiments of a given kind).

For example, say you have two memoized functions, `increment` and `add`, and run this:

with storage.run():

    for i in [1, 2, 3]:

        x = increment(i)

        y = add(i, x)

Running this code will memoize a bunch of calls to `increment` and `add` in a way that recognizes that the output of the call to `increment` is the second input to the call to `add` (behind the scenes, they point to the same saved object). Then, if you call `storage.similar(x, y)`, you'll see a message like this:

Pattern-matching to the following computational graph (all constraints apply):

    a0 = Q() # input to computation; can match anything

    a1 = Q() # input to computation; can match anything

    x = increment(a=a0)

    y = add(a=a1, b=x)

    result = storage.df(x, y)

and a table like this:

x y

2 3

3 5

4 7

which tells you all the values `x, y` have taken in executions of this kind of program. Since you've only ran this with `i = 1, 2, 3`, you get back just the result from this computation - but if you've been running similar things for weeks, you may not remember all the settings you've ran this with so far - and this query will reveal them.

How this works is that you look for values `a0, a1, x, y` in the storage that satisfy both the constraints `x = increment(a=a0), y = add(a=a1, b=x)`. Finding such values amounts to joining the memoization tables of `increment` and `add` in a certain way. Since this way is implicit in the computation itself, you can just point to the variables and get the result.

Hope that helps! I realize that this is a fairly non-standard way to query, but I believe it's well worth the benefits in terms of reducing the boilerplate required to interact with computational data. Let me know if you think of ways to improve the presentation or have any other questions!

amakelov

vor 3 Jahrendiscuss

Hi, author here. Sorry about the confusion - this blog post's intention was to give a more programming-language-themed introduction to the project (discussion on r/programminglanguages is here: https://www.reddit.com/r/ProgrammingLanguages/comments/12im9...). As such, it doesn't really talk very much about what you'd actually use this for - tracking computational experiments (for example). As for the part you highlighted, I was looking for a two-sentence summary that conveys the main features of the project... well, I guess I succeeded with the two-sentence part! :)

Let me try to unpack this in a hopefully saner format:

- the whole thing is based on memoization. You put a memoization decorator on all the functions whose results you want to persist/reuse, and you compose entire programs by calling such functions on the outputs of other functions. In between calls, you can also mix in data structure operations - so say a function returns a list, you can call another memoized function on just an element of this list.

- as such programs execute, some metadata is passed around that links together the inputs and outputs of each call, and the elements of each data structure. This dynamically builds a computational graph of the program behind the scenes.

- why would you need this graph? One reason is that there is a principled way to extract a SQL query from this graph that pattern-matches to all analogous computational graphs that have already been memoized. This gives you a flexible and natural interface to ask the queries you usually ask of your programs - "How do these outputs depend on these inputs across all the experiments of this kind?" - without writing extra code.

- since memoized results can go stale when you change your code, there is also a versioning system that tracks the dependencies that each memoized call accessed, and alerts you when a dependency changes. Tracking dependencies dynamically on a per-call basis, rather than statically on a per-function basis, gives you more opportunity to reuse computation automatically - for example, a function can have two branches that depend on different dependencies. If only one of the dependencies changes, you won't recompute the calls that used the other dependency. The "content-versioned" part refers to how the system recognizes which version a function is at: by looking at its code (content), instead of by you explicitly providing it with some arbitrary version name/number. This means that e.g. you can "go back in time" w.r.t a given function(s) by restoring the old code, and the storage will recognize that it's back in this world.

I hope that helps clarify things, and thanks a lot for bringing this up.

Submissions

Late Takes on OpenAI O1

Show HN: Mandala – Automatically save, query and version Python computations

Tidy Computations

Practical dependency tracking for Python function calls

comments