Agent params #792
Agent params #792
Conversation
alexmanle
requested review from
amaslenn,
jeffnvidia and
srivatsankrishnan
as code owners
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the
📝 WalkthroughWalkthroughAdded Pydantic agent configuration models, extended TestDefinition with an optional Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches🧪 Generate unit tests (beta)
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. Comment |
Greptile OverviewGreptile SummaryThis PR introduces a configuration override mechanism for agents, allowing users to customize agent hyperparameters through TOML configuration files. Key Changes:
Issues Found:
The implementation follows pydantic best practices with strict validation and clear error messages. Most previously reported issues have been addressed. Confidence Score: 4/5
Important Files Changed
|
![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
🤖 Fix all issues with AI agents
In `@src/cloudai/cli/handlers.py`:
- Around line 198-201: The code silently drops a user-supplied agent_config when
no validator exists; update the branch that checks config_class =
config_class_map.get(agent_type) to emit a logging.warning (not debug) when an
agent_config was provided and will be ignored, referencing agent_type and
agent_config in the message so callers understand overrides were ignored; keep
returning {} but ensure the warning clearly states the agent_type and that
agent_config will be ignored (optionally include a truncated/summary of
agent_config) to aid user visibility.
In `@src/cloudai/models/agent_config.py`:
- Around line 51-54: The algorithm Field currently accepts any string which
risks invalid runtime values; restrict and validate it by replacing the loose
Optional[str] with a strict set of allowed values (use a Python Enum or
typing.Literal for "ucb1", "ts" (thompson_sampling), "epsilon_greedy",
"softmax", "random") or add a Pydantic validator on the algorithm field in the
AgentConfig class to raise a clear validation error when an unsupported
algorithm is provided, and update the Field description to match the enforced
choices.
- Around line 42-45: The botorch_num_trials field currently allows any integer
but should only accept -1 or integers >= 1 (matching the semantic described and
the sibling sobol_num_trials). Add a Pydantic validator for botorch_num_trials
on the AgentConfig model (e.g., a method named validate_botorch_num_trials
decorated with `@validator`("botorch_num_trials")) that returns the value if it's
None or equals -1 or is >= 1, and raises a ValueError for other values (e.g., 0
or < -1).
src/cloudai/models/agent_config.py
Outdated
| sobol_num_trials: Optional[int] = Field(default=None, ge=1, description="Number of SOBOL initialization trials") | ||
| botorch_num_trials: Optional[int] = Field( | ||
| default=None, description="Number of BoTorch trials (-1 for unlimited until max_steps)" | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's examine the actual file structure and find agent_config.py
fd -t f "agent_config.py" | head -20Repository: NVIDIA/cloudai
Length of output: 92
🏁 Script executed:
# Read the agent_config.py file to see current state
cat -n src/cloudai/models/agent_config.pyRepository: NVIDIA/cloudai
Length of output: 3888
🏁 Script executed:
# Search for botorch_num_trials usage throughout the codebase
rg "botorch_num_trials" -B 2 -A 2Repository: NVIDIA/cloudai
Length of output: 499
🏁 Script executed:
# Search for files that use BayesianOptimizationConfig or botorch_num_trials
rg "BayesianOptimizationConfig|botorch_num_trials" -t py --files-with-matchesRepository: NVIDIA/cloudai
Length of output: 120
🏁 Script executed:
# Look for the BO agent implementation
fd -t f ".*agent.*" -path "*/bayesian*" -o -path "**/optimization*" | grep -i agentRepository: NVIDIA/cloudai
Length of output: 226
🏁 Script executed:
# Broader search for agent implementations
fd -t f "agent" -type fRepository: NVIDIA/cloudai
Length of output: 228
🏁 Script executed:
# Search more broadly for agent implementations
fd -e py "agent" | head -30Repository: NVIDIA/cloudai
Length of output: 152
🏁 Script executed:
# Also search for where BayesianOptimizationConfig is instantiated
rg "BayesianOptimizationConfig" -B 3 -A 10 -t pyRepository: NVIDIA/cloudai
Length of output: 2822
🏁 Script executed:
# Search for botorch usage in the codebase
rg "botorch" -i -B 3 -A 3Repository: NVIDIA/cloudai
Length of output: 673
🏁 Script executed:
# Find the base_agent.py file and search for BO implementation
cat src/cloudai/configurator/base_agent.py | head -100Repository: NVIDIA/cloudai
Length of output: 2224
🏁 Script executed:
# Search for where botorch_num_trials is actually used in the agent
rg "botorch_num_trials" -A 5 -B 5 src/cloudai/configurator/Repository: NVIDIA/cloudai
Length of output: 40
🏁 Script executed:
# Look for any BO agent or related files
fd -e py | xargs rg -l "class.*Bayesian|botorch_num_trials"Repository: NVIDIA/cloudai
Length of output: 92
🏁 Script executed:
# Search for actual agent implementations/classes
rg "class.*Agent.*:" -A 5 | head -100Repository: NVIDIA/cloudai
Length of output: 4155
🏁 Script executed:
# Look at test_agents.py to see how these configs are tested
cat tests/test_agents.pyRepository: NVIDIA/cloudai
Length of output: 2749
🏁 Script executed:
# Check if there are any actual BO agent implementations elsewhere
fd -e py -x rg -l "num_trials" {} \;Repository: NVIDIA/cloudai
Length of output: 94
🏁 Script executed:
# Look for any BO agent implementations or references
fd -e py | xargs rg -l "BayesianOptimization|BoTorch" 2>/dev/nullRepository: NVIDIA/cloudai
Length of output: 120
🏁 Script executed:
# Check if there's a bayesian agent file
fd -e py -path "*/agent*" -o -path "*/bayesian*"Repository: NVIDIA/cloudai
Length of output: 226
🏁 Script executed:
# Search for where the agent configs are actually used
rg "botorch_num_trials|sobol_num_trials" -B 5 -A 5Repository: NVIDIA/cloudai
Length of output: 905
🏁 Script executed:
# Check what agent types are actually implemented
ls -la src/cloudai/configurator/Repository: NVIDIA/cloudai
Length of output: 559
🏁 Script executed:
# Check if BO agent exists or is planned
fd "\.py$" | xargs rg -l "class.*Bayesian" 2>/dev/nullRepository: NVIDIA/cloudai
Length of output: 92
🏁 Script executed:
# Look at handlers.py more carefully to see how config is used
cat src/cloudai/cli/handlers.py | grep -A 20 "BayesianOptimizationConfig"Repository: NVIDIA/cloudai
Length of output: 1467
🏁 Script executed:
# Check if there's external documentation or plugin system for agents
rg "entry_points\|entrypoint" -A 5 -B 5Repository: NVIDIA/cloudai
Length of output: 40
🏁 Script executed:
# Search for any test or documentation mentioning BO agent
rg "bayesian|bo agent" -i -B 2 -A 2Repository: NVIDIA/cloudai
Length of output: 1069
Add validation constraint to botorch_num_trials.
The field currently accepts any integer, but the description specifies "(-1 for unlimited until max_steps)", implying only -1 or ≥1 are valid. Without validation, invalid values like 0 or -2 pass Pydantic validation. This is inconsistent with the sibling field sobol_num_trials (line 42), which has ge=1 constraint.
🛠 Proposed fix
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, field_validator
@@
class BayesianOptimizationConfig(AgentConfig):
"""Configuration overrides for Bayesian Optimization agent."""
@@
botorch_num_trials: Optional[int] = Field(
default=None, description="Number of BoTorch trials (-1 for unlimited until max_steps)"
)
+
+ `@field_validator`("botorch_num_trials")
+ `@classmethod`
+ def _validate_botorch_num_trials(cls, v: Optional[int]) -> Optional[int]:
+ if v is None:
+ return v
+ if v != -1 and v < 1:
+ raise ValueError("botorch_num_trials must be -1 or >= 1")
+ return v📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| sobol_num_trials: Optional[int] = Field(default=None, ge=1, description="Number of SOBOL initialization trials") | |
| botorch_num_trials: Optional[int] = Field( | |
| default=None, description="Number of BoTorch trials (-1 for unlimited until max_steps)" | |
| ) | |
| sobol_num_trials: Optional[int] = Field(default=None, ge=1, description="Number of SOBOL initialization trials") | |
| botorch_num_trials: Optional[int] = Field( | |
| default=None, description="Number of BoTorch trials (-1 for unlimited until max_steps)" | |
| ) | |
| `@field_validator`("botorch_num_trials") | |
| `@classmethod` | |
| def _validate_botorch_num_trials(cls, v: Optional[int]) -> Optional[int]: | |
| if v is None: | |
| return v | |
| if v != -1 and v < 1: | |
| raise ValueError("botorch_num_trials must be -1 or >= 1") | |
| return v |
🤖 Prompt for AI Agents
In `@src/cloudai/models/agent_config.py` around lines 42 - 45, The
botorch_num_trials field currently allows any integer but should only accept -1
or integers >= 1 (matching the semantic described and the sibling
sobol_num_trials). Add a Pydantic validator for botorch_num_trials on the
AgentConfig model (e.g., a method named validate_botorch_num_trials decorated
with `@validator`("botorch_num_trials")) that returns the value if it's None or
equals -1 or is >= 1, and raises a ValueError for other values (e.g., 0 or <
-1).
There was a problem hiding this comment.
@alexmanle would be good to consider this. wdyt?
There was a problem hiding this comment.
Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!
There was a problem hiding this comment.
Also would be good to test it with BO agent w/ AIConfigurator along with GA.
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| err = 1 | ||
| continue | ||
|
|
||
| agent = agent_class(env, **agent_overrides) if agent_overrides else agent_class(env) |
There was a problem hiding this comment.
Conditional passes kwargs only when non-empty, but empty dict is falsy. If agent_overrides = {}, the condition evaluates False and takes the else branch. Consider using is not None or explicit length check for clarity:
| agent = agent_class(env, **agent_overrides) if agent_overrides else agent_class(env) | |
| agent = agent_class(env, **agent_overrides) if agent_overrides is not None else agent_class(env) |
However, given validate_agent_overrides always returns a dict, the current logic works but is subtle.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| for error in e.errors(): | ||
| logging.error(f" - {'.'.join(str(var_name) for var_name in error['loc'])}: {error['msg']}") | ||
| logging.error("Valid overrides: ") | ||
| for item, desc in validate_agent_overrides(agent_type).items(): |
There was a problem hiding this comment.
calling validate_agent_overrides(agent_type) in error handler could raise another ValueError if agent doesn't support config, masking original validation error
| for item, desc in validate_agent_overrides(agent_type).items(): | |
| logging.error(f"Invalid agent_config for agent '{agent_type}': ") | |
| for error in e.errors(): | |
| logging.error(f" - {'.'.join(str(var_name) for var_name in error['loc'])}: {error['msg']}") |
3 participants
Summary
Adds a simple interface to override agents hyperparameters.
Test Plan
Tested using
AIConfiguratorworkload and GA agent.Test scenario
Example TOML:
New terminal output: