Skip to content

nixlbench: limit Neuron core allocation in SG mode for multi-processing #1913

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fengjica wants to merge 1 commit into
base: main
Choose a base branch
from
+29 −1
Open

nixlbench: limit Neuron core allocation in SG mode for multi-processing #1913

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 29 additions & 1 deletion benchmark/nixlbench/src/utils/neuron.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,14 @@
*/

#include "neuron.h"
#include "utils.h"

#include <cstdlib>
#include <dlfcn.h>

#include <cstddef>
#include <cstdint>
#include <iostream>
#include <memory>
#include <mutex>
#include <unordered_map>
Expand Down Expand Up @@ -142,6 +145,27 @@ getTensorFromVA(const void *va) {
int
neuronCoreCount() {
static const int core_count = []() {
// Scope this process to the minimum number of NeuronCores it needs,
// analogous to cudaSetDevice() scoping a CUDA process to one GPU.
// Without this, nrt_init() claims ALL visible cores by default,
// preventing sibling processes on the same host from initializing.
if (XFERBENCH_MODE_SG == xferBenchConfig::mode) {
const char *num_cores = getenv("NEURON_RT_NUM_CORES");
const char *visible_cores = getenv("NEURON_RT_VISIBLE_CORES");
if (!num_cores && !visible_cores) {
if (setenv("NEURON_RT_NUM_CORES", "1", 0) != 0) {
std::cerr << "nixlbench: failed to set NEURON_RT_NUM_CORES" << std::endl;
return -1;
}
std::cerr << "nixlbench: SG mode — set NEURON_RT_NUM_CORES=1" << std::endl;
} else {
std::cerr << "nixlbench: SG mode uses only the first core (vnc=0)"
<< " but NEURON_RT_NUM_CORES=" << (num_cores ? num_cores : "(unset)")
<< ", NEURON_RT_VISIBLE_CORES="
<< (visible_cores ? visible_cores : "(unset)") << std::endl;
}
}
Comment on lines +152 to +167

@coderabbitai coderabbitai Bot Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Gate SG core handling on pairwise SG
Use IS_PAIRWISE_AND_SG() here instead of checking xferBenchConfig::mode alone; otherwise non-pairwise SG runs still get forced to NEURON_RT_NUM_CORES=1 and vnc=0.

🧰 Tools
🪛 GitHub Actions: Clang Format Check / 0_clang-format.txt

[error] 154-160: clang-format check failed for modified C/C++ file. clang-format-diff-19 -p1 -style=file -v detected formatting differences (std::cerr output stream formatting/line wrapping).

🪛 GitHub Actions: Clang Format Check / clang-format

[error] 154-157: clang-format-diff-19 reported formatting differences (C/C++ formatting check failed). Run clang-format to apply the suggested changes.

[error] 159-164: clang-format-diff-19 reported formatting differences (C/C++ formatting check failed). The affected std::cerr statement formatting needs to match clang-format output.

🤖 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 `@benchmark/nixlbench/src/utils/neuron.cpp` around lines 152 - 165, Update the
SG core-handling condition in neuron.cpp to use IS_PAIRWISE_AND_SG() instead of
checking xferBenchConfig::mode alone. Keep the existing environment-variable
setup and warning behavior unchanged, but ensure it runs only for pairwise SG
configurations.


uint32_t vnc_count;
if (nrt_init(1 /* framework_type=NO_FW */, "nixl_bench", "nixl_bench") == 0 &&
nrt_get_visible_vnc_count(&vnc_count) == 0) {
Expand All @@ -158,7 +182,11 @@ neuronMalloc(void **addr, size_t buffer_size, int devid) {
nrt_tensor *tensor;
int status;

status = nrt_tensor_allocate(0 /* placement=device */, devid, buffer_size, nullptr, &tensor);
// VNC index is relative to the process's allocated core set.
// In SG mode each process owns 1 core, so always use vnc=0.
int vnc = (XFERBENCH_MODE_SG == xferBenchConfig::mode) ? 0 : devid;

status = nrt_tensor_allocate(0 /* placement=device */, vnc, buffer_size, nullptr, &tensor);
if (status != 0) return status;

NrtTensorPtr ptr{tensor};
Expand Down
Loading