Hang when using torch.multiprocessing after having called rr.init

[jleibs]

Minimal repro:

import gc

import rerun as rr
from torch import multiprocessing


def task() -> None:
    gc.collect()


def main() -> None:
    rr.init("foo")

    proc = multiprocessing.Process(
        target=task,
    )
    proc.start()
    proc.join


if __name__ == "__main__":
    main()

The hang only happens if:

1. We are using torch.multiprocessing
2. We are running on linux, where the default multiprocess behavior is fork
3. Something forces a gc to happen in the subprocess

Explanation

Although we already have a fork-handler that cleans up our global/thread-local recording streams, it's apparently still possible for an allocated PyRecordingStream to leak into the subprocess via fork (at least based on how pytorch multiprocessing works).

During __del__, we make one last call to a non-blocking flush.

While previously this was fine, we added an internal blocking batcher flush to our non-blocking sink flush, which still hangs for the same reason (the batcher processing thread is gone in the fork).

[jleibs]

Prior to 0.16, the easiest option to work around this is to call gc.collect() any time after rr.init() and before starting the multiprocessing job.