feat(core): make JobStore connection pooling configurable - #2150
Open
jonthedecepticon wants to merge 1 commit into
Open
feat(core): make JobStore connection pooling configurable#2150jonthedecepticon wants to merge 1 commit into
jonthedecepticon wants to merge 1 commit into
Conversation
|
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)
🚧 Files skipped from review as they are similar to previous changes (3)
WalkthroughThe job store engine now accepts SQLAlchemy engine keyword arguments and environment-based pool settings. It validates pool values, gives explicit arguments precedence, forwards settings to sync and async engines, adds tests, and documents the configuration. ChangesJob store pool configuration
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant EnvironmentVariables
participant get_db_engine
participant SQLAlchemyEngineFactory
Caller->>get_db_engine: Supply explicit engine kwargs
EnvironmentVariables->>get_db_engine: Provide pool settings
get_db_engine->>get_db_engine: Validate and merge settings
get_db_engine->>SQLAlchemyEngineFactory: Create sync or async engine
SQLAlchemyEngineFactory-->>Caller: Return configured engine
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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: 3
🤖 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 `@packages/nvidia_nat_core/src/nat/front_ends/fastapi/async_jobs/job_store.py`:
- Around line 648-650: Update get_db_engine and _pool_kwargs_from_env so each
pool environment variable is parsed only when its corresponding engine_kwargs
key is absent, allowing explicit arguments to bypass malformed environment
values. Preserve explicit-argument precedence and add regression tests covering
malformed overridden pool settings.
- Around line 613-623: Rewrite the complete get_db_engine docstring in Google
style, including a concise first-line summary that ends with a period and
Google-style Args, Returns, Raises, and relevant notes sections for the existing
API behavior. Preserve the documented engine_kwargs details and
environment-variable behavior while removing NumPy-style section syntax.
In `@packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_job_store.py`:
- Around line 734-736: Add the documented ("no", False) case to the
parameterization for the test covering _POOL_PRE_PING_FALSE_VALUES, preserving
the existing true and false value 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: e2993352-5a21-4c09-b135-e04bb8825603
📒 Files selected for processing (3)
docs/source/reference/rest-api/api-server-endpoints.mdpackages/nvidia_nat_core/src/nat/front_ends/fastapi/async_jobs/job_store.pypackages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_job_store.py
get_db_engine constructed the JobStore engine with no way to pass SQLAlchemy engine or pool options. Deployments behind managed Postgres could not set pool_pre_ping or pool_recycle, so the first async job submitted after an idle gap failed with "connection is closed". Forward **engine_kwargs to the engine factory, and read pool_pre_ping and pool_recycle from NAT_JOB_STORE_POOL_PRE_PING and NAT_JOB_STORE_POOL_RECYCLE when they are not passed explicitly. The environment is the only configuration that reaches the engines built in the Dask worker subprocesses. An unparseable value raises rather than silently falling back to the default, except when the caller supplies that setting explicitly, in which case the variable is left unread. Pooling defaults are unchanged when neither the keyword arguments nor the environment variables are set. Closes NVIDIA#2123 Signed-off-by: Jon Lambson <jonlambson@gmail.com>
jonthedecepticon
force-pushed
the
feat/2123-get-db-engine-kwargs
branch
from
July 31, 2026 17:06
30b514d to
f9d1fc0
Compare
Author
|
Thanks, addressed two of the three.
Skipping the Google-style docstring change. |
Author
|
Maintainers: this PR has no labels, so the Label Checker is failing with A vet for |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Closes #2123
get_db_enginebuilt the JobStore engine ascreate_engine_fn(db_url, echo=echo), so SQLAlchemy pool options could not be set. Behind managed PostgreSQL, idle pooled connections are dropped server side without the client noticing, and the first async job submitted after an idle gap fails withconnection is closed. It self-heals on the next request, so the cost is one failed submission per idle gap.Changes
get_db_engineaccepts**engine_kwargsand forwards them to the engine factory.pool_pre_pingandpool_recyclefall back toNAT_JOB_STORE_POOL_PRE_PINGandNAT_JOB_STORE_POOL_RECYCLEwhen not passed explicitly. Explicit keyword arguments take precedence.ValueErrornaming the variable, instead of silently leaving the default in place. A typo in the setting that exists to prevent this failure should not reproduce the failure.Pooling defaults are unchanged when neither the keyword arguments nor the environment variables are set.
Why environment variables rather than a config field
The JobStore engine is built in more than one process.
fastapi_front_end_plugin.pywritesNAT_JOB_STORE_DB_URLinto the environment so the Dask worker subprocesses inherit it,fastapi_front_end_plugin_worker.pyreads it back, andJobStore.__init__builds its own engine from it. AFastApiFrontEndConfigfield would configure the parent process only and leave the worker engines on the defaults, which is the wrong half of the problem. The environment is the existing transport for this subsystem, andget_db_enginealready resolvesdb_urlthe same way.I am happy to add a
FastApiFrontEndConfigfield and a matching CLI flag on top of this if you would prefer the YAML surface as well. It would be purely additive, and the environment fallback is still required for the workers either way.On defaulting
pool_pre_pingto trueLeft out. It is a behavior change for existing deployments and is not needed to resolve the issue now that the setting is reachable. Happy to default it for non-SQLite URLs if you would rather have it on out of the box.
How to verify
Requires the
async_endpointsextra.Covers forwarding on the sync and async paths, both environment variables, explicit keyword arguments overriding the environment, invalid values raising, and defaults staying unchanged when nothing is configured.
By Submitting this PR I confirm:
Summary by CodeRabbit