jaxdem.writers.async_base

Defines the base infrastructure for asynchronous data writing.

Classes

BaseAsyncWriter([directory, save_every, ...])

Infrastructure for non-blocking JAX data writing.

class jaxdem.writers.async_base.BaseAsyncWriter(directory: Path = PosixPath('frames'), save_every: int = 1, clean: bool = True, max_workers: int = 8, max_queue_size: int = 512)

Bases: object

Infrastructure for non-blocking JAX data writing.

This class provides a pool of background worker threads and a task queue to ensure that slow disk I/O operations and device-to-host transfers do not block the main simulation loop.

directory: Path

The root directory where simulation frames will be saved.

save_every: int

Frequency of saving. A frame is pushed to the queue every save_every calls to the save() method.

clean: bool

If True, the directory is deleted and recreated upon initialization. Basic safety checks are performed to prevent deleting the current working directory or the system root.

max_workers: int

The number of background worker threads to use for parallel I/O.

max_queue_size: int

Maximum number of pending tasks in the background queue. When the queue is full, submit() blocks until a worker frees a slot, providing backpressure so memory cannot grow without bound when the simulation outruns disk I/O. Set to 0 for an unbounded queue.

submit(func: Callable[[...], Any], *args: Any, **kwargs: Any) → None

Pushes a task to the background worker queue.

Blocks when the queue is full (see max_queue_size).

close() → None

Blocks the main thread until all pending writes are finished and shuts down the background threads.

block_until_ready() → None

Waits until all currently pending tasks in the queue are completed.