fix: prefer NVML v2 memory info for inference setup

[Rohithmatham12]

Summary

• try pynvml.nvmlDeviceGetMemoryInfo_v2() before the legacy v1 memory-info API in inference setup
• keep compatibility with older pynvml builds by falling back to nvmlDeviceGetMemoryInfo() when v2 is unavailable
• keep the existing torch/default memory fallback and ensure nvmlShutdown() runs even when NVML probing fails
• add unit coverage for v2 success when v1 would raise NVMLError_NotSupported, plus the older-pynvml fallback path

Why
DGX Spark / GB10 platforms can report pynvml.NVMLError_NotSupported from the legacy v1 nvmlDeviceGetMemoryInfo() call during Cosmos3 inference setup. The v2 NVML memory-info API is the supported path there, while older environments still need the v1 fallback.

Testing

• python3 -m py_compile cosmos_framework/inference/args.py cosmos_framework/inference/args_test.py
• git diff --check

Not run locally:

• python3 -m pytest cosmos_framework/inference/args_test.py -q because this local environment does not have pytest installed
• import smoke because this local environment is missing framework dependencies such as pydantic

Related to NVIDIA/cosmos#180

[Rohithmatham12]

Pushed a test-only follow-up for the unittest failure. CI's pynvml build does not expose nvmlDeviceGetMemoryInfo_v2, so the regression test now monkeypatches that symbol with raising=False before exercising the v2-preferred path.\n\nLocal validation after the change:\n- python3 -m py_compile cosmos_framework/inference/args.py cosmos_framework/inference/args_test.py\n- git diff --check\n\nI also tried the targeted pytest locally, but this environment does not have pytest installed.

[Rohithmatham12]

Sorry, the approval became stale because I pushed a one-line test-only follow-up to fix the failing unittest on CI. The implementation code is unchanged from the approved version.

[lfengad]

@Rohithmatham12 Thx for the check, let's wait for the PPL pass and we could do the merge

[foreverlms]

Here is a bit mis-understanding: we split the exception capture into two parts: one is in inside _get_nvml_device_memory_info, the other is outside this func. Should we keep the exception capture unified in one-place?

[foreverlms]

OK, this is trivial, this pr is good for me.