tweag / python-nix

Python-Nix FFI library using the new C API
MIT License
43 stars 3 forks source link

Memory management in expr #6

Open GuillaumeDesforges opened 8 months ago

GuillaumeDesforges commented 8 months ago

I need to make many evaluations (https://github.com/tweag/nixtract/tree/c-bindings) but memory usage explodes. Thus I want to eval a Value, then del it, expecting the memory to be freed. Unfortunately the memory usage does not seem to be freed even when Value is freed.

import gc
import nix.expr
import nix.expr_util
import nix.store
import psutil

store = nix.store.Store()
state = nix.expr.State([], store)

def get_nixpkgs_root_drvs():
    nix_builtins = state.eval_string("builtins", ".")
    nix_get_flake = nix_builtins["getFlake"]
    nix_nixpkgs_flake = nix_get_flake("nixpkgs")
    nix_nixpkgs_pkgs = nix_nixpkgs_flake["legacyPackages"]["x86_64-linux"]

    nix_nixpkgs_root_drvs = {}
    for k in nix_nixpkgs_pkgs.keys():
        try:
            v = nix_nixpkgs_pkgs[k]
        except:
            continue

        if (
            v.get_type() == nix.expr.Type.attrs
            and "type" in v
            and v["type"].force() == "derivation"
        ):
            nix_nixpkgs_root_drvs[k] = v

    return nix_nixpkgs_root_drvs

process = psutil.Process()

def _get_usage():
    return process.memory_info().rss

_base_usage = _get_usage()

def get_usage():
    return _get_usage() - _base_usage

def main():
    print("usage", get_usage())

    nixpkgs_root_drvs = get_nixpkgs_root_drvs()
    print("len(nixpkgs_root_drvs)", len(nixpkgs_root_drvs))

    print("usage", get_usage())

    del nixpkgs_root_drvs
    n_collect = gc.collect()
    print("collected", n_collect)
    nix.expr_util.lib.nix_gc_now()
    print("usage", get_usage())

    global store, state
    del store, state
    n_collect = gc.collect()
    print("collected", n_collect)
    nix.expr_util.lib.nix_gc_now()
    print("usage", get_usage())

main()

gives output

usage 0
len(nixpkgs_root_drvs) 18623
usage 2063163392
collected 18700
usage 2061344768
collected 0
usage 2061078528

Since the Python gc collects as many items as the length of the nixpkgs_root_drvs dict, I would assume that there is no Value that references to any C value in Python anymore, so the C values should be freed when calling nix.expr_util.lib.nix_gc_now() (which calls GC_gcollect under the hood AFAIK). However usage stays almost the same. It decreases only a little, probably because of some other Python objects that have been collected.

On a side note, I'm surprised that gc after del store, state doesn't collect anything.

How can I use python-nix in a way such that I can free the underlying C memory of a Value?

jlesquembre commented 8 months ago

I could be related to this: https://github.com/NixOS/nix/blob/b522b23e9e2ca4f700eb4a5d7d3dfeb2f206c7c6/src/libexpr/c/nix_api_value.h#L148