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

`mandala` is a framework I wrote to automate tracking ML experiments for my research. It differs from other experiment tracking tools by making persistence, query and versioning logic a generic part of the programming language itself, as opposed to an external logging tool you must learn and adapt to. The goal is to be able to write expressive computational code without thinking about persistence (like in an interactive session), and still have the full benefits of a versioned, queriable storage afterwards.

Surprisingly, it turns out that this vision can pretty much be achieved with two generic tools:

1. a memoization+versioning decorator, `@op`, which tracks inputs, outputs, code and runtime dependencies (other functions called, or global variables accessed) every time a function is called. Essentially, this makes function calls replace logging: if you want something saved, you write a function that returns it. Using (a lot of) hashing, `@op` ensures that the same version of the function is never executed twice on the same inputs.

Importantly, the decorator encourages/enforces composition. Before a call, `@op` functions wrap their inputs in special objects, `Ref`s, and return `Ref`s in turn. Furthermore, data structures can be made transparent to `@op`s, so that an `@op` can be called on a list of outputs of other `@op`s, or on an element of the output of another `@op`. This creates an expressive "web" of `@op` calls over time.

2. a data structure, `ComputationFrame`, can automatically organize any such web of `@op` calls into a high-level view, by grouping calls with a similar role into "operations", and their inputs/outputs into "variables". It can detect "imperative" patterns - like feedback loops, branching/merging, and grouping multiple results in a single object - and surface them in the graph.

`ComputationFrame`s are a "synthesis" of computation graphs and relational databases, and can be automatically "exported" as dataframes, where columns are variables and operations in the graph, and rows contain values and calls for (possibly partial) executions of the graph. The upshot is that you can query the relationships between any variables in a project in one line, even in the presence of very heterogeneous patterns in the graph.

I'm very excited about this project - which is still in an alpha version being actively developed - and especially about the `ComputationFrame` data structure. I'd love to hear the feedback of the HN community.

Colab quickstart: https://colab.research.google.com/github/amakelov/mandala/bl...

Blog post introducing `ComputationFrame`s (can be opened in Colab too): https://amakelov.github.io/mandala/blog/01_cf/

Docs: https://amakelov.github.io/mandala/

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

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...

Very cool concept, thanks for sharing this.

Is mandala designed for notebook-style interactive environments only or also when running python scripts more traditionally, in which case could this be integrated in a gitops like environment? (push to git, new run in CI/CD, export log metrics as an artifact with an easy way to compare to previous runs)

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).

Very cool! looking forward to trying it out - the graphs reminded me of a toy project I'd done a while back to better understand deterministic and reproducible execution in python as seen in marimo.io notebooks https://github.com/vrtnis/python-notebook-simulator