feat: add GPU filter (--gpus) to nvlink_monitor

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 <filename>`: 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 <id1,id2,...>`: 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.

monitor/arg_parser.cpp

@@ -1,5 +1,6 @@
 #include "arg_parser.h"
 
+#include <sstream>
 #include <string>
 
 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")) {

monitor/arg_parser.h

@@ -2,6 +2,7 @@
 #define NVLINK_MONITOR_ARG_PARSER_H
 
 #include <string>
+#include <vector>
 
 // 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<int>
+        gpuFilter;  // Empty = monitor all GPUs; else only listed ids
     std::string outputFilename;
     bool helpRequested = false;
     bool ok = true;

monitor/nvlink_monitor.cpp

@@ -35,10 +35,12 @@ void signal_handler(int signal) {
 
 // NvLinkMonitor constructor
 NvLinkMonitor::NvLinkMonitor(bool verbose, OutputFormat format,
+                             const std::vector<int>& 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<int>(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<GPUMonitorResult> NvLinkMonitor::getNvLinkData() {
     std::vector<GPUMonitorResult> 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 <id1,id2,...>      : 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);

monitor/nvlink_monitor.h

@@ -8,6 +8,7 @@
 #include <fstream>
 #include <iomanip>
 #include <iostream>
+#include <set>
 #include <string>
 #include <thread>
 #include <vector>
@@ -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<int> 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<int>& gpuFilter = {},
                   const std::string& outputFilename = "");
 
     /**

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<int>({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<int>({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<int>({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());