Expose Dask memory_limit config

[dagardner-nv]

Description

• Optionally allow the user to set the memory_limit value for a local Dask cluster.

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.

Summary by CodeRabbit

• New Features
  • Added a new configuration option to set Dask worker memory limits, supporting "auto", numeric values, and custom memory specifications.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Walkthrough

This change adds a configurable Dask worker memory limit to the FastAPI front-end. A new configuration field captures the memory limit setting, which is then normalized and passed to the LocalCluster initialization during plugin setup.

Changes

Cohort / File(s)	Summary
Configuration field addition src/nat/front_ends/fastapi/fastapi_front_end_config.py	Added dask_worker_memory_limit: str field with default value "auto" and documentation describing acceptable values for memory constraints.
Memory limit handling src/nat/front_ends/fastapi/fastapi_front_end_plugin.py	Introduced normalization logic that converts the config value to appropriate type (int for digits, float for numeric, otherwise preserves as string), then passes the resulting memory_limit to LocalCluster initialization.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'Expose Dask memory_limit config' is concise (33 chars), descriptive, uses imperative mood, and accurately reflects the main changes: adding a dask_worker_memory_limit field to configuration and passing it to LocalCluster initialization.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

src/nat/front_ends/fastapi/fastapi_front_end_plugin.py (1)

110-124: Narrow the exception type to ValueError.

The float() conversion raises ValueError for invalid input. Using a broad except Exception catches unrelated exceptions and is flagged by static analysis (BLE001, S110). Consider also adding a debug log for transparency when the value is kept as a string.

♻️ Suggested fix

                     # Try to convert to number if possible, otherwise leave as a string
                     try:
                         memory_limit = float(memory_limit)
-                    except Exception:
-                        pass  # Keep as string (e.g., "auto", "4GB")
+                    except ValueError:
+                        logger.debug("Keeping memory_limit as string: %s", memory_limit)

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 0a9b9b0 and 76031cf.

📒 Files selected for processing (2)

• src/nat/front_ends/fastapi/fastapi_front_end_config.py
• src/nat/front_ends/fastapi/fastapi_front_end_plugin.py

🧰 Additional context used

📓 Path-based instructions (6)

**/*.py

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.py: Follow PEP 20 and PEP 8 for Python style guidelines
Run yapf with PEP 8 base and 'column_limit = 120' for code formatting
Use 'ruff check --fix' for linting with configuration from 'pyproject.toml', fix warnings unless explicitly ignored
Use snake_case for functions and variables, PascalCase for classes, UPPER_CASE for constants
All public APIs require Python 3.11+ type hints on parameters and return values
Prefer 'collections.abc' / 'typing' abstractions (e.g., 'Sequence' over 'list') for type hints
Use 'typing.Annotated' for units or extra metadata when useful
Treat 'pyright' warnings (configured in 'pyproject.toml') as errors during development
Preserve stack traces and prevent duplicate logging when handling exceptions; use bare 'raise' statements when re-raising, and use 'logger.error()' for logging (not 'logger.exception()') to avoid duplicate stack trace output
When catching and logging exceptions without re-raising, always use 'logger.exception()' (equivalent to 'logger.error(exc_info=True)') to capture full stack trace information
Pydantic models using 'SecretStr', 'SerializableSecretStr', or 'OptionalSecretStr' should use 'default=None' for optional fields and 'default_factory=lambda: SerializableSecretStr("")' for non-optional fields to avoid initialization bugs
Provide Google-style docstrings for every public module, class, function and CLI command
The first line of docstrings must be a concise description ending with a period
Surround code entities in docstrings with backticks to avoid Vale false-positives
Validate and sanitise all user input, especially in web or CLI interfaces
Prefer 'httpx' with SSL verification enabled by default and follow OWASP Top-10 recommendations
Use 'async'/'await' for I/O-bound work (HTTP, DB, file reads)
Cache expensive computations with 'functools.lru_cache' or an external cache when appropriate
Leverage NumPy vectorised operations whenever beneficial and feasible

Files:

• src/nat/front_ends/fastapi/fastapi_front_end_config.py
• src/nat/front_ends/fastapi/fastapi_front_end_plugin.py

**/*.{py,yaml,yml,json,toml}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Indent with 4 spaces (never tabs) and ensure every file ends with a single newline

Files:

