From de0f9822c5b40ae71cb16624f37299dcb3b609fa Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Wed, 21 Oct 2026 21:47:03 +0800 Subject: [PATCH] chore(benchmark): delete device_info re-export shim, import unilab.utils.device directly (#2474) --- .../compute/benchmark_ane_inference.py | 2 +- .../benchmark/compute/benchmark_backends.py | 2 +- .../compute/benchmark_conversions.py | 2 +- .../compute/benchmark_mlp_inference.py | 2 +- .../compute/benchmark_mlx_compile.py | 2 +- scripts/benchmark/core/__init__.py | 3 --- scripts/benchmark/core/device_info.py | 7 ------ scripts/benchmark/env/benchmark_env_step.py | 5 ++-- .../benchmark_torch_env_shard_throughput.py | 2 +- .../benchmark_mujoco_backend_step_detail.py | 4 +--- .../benchmark_mujoco_rollout_model_sharing.py | 24 +++++-------------- .../physics/benchmark_physics_step_genesis.py | 4 +--- .../physics/benchmark_physics_step_mj_step.py | 4 +--- .../benchmark_physics_step_motrixsim.py | 4 +--- .../benchmark_physics_step_mujoco_warp.py | 4 +--- scripts/benchmark/physics/benchmark_sim.py | 2 +- .../benchmark_offpolicy_collector_active.py | 2 +- .../rl/benchmark_offpolicy_dp_scaling.py | 2 +- .../rl/benchmark_replay_buffer_placement.py | 2 +- .../benchmark_sac_replay_buffer_sampling.py | 2 +- scripts/benchmark/torch_env/run_benchmark.py | 2 +- 21 files changed, 25 insertions(+), 58 deletions(-) delete mode 100644 scripts/benchmark/core/device_info.py diff --git a/scripts/benchmark/compute/benchmark_ane_inference.py b/scripts/benchmark/compute/benchmark_ane_inference.py index dac695355..66cd38d7d 100644 --- a/scripts/benchmark/compute/benchmark_ane_inference.py +++ b/scripts/benchmark/compute/benchmark_ane_inference.py @@ -33,7 +33,7 @@ print_mlp_table, trimmed_mean, ) -from scripts.benchmark.core.device_info import get_device_info_dict, get_device_info_line +from unilab.utils.device import get_device_info_dict, get_device_info_line try: import coremltools as ct diff --git a/scripts/benchmark/compute/benchmark_backends.py b/scripts/benchmark/compute/benchmark_backends.py index 7ac2471e9..43ee8af69 100644 --- a/scripts/benchmark/compute/benchmark_backends.py +++ b/scripts/benchmark/compute/benchmark_backends.py @@ -22,7 +22,7 @@ print_table, torch_dtype, ) -from scripts.benchmark.core.device_info import get_device_info_dict +from unilab.utils.device import get_device_info_dict try: import numpy as np diff --git a/scripts/benchmark/compute/benchmark_conversions.py b/scripts/benchmark/compute/benchmark_conversions.py index a0a1e4969..647561262 100644 --- a/scripts/benchmark/compute/benchmark_conversions.py +++ b/scripts/benchmark/compute/benchmark_conversions.py @@ -34,7 +34,7 @@ parse_sizes, torch_dtype, ) -from scripts.benchmark.core.device_info import get_device_info_dict, get_device_info_line +from unilab.utils.device import get_device_info_dict, get_device_info_line _IS_MACOS = platform.system() == "Darwin" diff --git a/scripts/benchmark/compute/benchmark_mlp_inference.py b/scripts/benchmark/compute/benchmark_mlp_inference.py index 576f1e41b..386fea068 100644 --- a/scripts/benchmark/compute/benchmark_mlp_inference.py +++ b/scripts/benchmark/compute/benchmark_mlp_inference.py @@ -43,7 +43,7 @@ print_mlp_table, trimmed_mean, ) -from scripts.benchmark.core.device_info import get_device_info_dict, get_device_info_line +from unilab.utils.device import get_device_info_dict, get_device_info_line try: import numpy as np diff --git a/scripts/benchmark/compute/benchmark_mlx_compile.py b/scripts/benchmark/compute/benchmark_mlx_compile.py index 10411b37b..11a506978 100644 --- a/scripts/benchmark/compute/benchmark_mlx_compile.py +++ b/scripts/benchmark/compute/benchmark_mlx_compile.py @@ -38,7 +38,7 @@ import matplotlib -from scripts.benchmark.core.device_info import get_device_info_dict, get_device_info_line +from unilab.utils.device import get_device_info_dict, get_device_info_line matplotlib.use("Agg") import matplotlib.pyplot as plt diff --git a/scripts/benchmark/core/__init__.py b/scripts/benchmark/core/__init__.py index c5dbac369..334e2025b 100644 --- a/scripts/benchmark/core/__init__.py +++ b/scripts/benchmark/core/__init__.py @@ -1,7 +1,6 @@ """Core benchmark utilities for standardized benchmarking.""" from .backends import available_backends, mlx_dtype, numpy_dtype, sync_backend, torch_dtype -from .device_info import get_device_info_dict, get_device_info_line from .mlp_utils import MLPBenchRecord, env_nums_pow2, mlp_param_count, print_mlp_table, trimmed_mean from .output import print_table, save_json from .plotting import save_line_plot @@ -15,8 +14,6 @@ "torch_dtype", "mlx_dtype", "sync_backend", - "get_device_info_dict", - "get_device_info_line", "BenchRecord", "MLPBenchRecord", "bench_callable", diff --git a/scripts/benchmark/core/device_info.py b/scripts/benchmark/core/device_info.py deleted file mode 100644 index 21312bc67..000000000 --- a/scripts/benchmark/core/device_info.py +++ /dev/null @@ -1,7 +0,0 @@ -"""Thin re-export of the library device-info helpers (moved in issue #1240).""" - -from __future__ import annotations - -from unilab.utils.device import get_device_info_dict, get_device_info_line - -__all__ = ["get_device_info_dict", "get_device_info_line"] diff --git a/scripts/benchmark/env/benchmark_env_step.py b/scripts/benchmark/env/benchmark_env_step.py index c4a759fde..f5b3f2614 100644 --- a/scripts/benchmark/env/benchmark_env_step.py +++ b/scripts/benchmark/env/benchmark_env_step.py @@ -31,6 +31,8 @@ import numpy as np import torch +from unilab.utils.device import get_device_info_dict, get_device_info_line + plt: Any = None mpatches: Any = None @@ -61,12 +63,9 @@ def _load_helper_module(module_name: str, relative_path: str): return module -_DEVICE_INFO = _load_helper_module("bench_env_step_device_info", "device_info.py") _MEM_PROFILE = _load_helper_module("bench_env_step_mem_profile", "mem_profile.py") _OUTPUT = _load_helper_module("bench_env_step_output", "output.py") -get_device_info_dict = _DEVICE_INFO.get_device_info_dict -get_device_info_line = _DEVICE_INFO.get_device_info_line MEMORY_METRICS = _MEM_PROFILE.MEMORY_METRICS _build_memory_summary = _MEM_PROFILE.build_memory_summary _format_mb = _MEM_PROFILE.format_mb diff --git a/scripts/benchmark/env/benchmark_torch_env_shard_throughput.py b/scripts/benchmark/env/benchmark_torch_env_shard_throughput.py index 33199062b..4a7a63a40 100644 --- a/scripts/benchmark/env/benchmark_torch_env_shard_throughput.py +++ b/scripts/benchmark/env/benchmark_torch_env_shard_throughput.py @@ -43,8 +43,8 @@ if str(ROOT_DIR) not in sys.path: sys.path.append(str(ROOT_DIR)) -from scripts.benchmark.core.device_info import get_device_info_dict, get_device_info_line from scripts.benchmark.core.output import print_table, save_json +from unilab.utils.device import get_device_info_dict, get_device_info_line DEFAULT_OUTPUT_JSON = ( ROOT_DIR / "scripts" / "benchmark" / "outputs" / "torch_env_shard_throughput" / "results.json" diff --git a/scripts/benchmark/physics/benchmark_mujoco_backend_step_detail.py b/scripts/benchmark/physics/benchmark_mujoco_backend_step_detail.py index 9c917c99d..c44484c95 100644 --- a/scripts/benchmark/physics/benchmark_mujoco_backend_step_detail.py +++ b/scripts/benchmark/physics/benchmark_mujoco_backend_step_detail.py @@ -47,6 +47,7 @@ from unisim.backend.mujoco.xml import create_discardvisual_xml from unilab.dtype_config import get_global_dtype +from unilab.utils.device import get_device_info_dict, get_device_info_line matplotlib.use("Agg") import matplotlib.pyplot as plt @@ -65,12 +66,9 @@ def _load_helper_module(module_name: str, relative_path: str): return module -_DEVICE_INFO = _load_helper_module("bench_device_info", "device_info.py") _OUTPUT = _load_helper_module("bench_output", "output.py") _TASK_NAMES = _load_helper_module("bench_task_names", "task_names.py") -get_device_info_dict = _DEVICE_INFO.get_device_info_dict -get_device_info_line = _DEVICE_INFO.get_device_info_line save_json = _OUTPUT.save_json canonical_locomotion_task_ids = _TASK_NAMES.canonical_locomotion_task_ids locomotion_task_model_file = _TASK_NAMES.locomotion_task_model_file diff --git a/scripts/benchmark/physics/benchmark_mujoco_rollout_model_sharing.py b/scripts/benchmark/physics/benchmark_mujoco_rollout_model_sharing.py index 00f54da2f..1fdd03b5c 100644 --- a/scripts/benchmark/physics/benchmark_mujoco_rollout_model_sharing.py +++ b/scripts/benchmark/physics/benchmark_mujoco_rollout_model_sharing.py @@ -16,7 +16,6 @@ import argparse import gc -import importlib.util import json import os import resource @@ -32,6 +31,8 @@ import numpy as np +from unilab.utils.device import get_device_info_dict, get_device_info_line + try: import matplotlib @@ -67,19 +68,6 @@ PLOT_TEXT = "#0f172a" -def _load_device_info_helpers(): - device_info_path = Path(__file__).resolve().parent.parent / "core" / "device_info.py" - spec = importlib.util.spec_from_file_location("_local_benchmark_device_info", device_info_path) - if spec is None or spec.loader is None: - raise RuntimeError(f"Failed to load device info helpers from {device_info_path}") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module.get_device_info_dict, module.get_device_info_line - - -_get_device_info_dict, _get_device_info_line = _load_device_info_helpers() - - def _resolve_xml_path(xml_arg: str) -> str: xml_path = Path(xml_arg) if xml_path.is_absolute(): @@ -602,7 +590,7 @@ def _plot_summary_plots( if legend_handles: throughput_ax.legend(handles=legend_handles, frameon=False, loc="upper right") throughput_fig.suptitle( - f"MuJoCo Rollout Model Sharing\n{_get_device_info_line()}", + f"MuJoCo Rollout Model Sharing\n{get_device_info_line()}", y=0.98, fontsize=13, ) @@ -656,7 +644,7 @@ def _plot_summary_plots( _style_axis(axes[1]) memory_fig.suptitle( - f"MuJoCo Rollout Model Sharing Memory\n{_get_device_info_line()}", + f"MuJoCo Rollout Model Sharing Memory\n{get_device_info_line()}", y=0.985, fontsize=13, ) @@ -805,7 +793,7 @@ def main() -> None: out_dir / "results.json", all_results, { - "device_info": _get_device_info_dict(), + "device_info": get_device_info_dict(), "targets": [{"robot": name, "xml_path": xml_path} for name, xml_path in targets], "num_envs": args.num_envs, "steps": args.steps, @@ -821,7 +809,7 @@ def main() -> None: out_dir / f"{target_name}_results.json", results, { - "device_info": _get_device_info_dict(), + "device_info": get_device_info_dict(), "robot": target_name, "xml_path": xml_path, "num_envs": args.num_envs, diff --git a/scripts/benchmark/physics/benchmark_physics_step_genesis.py b/scripts/benchmark/physics/benchmark_physics_step_genesis.py index 3c61017ff..3310eb257 100644 --- a/scripts/benchmark/physics/benchmark_physics_step_genesis.py +++ b/scripts/benchmark/physics/benchmark_physics_step_genesis.py @@ -39,11 +39,9 @@ else: gs = _gs -from scripts.benchmark.core import device_info as _device_info from scripts.benchmark.core import task_names as _task_names +from unilab.utils.device import get_device_info_dict, get_device_info_line -get_device_info_dict = _device_info.get_device_info_dict -get_device_info_line = _device_info.get_device_info_line canonical_locomotion_task_ids = _task_names.canonical_locomotion_task_ids locomotion_task_spec = _task_names.locomotion_task_spec locomotion_task_model_file = _task_names.locomotion_task_model_file diff --git a/scripts/benchmark/physics/benchmark_physics_step_mj_step.py b/scripts/benchmark/physics/benchmark_physics_step_mj_step.py index c668926a2..db854fb96 100644 --- a/scripts/benchmark/physics/benchmark_physics_step_mj_step.py +++ b/scripts/benchmark/physics/benchmark_physics_step_mj_step.py @@ -33,11 +33,9 @@ import numpy as np from matplotlib.patches import Patch -from scripts.benchmark.core import device_info as _device_info from scripts.benchmark.core import task_names as _task_names +from unilab.utils.device import get_device_info_dict, get_device_info_line -get_device_info_dict = _device_info.get_device_info_dict -get_device_info_line = _device_info.get_device_info_line canonical_locomotion_task_ids = _task_names.canonical_locomotion_task_ids locomotion_task_spec = _task_names.locomotion_task_spec locomotion_task_model_file = _task_names.locomotion_task_model_file diff --git a/scripts/benchmark/physics/benchmark_physics_step_motrixsim.py b/scripts/benchmark/physics/benchmark_physics_step_motrixsim.py index 0c2d42655..c0a231bf1 100644 --- a/scripts/benchmark/physics/benchmark_physics_step_motrixsim.py +++ b/scripts/benchmark/physics/benchmark_physics_step_motrixsim.py @@ -42,11 +42,9 @@ matplotlib.use("Agg") import matplotlib.pyplot as plt -from scripts.benchmark.core import device_info as _benchmark_device_info from scripts.benchmark.core import task_names as _benchmark_task_names +from unilab.utils.device import get_device_info_dict, get_device_info_line -get_device_info_dict = _benchmark_device_info.get_device_info_dict -get_device_info_line = _benchmark_device_info.get_device_info_line canonical_locomotion_task_ids = _benchmark_task_names.canonical_locomotion_task_ids locomotion_task_spec = _benchmark_task_names.locomotion_task_spec locomotion_task_model_file = _benchmark_task_names.locomotion_task_model_file diff --git a/scripts/benchmark/physics/benchmark_physics_step_mujoco_warp.py b/scripts/benchmark/physics/benchmark_physics_step_mujoco_warp.py index 5cae4fa3c..296f7ad12 100644 --- a/scripts/benchmark/physics/benchmark_physics_step_mujoco_warp.py +++ b/scripts/benchmark/physics/benchmark_physics_step_mujoco_warp.py @@ -55,11 +55,9 @@ except ImportError: warp = None -from scripts.benchmark.core import device_info as _device_info from scripts.benchmark.core import task_names as _task_names +from unilab.utils.device import get_device_info_dict, get_device_info_line -get_device_info_dict = _device_info.get_device_info_dict -get_device_info_line = _device_info.get_device_info_line canonical_locomotion_task_ids = _task_names.canonical_locomotion_task_ids locomotion_task_spec = _task_names.locomotion_task_spec locomotion_task_model_file = _task_names.locomotion_task_model_file diff --git a/scripts/benchmark/physics/benchmark_sim.py b/scripts/benchmark/physics/benchmark_sim.py index f4e698fc9..1ed84afb2 100644 --- a/scripts/benchmark/physics/benchmark_sim.py +++ b/scripts/benchmark/physics/benchmark_sim.py @@ -35,7 +35,7 @@ import matplotlib -from scripts.benchmark.core.device_info import get_device_info_dict, get_device_info_line +from unilab.utils.device import get_device_info_dict, get_device_info_line matplotlib.use("Agg") import matplotlib.pyplot as plt diff --git a/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py b/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py index e11c194b3..06b964970 100644 --- a/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py +++ b/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py @@ -58,8 +58,8 @@ if str(ROOT_DIR) not in sys.path: sys.path.append(str(ROOT_DIR)) -from scripts.benchmark.core.device_info import get_device_info_dict, get_device_info_line from unilab.base.backend_timing import UPDATE_STATE_DETAIL_TIMING_KEYS +from unilab.utils.device import get_device_info_dict, get_device_info_line DEFAULT_OUTPUT_JSON = ( ROOT_DIR / "scripts" / "benchmark" / "outputs" / "offpolicy_collector_active" / "results.json" diff --git a/scripts/benchmark/rl/benchmark_offpolicy_dp_scaling.py b/scripts/benchmark/rl/benchmark_offpolicy_dp_scaling.py index c09acfa36..169d241bd 100644 --- a/scripts/benchmark/rl/benchmark_offpolicy_dp_scaling.py +++ b/scripts/benchmark/rl/benchmark_offpolicy_dp_scaling.py @@ -55,8 +55,8 @@ if str(ROOT_DIR) not in sys.path: sys.path.append(str(ROOT_DIR)) -from scripts.benchmark.core.device_info import get_device_info_dict, get_device_info_line from scripts.benchmark.core.output import print_table, save_json +from unilab.utils.device import get_device_info_dict, get_device_info_line DEFAULT_OUTPUT_JSON = ( ROOT_DIR / "scripts" / "benchmark" / "outputs" / "offpolicy_dp_scaling" / "results.json" diff --git a/scripts/benchmark/rl/benchmark_replay_buffer_placement.py b/scripts/benchmark/rl/benchmark_replay_buffer_placement.py index 953c254e3..12b1df6a6 100644 --- a/scripts/benchmark/rl/benchmark_replay_buffer_placement.py +++ b/scripts/benchmark/rl/benchmark_replay_buffer_placement.py @@ -39,7 +39,7 @@ if str(ROOT_DIR) not in sys.path: sys.path.append(str(ROOT_DIR)) -from scripts.benchmark.core.device_info import get_device_info_dict, get_device_info_line +from unilab.utils.device import get_device_info_dict, get_device_info_line plt: Any | None = None try: diff --git a/scripts/benchmark/rl/benchmark_sac_replay_buffer_sampling.py b/scripts/benchmark/rl/benchmark_sac_replay_buffer_sampling.py index 55238f9ba..302cac01d 100644 --- a/scripts/benchmark/rl/benchmark_sac_replay_buffer_sampling.py +++ b/scripts/benchmark/rl/benchmark_sac_replay_buffer_sampling.py @@ -49,7 +49,7 @@ if str(ROOT_DIR) not in sys.path: sys.path.append(str(ROOT_DIR)) -from scripts.benchmark.core.device_info import get_device_info_dict, get_device_info_line +from unilab.utils.device import get_device_info_dict, get_device_info_line plt: Any | None = None try: diff --git a/scripts/benchmark/torch_env/run_benchmark.py b/scripts/benchmark/torch_env/run_benchmark.py index 91d21a3c7..b728dfbb0 100644 --- a/scripts/benchmark/torch_env/run_benchmark.py +++ b/scripts/benchmark/torch_env/run_benchmark.py @@ -38,7 +38,6 @@ if str(ROOT_DIR) not in sys.path: sys.path.append(str(ROOT_DIR)) -from scripts.benchmark.core.device_info import get_device_info_dict from scripts.benchmark.torch_env import motion_tracking, walk_flat from scripts.benchmark.torch_env.xp import ( NumpyBackend, @@ -48,6 +47,7 @@ TorchBackend, TorchRng, ) +from unilab.utils.device import get_device_info_dict DEFAULT_OUTPUT_JSON = ROOT_DIR / "scripts" / "benchmark" / "outputs" / "torch_env" / "results.json" RESET_SIZES = (16, 256, 2048)