feat(core): make JobStore connection pooling configurable

[jonthedecepticon]

Description

Closes #2123

get_db_engine built the JobStore engine as create_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 with connection is closed. It self-heals on the next request, so the cost is one failed submission per idle gap.

Changes

• get_db_engine accepts **engine_kwargs and forwards them to the engine factory.
• pool_pre_ping and pool_recycle fall back to NAT_JOB_STORE_POOL_PRE_PING and NAT_JOB_STORE_POOL_RECYCLE when not passed explicitly. Explicit keyword arguments take precedence.
• Unparseable values raise ValueError naming 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.
• Both variables are documented under Asynchronous Generate in the REST API reference.

Pooling defaults are unchanged when neither the keyword arguments nor the environment variables are set.

Environment variables now, config field is a maintainer's call

An earlier revision of this description claimed a FastApiFrontEndConfig field would reach only the parent process. That was wrong, and @chuenchen309 was right to push on it. db_url is already such a field, and it reaches other processes because the parent resolves it and writes it into os.environ. Pool settings could ride the same path.

Checking that claim turned up two further errors of mine, both now fixed in the docstring, the docs, and the commit message:

• The os.environ.update() in fastapi_front_end_plugin.py runs at line 159, after LocalCluster is created at line 122 and before uvicorn or gunicorn starts at line 179. It serves the API server workers, not the Dask workers. The Dask side receives db_url as an explicit dask_client.submit() argument.
• Because of that, the claim that the environment is "the only configuration that reaches the engines built in the Dask worker subprocesses" was false. Pool settings reach a local Dask cluster only because the operator exports them before nat serve and the workers inherit the environment at spawn.

That last point exposes a real limitation, now documented: with a separately managed scheduler_address, those workers inherit nothing from the API server host and need the variables set in their own environment. Threading the settings explicitly, the way db_url already is, would close that gap.

So the honest framing is that this PR is the smaller change and it resolves the reported failure, but the choice between it and a YAML surface is a preference about the config surface rather than a constraint. I am happy to add a FastApiFrontEndConfig field, a matching CLI flag, and explicit threading through periodic_cleanup and run_generation if a maintainer wants the wider fix. The environment fallback stays useful either way for processes that receive no explicit arguments.

On defaulting pool_pre_ping to true

Left 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_endpoints extra.

uv run pytest packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_job_store.py -k get_db_engine

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:

• I am familiar with the Contributing Guidelines.
• We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  • Any contribution which contains commits that are not Signed-Off will not be accepted.
• When the PR is ready for review, new or existing tests cover these changes.
• When the PR is ready for review, the documentation is up to date with these changes.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Walkthrough

The 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.

Changes

Job store pool configuration

Layer / File(s)	Summary
Engine pool configuration and validation packages/nvidia_nat_core/src/nat/front_ends/fastapi/async_jobs/job_store.py, packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_job_store.py, docs/source/reference/rest-api/api-server-endpoints.md	get_db_engine accepts engine keyword arguments and merges them with validated NAT_JOB_STORE_POOL_PRE_PING and NAT_JOB_STORE_POOL_RECYCLE values. Explicit arguments take precedence. Tests cover sync and async engines, defaults, parsing, precedence, and invalid values. Documentation describes the settings and affected processes.

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

Loading

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The changes satisfy issue #2123 by forwarding engine and pool options and supporting pool_pre_ping and pool_recycle configuration.
Out of Scope Changes check	✅ Passed	The implementation, tests, and documentation are directly related to configurable JobStore connection pooling.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title is concise, descriptive, uses imperative mood, and accurately summarizes the configurable JobStore connection pooling changes.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

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

📥 Commits

Reviewing files that changed from the base of the PR and between c10dc4d and 30b514d.

📒 Files selected for processing (3)

• docs/source/reference/rest-api/api-server-endpoints.md
• packages/nvidia_nat_core/src/nat/front_ends/fastapi/async_jobs/job_store.py
• packages/nvidia_nat_core/tests/nat/front_ends/fastapi/test_job_store.py

[jonthedecepticon]

Thanks, addressed two of the three.

• Environment validation on an overridden setting: fixed. _apply_pool_env_defaults now reads a variable only when the matching key is absent from engine_kwargs, so an explicit argument leaves the variable unread rather than validated. Added a regression test covering malformed values for both variables alongside explicit overrides.
• no value: added to the parameter table, along with a surrounding-whitespace case.

