Assign gRPC workers to distinct GPUs via cudaSetDevice - #1449

tmckayus · 2026-06-22T20:25:08Z

Round-robin workers across visible CUDA devices at worker startup and initialize a per-process RMM pool (CUOPT_GIGABYTES_PER_PROC) so multiple workers can share a GPU without exhausting memory.

The logic to actually map a worker process to a specific GPU was missing previously (oversight).

Round-robin workers across visible CUDA devices at worker startup and initialize a per-process RMM pool (CUOPT_GIGABYTES_PER_PROC) so multiple workers can share a GPU without exhausting memory.

tmckayus · 2026-06-22T20:28:10Z

 #include "grpc_pipe_serialization.hpp"
 #include "grpc_server_types.hpp"

+#include <cctype>


Note that these routines below are simply for logging in the main server, so that we can log information about the number of workers and gpus being used without actually initializing CUDA before forking the worker processes (which causes issues).

The workers themselves simply ask CUDA for available devices and apply a mod on worker id

coderabbitai · 2026-06-22T20:31:00Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c1604136-0333-4501-807c-bb433178a4d2

📥 Commits

Reviewing files that changed from the base of the PR and between 25fbe56 and ec2a5b7.

📒 Files selected for processing (3)

cpp/src/grpc/server/grpc_server_types.hpp
cpp/src/grpc/server/grpc_worker.cpp
cpp/src/grpc/server/grpc_worker_infra.cpp

🚧 Files skipped from review as they are similar to previous changes (2)

cpp/src/grpc/server/grpc_server_types.hpp
cpp/src/grpc/server/grpc_worker_infra.cpp

📝 Walkthrough

Walkthrough

Adds GPU layout detection and logging to the gRPC server startup sequence, and initializes per-worker CUDA device selection and RMM memory pool before workers begin processing jobs. GPU count is determined from CUDA_VISIBLE_DEVICES with a fallback to nvidia-smi -L.

Changes

GPU Layout Logging and Worker CUDA/RMM Initialization

Layer / File(s)	Summary
GPU discovery helpers and startup layout logging `cpp/src/grpc/server/grpc_server_types.hpp`, `cpp/src/grpc/server/grpc_worker_infra.cpp`, `cpp/src/grpc/server/grpc_server_main.cpp`	Adds `visible_gpu_count` global and forward declarations to the shared header. Implements `count_cuda_visible_devices()` (parses `CUDA_VISIBLE_DEVICES`) and `discover_gpu_count_via_nvidia_smi()` (runs `nvidia-smi -L` via `popen`) as anonymous-namespace helpers, and `log_worker_gpu_layout()` which uses them with fallback logic and warns when worker count exceeds GPU count. Calls `log_worker_gpu_layout()` in the server startup sequence before `spawn_workers()`.
Per-worker CUDA device selection and RMM pool initialization `cpp/src/grpc/server/grpc_worker.cpp`	Implements RMM pool initialization via one-time `pool_memory_resource` setup sized by `CUOPT_GIGABYTES_PER_PROC` with validation and warnings on invalid values. Adds `init_worker_cuda_environment(worker_id)` to select the CUDA device via `worker_id % device_count`, call `cudaSetDevice`, and initialize the RMM pool. Wires `init_worker_cuda_environment(worker_id)` into `worker_process` after the startup log and before shared-memory bookkeeping; on failure, logs an error and terminates the worker via `_exit(1)`.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 44.44% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately captures the primary change: assigning gRPC workers to distinct GPUs via cudaSetDevice, which is the core feature implemented across the files.
Description check	✅ Passed	The description is directly related to the changeset, explaining the round-robin worker-to-GPU assignment and per-process RMM pool initialization, both clearly evident in the file modifications.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (1)

cpp/src/grpc/server/grpc_worker.cpp (1)

21-25: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick win

Harden CUOPT_GIGABYTES_PER_PROC parsing and bounds checks.

Line 23 uses std::atoi, which silently accepts partial strings and has undefined overflow behavior; a malformed env value can produce unintended pool sizing and unstable startup behavior.

💡 Proposed refactor

+ `#include` <cerrno>
+ `#include` <limits>
@@
   int pool_gigs = 1;
   if (const char* env = std::getenv("CUOPT_GIGABYTES_PER_PROC")) {
-    const int parsed = std::atoi(env);
-    if (parsed > 0) { pool_gigs = parsed; }
+    char* end         = nullptr;
+    errno             = 0;
+    long long parsed  = std::strtoll(env, &end, 10);
+    if (errno == 0 && end != env && *end == '\0' && parsed > 0 &&
+        parsed <= std::numeric_limits<int>::max()) {
+      pool_gigs = static_cast<int>(parsed);
+    } else {
+      SERVER_LOG_WARN("[Worker] Ignoring invalid CUOPT_GIGABYTES_PER_PROC='%s'", env);
+    }
   }