• src/nat/front_ends/fastapi/fastapi_front_end_config.py
• src/nat/front_ends/fastapi/fastapi_front_end_plugin.py

**/*.{py,js,ts,tsx,jsx,sh,yaml,yml,json,toml,md,mdx,rst}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

**/*.{py,js,ts,tsx,jsx,sh,yaml,yml,json,toml,md,mdx,rst}: Every file must start with the standard SPDX Apache-2.0 header
Confirm that copyright years are up-to-date whenever a file is changed
All source files must include the SPDX Apache-2.0 header template

Files:

• src/nat/front_ends/fastapi/fastapi_front_end_config.py
• src/nat/front_ends/fastapi/fastapi_front_end_plugin.py

**/*.{py,md,mdx,rst}

📄 CodeRabbit inference engine (.cursor/rules/general.mdc)

Version numbers are derived automatically by 'setuptools-scm'; never hard-code them in code or docs

Files:

• src/nat/front_ends/fastapi/fastapi_front_end_config.py
• src/nat/front_ends/fastapi/fastapi_front_end_plugin.py

**/*

⚙️ CodeRabbit configuration file

**/*: # Code Review Instructions

• Ensure the code follows best practices and coding standards. - For Python code, follow
  PEP 20 and
  PEP 8 for style guidelines.
• Check for security vulnerabilities and potential issues. - Python methods should use type hints for all parameters and return values (except for return values of None,
  in that situation no return type hint is needed).
  Example:

  def my_function(param1: int, param2: str) -> bool:
      pass

• For Python exception handling, ensure proper stack trace preservation:
  • When re-raising exceptions: use bare raise statements to maintain the original stack trace,
    and use logger.error() (not logger.exception()) to avoid duplicate stack trace output.
  • When catching and logging exceptions without re-raising: always use logger.exception()
    to capture the full stack trace information.

Documentation Review Instructions - Verify that documentation and comments are clear and comprehensive. - Verify that the documentation doesn't contain any TODOs, FIXMEs or placeholder text like "lorem ipsum". - Verify that the documentation doesn't contain any offensive or outdated terms. - Verify that documentation and comments are free of spelling mistakes, ensure the documentation doesn't contain any

words listed in the ci/vale/styles/config/vocabularies/nat/reject.txt file, words that might appear to be
spelling mistakes but are listed in the ci/vale/styles/config/vocabularies/nat/accept.txt file are OK.

• Documentation in Markdown files should not contain usage of a possessive 's with inanimate objects
  (ex: "the system's performance" should be "the performance of the system").
• Documentation in Markdown files should not use NAT as an acronym, always spell out NeMo Agent Toolkit.
  The exception to this rule is when referring to package names or code identifiers that contain "nat", th...

Files:

• src/nat/front_ends/fastapi/fastapi_front_end_config.py
• src/nat/front_ends/fastapi/fastapi_front_end_plugin.py

src/nat/**/*

⚙️ CodeRabbit configuration file

This directory contains the core functionality of the toolkit. Changes should prioritize backward compatibility.

Files:

• src/nat/front_ends/fastapi/fastapi_front_end_config.py
• src/nat/front_ends/fastapi/fastapi_front_end_plugin.py

🧬 Code graph analysis (1)

src/nat/front_ends/fastapi/fastapi_front_end_plugin.py (1)

src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py (1)

• front_end_config (137-138)

🪛 Ruff (0.14.11)

src/nat/front_ends/fastapi/fastapi_front_end_plugin.py

123-124: try-except-pass detected, consider logging the exception

(S110)

---

123-123: Do not catch blind exception: Exception

(BLE001)

🔇 Additional comments (2)

src/nat/front_ends/fastapi/fastapi_front_end_config.py (1)

240-244: LGTM!

The new dask_worker_memory_limit configuration field is well-defined with a sensible default of "auto" and a clear description that documents the acceptable values and provides a reference to Dask documentation. Using str type is appropriate since the parsing/conversion logic is handled in the plugin.

src/nat/front_ends/fastapi/fastapi_front_end_plugin.py (1)

126-131: LGTM!

The memory_limit parameter is correctly passed to LocalCluster. The conversion logic appropriately handles all documented Dask memory_limit formats: "auto", integer bytes, float fraction, and memory strings like "4GB".

[dagardner-nv]

/merge