opentau.utils.accelerate_utils

Utilities for managing Accelerate accelerator instances.

This module provides functions for setting and getting a global accelerator instance, which is useful for accessing accelerator state throughout the codebase.

Functions

acc_print(*args, **kwargs)	Print with process index prefix when using accelerate.
get_proc_accelerator()	Get the global accelerator instance for the current process.
set_proc_accelerator(accelerator[, allow_reset])	Set the global accelerator instance for the current process.

opentau.utils.accelerate_utils.acc_print(*args, **kwargs) → None

Print with process index prefix when using accelerate.

If an accelerator is set, prints with a prefix showing the process index. Otherwise, prints normally.

Parameters:

• *args – Positional arguments to pass to print.

• **kwargs – Keyword arguments to pass to print.

opentau.utils.accelerate_utils.get_proc_accelerator() → Accelerator

Get the global accelerator instance for the current process.

Returns:

The accelerator instance, or None if not set.

opentau.utils.accelerate_utils.set_proc_accelerator(accelerator: Accelerator, allow_reset: bool = False) → None

Set the global accelerator instance for the current process.

Parameters:

• accelerator – Accelerator instance to set.

• allow_reset – If True, allow resetting an already-set accelerator. Defaults to False.

Raises:

• AssertionError – If accelerator is not an Accelerator instance.

• RuntimeError – If accelerator is already set and allow_reset is False.