diff --git a/src/cloudai/cli/handlers.py b/src/cloudai/cli/handlers.py index d474ff421..c8ff8a74f 100644 --- a/src/cloudai/cli/handlers.py +++ b/src/cloudai/cli/handlers.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -21,11 +21,12 @@ import signal from contextlib import contextmanager from pathlib import Path -from typing import Callable, List, Optional +from typing import Any, Callable, List, Optional from unittest.mock import Mock import toml import yaml +from pydantic import ValidationError from cloudai.core import ( BaseInstaller, @@ -40,6 +41,11 @@ TestParser, TestScenario, ) +from cloudai.models.agent_config import ( + BayesianOptimizationConfig, + GeneticAlgorithmConfig, + MultiArmedBanditConfig, +) from cloudai.models.scenario import ReportConfig from cloudai.models.workload import TestDefinition from cloudai.parser import HOOK_ROOT @@ -145,7 +151,19 @@ def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int: continue env = CloudAIGymEnv(test_run=test_run, runner=runner.runner) - agent = agent_class(env) + + try: + agent_overrides = validate_agent_overrides(agent_type, test_run.test.agent_config) + except ValidationError as e: + logging.error(f"Invalid agent_config for agent '{agent_type}':") + for error in e.errors(): + field = ".".join(str(loc) for loc in error["loc"]) + logging.error(f" - {field}: {error['msg']}") + err = 1 + continue + + agent = agent_class(env, **agent_overrides) if agent_overrides else agent_class(env) + for step in range(agent.max_steps): result = agent.select_action() if result is None: @@ -166,6 +184,31 @@ def handle_dse_job(runner: Runner, args: argparse.Namespace) -> int: return err +def validate_agent_overrides(agent_type: str, agent_config: Optional[dict[str, Any]]) -> dict[str, Any]: + """Validate and process agent configuration overrides.""" + if not agent_config: + return {} + + config_class_map = { + "ga": GeneticAlgorithmConfig, + "bo": BayesianOptimizationConfig, + "mab": MultiArmedBanditConfig, + } + + config_class = config_class_map.get(agent_type) + if not config_class: + logging.debug(f"No config validation available for agent type '{agent_type}', using defaults.") + return {} + + validated_config = config_class.model_validate(agent_config) + agent_kwargs = validated_config.model_dump(exclude_none=True) + + if agent_kwargs: + logging.info(f"Applying agent config overrides for '{agent_type}': {agent_kwargs}") + + return agent_kwargs + + def generate_reports(system: System, test_scenario: TestScenario, result_dir: Path) -> None: registry = Registry() diff --git a/src/cloudai/models/agent_config.py b/src/cloudai/models/agent_config.py new file mode 100644 index 000000000..3e090a622 --- /dev/null +++ b/src/cloudai/models/agent_config.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from abc import ABC +from typing import Any, Optional + +from pydantic import BaseModel, ConfigDict, Field + + +class AgentConfig(BaseModel, ABC): + """Base configuration for agent overrides.""" + + model_config = ConfigDict(extra="forbid") + random_seed: Optional[int] = Field(default=None, description="Random seed for reproducibility") + + +class GeneticAlgorithmConfig(AgentConfig): + """Configuration overrides for Genetic Algorithm agent.""" + + population_size: Optional[int] = Field(default=None, ge=2, description="Population size for the genetic algorithm") + n_offsprings: Optional[int] = Field(default=None, ge=1, description="Number of offsprings per generation") + crossover_prob: Optional[float] = Field(default=None, ge=0.0, le=1.0, description="Crossover probability") + mutation_prob: Optional[float] = Field(default=None, ge=0.0, le=1.0, description="Mutation probability") + + +class BayesianOptimizationConfig(AgentConfig): + """Configuration overrides for Bayesian Optimization agent.""" + + 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)" + ) + + +class MultiArmedBanditConfig(AgentConfig): + """Configuration overrides for Multi-Armed Bandit agent.""" + + algorithm: Optional[str] = Field( + default=None, + description="MAB algorithm: ucb1, ts (thompson_sampling), epsilon_greedy, softmax, or random", + ) + algorithm_params: Optional[dict[str, Any]] = Field( + default=None, description="Algorithm-specific parameters (e.g., alpha for UCB1, epsilon for epsilon_greedy)" + ) + seed_parameters: Optional[dict[str, Any]] = Field( + default=None, description="Initial seed configuration to evaluate first" + ) + max_arms: Optional[int] = Field(default=None, ge=1, description="Maximum number of arms in the action space") + warm_start_size: Optional[int] = Field( + default=None, ge=0, description="Number of arms to randomly explore initially" + ) + epsilon_override: Optional[float] = Field( + default=None, ge=0.0, le=1.0, description="Epsilon value for exploration (overrides algorithm epsilon)" + ) + max_explore_steps: Optional[int] = Field( + default=None, ge=0, description="Maximum steps for epsilon exploration (None for unlimited)" + ) + prefer_unseen_random: Optional[bool] = Field( + default=None, description="Prefer unseen arms during random exploration (epsilon)" + ) diff --git a/src/cloudai/models/workload.py b/src/cloudai/models/workload.py index 1745ae734..0a962cf59 100644 --- a/src/cloudai/models/workload.py +++ b/src/cloudai/models/workload.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -107,6 +107,7 @@ class TestDefinition(BaseModel, ABC): agent_steps: int = 1 agent_metrics: list[str] = Field(default=["default"]) agent_reward_function: str = "inverse" + agent_config: Optional[dict[str, Any]] = None @property def cmd_args_dict(self) -> Dict[str, Union[str, List[str]]]: