[https://nvbugs/6525008][fix] Isolate FlashInfer JIT workspaces for MPI workers - #17233
[https://nvbugs/6525008][fix] Isolate FlashInfer JIT workspaces for MPI workers#17233VALLIS-NERIA wants to merge 6 commits into
Conversation
…PI workers (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
|
/bot run |
|
PR_Github #63608 [ run ] triggered by Bot. Commit: |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
WalkthroughThe change adds FlashInfer workspace-slot locking for MPI workers. Multi-worker pools use the bootstrap and forward FlashInfer environment variables. Explicit workspaces remain unchanged. Tests and documentation cover workspace selection, lock release, pool configuration, and isolation settings. ChangesFlashInfer workspace isolation
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant MpiPoolSession
participant MPIPoolExecutor
participant FlashInferBootstrap
participant WorkspaceSlot
participant MPIWorkerServer
MpiPoolSession->>MPIPoolExecutor: configure worker launch and FlashInfer variables
MPIPoolExecutor->>FlashInferBootstrap: start multi-worker bootstrap
FlashInferBootstrap->>WorkspaceSlot: acquire available rank-based slot
WorkspaceSlot-->>FlashInferBootstrap: return workspace path
FlashInferBootstrap->>MPIWorkerServer: set workspace and cubin environment
MPIWorkerServer-->>FlashInferBootstrap: exit server
FlashInferBootstrap->>WorkspaceSlot: release lock
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/unittest/llmapi/test_flashinfer_workspace.py`:
- Around line 27-98: Add the six test cases from
test_configure_flashinfer_workspace_is_opt_in,
test_configure_flashinfer_workspace_is_per_process,
test_configure_flashinfer_workspace_preserves_explicit_value, and the three
parameterized cases of test_mpi_pool_configures_worker_workspace_isolation to an
applicable test-db YAML file such as l0_a100.yml, and add
tests/unittest/llmapi/test_flashinfer_workspace.py to the QA test list when
required for coverage.
🪄 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: 843218a1-8531-4445-8842-170cd9c3cdbf
📒 Files selected for processing (4)
tensorrt_llm/__init__.pytensorrt_llm/_flashinfer_workaround.pytensorrt_llm/llmapi/mpi_session.pytests/unittest/llmapi/test_flashinfer_workspace.py
| def test_configure_flashinfer_workspace_is_opt_in(monkeypatch: pytest.MonkeyPatch) -> None: | ||
| monkeypatch.delenv(workaround._FLASHINFER_WORKSPACE_ENV, raising=False) | ||
| monkeypatch.delenv(workaround._FLASHINFER_WORKSPACE_ISOLATION_ENV, raising=False) | ||
|
|
||
| workaround._configure_flashinfer_workspace() | ||
|
|
||
| assert workaround._FLASHINFER_WORKSPACE_ENV not in os.environ | ||
|
|
||
|
|
||
| def test_configure_flashinfer_workspace_is_per_process( | ||
| monkeypatch: pytest.MonkeyPatch, tmp_path: Path | ||
| ) -> None: | ||
| monkeypatch.delenv(workaround._FLASHINFER_WORKSPACE_ENV, raising=False) | ||
| monkeypatch.setenv(workaround._FLASHINFER_WORKSPACE_ISOLATION_ENV, "1") | ||
| monkeypatch.setattr(workaround.tempfile, "gettempdir", lambda: str(tmp_path)) | ||
|
|
||
| workaround._configure_flashinfer_workspace() | ||
|
|
||
| get_user_id = getattr(os, "getuid", lambda: 0) | ||
| expected = tmp_path / f"trtllm-flashinfer-{get_user_id()}-{os.getpid()}" | ||
| assert os.environ[workaround._FLASHINFER_WORKSPACE_ENV] == str(expected) | ||
|
|
||
|
|
||
| def test_configure_flashinfer_workspace_preserves_explicit_value( | ||
| monkeypatch: pytest.MonkeyPatch, tmp_path: Path | ||
| ) -> None: | ||
| explicit_workspace = tmp_path / "explicit" | ||
| monkeypatch.setenv(workaround._FLASHINFER_WORKSPACE_ISOLATION_ENV, "1") | ||
| monkeypatch.setenv(workaround._FLASHINFER_WORKSPACE_ENV, str(explicit_workspace)) | ||
|
|
||
| workaround._configure_flashinfer_workspace() | ||
|
|
||
| assert os.environ[workaround._FLASHINFER_WORKSPACE_ENV] == str(explicit_workspace) | ||
|
|
||
|
|
||
| @pytest.mark.skipif(not ENABLE_MULTI_DEVICE, reason="multi-device required") | ||
| @pytest.mark.parametrize( | ||
| "n_workers, override, expected", | ||
| [ | ||
| (1, None, None), | ||
| (4, None, "1"), | ||
| (4, "0", "0"), | ||
| ], | ||
| ) | ||
| def test_mpi_pool_configures_worker_workspace_isolation( | ||
| monkeypatch: pytest.MonkeyPatch, | ||
| n_workers: int, | ||
| override: str | None, | ||
| expected: str | None, | ||
| ) -> None: | ||
| captured: dict[str, object] = {} | ||
|
|
||
| class FakeMpiPoolExecutor: | ||
| def __init__(self, **kwargs: object) -> None: | ||
| captured.update(kwargs) | ||
|
|
||
| monkeypatch.delenv(workaround._FLASHINFER_WORKSPACE_ISOLATION_ENV, raising=False) | ||
| monkeypatch.setattr(mpi_session, "MPIPoolExecutor", FakeMpiPoolExecutor) | ||
| env_overrides = ( | ||
| {workaround._FLASHINFER_WORKSPACE_ISOLATION_ENV: override} if override is not None else {} | ||
| ) | ||
| session = SimpleNamespace( | ||
| n_workers=n_workers, | ||
| _env_overrides=env_overrides, | ||
| mpi_pool=None, | ||
| ) | ||
|
|
||
| mpi_session.MpiPoolSession._start_mpi_pool(session) | ||
|
|
||
| env = captured["env"] | ||
| assert isinstance(env, dict) | ||
| assert env.get(workaround._FLASHINFER_WORKSPACE_ISOLATION_ENV) == expected |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -F 'tests/unittest/llmapi/test_flashinfer_workspace.py' \
tests/integration/test_lists || true
rg -n \
'test_configure_flashinfer_workspace|test_mpi_pool_configures_worker_workspace_isolation' \
tests/integration/test_lists || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- changed-file status ---'
git status --short -- tests/unittest/llmapi/test_flashinfer_workspace.py
git diff --stat -- tests/unittest/llmapi/test_flashinfer_workspace.py
printf '%s\n' '--- test-list files ---'
git ls-files tests/integration/test_lists | sed -n '1,120p'
printf '%s\n' '--- matching module and test names ---'
rg -n -F 'test_flashinfer_workspace.py' tests/integration/test_lists || true
rg -n 'test_configure_flashinfer_workspace|test_mpi_pool_configures_worker_workspace_isolation' \
tests/integration/test_lists || true
printf '%s\n' '--- nearby test-list entries for unittest modules ---'
rg -n 'tests/unittest/|unittest/llmapi|llmapi' tests/integration/test_lists | sed -n '1,120p' || true
printf '%s\n' '--- test module outline ---'
ast-grep outline tests/unittest/llmapi/test_flashinfer_workspace.py --view expandedRepository: NVIDIA/TensorRT-LLM
Length of output: 23426
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- test-list documentation ---'
cat -n tests/integration/test_lists/qa/README.md | sed -n '1,180p'
cat -n tests/integration/test_lists/test-db/README.md | sed -n '1,180p'
printf '%s\n' '--- CBTS artifacts ---'
git ls-files | rg -i 'cbts|touchmap|coverage' | sed -n '1,160p' || true
find . -maxdepth 5 -type f \( -iname '*cbts*' -o -iname '*touchmap*' -o -iname '*coverage*' \) -print | sed -n '1,160p'
printf '%s\n' '--- all list references by basename ---'
rg -n -F 'test_flashinfer_workspace' tests/integration/test_lists || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 9667
🏁 Script executed:
#!/bin/bash
set -euo pipefail
cat -n tests/unittest/llmapi/test_flashinfer_workspace.py | sed -n '1,110p'
rg -n 'test_mpi_session.py|flashinfer|workspace_isolation|ENABLE_MULTI_DEVICE' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa \
tests/unittest/llmapi/test_flashinfer_workspace.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 8761
Add the new tests to a CI test list.
tests/unittest/llmapi/test_flashinfer_workspace.py is absent from tests/integration/test_lists/test-db/ and tests/integration/test_lists/qa/. Add the six test cases to the applicable test-db YAML file, such as l0_a100.yml, and include the module in QA coverage if required.
Test coverage summary: three workspace tests and three MPI parameter cases cover the new behavior. Coverage verdict: insufficient because the tests are not scheduled by a test list.
🤖 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 `@tests/unittest/llmapi/test_flashinfer_workspace.py` around lines 27 - 98, Add
the six test cases from test_configure_flashinfer_workspace_is_opt_in,
test_configure_flashinfer_workspace_is_per_process,
test_configure_flashinfer_workspace_preserves_explicit_value, and the three
parameterized cases of test_mpi_pool_configures_worker_workspace_isolation to an
applicable test-db YAML file such as l0_a100.yml, and add
tests/unittest/llmapi/test_flashinfer_workspace.py to the QA test list when
required for coverage.
Source: Path instructions
…ound (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tensorrt_llm/__init__.py`:
- Around line 25-29: Update the workspace initialization around
_FLASHINFER_WORKSPACE_ENV to distinguish explicit caller-provided values from
parent-generated values: preserve explicit values, but clear inherited generated
workspace values before worker import and assign a worker-local workspace when
needed. Add a regression test covering MPIPoolExecutor(env=...) environment
propagation and verifying workers do not reuse the parent-generated
FLASHINFER_WORKSPACE_BASE.
🪄 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: c35474ba-5f63-4991-b98b-0c6a8d1bef92
📒 Files selected for processing (3)
tensorrt_llm/__init__.pytensorrt_llm/llmapi/mpi_session.pytests/unittest/llmapi/test_flashinfer_workspace.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tensorrt_llm/llmapi/mpi_session.py
- tests/unittest/llmapi/test_flashinfer_workspace.py
|
PR_Github #63654 [ run ] triggered by Bot. Commit: |
|
PR_Github #63608 [ run ] completed with state |
…MPI workers (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@tensorrt_llm/llmapi/mpi_session.py`:
- Around line 34-40: Update the generated _FLASHINFER_WORKER_BOOTSTRAP so it
checks TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS before assigning
FLASHINFER_WORKSPACE_BASE, while preserving setdefault behavior for explicit
workspace values and continuing to launch the MPI futures server. Add a
regression test covering the disabled flag and confirming the workspace is not
changed.
- Around line 448-454: Update the environment construction used by the
MPIPoolExecutor call in the MPI session setup to explicitly propagate
FLASHINFER_WORKSPACE_BASE, sourcing it from the configured environment or
requiring it through env_overrides. Preserve existing environment handling and
add a spawn-level test verifying that a value present only in the parent process
reaches workers.
🪄 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: 92b9d83f-14d3-4cab-9703-a6da5d2f56cf
📒 Files selected for processing (2)
tensorrt_llm/llmapi/mpi_session.pytests/unittest/llmapi/test_flashinfer_workspace.py
|
PR_Github #63670 [ run ] triggered by Bot. Commit: |
|
PR_Github #63654 [ run ] completed with state |
|
The mechanism itself checks out. The cache relocation is wider than "separate temporary caches."
Test collection on single-device builds. Minor: is this still needed once #17175 moves us to flashinfer 0.6.16? If that release fixes the lock ordering, it'd be better to land this with an explicit removal condition than as a permanent workaround. Holding on the CI signal for now — 63608 and 63654 both ABORTED and 63670 is still running, so no pipeline has completed on any head yet. |
|
PR_Github #63670 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63686 [ run ] triggered by Bot. Commit: |
|
PR_Github #63686 [ run ] completed with state
|
|
/bot run |
|
PR_Github #63697 [ run ] triggered by Bot. Commit: |
|
PR_Github #63697 [ run ] completed with state
|
…y Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/llmapi/test_flashinfer_workspace.py (1)
75-102: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winClear inherited FlashInfer settings before the parameterized test.
_start_mpi_pool()copiesos.environbefore it appliesenv_overrides. An inheritedTRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0orFLASHINFER_WORKSPACE_BASEmakes the(4, {}, ...)case select no bootstrap. An inheritedFLASHINFER_CUBIN_DIRalso makes the default-cubin assertion fail. Clear these variables before creatingsession.Proposed fix
def test_mpi_pool_configures_worker_bootstrap( monkeypatch: pytest.MonkeyPatch, n_workers: int, env_overrides: dict[str, str], expected: list[str] | None, ) -> None: + monkeypatch.delenv(_FLASHINFER_WORKSPACE_ENV, raising=False) + monkeypatch.delenv(_FLASHINFER_CUBIN_ENV, raising=False) + monkeypatch.delenv( + "TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS", raising=False + ) captured: dict[str, object] = {}🤖 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 `@tests/unittest/llmapi/test_flashinfer_workspace.py` around lines 75 - 102, Clear the inherited FlashInfer environment variables before creating the session in test_mpi_pool_configures_worker_bootstrap: TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS, FLASHINFER_WORKSPACE_BASE, and FLASHINFER_CUBIN_DIR. Preserve env_overrides so each parameterized case still controls its intended value.
🤖 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.
Outside diff comments:
In `@tests/unittest/llmapi/test_flashinfer_workspace.py`:
- Around line 75-102: Clear the inherited FlashInfer environment variables
before creating the session in test_mpi_pool_configures_worker_bootstrap:
TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS, FLASHINFER_WORKSPACE_BASE, and
FLASHINFER_CUBIN_DIR. Preserve env_overrides so each parameterized case still
controls its intended value.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d0d3e63c-b6ef-4bc0-96e6-cc80adf3980c
📒 Files selected for processing (3)
tensorrt_llm/llmapi/mpi_session.pytests/integration/test_lists/test-db/l0_a100.ymltests/unittest/llmapi/test_flashinfer_workspace.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/llmapi/mpi_session.py
|
/bot run |
|
PR_Github #63723 [ run ] triggered by Bot. Commit: |
|
What remains is the cost of the approach, and it got sharper this week.
And #17175 makes that worse. It bumps flashinfer to 0.6.16, which adds an on-disk JIT cache for CuTe-DSL kernels ( I also looked at whether 0.6.16 simply fixes what you are working around: it does not, as far as I can tell. Given that, is a persistent per-rank directory ( |
|
PR_Github #63723 [ run ] completed with state
|
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
The race analysis and the gating (n_workers > 1, explicit-workspace passthrough, kill switch) look right, and I appreciate that the two real-spawn tests actually exercise the python_args path rather than just asserting on the argument list — that's the part most likely to break silently on an mpi4py bump.
Main concern is the cache cost, raised inline: a fresh mkdtemp per worker per run makes every FlashInfer JIT build cold on every launch, for every worker, forever. That's trading a rare truncated-file failure for a guaranteed multi-minute startup regression on any multi-GPU run using the FlashInfer backend. A rank-keyed stable directory (e.g. $FLASHINFER_WORKSPACE_BASE_default/rank<N>) removes the same cross-rank race while keeping the cache warm across runs, at the cost of N copies of the build tree. If there's a reason that doesn't work (concurrent jobs on one node reusing rank IDs?), say so in the commit message — it's the obvious question a future reader will have.
Also worth stating in the code or the NVBug: this is a workaround for a specific upstream bug in flashinfer-python 0.6.15 (requirements.txt:61). A short comment naming the version and the upstream issue makes it removable later instead of permanent.
Secondary: TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS is a new user-visible knob with no documentation anywhere. One line somewhere discoverable would help support triage.
| _FLASHINFER_WORKER_BOOTSTRAP = ( | ||
| "import os,shutil,tempfile;" | ||
| "from mpi4py import MPI;" | ||
| "workspace=tempfile.mkdtemp(" |
There was a problem hiding this comment.
FLASHINFER_WORKSPACE_BASE is where FlashInfer caches JIT-compiled .so files, not just scratch — see the comment at tests/unittest/auto_deploy/standalone/test_standalone_package.py:250. A fresh mkdtemp per worker plus rmtree on exit means every worker recompiles every kernel on every launch and throws the artifacts away. On a cold cache that's minutes of added startup per run, permanently, for all multi-GPU FlashInfer-backend flows.
A rank-keyed stable path (e.g. ~/.cache/flashinfer-rank{rank} or <default_base>/rank{rank}) removes the same cross-rank race and keeps the cache across runs. Cost is N copies of the build tree, which seems much cheaper than N cold builds per run.
| "prefix=f'trtllm-flashinfer-{MPI.COMM_WORLD.Get_rank()}-{os.getpid()}-');" | ||
| "os.environ.setdefault('FLASHINFER_WORKSPACE_BASE',workspace);" | ||
| "from mpi4py.futures.server import main\n" | ||
| "try:main()\n" |
There was a problem hiding this comment.
The finally: shutil.rmtree(...) only runs on a clean worker exit. shutdown_abort() in this same file SIGKILLs worker PIDs, and MPI_Abort paths don't unwind either — so on any crash/abort the temp workspace (a full JIT build tree, easily hundreds of MB) is orphaned in /tmp with a random suffix that nothing will ever collect. On a shared node with repeated failing runs that fills /tmp.
If you keep the temp-dir approach, consider mkdtemp(dir=<stable parent>) and having the parent session clean the parent dir on shutdown, so the abort path is recoverable.
| and "FLASHINFER_WORKSPACE_BASE" not in env) | ||
| if isolate_workspace: | ||
| # Keep downloaded cubins shared; only generated JIT sources race. | ||
| env.setdefault( |
There was a problem hiding this comment.
Hardcoding ~/.cache/flashinfer/cubins duplicates FlashInfer's own default-location logic, which is version-dependent and derived from the workspace base upstream. Two risks: (1) if upstream changes the layout (e.g. adds a version subdir), workers silently get a different cubin dir than a non-MPI run, defeating the sharing this is meant to preserve; (2) expanduser resolves the parent's $HOME and passes it as an absolute path — on a multi-node spawn without a shared home, that path may not exist or be writable on the worker node.
Preferably read the default from the installed flashinfer jit env module rather than reconstructing it here.
| from types import SimpleNamespace | ||
|
|
||
| import pytest | ||
| from mpi4py.futures import server |
There was a problem hiding this comment.
This import runs at collection time, before pytestmark on line 29 is consulted. On a build with ENABLE_MULTI_DEVICE=0 (where mpi4py isn't necessarily installed) this is a collection error, not a skip — and it takes the whole file down. Note test_mpi_session.py deliberately avoids importing mpi4py directly for this reason.
Move it into the helper that needs it, or guard with server = pytest.importorskip("mpi4py.futures.server").
| rank = mpi_session.mpi4py.MPI.COMM_WORLD.Get_rank() | ||
| assert prefixes == [f"trtllm-flashinfer-{rank}-{os.getpid()}-"] | ||
| assert os.environ[_FLASHINFER_WORKSPACE_ENV] == str(workspace) | ||
| assert not workspace.exists() |
There was a problem hiding this comment.
This assertion is vacuous: mkdtemp is faked on line 53 and never creates workspace, so it can't exist regardless of whether the finally: rmtree ran. If the intent is to prove cleanup happens, have fake_mkdtemp actually workspace.mkdir() (and drop a file in it) first — then the assertion tests something.
…Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
|
/bot run |
|
PR_Github #64033 [ run ] triggered by Bot. Commit: |
|
PR_Github #64033 [ run ] completed with state
|
5 participants
Description
FlashInfer 0.6.15 writes generated JIT configuration sources before acquiring its compilation lock. Concurrent MPI workers using the shared default workspace can therefore observe a temporarily truncated file and fail compilation with undefined generated types and constants.
For multi-worker
MpiPoolSessioninstances without an explicit workspace, TensorRT-LLM now starts each worker with a rank- and PID-specific temporaryFLASHINFER_WORKSPACE_BASEbeforempi4py.futures.serverimports the user main module. Downloaded/precompiled cubins remain in the sharedFLASHINFER_CUBIN_DIR; explicit workspace and cubin paths are propagated unchanged.TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0disables the workaround, and normal worker shutdown removes the temporary workspace.The workaround intentionally covers dynamically spawned
MpiPoolSessionworkers. Externalmpirun/MpiCommSessionranks must set their workspace before importing TensorRT-LLM.Test Coverage
MPIPoolExecutortests verify distinct JIT workspaces, a shared cubin directory, and parent-only explicit workspace propagation.tests/unittest/llmapi/test_flashinfer_workspace.pyin the A100 isolation stage.selective_state_updateJIT with distinct workspaces.The original end-to-end command on NSC/Pyxis could not reach model initialization because dynamic MPI spawn failed with
MPI_ERR_SPAWN; the worker bootstrap itself is covered by direct spawn tests.PR Checklist
GitHub Bot Help
To see a list of available CI bot commands, comment
/bot help.Dev Engineer Review
Code Changes:
MpiPoolSessioninstances without explicit workspaces.TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0.MpiCommSession.Configuration and Test List:
unittest/llmapi/test_flashinfer_workspace.py ISOLATIONtotests/integration/test_lists/test-db/l0_a100.yml.QA Engineer Review
Test Changes:
Test Coverage:
tests/unittest/llmapi/test_flashinfer_workspace.pyis included intests/integration/test_lists/test-db/l0_a100.ymlwith theISOLATIONtag.MPI_ERR_SPAWN.