jaxdem.writers.async_base

Defines the base infrastructure for asynchronous data writing.

Classes

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

Infrastructure for non-blocking JAX data writing.

Exceptions

AsyncWriterError(failures)

Raised when one or more background writer tasks fail.

exception jaxdem.writers.async_base.AsyncWriterError(failures: tuple[tuple[str, Exception], ...])

Bases: RuntimeError

Raised when one or more background writer tasks fail.

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

Bases: object

Infrastructure for non-blocking JAX data writing.

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

directory: Path

The root directory where the writer saves simulation frames.

save_every: int

Save frequency. The writer pushes a frame to the queue on the first call and on every save_every-th call to the subclass’s save method.

clean: bool

If True, the writer deletes and recreates directory on initialization. Safety checks 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. This backpressure keeps memory bounded 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.

Raises AsyncWriterError after shutdown if any task failed. Reported failures are consumed, so a later call does not report the same failures again.

block_until_ready() → None

Waits until all pending tasks in the queue complete.

Raises AsyncWriterError after the queue drains if any task failed. Reported failures are consumed, so a later completion boundary does not report the same failures again.