From ae76c586fb8b35285c266c943d2ecfab74cebf6e Mon Sep 17 00:00:00 2001 From: Yuxin Chen Date: Fri, 10 Jul 2026 08:53:20 +0000 Subject: [PATCH] feat: add GPU filter (--gpus) to nvlink_monitor Add `-g, --gpus ` to nvlink_monitor so the user can restrict monitoring to a subset of GPUs. On an 8-GPU system where a bandwidth test only touches 2 GPUs, this cuts the irrelevant output and lets the user focus on the links that matter. Behavior: - Empty filter (default) => monitor all GPUs (unchanged). - Non-empty filter => getNvLinkData skips GPUs whose index is not listed, so CSV/JSON/text output and bandwidth calc only cover the selected GPUs. Snapshot indices stay aligned across samples (filter is stable for the monitor's lifetime), so calculateBandwidth's index-based matching is unaffected. - The filter is validated against discovered GPUs in the constructor after discovery: an out-of-range id throws std::runtime_error (surfaced as "Error: GPU id N out of range (0-M)" by main's catch) so a typo does not silently monitor nothing. A "Monitoring K of N GPU(s) (filter applied)" notice is printed to stderr when a filter is active. Parsing (monitor/arg_parser.{h,cpp}): comma-separated integer list, e.g. "0,1,3". Whitespace around tokens is tolerated. Negative ids, non-numeric tokens, missing value, and empty lists are all rejected with ok=false. Tests: 72 pass (was 66; +6 for --gpus: single short, multiple long, whitespace tolerated, invalid token rejected, negative rejected, missing value rejected; Defaults asserts gpuFilter empty). make check-format clean. make monitor builds with -Wall -Wextra, no warnings. Verified on an 8x H20 system: -g 0,1 limits CSV output to GPUs 0 and 1; -g 9 errors with "out of range (0-7)". Signed-off-by: staryxchen --- README.md | 1 + monitor/arg_parser.cpp | 38 +++++++++++++++++++++++++++++++++++ monitor/arg_parser.h | 3 +++ monitor/nvlink_monitor.cpp | 39 +++++++++++++++++++++++++++++++++++- monitor/nvlink_monitor.h | 12 +++++++++-- test/test_arg_parser.cpp | 41 ++++++++++++++++++++++++++++++++++++++ 6 files changed, 131 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index e873ac6..724c63e 100644 --- a/README.md +++ b/README.md @@ -160,6 +160,7 @@ pipe or redirect stdout directly for analysis: - `-v, --verbose`: Enable detailed NvLink output (shows individual link bandwidth) - `-o, --output `: Redirect output to file - `-f, --format text|csv|json`: Output format (default: text). CSV and JSON are machine-readable; `--verbose` only affects the text path +- `-g, --gpus `: Only monitor the listed GPU indices (e.g. `-g 0,1,3`); default = all GPUs. Out-of-range ids are rejected at startup - `-h, --help`: Show help information **Note:** The interval parameter supports decimal values (e.g., 0.5 for 500ms, 0.1 for 100ms). The minimum practical interval is 1 microsecond (0.000001s), but very small intervals may affect system performance. diff --git a/monitor/arg_parser.cpp b/monitor/arg_parser.cpp index be90ef9..f9425cf 100644 --- a/monitor/arg_parser.cpp +++ b/monitor/arg_parser.cpp @@ -1,5 +1,6 @@ #include "arg_parser.h" +#include #include MonitorCliArgs parseMonitorArgs(int argc, char* argv[]) { @@ -58,6 +59,43 @@ MonitorCliArgs parseMonitorArgs(int argc, char* argv[]) { "Invalid format: " + f + " (expected text, csv, or json)"; return args; } + } else if (arg == "-g" || arg == "--gpus") { + if (i + 1 >= argc) { + args.ok = false; + args.errorMessage = "Missing value for " + arg; + return args; + } + // Parse a comma-separated list of GPU indices, e.g. "0,1,3". + // Whitespace around tokens is tolerated. + std::string val = argv[++i]; + std::string token; + std::istringstream iss(val); + while (std::getline(iss, token, ',')) { + // Trim leading/trailing whitespace. + size_t a = token.find_first_not_of(" \t"); + size_t b = token.find_last_not_of(" \t"); + if (a == std::string::npos) continue; // skip empty token + token = token.substr(a, b - a + 1); + try { + int id = std::stoi(token); + if (id < 0) { + args.ok = false; + args.errorMessage = + "Invalid GPU id (negative): " + token; + return args; + } + args.gpuFilter.push_back(id); + } catch (const std::exception&) { + args.ok = false; + args.errorMessage = "Invalid GPU id: " + token; + return args; + } + } + if (args.gpuFilter.empty()) { + args.ok = false; + args.errorMessage = "No valid GPU ids in: " + val; + return args; + } } else if (arg == "-c" || arg == "--continuous") { if (i + 1 < argc && (std::string(argv[i + 1]) == "true" || std::string(argv[i + 1]) == "false")) { diff --git a/monitor/arg_parser.h b/monitor/arg_parser.h index d245239..b6b97a5 100644 --- a/monitor/arg_parser.h +++ b/monitor/arg_parser.h @@ -2,6 +2,7 @@ #define NVLINK_MONITOR_ARG_PARSER_H #include +#include // Output format for nvlink_monitor. // Text - human-readable (default) @@ -16,6 +17,8 @@ struct MonitorCliArgs { bool continuous = true; bool verbose = false; OutputFormat format = OutputFormat::Text; + std::vector + gpuFilter; // Empty = monitor all GPUs; else only listed ids std::string outputFilename; bool helpRequested = false; bool ok = true; diff --git a/monitor/nvlink_monitor.cpp b/monitor/nvlink_monitor.cpp index 304f966..a2dc133 100644 --- a/monitor/nvlink_monitor.cpp +++ b/monitor/nvlink_monitor.cpp @@ -35,10 +35,12 @@ void signal_handler(int signal) { // NvLinkMonitor constructor NvLinkMonitor::NvLinkMonitor(bool verbose, OutputFormat format, + const std::vector& gpuFilterVec, const std::string& outputFilename) : verboseOutput(verbose), outputFormat(format), csvHeaderPrinted(false), + gpuFilter(gpuFilterVec.begin(), gpuFilterVec.end()), fileOutput(!outputFilename.empty()) { // Initialize NVML nvmlReturn_t result = nvmlInit(); @@ -60,6 +62,31 @@ NvLinkMonitor::NvLinkMonitor(bool verbose, OutputFormat format, // Discover GPUs discoverGPUs(); + + // Validate the GPU filter against discovered GPUs (after discovery so we + // know the valid id range). Throw on out-of-range ids so the user gets a + // clear error instead of silently monitoring nothing. + if (!gpuFilter.empty() && !gpus.empty()) { + int max_id = static_cast(gpus.size()) - 1; + for (int id : gpuFilter) { + if (id < 0 || id > max_id) { + throw std::runtime_error("GPU id " + std::to_string(id) + + " out of range (0-" + + std::to_string(max_id) + ")"); + } + } + std::cerr << "Monitoring " << gpuFilter.size() << " of " << gpus.size() + << " GPU(s) (filter applied)" << std::endl; + } +} + +bool NvLinkMonitor::isGpuSelected(const std::string& id) const { + if (gpuFilter.empty()) return true; // no filter => all GPUs + try { + return gpuFilter.count(std::stoi(id)) > 0; + } catch (const std::exception&) { + return false; // non-numeric id not in filter + } } // NvLinkMonitor destructor @@ -133,6 +160,9 @@ std::vector NvLinkMonitor::getNvLinkData() { std::vector results; for (const auto& gpu : gpus) { + // Skip GPUs not in the --gpus filter (empty filter = all). + if (!isGpuSelected(gpu.id)) continue; + GPUMonitorResult result; result.gpuId = gpu.id; result.nvLinkCount = gpu.nvLinkCount; @@ -468,6 +498,9 @@ void printHelp(const char* programName) { std::cout << " -f, --format text|csv|json : Output format (default: text)" << std::endl; + std::cout + << " -g, --gpus : Only monitor listed GPU indices" + << std::endl; std::cout << " -h, --help : Show this help message" << std::endl; std::cout << std::endl; @@ -497,6 +530,9 @@ void printHelp(const char* programName) { std::cout << " " << programName << " --format json -o data.jsonl # JSONL output to file" << std::endl; + std::cout << " " << programName + << " -g 0,1 # Only monitor GPUs 0 and 1" + << std::endl; } int main(int argc, char* argv[]) { @@ -517,7 +553,8 @@ int main(int argc, char* argv[]) { signal(SIGTERM, signal_handler); try { - NvLinkMonitor monitor(args.verbose, args.format, args.outputFilename); + NvLinkMonitor monitor(args.verbose, args.format, args.gpuFilter, + args.outputFilename); if (args.continuous) { monitor.runContinuousMonitoring(args.interval); diff --git a/monitor/nvlink_monitor.h b/monitor/nvlink_monitor.h index 847f29c..c457e92 100644 --- a/monitor/nvlink_monitor.h +++ b/monitor/nvlink_monitor.h @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -64,21 +65,28 @@ class NvLinkMonitor { bool verboseOutput; // Flag for detailed NvLink output OutputFormat outputFormat; // Output format (text/csv/json) bool csvHeaderPrinted; // Whether the CSV header has been emitted + std::set gpuFilter; // Empty = monitor all; else only listed indices std::ofstream outputFile; // Output file stream bool fileOutput; // Flag for file output + // Returns true if the GPU with the given id string should be monitored. + // Empty filter => all GPUs; otherwise the id must be in the filter set. + bool isGpuSelected(const std::string& id) const; + public: /** * @brief Constructor - initializes NVML and discovers GPUs * @param verbose Enable detailed NvLink output * @param format Output format (text/csv/json) + * @param gpuFilter Optional list of GPU indices to monitor (empty = all) * @param outputFilename Optional output file name (empty for console * output) - * @throws std::runtime_error if NVML initialization fails or file cannot be - * opened + * @throws std::runtime_error if NVML initialization fails, the file cannot + * be opened, or a filtered GPU id is out of range */ NvLinkMonitor(bool verbose = false, OutputFormat format = OutputFormat::Text, + const std::vector& gpuFilter = {}, const std::string& outputFilename = ""); /** diff --git a/test/test_arg_parser.cpp b/test/test_arg_parser.cpp index 105678c..9d80680 100644 --- a/test/test_arg_parser.cpp +++ b/test/test_arg_parser.cpp @@ -37,6 +37,7 @@ TEST(MonitorArgs, Defaults) { EXPECT_TRUE(args.continuous); EXPECT_FALSE(args.verbose); EXPECT_EQ(args.format, OutputFormat::Text); + EXPECT_TRUE(args.gpuFilter.empty()); EXPECT_TRUE(args.outputFilename.empty()); EXPECT_FALSE(args.helpRequested); } @@ -201,6 +202,46 @@ TEST(MonitorArgs, FormatMissingValueRejected) { EXPECT_FALSE(args.ok); } +TEST(MonitorArgs, GpusSingleShort) { + Argv a{"nvlink_monitor", "-g", "3"}; + auto args = parseMonitorArgs(a.argc(), a.argv()); + EXPECT_TRUE(args.ok); + EXPECT_EQ(args.gpuFilter, std::vector({3})); +} + +TEST(MonitorArgs, GpusMultipleLong) { + Argv a{"nvlink_monitor", "--gpus", "0,1,3"}; + auto args = parseMonitorArgs(a.argc(), a.argv()); + EXPECT_TRUE(args.ok); + EXPECT_EQ(args.gpuFilter, std::vector({0, 1, 3})); +} + +TEST(MonitorArgs, GpusToleratesWhitespace) { + Argv a{"nvlink_monitor", "--gpus", "0, 1, 3"}; + auto args = parseMonitorArgs(a.argc(), a.argv()); + EXPECT_TRUE(args.ok); + EXPECT_EQ(args.gpuFilter, std::vector({0, 1, 3})); +} + +TEST(MonitorArgs, GpusInvalidTokenRejected) { + Argv a{"nvlink_monitor", "--gpus", "0,a,3"}; + auto args = parseMonitorArgs(a.argc(), a.argv()); + EXPECT_FALSE(args.ok); + EXPECT_FALSE(args.errorMessage.empty()); +} + +TEST(MonitorArgs, GpusNegativeRejected) { + Argv a{"nvlink_monitor", "--gpus", "0,-1"}; + auto args = parseMonitorArgs(a.argc(), a.argv()); + EXPECT_FALSE(args.ok); +} + +TEST(MonitorArgs, GpusMissingValueRejected) { + Argv a{"nvlink_monitor", "--gpus"}; + auto args = parseMonitorArgs(a.argc(), a.argv()); + EXPECT_FALSE(args.ok); +} + TEST(MonitorArgs, UnknownOptionRejected) { Argv a{"nvlink_monitor", "--bogus"}; auto args = parseMonitorArgs(a.argc(), a.argv());