utils

olmo_core.utils.generate_uuid()

Generate a unique ID.

Return type:

str

olmo_core.utils.get_default_thread_count()

Get the default maximum number of threads allowed.

Return type:

int

olmo_core.utils.wait_for(condition, description, timeout=10.0)

Wait for the condition function to return True.

olmo_core.utils.apply_to_tensors(fn, container)

Recursively apply fn to all tensors in a container.

Return type:

None

olmo_core.utils.move_to_device(o, device, non_blocking=None)

Move a tensor or container of tensors to the given device.

Parameters:

• o (TypeVar(T)) – The object to move.

• device (device) – The device to move to.

Return type:

TypeVar(T)

olmo_core.utils.mark_dynamic(x, dim, strict=True)

Mark a tensor as having dynamic sizes for torch.compile().

olmo_core.utils.get_default_device()

Get the default device.

Return type:

device

olmo_core.utils.seed_all(seed)

Seed all RNG states.

olmo_core.utils.same_storage(x, y)

Check if two tensors share the same storage.

Return type:

bool

olmo_core.utils.gc_cuda()

Run garbage collection, including emptying the CUDA cache.

olmo_core.utils.has_flash_attn()

Check if flash-attn is available.

Return type:

bool

class olmo_core.utils.LogFilterType(value)

Bases: StrEnum

Determines which ranks are allowed to emit INFO messages.

rank0_only = 'rank0_only'

INFO messages are only emitted from the global rank 0.

local_rank0_only = 'local_rank0_only'

INFO messages are only emitted from the local (node) rank 0.

all_ranks = 'all_ranks'

All ranks emit INFO messages.

olmo_core.utils.log_extra_field(field_name, field_value)

Add an additional field to each log record.

Note

For these fields to actually show up in the logs you need to use a formatter/handler that displays them.

Parameters:

• field_name (str) – The name of the field to attach.

• field_value (Any) – The value of the field to attach.

Return type:

None

olmo_core.utils.setup_logging(log_filter_type='rank0_only', force=False)

Configure logging.

See also

prepare_cli_environment()

Parameters:

• log_filter_type (LogFilterType, default: 'rank0_only') – Which ranks emit INFO and below messages.

• force (bool, default: False) – Force configuring logging even if it was already configured.

Return type:

None

olmo_core.utils.logging_configured()

Returns True if logging has been configured (like with setup_logging()), otherwise returns False.

Return type:

bool

olmo_core.utils.excepthook(exctype, value, traceback)

Used to patch sys.excepthook in order to log exceptions. Use install_excepthook() to install this.

olmo_core.utils.install_excepthook()

Install the custom excepthook().

See also

prepare_cli_environment()

olmo_core.utils.filter_warnings()

Configure warning filters for warnings we don’t need to see.

See also

prepare_cli_environment()

olmo_core.utils.set_env_variables()

Set common needed env vars if they’re not already set.

See also

prepare_cli_environment()

olmo_core.utils.prepare_cli_environment(log_filter_type=None)

Prepare the environment for a script/CLI. This should be called at the very beginning of the script/command, like at the top of the if __name__ == "__main__": ... block.

Internally this calls:

• setup_logging()

• install_excepthook()

• filter_warnings()

• set_env_variables()

• add_cached_path_clients()

Tip

If you’re looking to setup the environment specifically for distributed training, see prepare_training_environment() instead.

Parameters:

log_filter_type (Optional[LogFilterType], default: None) –

Determines which ranks are allowed to emit log messages below the WARNING level. You can also configure this through the env var LOG_FILTER_TYPE. If neither are set, this defaults to “rank0_only”.

Note

All ranks will always emit messages at the WARNING level or higher.

olmo_core.utils.threaded_generator(g, maxsize=16, thread_name=None)

Wraps a generator g and runs it in a thread.

olmo_core.utils.roundrobin(*iterables)

Call the given iterables in a round-robin fashion. For example: roundrobin('ABC', 'D', 'EF') --> A D E B F C

olmo_core.utils.format_int(count)

Format a large integer into a more human-readable string.

Return type:

str

olmo_core.utils.flatten_dict(d)

Flatten a nested dictionary with strings keys using dot notation.

Return type:

Dict[str, Any]

olmo_core.utils.cuda_sync_debug_mode(debug_mode)

A context manager for temporarily setting the CUDA sync debug mode.

olmo_core.utils.get_element_size(dtype)

Get the size in bytes of element of the given PyTorch dtype.

Return type:

int

olmo_core.utils.info_value_of_dtype(dtype)

Returns the finfo or iinfo object of a given PyTorch data type. Does not allow torch.bool.

olmo_core.utils.min_value_of_dtype(dtype)

Returns the minimum value of a given PyTorch data type. Does not allow torch.bool.