Skipping the Google-style docstring change. job_store.py is entirely NumPy-style (nine Parameters sections, no Args: blocks), and this change extends the existing get_db_engine docstring rather than introducing a new one. Converting only this function would leave the file inconsistent. Happy to convert the whole module in a separate PR if that is the direction you want.

[jonthedecepticon]

Maintainers: this PR has no labels, so the Label Checker is failing with Missing category & breaking labels. I cannot set them as an external contributor. Could someone add feature request and non-breaking?

A vet for copy-pr-bot would also be appreciated, since the CI pipeline has not run on this PR yet. In the meantime it passes locally: the get_db_engine tests, copyright.py --verify-apache-v2, path_checks.py, pre-commit on the changed files, Vale, and the docs HTML build. The branch is current with develop.

[chuenchen309]

Checked the multi-process argument against develop, since that's what the design turns on.

The premise holds. fastapi_front_end_plugin.py writes the resolved URL into the process environment before the workers spawn:

# Set environment variabls such that the worker subprocesses will know how to connect to dask and to
# the database
os.environ.update({
    "NAT_DASK_SCHEDULER_ADDRESS": self._scheduler_address,
    "NAT_JOB_STORE_DB_URL": db_url,
    ...
})

and fastapi_front_end_plugin_worker.py reads NAT_JOB_STORE_DB_URL back out. So an environment variable does reach the worker engines, which is what you need.

But that same code undercuts the reason given for not adding a config field. db_url already is a FastApiFrontEndConfig field — the parent builds its own engine from self.front_end_config.db_url a few lines above — and it reaches the workers only because the parent resolves it and writes it into os.environ there. "A FastApiFrontEndConfig field would configure the parent process only" is therefore a description of not adding those two dictionary entries, not a property of the system. pool_pre_ping and pool_recycle would ride the same path, and more simply than db_url does: db_url needs str(db_engine.url) because it may be None and get defaulted by the engine, whereas pool settings have no equivalent resolution step.

I don't think that makes the environment-variable version wrong — it's the smaller change and it works. It does mean the choice between it and a YAML surface is a maintainer's preference about the config surface, not a constraint, and the PR currently presents it as a constraint. Worth stating plainly so whoever reviews it decides the actual question.

One thing outside the diff: this has been open since July 31 with the Label Checker red, and the author has said they cannot set labels as an external contributor. That needs a maintainer regardless of the design question above.

AI-assisted review; the plugin and worker files were read on develop rather than inferred from the PR description.

[jonthedecepticon]

@chuenchen309 You are right, and thank you for reading the actual code rather than the description. I have corrected the PR description, the docstring, the docs, and the commit message.

Your point stands as written: db_url is already a FastApiFrontEndConfig field and it reaches other processes only because the parent resolves it and writes it into os.environ. Calling that a constraint was wrong.

Checking it turned up two more errors of mine that you did not claim, both worse than the one you did:

• The os.environ.update() runs at fastapi_front_end_plugin.py:159, after LocalCluster is created at line 122 and before uvicorn starts at line 179. So it serves the API server workers, not the Dask workers. The Dask side gets db_url as an explicit dask_client.submit() argument (async_job.py:72,130). My description had the wrong set of workers.
• That made the docstring claim about the environment being the only configuration reaching the engines built in the Dask worker subprocesses false. Pool settings reach a local cluster only because the operator exports them before nat serve and the workers inherit the environment at spawn.

The second one exposes a real limitation I had not documented: with a separately managed scheduler_address, those workers inherit nothing from the API server host, so the variables have to be set in their environment too. That is now stated in the docs. It also strengthens your implicit argument, since the Dask path already threads db_url explicitly and pool settings could join it with no environment variable at all.

I have left the implementation as is, since it is the smaller change and it resolves the reported failure, but the description now presents the choice as a preference for a maintainer to settle rather than as a constraint. If the YAML surface is wanted, I will add the config field, the CLI flag, and explicit threading through periodic_cleanup and run_job.

[jonthedecepticon]

@willkill07 Sorry to ping directly. This has been open since July 31 with Label Checker red, and as an external contributor I cannot set labels. Could you add feature request and non-breaking, and vet for CI when you get a chance? Latest commit is 32cb07a.

The CI pipeline has not run on this PR yet. It passes locally: the get_db_engine tests, copyright.py --verify-apache-v2, path_checks.py, pre-commit on the changed files, Vale, and the docs HTML build. Rebased on develop as of 2618705.