As per coding guidelines, “Prevent numerical instability (overflow, underflow, precision loss) producing wrong results.”

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@cpp/src/grpc/server/grpc_worker.cpp` around lines 21 - 25, The
CUOPT_GIGABYTES_PER_PROC environment variable parsing on line 23 uses std::atoi
which silently accepts partial strings and has undefined overflow behavior,
creating numerical instability. Replace std::atoi with std::stoi wrapped in a
try-catch block to handle parsing errors gracefully, and add upper bounds
checking to ensure the parsed value does not exceed a reasonable maximum (in
addition to the existing positive check). This will prevent malformed
environment values from producing unintended pool sizing and unstable startup
behavior.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@cpp/src/grpc/server/grpc_worker.cpp`:
- Around line 42-50: The init_worker_cuda_environment function currently logs
errors but returns void, making it impossible for the caller to detect
initialization failures. Change the function signature from void to return a
bool or status code, then modify the error handling at lines 46-50 and the other
CUDA error cases to return false instead of just logging and returning. Update
the caller at line 582 where init_worker_cuda_environment is invoked to check
the return status and fail worker startup if the CUDA environment initialization
returns false, preventing the worker from becoming active with a failed GPU
state.

---

Nitpick comments:
In `@cpp/src/grpc/server/grpc_worker.cpp`:
- Around line 21-25: The CUOPT_GIGABYTES_PER_PROC environment variable parsing
on line 23 uses std::atoi which silently accepts partial strings and has
undefined overflow behavior, creating numerical instability. Replace std::atoi
with std::stoi wrapped in a try-catch block to handle parsing errors gracefully,
and add upper bounds checking to ensure the parsed value does not exceed a
reasonable maximum (in addition to the existing positive check). This will
prevent malformed environment values from producing unintended pool sizing and
unstable startup behavior.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8c7410cd-ca08-4ff6-a3a0-8e77c3c7f64a

📥 Commits

Reviewing files that changed from the base of the PR and between 8c79892 and 25fbe56.

📒 Files selected for processing (4)

cpp/src/grpc/server/grpc_server_main.cpp
cpp/src/grpc/server/grpc_server_types.hpp
cpp/src/grpc/server/grpc_worker.cpp
cpp/src/grpc/server/grpc_worker_infra.cpp

ramakrishnap-nv · 2026-06-22T20:43:07Z

+
+void init_worker_rmm_pool()
+{
+  int pool_gigs = 1;


Are we using 1GB by default as rmm pool memory?

I think so, let me double check that against the Python server

@ramakrishnap-nv yes the default in the Python server is 1GB, can be overridden with environment variable

akifcorduk · 2026-06-24T12:31:17Z

+  const int pool_gigs = parse_pool_gigs_env();
+
+  // Keep the pool alive for the lifetime of this worker process.
+  static std::unique_ptr<rmm::mr::pool_memory_resource> pool_mr;


Do you think we will ever do multi-threading or is it always going to be multi-processing? If we plan to do multi threading, static variables might cause issues.

I think we'll always do multi-process, because it is the easiest way to cancel a job (we don't have cancel hooks in the solver) or restart a worker if it has CUDA errors or something. And it makes it easy to map a GPU to an isolated entity

akifcorduk · 2026-06-24T12:35:23Z

+    return false;
+  }
+
+  init_worker_rmm_pool();


Shouldn't this be done once per device instead of once per worker, otherwise we are allocating a new memory resource and assign it to the device again which might cause issues when there are multiple workers using the same device?

good feedback, I'll take a look at that

@akifcorduk
This is the way we did it in the Python server -- each process gets it own pool on a device, and that only happens once per process. Not sure how we would share a pool across processes -- is there a simple way to do that?

If it might be a problem, the simple fix is to cap the number of workers to the number of GPUs available. This would guarantee only 1 pool per device (held by 1 process) unless you ran multiple servers. It would just prevent oversubscribing.

Oh worker is a process. Sorry for the confusion, in cpp OMP worker keyword is per thread. Yes, one pool per process is needed, it won't work otherwise, a different process cannot access cuda memory.

akifcorduk

LGTM, thanks Trevor!

tmckayus · 2026-06-25T15:46:18Z

/merge

Assign gRPC workers to distinct GPUs via cudaSetDevice

25fbe56

Round-robin workers across visible CUDA devices at worker startup and initialize a per-process RMM pool (CUOPT_GIGABYTES_PER_PROC) so multiple workers can share a GPU without exhausting memory.

tmckayus requested a review from a team as a code owner June 22, 2026 20:25

tmckayus requested review from hlinsen and rg20 June 22, 2026 20:25

tmckayus added bug Something isn't working non-breaking Introduces a non-breaking change labels Jun 22, 2026

tmckayus commented Jun 22, 2026

View reviewed changes

tmckayus requested a review from ramakrishnap-nv June 22, 2026 20:29

coderabbitai Bot reviewed Jun 22, 2026

View reviewed changes

Comment thread cpp/src/grpc/server/grpc_worker.cpp Outdated

ramakrishnap-nv approved these changes Jun 22, 2026

View reviewed changes

ramakrishnap-nv self-requested a review June 22, 2026 20:38

ramakrishnap-nv reviewed Jun 22, 2026

View reviewed changes

address coderabbit feedback on grpc server per-gpu worker gpu

ec2a5b7

akifcorduk reviewed Jun 24, 2026

View reviewed changes

akifcorduk approved these changes Jun 25, 2026

View reviewed changes

rapids-bot Bot merged commit 2d6569b into NVIDIA:main Jun 25, 2026
98 checks passed

Uh oh!

Conversation

tmckayus commented Jun 22, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot commented Jun 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

akifcorduk left a comment

Choose a reason for hiding this comment

Uh oh!

tmckayus commented Jun 25, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

coderabbitai Bot commented Jun 22, 2026 •

edited

Loading