chore(benchmark): delete device_info re-export shim, import unilab.utils.device directly (#2474)

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

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

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"
 

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

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

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",

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

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"

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

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,

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

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

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

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

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

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"

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"

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:

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:

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)