#!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
# Copyright 2026 Tensor Auto Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""General utility functions for device management, logging, and common operations.
This module provides utilities for device selection, logging initialization,
number formatting, platform-specific operations, and various helper functions
used throughout the OpenTau codebase.
"""
import enum
import inspect
import logging
import os
import platform
import warnings
from copy import copy
from dataclasses import fields, is_dataclass
from datetime import datetime, timezone
from functools import wraps
from typing import Any, Callable
import accelerate
import numpy as np
import torch
[docs]
def inside_slurm() -> bool:
"""Check whether the Python process was launched through SLURM.
Returns:
True if the process is running in a SLURM environment, False otherwise.
"""
# TODO(rcadene): return False for interactive mode `--pty bash`
return "SLURM_JOB_ID" in os.environ
[docs]
def auto_torch_device() -> torch.device:
"""Automatically select the best available torch device.
Returns:
torch.device instance. Priority: CUDA > MPS > CPU.
"""
if torch.cuda.is_available():
return torch.device("cuda")
if torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")
[docs]
def get_safe_torch_device(try_device: str, log: bool = False, accelerator: Callable = None) -> torch.device:
"""Given a string, return a torch.device with checks on whether the device is available."""
match try_device:
case "cuda":
assert torch.cuda.is_available()
device = accelerator.device if accelerator else torch.device("cuda")
case "mps":
assert torch.backends.mps.is_available()
device = torch.device("mps")
case "cpu":
device = torch.device("cpu")
if log:
logging.warning("Using CPU, this will be slow.")
case _:
device = torch.device(try_device)
if log:
logging.warning(f"Using custom {try_device} device.")
return device
[docs]
def get_safe_dtype(dtype: torch.dtype, device: str | torch.device) -> torch.dtype:
"""Get a dtype that is compatible with the given device.
MPS is currently not compatible with float64, so this function converts
float64 to float32 for MPS devices.
Args:
dtype: Desired dtype.
device: Target device (string or torch.device).
Returns:
Compatible dtype for the device.
"""
if isinstance(device, torch.device):
device = device.type
if device == "mps" and dtype == torch.float64:
return torch.float32
else:
return dtype
[docs]
def is_torch_device_available(try_device: str) -> bool:
"""Check if a torch device is available.
Args:
try_device: Device name to check ('cuda', 'mps', or 'cpu').
Returns:
True if the device is available, False otherwise.
Raises:
ValueError: If try_device is not one of the recognized device types.
"""
if try_device == "cuda":
return torch.cuda.is_available()
elif try_device == "mps":
return torch.backends.mps.is_available()
elif try_device == "cpu":
return True
else:
raise ValueError(f"Unknown device '{try_device}.")
[docs]
def is_amp_available(device: str) -> bool:
"""Check if automatic mixed precision (AMP) is available for a device.
Args:
device: Device name to check ('cuda', 'mps', or 'cpu').
Returns:
True if AMP is available for the device, False otherwise.
Raises:
ValueError: If device is not one of the recognized device types.
"""
if device in ["cuda", "cpu"]:
return True
elif device == "mps":
return False
else:
raise ValueError(f"Unknown device '{device}.")
# Global variable to ensure logging is initialized only once
_logging_init_stack = ""
# accelerate's save_accelerator_state emits one INFO line per dataloader; a
# WeightedDatasetMixture run can have hundreds. The matching load path already
# emits a single summary line, so silence the save side to match. The exact
# prefix string was verified at accelerate>=1.4.0 (our floor in pyproject.toml)
# through current main; re-check on the next accelerate bump.
class _DropDataloaderSamplerSaveSpam(logging.Filter):
def filter(self, record: logging.LogRecord) -> bool:
msg = record.getMessage()
return not (msg.startswith("Sampler state for dataloader ") and " saved in " in msg)
def _format_stack(stack: list[inspect.FrameInfo]) -> str:
return "\n".join(
f" File '{frame.filename}', line {frame.lineno}, in {frame.function}"
for frame in stack[1:] # skip the current frame
)
[docs]
def init_logging(accelerator: accelerate.Accelerator | None = None, level=logging.INFO) -> None:
"""Initialize logging configuration with custom formatter.
This function sets up logging with a custom formatter that includes
timestamp, filename, and line number. It can only be initialized once
per process and will warn if called multiple times.
Args:
accelerator: Optional Accelerator instance. If provided, logging level
is set to WARNING on non-main processes to avoid duplicate logs.
level: Logging level to use. Defaults to logging.INFO.
"""
global _logging_init_stack
stack = inspect.stack()
if _logging_init_stack:
warnings.warn(
f"""Logging was already initialized through the following stack:
{_logging_init_stack}
Not initializing again through the following stack:
{_format_stack(stack)}""",
stacklevel=2,
)
else:
_logging_init_stack = _format_stack(stack)
class CustomFormatter(logging.Formatter):
def format(self, record):
dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
fnameline = f"{record.pathname}:{record.lineno}"
return f"{record.levelname} {dt} {fnameline[-15:]:>15} {record.getMessage()}"
console_handler = logging.StreamHandler()
console_handler.setFormatter(CustomFormatter())
logging.basicConfig(level=level, force=True, handlers=[console_handler])
accel_ckpt_logger = logging.getLogger("accelerate.checkpointing")
if not any(isinstance(f, _DropDataloaderSamplerSaveSpam) for f in accel_ckpt_logger.filters):
accel_ckpt_logger.addFilter(_DropDataloaderSamplerSaveSpam())
if accelerator and not accelerator.is_main_process:
# Disable duplicate logging on non-main processes
logging.info(f"Setting logging level on non-main process {accelerator.process_index} to WARNING.")
logging.getLogger().setLevel(logging.WARNING)
[docs]
def capture_timestamp_utc() -> datetime:
"""Capture the current UTC timestamp.
Returns:
datetime object representing the current UTC time.
"""
return datetime.now(timezone.utc)
[docs]
def say(text: str, blocking: bool = False) -> None:
"""Use text-to-speech to speak text (platform-specific).
Args:
text: Text to speak.
blocking: If True, wait for speech to complete before returning.
Defaults to False.
"""
# Check if mac, linux, or windows.
if platform.system() == "Darwin":
cmd = f'say "{text}"'
if not blocking:
cmd += " &"
elif platform.system() == "Linux":
cmd = f'spd-say "{text}"'
if blocking:
cmd += " --wait"
elif platform.system() == "Windows":
# TODO(rcadene): Make blocking option work for Windows
cmd = (
'PowerShell -Command "Add-Type -AssemblyName System.Speech; '
f"(New-Object System.Speech.Synthesis.SpeechSynthesizer).Speak('{text}')\""
)
os.system(cmd) # nosec: B605
[docs]
def log_say(text: str, play_sounds: bool, blocking: bool = False) -> None:
"""Log text and optionally speak it using text-to-speech.
Args:
text: Text to log and optionally speak.
play_sounds: If True, also speak the text using text-to-speech.
blocking: If True, wait for speech to complete. Defaults to False.
"""
logging.info(text)
if play_sounds:
say(text, blocking)
[docs]
def get_channel_first_image_shape(image_shape: tuple) -> tuple:
"""Convert image shape from HWC to CHW format if needed.
Args:
image_shape: Image shape tuple, either (H, W, C) or (C, H, W).
Returns:
Image shape in CHW format (C, H, W).
Raises:
ValueError: If the input shape is not in a recognized format.
"""
shape = copy(image_shape)
if shape[2] < shape[0] and shape[2] < shape[1]: # (h, w, c) -> (c, h, w)
shape = (shape[2], shape[0], shape[1])
elif not (shape[0] < shape[1] and shape[0] < shape[2]):
raise ValueError(image_shape)
return shape
[docs]
def has_method(cls: object, method_name: str) -> bool:
"""Check if a class or object has a specific method.
Args:
cls: Class or object to check.
method_name: Name of the method to check for.
Returns:
True if the method exists and is callable, False otherwise.
"""
return hasattr(cls, method_name) and callable(getattr(cls, method_name))
[docs]
def is_valid_numpy_dtype_string(dtype_str: str) -> bool:
"""Check if a string can be converted to a numpy dtype.
Args:
dtype_str: String to check.
Returns:
True if the string is a valid numpy dtype, False otherwise.
"""
try:
# Attempt to convert the string to a numpy dtype
np.dtype(dtype_str)
return True
except TypeError:
# If a TypeError is raised, the string is not a valid dtype
return False
[docs]
def is_launched_with_accelerate() -> bool:
"""Check if the process was launched with accelerate.
Returns:
True if ACCELERATE_MIXED_PRECISION is in the environment, False otherwise.
"""
return "ACCELERATE_MIXED_PRECISION" in os.environ
[docs]
def attempt_torch_compile(fn: callable, device_hint=None) -> callable:
"""Attempt to compile a PyTorch function using torch.compile.
The argument device_hint is used to check if torch.compile works reliably
on the device. Compilation is skipped if the device is MPS (Metal Performance
Shaders) as it is experimental.
Args:
fn: Function to compile.
device_hint: Optional device hint to check compatibility.
Returns:
Compiled function if compilation succeeds, otherwise the original function.
"""
if device_hint and "mps" in str(device_hint):
logging.warning("torch.compile is experimental on MPS devices. Compilation skipped.")
return fn
if hasattr(torch, "compile"):
logging.info("Attempting to compile the policy with torch.compile()...")
try:
fn = torch.compile(fn, mode="max-autotune")
logging.info("Policy compiled successfully.")
except Exception as e:
logging.warning(f"torch.compile failed with error: {e}. Proceeding without compilation.")
else:
logging.warning(
"torch.compile is not available. Requires PyTorch 2.0+. Proceeding without compilation."
)
return fn
[docs]
def create_dummy_observation(cfg, device, dtype=torch.bfloat16) -> dict:
"""Create a dummy observation dictionary for testing or initialization.
Args:
cfg: Configuration object with num_cams, resolution, max_state_dim,
and action_chunk attributes.
device: Device to create tensors on.
dtype: Data type for tensors. Defaults to torch.bfloat16.
Returns:
Dictionary containing dummy camera observations, state, prompt, and
padding flags.
"""
camera_observations = {
f"camera{i}": torch.zeros((1, 3, *cfg.resolution), dtype=dtype, device=device)
for i in range(cfg.num_cams)
}
return {
**camera_observations,
"state": torch.zeros((1, cfg.max_state_dim), dtype=dtype, device=device),
"prompt": ["Pick up yellow lego block and put it in the bin"],
"img_is_pad": torch.zeros((1, cfg.num_cams), dtype=torch.bool, device=device),
"action_is_pad": torch.zeros((1, cfg.action_chunk), dtype=torch.bool, device=device),
# Pin per-sample normalization to dataset row 0 so checkpoints trained
# with per-dataset stats on >1 datasets don't trip
# `_resolve_dataset_index`'s "missing selector" guard. A no-op for
# single-dataset policies (the fallback would resolve to the same).
"dataset_index": torch.zeros((1,), dtype=torch.long, device=device),
}
[docs]
def encode_accelerator_state_dict(obj) -> Any:
"""Encode an object into a JSON/YAML-compatible primitive type.
Args:
obj: Object to encode (can be Enum, dict, list, tuple, dataclass, etc.).
Returns:
Encoded object with all nested structures converted to primitives.
"""
if isinstance(obj, enum.Enum):
return encode_accelerator_state_dict(obj.value)
elif isinstance(obj, (str, int, float, bool)) or obj is None:
return obj
elif isinstance(obj, (list, tuple)):
return [encode_accelerator_state_dict(item) for item in obj]
elif isinstance(obj, dict):
return {key.replace(".", "_"): encode_accelerator_state_dict(value) for key, value in obj.items()}
elif is_dataclass(obj):
return {f.name: encode_accelerator_state_dict(getattr(obj, f.name)) for f in fields(obj)}
else:
return str(obj) # Fallback to string representation for unsupported types
[docs]
def on_accelerate_main_proc(*, local=False, _sync=False):
r"""Returns a decorator to run a function only on the main process when using `accelerate`.
If `local` is True (defaults to False), the function will run on the main process of each node
(useful for multi-node setups).
If `_sync` is True (defaults to False), the output of the function will be broadcasted to all processes.
If `_sync` is True, you must ensure that all processes call the decorated function, otherwise it will deadlock.
YOU SHOULD BE EXTREMELY CAREFUL WHEN USING THIS DECORATOR with _sync=True. Consider the following example::
@on_accelerate_main_proc()
def f():
return g()
@on_accelerate_main_proc(_sync=True)
def g():
return 42
In this case, if f() is called on all processes, they will deadlock at g() because child processes don't even
enter f(), hence never call g(), and thus won't reach the broadcast.
Another example::
@on_accelerate_main_proc(_sync=cond())
def f():
print("hi")
"""
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
state = accelerate.state.PartialState()
if not is_launched_with_accelerate() or not state.use_distributed:
return func(*args, **kwargs)
output, exception = None, None
flag = state.is_local_main_process if local else state.is_main_process
if flag:
try:
output = func(*args, **kwargs)
except Exception as e:
exception = e
if _sync:
payload = [output, exception]
accelerate.utils.broadcast_object_list(payload, from_process=0)
output, exception = payload
if exception is not None:
raise RuntimeError("An exception occurred in the main process.") from exception
return output
return wrapper
return decorator