补充基准测试设备元数据采集#10
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces environment metadata collection (including PyTorch version and CUDA status/device details) into the benchmark records, along with a new unit test to verify the behavior when CUDA is unavailable. The feedback suggests caching the metadata to avoid redundant calls during benchmark runs, adding exception handling to prevent CUDA-related errors from crashing the benchmarks, and resetting the cache in the unit test to ensure proper mock isolation.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def _collect_environment_metadata(): | ||
| metadata = { | ||
| 'torch_version': torch.__version__, | ||
| 'cuda_available': torch.cuda.is_available(), | ||
| 'cuda_device_count': torch.cuda.device_count(), | ||
| } | ||
| if metadata['cuda_available'] and metadata['cuda_device_count'] > 0: | ||
| current_device = torch.cuda.current_device() | ||
| metadata.update({ | ||
| 'cuda_current_device': current_device, | ||
| 'cuda_device_name': torch.cuda.get_device_name(current_device), | ||
| }) | ||
| return metadata |
There was a problem hiding this comment.
性能与鲁棒性改进建议
-
性能优化(缓存元数据):
_collect_environment_metadata会在每次记录 benchmark 时(即_record中)被调用。由于测试运行期间的 PyTorch 版本和 GPU 硬件环境是静态不变的,重复调用torch.cuda相关 API 会带来不必要的开销。建议引入一个模块级的缓存变量(例如_metadata_cache)来避免重复计算。 -
鲁棒性增强(异常处理):
在某些特殊或配置不当的环境中(例如 CUDA 驱动异常、部分国产 GPU 适配环境等),调用torch.cuda.current_device()或torch.cuda.get_device_name()可能会抛出RuntimeError或AssertionError。为了防止元数据收集失败导致整个 Benchmark 测试中断,建议对 CUDA 相关的属性获取加上try...except保护。
_metadata_cache = None
def _collect_environment_metadata():
global _metadata_cache
if _metadata_cache is not None:
return _metadata_cache
metadata = {
'torch_version': torch.__version__,
'cuda_available': torch.cuda.is_available(),
'cuda_device_count': torch.cuda.device_count(),
}
if metadata['cuda_available'] and metadata['cuda_device_count'] > 0:
try:
current_device = torch.cuda.current_device()
metadata.update({
'cuda_current_device': current_device,
'cuda_device_name': torch.cuda.get_device_name(current_device),
})
except Exception:
pass
_metadata_cache = metadata
return metadata| def test_collect_environment_metadata_without_cuda(monkeypatch): | ||
| import tests.pytest_benchmark_plugin as plugin | ||
|
|
||
| monkeypatch.setattr(plugin.torch.cuda, 'is_available', lambda: False) |
There was a problem hiding this comment.
为了配合 _collect_environment_metadata 的缓存机制,在单测中需要显式重置 _metadata_cache,以确保 mock 的 is_available 和 device_count 能够正确生效,避免受到其他测试用例或先前调用的缓存干扰。
| monkeypatch.setattr(plugin.torch.cuda, 'is_available', lambda: False) | |
| plugin._metadata_cache = None | |
| monkeypatch.setattr(plugin.torch.cuda, 'is_available', lambda: False) |
?? API ?? 2 ??????
这个改动补全了基准测试插件的设备元数据采集逻辑,在重复收集和设备查询异常时都能保持稳定行为,避免 benchmark 流程因为环境探测波动而失败。同时为元数据缓存增加了可控重置,保证测试场景之间互不污染。
相关验证已经完成,新增与调整后的测试覆盖了缓存命中、异常回退和测试隔离场景,并通过基于插件加载的本地校验确认行为符合预期。