diff --git a/server/src/common/backend_ipc.cpp b/server/src/common/backend_ipc.cpp index 0403a7d73..e8e99f128 100644 --- a/server/src/common/backend_ipc.cpp +++ b/server/src/common/backend_ipc.cpp @@ -179,6 +179,13 @@ bool BackendIpcProcess::start(const BackendIpcLaunchConfig & cfg) { } argv_storage.emplace_back("--stream-fd=" + std::to_string(stream_pipe[1])); + // Apply any per-daemon environment overrides (e.g. GPU visibility + // pinning) in the child before exec, so the parent's own environment is + // untouched. + for (const auto & kv : cfg.child_env) { + ::setenv(kv.first.c_str(), kv.second.c_str(), 1); + } + std::vector argv; argv.reserve(argv_storage.size() + 1); for (std::string & arg : argv_storage) argv.push_back(arg.data()); diff --git a/server/src/common/backend_ipc.h b/server/src/common/backend_ipc.h index 62ea6f493..22eeb2638 100644 --- a/server/src/common/backend_ipc.h +++ b/server/src/common/backend_ipc.h @@ -12,6 +12,7 @@ #include #include #include +#include #include #if !defined(_WIN32) @@ -81,6 +82,10 @@ struct BackendIpcLaunchConfig { std::string work_dir; BackendIpcPayloadTransport payload_transport = BackendIpcPayloadTransport::Auto; size_t shared_payload_bytes = 0; + // Extra environment variables to set in the child *before* exec (e.g. to pin + // the daemon to a subset of GPUs via ROCR_VISIBLE_DEVICES/CUDA_VISIBLE_DEVICES). + // Empty by default: the child inherits the parent environment unchanged. + std::vector> child_env; }; struct BackendIpcPayloadSegment { diff --git a/server/src/common/target_shard_ipc.cpp b/server/src/common/target_shard_ipc.cpp index 776aa6b67..549581f47 100644 --- a/server/src/common/target_shard_ipc.cpp +++ b/server/src/common/target_shard_ipc.cpp @@ -171,7 +171,26 @@ bool TargetShardIpcSession::start(const TargetShardIpcLaunchConfig & cfg) { layer_begin_list += std::to_string(cfg.layer_begins[i]); layer_end_list += std::to_string(cfg.layer_ends[i]); } - launch.args.push_back("--target-gpus=" + gpu_list); + // Optional GPU isolation for the target-shard daemon. By default the main + // process and the daemon both enumerate every GPU; for a cross-backend split + // (CUDA main + HIP daemon) that is harmless, but for a same-backend split + // (e.g. HIP main + HIP daemon across two AMD GPUs) two ROCr runtimes each + // initializing the *other* process's device can hard-fault the machine. + // When DFLASH_TARGET_SHARD_ISOLATE_GPUS is set, pin the daemon to see only + // its assigned GPUs via {CUDA,ROCR}_VISIBLE_DEVICES; under that pinned view + // the visible GPUs re-index from 0, so remap --target-gpus to match. + std::string target_gpu_list = gpu_list; + if (std::getenv("DFLASH_TARGET_SHARD_ISOLATE_GPUS") != nullptr) { + launch.child_env.emplace_back("ROCR_VISIBLE_DEVICES", gpu_list); + launch.child_env.emplace_back("CUDA_VISIBLE_DEVICES", gpu_list); + std::string remapped; + for (size_t i = 0; i < cfg.gpus.size(); ++i) { + if (i > 0) remapped += ","; + remapped += std::to_string(i); + } + target_gpu_list = remapped; + } + launch.args.push_back("--target-gpus=" + target_gpu_list); launch.args.push_back("--layer-begins=" + layer_begin_list); launch.args.push_back("--layer-ends=" + layer_end_list); launch.args.push_back("--max-ctx=" + std::to_string(cfg.max_ctx));