diff --git a/.semgrepignore b/.semgrepignore index 0d3f1ac..29b6684 100644 --- a/.semgrepignore +++ b/.semgrepignore @@ -8,3 +8,4 @@ src/amzn_nova_forge/model/result/eval_result.py # Safe exec usage in controlled environment for reward verification src/amzn_nova_forge/util/reward_verifier.py +tests/unit/manager/test_smtj_dataprep_runtime_manager.py \ No newline at end of file diff --git a/docs/spec/model-customizer.md b/docs/spec/model-customizer.md index 9285b23..9185f9f 100644 --- a/docs/spec/model-customizer.md +++ b/docs/spec/model-customizer.md @@ -248,6 +248,36 @@ print(f"Training job started: {result.job_id}") print(f"Checkpoint path: {result.model_artifacts.checkpoint_s3_path}") ``` --- + +#### `get_config()` +Returns the overridable training parameters. Delegates to `ForgeTrainer.get_config()`. + +> **Deprecated**: Use `ForgeTrainer.get_config()` instead. + +**Signature:** +```python +def get_config( + self, + overrides: Optional[Dict[str, Any]] = None, +) -> RecipeConfig +``` + +**Parameters:** +- `overrides` (Optional[Dict[str, Any]]): Optional overrides to merge with defaults + +**Returns:** +- `RecipeConfig`: Frozen dataclass containing all overridable parameters with defaults and constraints. See `ForgeTrainer.get_config()` for details. + +**Example:** +```python +config = customizer.get_config() +print(config) + +# With overrides +config = customizer.get_config(overrides={"lr": 1e-5}) +print(config.to_dict()) +``` +--- #### `evaluate()` Generates the recipe YAML, configures runtime, and launches an evaluation job. diff --git a/docs/spec/service-classes.md b/docs/spec/service-classes.md index 5929827..39e3fd6 100644 --- a/docs/spec/service-classes.md +++ b/docs/spec/service-classes.md @@ -164,6 +164,49 @@ print(f"Checkpoint path: {result.model_artifacts.checkpoint_s3_path}") ``` --- +##### `get_config()` +Returns the overridable training parameters for the current model/method/platform without starting a job or running validation. Use this to inspect defaults, valid ranges, and available parameters before calling `train()`. + +**Signature:** +```python +def get_config( + self, + overrides: Optional[TrainingOverrides] = None, +) -> RecipeConfig +``` + +**Parameters:** +- `overrides` (Optional[TrainingOverrides]): Optional overrides to merge with defaults. When provided, `ConfigParameter.default` values in the result reflect the merged values. Plain `dict` values are also accepted + +**Returns:** +- `RecipeConfig`: Frozen dataclass with: + - `model` (Model): The model + - `method` (TrainingMethod): The training method + - `platform` (Platform): The infrastructure platform + - `parameters` (tuple[ConfigParameter, ...]): All overridable parameters, each with `name`, `type`, `default`, `description`, `min`, `max`, `enum`, `required` + - `to_dict()`: Returns `{name: default}` mapping directly passable to `train(overrides=...)` + +**Example:** +```python +# Inspect default configuration +config = trainer.get_config() +print(config) +# RecipeConfig(model=NOVA_MICRO, method=SFT_LORA, platform=SMTJ) +# lr: float = 5e-06 [1e-06..0.0001] +# max_epochs: integer = 2 [1..5] +# warmup_steps: integer = 10 [0..100] +# global_batch_size: integer = 64 [1..512] + +# Preview with your overrides applied +config = trainer.get_config(overrides={"lr": 1e-5, "max_epochs": 4}) +print(config.to_dict()) +# {"lr": 1e-5, "max_epochs": 4, "warmup_steps": 10, "global_batch_size": 64, ...} +``` + +**Note:** `get_config()` is not supported for `RFT_MULTITURN_*` training methods. For those methods, use `train(dry_run=True)` to inspect the recipe. + +--- + ##### `get_logs()` Retrieves and displays CloudWatch logs for a training job. diff --git a/pyproject.toml b/pyproject.toml index 0f09909..037792f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ dependencies = [ "filetype", "matplotlib", "pydantic>=2.0", - "pyarrow", + "pyarrow>=14.0.1", "sagemaker>=3.5.0", ] [project.optional-dependencies] diff --git a/samples/nova_quickstart.ipynb b/samples/nova_quickstart.ipynb index 32683dc..9904bb1 100644 --- a/samples/nova_quickstart.ipynb +++ b/samples/nova_quickstart.ipynb @@ -480,6 +480,18 @@ "print(f\" Method: SFT with LoRA\")" ] }, + { + "cell_type": "markdown", + "source": "### Inspect Available Training Parameters\n\nBefore starting training, you can use `get_config()` to see what parameters are available, their defaults, and valid ranges. This is more convenient than running `train(dry_run=True)` just to inspect the recipe.", + "metadata": {} + }, + { + "cell_type": "code", + "source": "# View all overridable parameters with their defaults and valid ranges\nconfig = trainer.get_config()\nprint(config)\n\n# Preview what the config looks like with your planned overrides\nconfig_with_overrides = trainer.get_config(overrides={\"lr\": 5e-6, \"warmup_steps\": 100})\nprint(\"\\nWith overrides applied:\")\nprint(config_with_overrides.to_dict())", + "metadata": {}, + "execution_count": null, + "outputs": [] + }, { "cell_type": "markdown", "metadata": {}, @@ -1214,4 +1226,4 @@ }, "nbformat": 4, "nbformat_minor": 4 -} +} \ No newline at end of file diff --git a/samples/rft_multiturn_quickstart.ipynb b/samples/rft_multiturn_quickstart.ipynb index a7fa4cd..10a49f9 100644 --- a/samples/rft_multiturn_quickstart.ipynb +++ b/samples/rft_multiturn_quickstart.ipynb @@ -1025,7 +1025,7 @@ " \"max_new_tokens\": 4096,\n", " \"max_length\": 16384,\n", " \"global_batch_size\": 64,\n", - " \"reasoning_effort\": \"null\",\n", + " \"reasoning_effort\": None,\n", " \"timeout\": 1800,\n", " },\n", " rft_multiturn_infra=rft_infra, # Pass RFT infrastructure\n", diff --git a/src/amzn_nova_forge/__init__.py b/src/amzn_nova_forge/__init__.py index 1bb9829..679b318 100644 --- a/src/amzn_nova_forge/__init__.py +++ b/src/amzn_nova_forge/__init__.py @@ -11,6 +11,7 @@ # 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 .core.data_mixing_config import DataMixingConfig from .core.enums import ( DeploymentMode, DeployPlatform, @@ -29,7 +30,15 @@ BaseJobResult, JobStatus, ) -from .core.types import DeploymentResult, EndpointInfo, ForgeConfig, ModelArtifacts +from .core.training_overrides import TrainingOverrides +from .core.types import ( + ConfigParameter, + DeploymentResult, + EndpointInfo, + ForgeConfig, + ModelArtifacts, + RecipeConfig, +) from .dataset import ( ArrowDatasetLoader, CSVDatasetLoader, @@ -121,6 +130,10 @@ def __getattr__(name: str): "DeployPlatform", "Platform", "NovaModelCustomizer", + "ConfigParameter", + "DataMixingConfig", + "RecipeConfig", + "TrainingOverrides", "ForgeConfig", "ForgeTrainer", "ForgeEvaluator", diff --git a/src/amzn_nova_forge/__version__.py b/src/amzn_nova_forge/__version__.py index e561852..9f90d55 100644 --- a/src/amzn_nova_forge/__version__.py +++ b/src/amzn_nova_forge/__version__.py @@ -11,4 +11,4 @@ # 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. -VERSION = "1.4.6" # pragma: no cover +VERSION = "1.4.7" # pragma: no cover diff --git a/src/amzn_nova_forge/core/__init__.py b/src/amzn_nova_forge/core/__init__.py index 789fc03..fec2e55 100644 --- a/src/amzn_nova_forge/core/__init__.py +++ b/src/amzn_nova_forge/core/__init__.py @@ -52,12 +52,14 @@ ) from amzn_nova_forge.core.runtime import RuntimeManager from amzn_nova_forge.core.types import ( + ConfigParameter, DeploymentResult, EndpointInfo, ForgeConfig, JobConfig, ModelArtifacts, ModelConfigDict, + RecipeConfig, validate_region, ) @@ -94,12 +96,14 @@ # runtime "RuntimeManager", # types + "ConfigParameter", "DeploymentResult", "EndpointInfo", "ForgeConfig", "JobConfig", "ModelArtifacts", "ModelConfigDict", + "RecipeConfig", # result classes (lazy-loaded) "BaseJobResult", "BedrockStatusManager", diff --git a/src/amzn_nova_forge/core/constants.py b/src/amzn_nova_forge/core/constants.py index 62ebd2c..0af06d8 100644 --- a/src/amzn_nova_forge/core/constants.py +++ b/src/amzn_nova_forge/core/constants.py @@ -31,6 +31,12 @@ DEFAULT_BATCH_TRACE_CACHE_DIR = "~/.nova-forge/batch_trace_cache" BATCH_TRACE_LOG_SUBDIR = "batch_tracing" +# Data mixing field name patterns +DATAMIX_NOVA_PREFIX = "nova_" +DATAMIX_PERCENT_SUFFIX = "_percent" +DATAMIX_CUSTOMER_DATA_FIELD = "customer_data_percent" +DATAMIX_DATASET_CATALOG_FIELD = "dataset_catalog" + REGION_TO_ESCROW_ACCOUNT_MAPPING = { "us-east-1": "708977205387", "us-west-2": "176779409107", diff --git a/src/amzn_nova_forge/core/data_mixing_config.py b/src/amzn_nova_forge/core/data_mixing_config.py new file mode 100644 index 0000000..872da12 --- /dev/null +++ b/src/amzn_nova_forge/core/data_mixing_config.py @@ -0,0 +1,42 @@ +# Copyright Amazon.com, Inc. or its affiliates + +# 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. +"""Typed configuration for data mixing percentages.""" + +from typing_extensions import TypedDict + + +class DataMixingConfig(TypedDict, total=False): + """Typed configuration for data mixing between customer and Nova curated data. + + All percent fields must be 0-100. When customer_data_percent < 100, the + nova_*_percent fields must sum to 100. Unrecognized nova_*_percent fields + are accepted since TypedDict does not reject unknown keys at runtime. + + Can be passed directly to ``DataMixing.set_config()`` — plain dicts + satisfy this annotation without construction. + + Example:: + + config: DataMixingConfig = { + "customer_data_percent": 50.0, + "nova_code_percent": 30.0, + "nova_general_percent": 70.0, + } + """ + + customer_data_percent: float + nova_code_percent: float + nova_general_percent: float + nova_math_percent: float + nova_image_percent: float diff --git a/src/amzn_nova_forge/core/training_overrides.py b/src/amzn_nova_forge/core/training_overrides.py new file mode 100644 index 0000000..2881228 --- /dev/null +++ b/src/amzn_nova_forge/core/training_overrides.py @@ -0,0 +1,67 @@ +# Copyright Amazon.com, Inc. or its affiliates + +# 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. +"""TypedDict for training recipe overrides.""" + +from typing_extensions import TypedDict + + +class TrainingOverrides(TypedDict, total=False): + """Typed hints for training recipe override keys. + + Provides IDE autocomplete and static type checking for common override + parameters passed to ``ForgeTrainer.train()``, ``ForgeEvaluator.evaluate()``, + and ``RecipeBuilder.build_and_validate()``. All keys are optional + (``total=False``). + + Plain dicts remain fully supported — this is structural typing, so + ``{"lr": 0.001}`` satisfies the annotation without constructing a + ``TrainingOverrides`` instance. Unknown keys from new recipe templates + pass through at runtime without restriction. + + Example:: + + overrides: TrainingOverrides = { + "lr": 5e-6, + "max_steps": 1000, + "global_batch_size": 128, + } + """ + + lr: float + loraplus_lr_ratio: float + lora_plus_lr_ratio: float + + max_steps: int + max_epochs: int + save_steps: int + warmup_steps: int + + global_batch_size: int + max_length: int + + alpha: float + + beta: float + + reasoning_enabled: bool + reasoning_effort: str | None + + val_check_interval: int + validation_s3_path: str + validation_data_s3_path: str + + top_logprobs: int + + +NULLABLE_OVERRIDE_FIELDS: frozenset[str] = frozenset({"reasoning_effort"}) diff --git a/src/amzn_nova_forge/core/types.py b/src/amzn_nova_forge/core/types.py index 416a5c9..05a1b8f 100644 --- a/src/amzn_nova_forge/core/types.py +++ b/src/amzn_nova_forge/core/types.py @@ -27,7 +27,7 @@ DEFAULT_JOB_CACHE_DIR, REGION_TO_ESCROW_ACCOUNT_MAPPING, ) -from amzn_nova_forge.core.enums import DeployPlatform, Platform, TrainingMethod +from amzn_nova_forge.core.enums import DeployPlatform, Model, Platform, TrainingMethod if TYPE_CHECKING: from amzn_nova_forge.core.job_cache import JobCachingConfig @@ -157,3 +157,47 @@ class JobConfig: method: Optional[TrainingMethod] = None # Training method (required for Bedrock) data_mixing_config: Optional[Dict[str, Any]] = None # Datamix percent fields (SMTJServerless) # TODO: The mlflow config is populated in recipe for both SMTJ and SMHP but will only work for SMHP as SMTJ support for mlflow is only through boto3, fix this with sagemaker 3 update + + +@dataclass(frozen=True) +class ConfigParameter: + """A single overridable parameter in a training recipe.""" + + name: str + type: str + default: Any + description: Optional[str] = None + min: Optional[float] = None + max: Optional[float] = None + enum: Optional[tuple] = None + required: bool = False + + +@dataclass(frozen=True) +class RecipeConfig: + """Frozen snapshot of the overridable configuration for a training recipe.""" + + model: Model + method: TrainingMethod + platform: Platform + parameters: tuple[ConfigParameter, ...] + + def to_dict(self) -> Dict[str, Any]: + return {p.name: p.default for p in self.parameters} + + def __repr__(self) -> str: + lines = [ + f"RecipeConfig(model={self.model.name}, method={self.method.name}, platform={self.platform.name})" + ] + for p in self.parameters: + constraint = "" + if p.min is not None and p.max is not None: + constraint = f" [{p.min}..{p.max}]" + elif p.min is not None: + constraint = f" [>={p.min}]" + elif p.max is not None: + constraint = f" [<={p.max}]" + if p.enum is not None: + constraint = f" enum={list(p.enum)}" + lines.append(f" {p.name}: {p.type} = {p.default}{constraint}") + return "\n".join(lines) diff --git a/src/amzn_nova_forge/deployer/forge_deployer.py b/src/amzn_nova_forge/deployer/forge_deployer.py index fbe8e12..aacb7cd 100644 --- a/src/amzn_nova_forge/deployer/forge_deployer.py +++ b/src/amzn_nova_forge/deployer/forge_deployer.py @@ -224,12 +224,17 @@ def get_logs( endpoint_arn: Optional[str] = None, platform: Optional[DeployPlatform] = None, ) -> None: - """Log deployment status information.""" + """Log deployment status information. + + Raises: + ValueError: If neither ``job_result`` nor ``endpoint_arn`` is provided. + """ arn = endpoint_arn or (job_result.endpoint.uri if job_result else None) if not arn: - logger.info("No endpoint ARN available. Call deploy() first.") - return + raise ValueError( + "No endpoint ARN available. Pass either a job_result from deploy() or an explicit endpoint_arn." + ) if job_result is not None: platform = job_result.endpoint.platform diff --git a/src/amzn_nova_forge/evaluator/forge_evaluator.py b/src/amzn_nova_forge/evaluator/forge_evaluator.py index b7c83b7..124b356 100644 --- a/src/amzn_nova_forge/evaluator/forge_evaluator.py +++ b/src/amzn_nova_forge/evaluator/forge_evaluator.py @@ -346,13 +346,19 @@ def get_logs( start_from_head: bool = False, end_time: Optional[int] = None, ) -> None: - """Stream CloudWatch logs for an evaluation job.""" + """Stream CloudWatch logs for an evaluation job. + + Raises: + ValueError: If neither ``job_result`` nor both ``job_id`` and + ``started_time`` are provided. + """ resolved_job_id = job_result.job_id if job_result else job_id resolved_started = job_result.started_time if job_result else started_time if not resolved_job_id or not resolved_started: - logger.info("Provide either a job_result or explicit job_id and started_time.") - return + raise ValueError( + "No job reference provided. Pass either a job_result or explicit job_id and started_time." + ) kwargs: Dict[str, Any] = {} if self._platform == Platform.SMHP: diff --git a/src/amzn_nova_forge/inference/forge_inference.py b/src/amzn_nova_forge/inference/forge_inference.py index d44ee9a..7838bda 100644 --- a/src/amzn_nova_forge/inference/forge_inference.py +++ b/src/amzn_nova_forge/inference/forge_inference.py @@ -327,17 +327,24 @@ def get_logs( """Stream CloudWatch logs for a batch inference job. Provide either a ``job_result`` or explicit ``job_id`` + ``started_time``. + + Raises: + ValueError: If neither ``job_result`` nor both ``job_id`` and + ``started_time`` are provided, or if the platform cannot be + determined (no ``infra`` in constructor). """ resolved_job_id = job_result.job_id if job_result else job_id resolved_started = job_result.started_time if job_result else started_time if not resolved_job_id or not resolved_started: - logger.info("Provide either a job_result or explicit job_id and started_time.") - return + raise ValueError( + "No job reference provided. Pass either a job_result or explicit job_id and started_time." + ) if self._platform is None: - logger.info("Cannot determine platform — provide infra in the constructor.") - return + raise ValueError( + "Cannot determine platform. Provide infra in the ForgeInference constructor." + ) platform: Platform = self._platform kwargs: Dict[str, Any] = {} diff --git a/src/amzn_nova_forge/manager/glue_runtime_manager.py b/src/amzn_nova_forge/manager/glue_runtime_manager.py index f687c74..845fff0 100644 --- a/src/amzn_nova_forge/manager/glue_runtime_manager.py +++ b/src/amzn_nova_forge/manager/glue_runtime_manager.py @@ -15,6 +15,7 @@ from __future__ import annotations +import hashlib import io import json import os @@ -43,6 +44,7 @@ GLUE_IAM_ROLE_NAME = "GlueDataPrepExecutionRole" _WHL_FILENAME = "agi_data_curator-1.0.0-py3-none-any.whl" +_WHL_SHA256 = "18d631f2a367dc744d34ac3b3e1cd8718a4fe3a859a79d63c29472169e7abba1" _TERMINAL_STATES = {"SUCCEEDED", "FAILED", "STOPPED", "TIMEOUT", "ERROR"} # Glue entry-point script, uploaded to S3 and executed inside the Glue Ray worker. @@ -271,6 +273,12 @@ def _upload_bundled_module(self) -> str: """Upload the bundled agi_data_curator wheel and return its S3 path.""" whl_ref = resources.files("amzn_nova_forge.dataset.bundled").joinpath(_WHL_FILENAME) whl_bytes = whl_ref.read_bytes() + actual_sha = hashlib.sha256(whl_bytes).hexdigest() + if actual_sha != _WHL_SHA256: + raise RuntimeError( + f"Integrity check failed for bundled wheel {_WHL_FILENAME}: " + f"expected sha256 {_WHL_SHA256}, got {actual_sha}" + ) return self._upload_whl_as_zip(_WHL_FILENAME, whl_bytes) def _upload_module(self, module_path: str) -> str: diff --git a/src/amzn_nova_forge/manager/runtime_manager.py b/src/amzn_nova_forge/manager/runtime_manager.py index da8dbe5..ac1b314 100644 --- a/src/amzn_nova_forge/manager/runtime_manager.py +++ b/src/amzn_nova_forge/manager/runtime_manager.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import enum +import hashlib import io import json import os @@ -76,6 +77,10 @@ DEFAULT_SMTJ_JOB_MAX_RUNTIME = 86400 # 1 day DEFAULT_SMTJ_JOB_SUBMIT_POLL_TIMEOUT = 30 # seconds to poll for job after submission +_BUNDLED_WHL_SHA256 = { + "agi_data_curator-1.0.0-py3-none-any.whl": "18d631f2a367dc744d34ac3b3e1cd8718a4fe3a859a79d63c29472169e7abba1", +} + DEFAULT_DATAPREP_MAX_JOB_RUNTIME = 3600 DEFAULT_VOLUME_SIZE_GB = 100 @@ -822,9 +827,20 @@ def _install_deps(pip_install): site_dirs = site.getsitepackages() target = site_dirs[0] if site_dirs else site.getusersitepackages() logger.info("Extracting bundled wheels to %s: %s", target, whls) + target_abs = os.path.realpath(target) for whl_path in whls: with zipfile.ZipFile(whl_path, "r") as zf: - zf.extractall(target) + for name in zf.namelist(): + # Reject absolute paths and path traversal + if os.path.isabs(name) or name.startswith("/"): + raise RuntimeError(f"Unsafe absolute path: {name!r} in {whl_path}") + resolved = os.path.realpath(os.path.join(target_abs, name)) + if ( + resolved != target_abs + and not resolved.startswith(target_abs + os.sep) + ): + raise RuntimeError(f"Unsafe path traversal: {name!r} in {whl_path}") + zf.extractall(target_abs) if pip_install: logger.info("Installing additional packages: %s", pip_install) @@ -1125,6 +1141,14 @@ def _upload_bundled_whls(self) -> None: deps_prefix = f"{self.s3_artifact_prefix}/deps" for whl_name in whl_names: whl_bytes = bundled.joinpath(whl_name).read_bytes() + expected_sha = _BUNDLED_WHL_SHA256.get(whl_name) + if expected_sha: + actual_sha = hashlib.sha256(whl_bytes).hexdigest() + if actual_sha != expected_sha: + raise RuntimeError( + f"Integrity check failed for bundled wheel {whl_name}: " + f"expected sha256 {expected_sha}, got {actual_sha}" + ) whl_key = f"{deps_prefix}/{whl_name}" self.s3_client.put_object( Bucket=self.s3_artifact_bucket, diff --git a/src/amzn_nova_forge/model/nova_model_customizer.py b/src/amzn_nova_forge/model/nova_model_customizer.py index ded7c8e..1b2f1dd 100644 --- a/src/amzn_nova_forge/model/nova_model_customizer.py +++ b/src/amzn_nova_forge/model/nova_model_customizer.py @@ -41,11 +41,11 @@ TrainingResult, ) from amzn_nova_forge.core.result.inference_result import InferenceResult -from amzn_nova_forge.core.result.job_result import JobStatus from amzn_nova_forge.core.types import ( DeploymentResult, EndpointInfo, ForgeConfig, + RecipeConfig, ValidationConfig, ) from amzn_nova_forge.manager.runtime_manager import ( @@ -72,19 +72,6 @@ from amzn_nova_forge.telemetry.telemetry_logging import ( _telemetry_emitter, ) -from amzn_nova_forge.util.bedrock import ( - BEDROCK_EXECUTION_ROLE_NAME, - DEPLOYMENT_ARN_NAME, - check_existing_deployment, - delete_existing_deployment, - find_bedrock_model_by_tag, - get_required_bedrock_deletion_permissions, - get_required_bedrock_update_permissions, - invoke_model, - monitor_model_create, - update_provisioned_throughput_model, - wait_for_model_ready, -) from amzn_nova_forge.util.checkpoint_util import extract_checkpoint_path_from_job_output from amzn_nova_forge.util.data_mixing import DataMixing from amzn_nova_forge.util.data_utils import is_multimodal_data @@ -515,7 +502,7 @@ def set_data_mixing_config(self, config: Dict[str, Any]) -> None: "Data mixing is not enabled for this customizer. Set data_mixing = True in 'NovaModelCustomizer' object." ) - self.data_mixing.set_config(config, normalize=True) + self.data_mixing.set_config(config, normalize=True) # type: ignore[arg-type] @_telemetry_emitter( Feature.TRAINING, @@ -617,7 +604,7 @@ def train( training_result = trainer.train( job_name=job_name, recipe_path=recipe_path, - overrides=overrides, + overrides=overrides, # type: ignore[arg-type] rft_lambda_arn=rft_lambda_arn, dry_run=dry_run, rft_multiturn_infra=rft_multiturn_infra, @@ -645,6 +632,51 @@ def train( return training_result + @_telemetry_emitter( + Feature.TRAINING, + "customizer.get_config", + extra_info_fn=lambda self, *args, **kwargs: { + "method": self.method, + "model": self.model.value, + "platform": self.platform, + }, + ) + def get_config( + self, + overrides: Optional[Dict[str, Any]] = None, + ) -> RecipeConfig: + """Return the overridable training configuration. + + .. deprecated:: + Use ``ForgeTrainer.get_config()`` instead. + """ + warnings.warn( + "NovaModelCustomizer.get_config() is deprecated and will be removed in a future version. " + "Use ForgeTrainer.get_config() instead.", + DeprecationWarning, + stacklevel=2, + ) + + from amzn_nova_forge.trainer.forge_trainer import ForgeTrainer + + trainer = ForgeTrainer( + model=self.model, + method=self.method, + infra=self.infra, + training_data_s3_path=self.data_s3_path, + model_s3_path=self.model_path, + data_mixing_enabled=self.data_mixing_enabled, + config=self._build_forge_config(), + region=self.region, + is_multimodal=self.is_multimodal, + hub_content_version=self.hub_content_version, + ) + + if self.data_mixing is not None: + trainer.data_mixing = self.data_mixing + + return trainer.get_config(overrides=overrides) # type: ignore[arg-type] + @_telemetry_emitter( Feature.EVAL, "evaluate", @@ -1242,8 +1274,8 @@ def get_logs( limit=limit, start_from_head=start_from_head, end_time=end_time ) else: - logger.info( - "No job_id and job_started_time found for this model, please call .train() or .evaluate() first." + raise ValueError( + "No job_id and job_started_time found. Call .train() or .evaluate() before calling .get_logs()." ) @_telemetry_emitter( diff --git a/src/amzn_nova_forge/recipe/recipe_builder.py b/src/amzn_nova_forge/recipe/recipe_builder.py index 719220b..664c9e7 100644 --- a/src/amzn_nova_forge/recipe/recipe_builder.py +++ b/src/amzn_nova_forge/recipe/recipe_builder.py @@ -24,6 +24,10 @@ import yaml from amzn_nova_forge.core.constants import ( + DATAMIX_CUSTOMER_DATA_FIELD, + DATAMIX_DATASET_CATALOG_FIELD, + DATAMIX_NOVA_PREFIX, + DATAMIX_PERCENT_SUFFIX, EVAL_TASK_METRIC_MAP, EVAL_TASK_STRATEGY_MAP, HYPERPOD_RECIPE_PATH, @@ -36,7 +40,8 @@ Version, ) from amzn_nova_forge.core.runtime import RuntimeManager -from amzn_nova_forge.core.types import ValidationConfig +from amzn_nova_forge.core.training_overrides import NULLABLE_OVERRIDE_FIELDS +from amzn_nova_forge.core.types import ConfigParameter, RecipeConfig, ValidationConfig from amzn_nova_forge.monitor import MLflowMonitor from amzn_nova_forge.rft_multiturn import RFTMultiturnInfrastructure from amzn_nova_forge.util.checkpoint_util import validate_checkpoint_uri @@ -588,7 +593,10 @@ def update_overrides_template(recipe_template: Dict[str, Any]): overrides_template[recipe_template_key] = { "default": overrides[recipe_template_key], "type": type(recipe_template_value).__name__, - "required": True, + "required": not ( + overrides[recipe_template_key] is None + and recipe_template_key in NULLABLE_OVERRIDE_FIELDS + ), } if existing_enum is not None: overrides_template[recipe_template_key]["enum"] = existing_enum @@ -691,13 +699,13 @@ def update_overrides_template(recipe_template: Dict[str, Any]): # TODO Investigate some percent params are integer type but default value is float # Changing all to float right now for key, value in data_mixing_config.items(): - if key in overrides_template and key != DataMixing.DATASET_CATALOG_FIELD: - if key == DataMixing.CUSTOMER_DATA_FIELD: + if key in overrides_template and key != DATAMIX_DATASET_CATALOG_FIELD: + if key == DATAMIX_CUSTOMER_DATA_FIELD: data_mixing_recipe_key = "percent" else: - data_mixing_recipe_key = key.removeprefix(DataMixing.NOVA_PREFIX) + data_mixing_recipe_key = key.removeprefix(DATAMIX_NOVA_PREFIX) data_mixing_recipe_key = data_mixing_recipe_key.removesuffix( - DataMixing.PERCENT_SUFFIX + DATAMIX_PERCENT_SUFFIX ) overrides_template[data_mixing_recipe_key] = overrides_template[key] @@ -1008,3 +1016,80 @@ def build_and_validate( else None, image_uri, ) + + _NON_CONFIG_KEYS = frozenset( + { + "name", + "model_type", + "model_name_or_path", + "data_s3_path", + "output_s3_path", + "peft_scheme", + "task", + } + ) + + def get_overridable_config( + self, + overrides: Optional[Dict[str, Any]] = None, + input_recipe_path: Optional[str] = None, + ) -> RecipeConfig: + """Return overridable recipe parameters without validation or file I/O.""" + recipe_metadata, recipe_template, overrides_template, _image_uri = load_recipe_templates( + model=self.model, + method=self.method, + platform=self.platform, + region=self.region, + instance_type=self.instance_type, + data_mixing_enabled=True if self.data_mixing_instance else False, + eval_task=getattr(self, "eval_task", None), + image_uri_override=self.image_uri_override, + rft_multiturn_infra=self.rft_multiturn_infra, + is_multimodal=self.is_multimodal, + hub_content_version=self.hub_content_version, + ) + + if not overrides_template: + return RecipeConfig( + model=self.model, + method=self.method, + platform=self.platform, + parameters=(), + ) + + self._resolve_user_inputs( + recipe_template=recipe_template, + overrides_template=overrides_template, + overrides=overrides or {}, + input_recipe_path=input_recipe_path, + allowed_instance_count=recipe_metadata.get("InstanceCount") + if recipe_metadata + else None, + allowed_instance_types=recipe_metadata.get("SupportedInstanceTypes") + if recipe_metadata + else None, + ) + + parameters = [] + for key, meta in overrides_template.items(): + if key in self._NON_CONFIG_KEYS or not isinstance(meta, dict): + continue + parameters.append( + ConfigParameter( + name=key, + type=meta.get("type", "unknown"), + default=meta.get("default"), + description=meta.get("description"), + min=meta.get("min"), + max=meta.get("max"), + enum=tuple(meta["enum"]) if "enum" in meta else None, + required=meta.get("required", False), + ) + ) + + return RecipeConfig( + model=self.model, + method=self.method, + platform=self.platform, + parameters=tuple(parameters), + ) diff --git a/src/amzn_nova_forge/trainer/forge_trainer.py b/src/amzn_nova_forge/trainer/forge_trainer.py index 43ab938..d31cf3a 100644 --- a/src/amzn_nova_forge/trainer/forge_trainer.py +++ b/src/amzn_nova_forge/trainer/forge_trainer.py @@ -41,10 +41,12 @@ TrainingResult, ) from amzn_nova_forge.core.runtime import RuntimeManager +from amzn_nova_forge.core.training_overrides import TrainingOverrides from amzn_nova_forge.core.types import ( ForgeConfig, JobConfig, ModelArtifacts, + RecipeConfig, validate_region, ) from amzn_nova_forge.manager.runtime_manager import SMHPRuntimeManager @@ -226,7 +228,7 @@ def train( self, job_name: str, recipe_path: Optional[str] = None, - overrides: Optional[Dict[str, Any]] = None, + overrides: Optional[TrainingOverrides] = None, rft_lambda_arn: Optional[str] = None, dry_run: bool = False, rft_multiturn_infra: Optional[RFTMultiturnInfrastructure] = None, @@ -293,7 +295,7 @@ def train( resolved_data_s3_path, resolved_image_uri, ) = recipe_builder.build_and_validate( - overrides=overrides, + overrides=dict(overrides) if overrides else {}, input_recipe_path=recipe_path, output_recipe_path=self._config.generated_recipe_dir, validation_config=self._config.validation_config, @@ -397,6 +399,52 @@ def train( return training_result + @_telemetry_emitter( + Feature.TRAINING, + "forgetrainer.get_config", + extra_info_fn=lambda self, *args, **kwargs: { + "method": self.method, + "model": self.model, + "platform": self._platform, + }, + ) + def get_config( + self, + overrides: Optional[TrainingOverrides] = None, + ) -> RecipeConfig: + """Return overridable training parameters without starting a job. + + Not supported for RFT_MULTITURN methods (rft_multiturn_infra + is not available outside of ``train()``). + + Args: + overrides: Optional overrides to merge with defaults. + + Returns: + Frozen RecipeConfig with all overridable parameters. + """ + recipe_builder = RecipeBuilder( + region=self.region, + job_name="__config_preview__", + platform=self._platform, + model=self.model, + method=self.method, + instance_type=self.infra.instance_type, + instance_count=self.infra.instance_count, + infra=self.infra, + data_s3_path=self.training_data_s3_path, + output_s3_path=self.output_s3_path, + model_path=self.model_s3_path, + validation_data_s3_path=self.holdout_data_s3_path, + val_check_interval=self.val_check_interval, + data_mixing_instance=self.data_mixing, + image_uri_override=self._config.image_uri, + is_multimodal=self._is_multimodal, + mlflow_monitor=self._config.mlflow_monitor, + hub_content_version=self.hub_content_version, + ) + return recipe_builder.get_overridable_config(overrides=overrides) # type: ignore[arg-type] + @_telemetry_emitter(Feature.TRAINING, "get_logs") def get_logs( self, @@ -410,13 +458,18 @@ def get_logs( """Stream CloudWatch logs for a training job. Provide either a ``job_result`` or explicit ``job_id`` + ``started_time``. + + Raises: + ValueError: If neither ``job_result`` nor both ``job_id`` and + ``started_time`` are provided. """ resolved_job_id = job_result.job_id if job_result else job_id resolved_started = job_result.started_time if job_result else started_time if not resolved_job_id or not resolved_started: - logger.info("Provide either a job_result or explicit job_id and started_time.") - return + raise ValueError( + "No job reference provided. Pass either a job_result or explicit job_id and started_time." + ) kwargs: Dict[str, Any] = {} if self._platform == Platform.SMHP: diff --git a/src/amzn_nova_forge/util/data_mixing.py b/src/amzn_nova_forge/util/data_mixing.py index 24c546e..8347b69 100644 --- a/src/amzn_nova_forge/util/data_mixing.py +++ b/src/amzn_nova_forge/util/data_mixing.py @@ -21,6 +21,13 @@ from dataclasses import dataclass, field from typing import Any, Dict, Optional +from amzn_nova_forge.core.constants import ( + DATAMIX_CUSTOMER_DATA_FIELD, + DATAMIX_DATASET_CATALOG_FIELD, + DATAMIX_NOVA_PREFIX, + DATAMIX_PERCENT_SUFFIX, +) +from amzn_nova_forge.core.data_mixing_config import DataMixingConfig from amzn_nova_forge.util.logging import logger from amzn_nova_forge.validation.validator import Validator @@ -41,12 +48,6 @@ class DataMixing: _dataset_catalog: Stores the dataset catalog value (read-only from template) """ - # Constants for field name patterns - NOVA_PREFIX = "nova_" - PERCENT_SUFFIX = "_percent" - CUSTOMER_DATA_FIELD = "customer_data_percent" - DATASET_CATALOG_FIELD = "dataset_catalog" - config: Dict[str, Any] = field(default_factory=dict) _default_nova_fields: set = field(default_factory=lambda: set(), init=False, repr=False) _dataset_catalog: Optional[Any] = field(default=None, init=False, repr=False) @@ -60,15 +61,16 @@ def get_config(self) -> Dict[str, Any]: """ result = self.config.copy() if self._dataset_catalog is not None: - result[self.DATASET_CATALOG_FIELD] = self._dataset_catalog + result[DATAMIX_DATASET_CATALOG_FIELD] = self._dataset_catalog return result - def set_config(self, config: Dict[str, Any], normalize: bool = True) -> None: + def set_config(self, config: DataMixingConfig, normalize: bool = True) -> None: """ Set the data mixing configuration. Args: - config: Dictionary containing the data mixing configuration. + config: A DataMixingConfig TypedDict (or any dict matching that shape) + containing the data mixing configuration. Keys should include nova_*_percent fields and customer_data_percent. Any nova_*_percent fields not specified will be set to 0 if normalize=True. normalize: If True, unspecified nova fields will be set to 0. Default is True. @@ -76,37 +78,32 @@ def set_config(self, config: Dict[str, Any], normalize: bool = True) -> None: Raises: ValueError: If configuration is invalid or contains unknown nova fields """ - if self.DATASET_CATALOG_FIELD in config: + config_dict = dict(config) + + # Strip dataset_catalog with warning + if DATAMIX_DATASET_CATALOG_FIELD in config_dict: logger.warning( - f"{self.DATASET_CATALOG_FIELD} cannot be set in data mixing configuration. " - f"Ignoring value: {config[self.DATASET_CATALOG_FIELD]}" + "%s cannot be set in data mixing configuration. Ignoring value: %s", + DATAMIX_DATASET_CATALOG_FIELD, + config_dict[DATAMIX_DATASET_CATALOG_FIELD], ) + del config_dict[DATAMIX_DATASET_CATALOG_FIELD] + + # Filter None values + config_dict = {k: v for k, v in config_dict.items() if v is not None} # Create new config with normalization if needed - new_config = {} - - if normalize: - # Start with all known fields set to 0 - if self._default_nova_fields or self.CUSTOMER_DATA_FIELD in self._default_nova_fields: - # Include nova fields - for field in self._default_nova_fields: - new_config[field] = 0 - # Include customer_data_percent if it's in defaults - if self.CUSTOMER_DATA_FIELD in self._default_nova_fields: - new_config[self.CUSTOMER_DATA_FIELD] = 0 + new_config: Dict[str, Any] = {} + + if normalize and self._default_nova_fields: + for f in self._default_nova_fields: + new_config[f] = 0 # Update with provided config - new_config.update(config) - - # Validate the configuration - Validator.validate_data_mixing_config( - new_config, - self.NOVA_PREFIX, - self.PERCENT_SUFFIX, - self.CUSTOMER_DATA_FIELD, - self.DATASET_CATALOG_FIELD, - self._default_nova_fields, - ) + new_config.update(config_dict) + + # Validate + Validator.validate_data_mixing_config(new_config) # Store the new config self.config = new_config @@ -133,16 +130,16 @@ def _load_defaults_from_template(self, overrides_template: Dict[str, Any]) -> No default_config = {} for key, value in overrides_template.items(): - if key.startswith(self.NOVA_PREFIX) and key.endswith(self.PERCENT_SUFFIX): + if key.startswith(DATAMIX_NOVA_PREFIX) and key.endswith(DATAMIX_PERCENT_SUFFIX): self._default_nova_fields.add(key) if isinstance(value, dict) and "default" in value: default_config[key] = value["default"] - elif key == self.CUSTOMER_DATA_FIELD: + elif key == DATAMIX_CUSTOMER_DATA_FIELD: # Add customer_data_percent to default fields self._default_nova_fields.add(key) if isinstance(value, dict) and "default" in value: default_config[key] = value["default"] - elif key == self.DATASET_CATALOG_FIELD: + elif key == DATAMIX_DATASET_CATALOG_FIELD: # Store dataset_catalog separately as it's read-only if isinstance(value, dict) and "default" in value: self._dataset_catalog = value["default"] @@ -151,8 +148,8 @@ def _load_defaults_from_template(self, overrides_template: Dict[str, Any]) -> No def _is_data_mixing_field(self, key: str) -> bool: return ( - (key.startswith(DataMixing.NOVA_PREFIX) and key.endswith(DataMixing.PERCENT_SUFFIX)) - or key == DataMixing.CUSTOMER_DATA_FIELD - or DataMixing.NOVA_PREFIX + key + DataMixing.PERCENT_SUFFIX in self.get_config() - or key == DataMixing.DATASET_CATALOG_FIELD + (key.startswith(DATAMIX_NOVA_PREFIX) and key.endswith(DATAMIX_PERCENT_SUFFIX)) + or key == DATAMIX_CUSTOMER_DATA_FIELD + or DATAMIX_NOVA_PREFIX + key + DATAMIX_PERCENT_SUFFIX in self.get_config() + or key == DATAMIX_DATASET_CATALOG_FIELD ) diff --git a/src/amzn_nova_forge/validation/validator.py b/src/amzn_nova_forge/validation/validator.py index eb7c4a0..b47d6ec 100644 --- a/src/amzn_nova_forge/validation/validator.py +++ b/src/amzn_nova_forge/validation/validator.py @@ -20,6 +20,9 @@ from amzn_nova_forge.core.constants import ( BYOD_AVAILABLE_EVAL_TASKS, + DATAMIX_CUSTOMER_DATA_FIELD, + DATAMIX_NOVA_PREFIX, + DATAMIX_PERCENT_SUFFIX, EVAL_TASK_STRATEGY_MAP, get_available_subtasks, ) @@ -754,71 +757,52 @@ def _validate_infrastructure(infra: Any, errors: List[str]) -> None: errors.append(f"Failed to validate cluster infrastructure: {str(e)}") @staticmethod - def validate_data_mixing_config( - config: dict, - nova_prefix: str, - percent_suffix: str, - customer_data_field: str, - dataset_catalog_fields: str, - expected_nova_fields: set, - ) -> None: - """ - Validate the data mixing configuration. The validation rules are as follows: - - The datamix config should have valid fields. - - Customer data can be between 0-100. - - If customer data is 100, then no nova data is used - - If customer data < 100, then sum of nova data percent fields should be 100. + def validate_data_mixing_config(config: Dict[str, Any]) -> None: + """Validate a data mixing config dict. + + Checks: + - Each percent field is 0-100 and numeric. + - When ``customer_data_percent`` is 100, all nova data should sum to 0. + - When ``customer_data_percent`` < 100, nova data percentages must sum to 100. Raises: - ValueError: If configuration is invalid + ValueError: If configuration is invalid. """ + range_errors: List[str] = [] + nova_fields: Dict[str, Any] = {} + customer_percent = config.get(DATAMIX_CUSTOMER_DATA_FIELD) - nova_fields = {} - customer_percent = 0 - total = 0 + for key, value in config.items(): + is_percent_field = key == DATAMIX_CUSTOMER_DATA_FIELD or ( + key.startswith(DATAMIX_NOVA_PREFIX) and key.endswith(DATAMIX_PERCENT_SUFFIX) + ) + if not is_percent_field: + continue - if expected_nova_fields: - for key in config.keys(): - # Skip customer_data_percent - if key == dataset_catalog_fields: - continue - # Check if key is a nova field that's not in the known defaults - if key not in expected_nova_fields: - raise ValueError( - f"Invalid nova field '{key}'. Valid fields are: {sorted(expected_nova_fields)}" - ) + if value is not None and not isinstance(value, (int, float)): + range_errors.append( + f"{key} must be a number, got {type(value).__name__}: {value!r}" + ) + elif value is not None and not 0 <= value <= 100: + range_errors.append(f"{key} must be between 0 and 100, got {value}") - for key, value in config.items(): - if key == customer_data_field: - customer_percent = value - if value is not None and not 0 <= value <= 100: - raise ValueError( - f"{customer_data_field} must be between 0 and 100, got {value}" - ) - elif key.startswith(nova_prefix) and key.endswith(percent_suffix): + if key != DATAMIX_CUSTOMER_DATA_FIELD: nova_fields[key] = value - # Each nova field must be between 0 and 100 - if value is not None and not 0 <= value <= 100: - raise ValueError(f"{key} must be between 0 and 100, got {value}") - - if nova_fields: - non_none_values = [v for v in nova_fields.values() if v is not None] - if non_none_values: - total = sum(non_none_values) - if abs(total - 100.0) > 0.01: # Allow small floating point errors - raise ValueError( - f"Nova data percentages must sum to 100, got {total}. Fields: {nova_fields}" - ) - if customer_percent == 100 and total > 0: - raise ValueError( - f"Since {customer_data_field} is 100 %, all nova data should sum to 0 %" - ) + if range_errors: + raise ValueError("; ".join(range_errors)) - if customer_percent < 100 and total == 0: + nova_values = [v for v in nova_fields.values() if v is not None] + total = sum(nova_values) if nova_values else 0 + + if customer_percent == 100: + if total > 0: + raise ValueError( + f"Since {DATAMIX_CUSTOMER_DATA_FIELD} is 100 %, all nova data should sum to 0 %" + ) + elif nova_values and abs(total - 100.0) > 0.01: raise ValueError( - f"Since {customer_data_field} is less than 100 % {customer_percent}%, all nova data cannot be 0" - f"Fields: {nova_fields} should sum to 100 %" + f"Nova data percentages must sum to 100, got {total}. Fields: {nova_fields}" ) @staticmethod @@ -1040,7 +1024,7 @@ def get_recipe_value(data: Dict[str, Any], key_to_find: str) -> Any: continue # Validate proper types are used - if "type" in override_metadata: + if "type" in override_metadata and key != "temperature": # None is allowed for optional fields (recipe key exists but has no value set) if recipe_value is None and not override_metadata.get("required", False): continue @@ -1055,6 +1039,15 @@ def get_recipe_value(data: Dict[str, Any], key_to_find: str) -> Any: f"'{key}' expects {override_metadata['type']}. You provided {type(recipe_value).__name__}." ) continue # If wrong type is used, continue to next key to prevent type exceptions + + # Special validation for temperature: accepts both int and float + if key == "temperature": + if not isinstance(recipe_value, (int, float)): + errors.append( + f"'temperature' expects int or float. You provided {type(recipe_value).__name__}." + ) + continue + # Validate enum constraints are met if "enum" in override_metadata: if recipe_value not in override_metadata["enum"] and recipe_value != "": diff --git a/tests/unit/core/test_config_types.py b/tests/unit/core/test_config_types.py new file mode 100644 index 0000000..72602a8 --- /dev/null +++ b/tests/unit/core/test_config_types.py @@ -0,0 +1,125 @@ +# Copyright Amazon.com, Inc. or its affiliates + +# 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. +"""Unit tests for ConfigParameter and RecipeConfig.""" + +import unittest + +from amzn_nova_forge.core.enums import Model, Platform, TrainingMethod +from amzn_nova_forge.core.types import ConfigParameter, RecipeConfig + + +class TestConfigParameter(unittest.TestCase): + def test_frozen(self): + param = ConfigParameter(name="lr", type="float", default=5e-6) + with self.assertRaises(AttributeError): + param.name = "other" + + def test_defaults(self): + param = ConfigParameter(name="lr", type="float", default=5e-6) + self.assertIsNone(param.description) + self.assertIsNone(param.min) + self.assertIsNone(param.max) + self.assertIsNone(param.enum) + self.assertFalse(param.required) + + def test_with_constraints(self): + param = ConfigParameter( + name="lr", + type="float", + default=5e-6, + description="Learning rate", + min=1e-6, + max=1e-4, + ) + self.assertEqual(param.min, 1e-6) + self.assertEqual(param.max, 1e-4) + self.assertEqual(param.description, "Learning rate") + + def test_with_enum(self): + param = ConfigParameter( + name="reasoning_effort", + type="string", + default="medium", + enum=("low", "medium", "high"), + ) + self.assertEqual(param.enum, ("low", "medium", "high")) + + def test_equality(self): + a = ConfigParameter(name="lr", type="float", default=5e-6) + b = ConfigParameter(name="lr", type="float", default=5e-6) + self.assertEqual(a, b) + + def test_hashable(self): + param = ConfigParameter(name="lr", type="float", default=5e-6) + self.assertIsInstance(hash(param), int) + + +class TestRecipeConfig(unittest.TestCase): + def setUp(self): + self.params = ( + ConfigParameter(name="lr", type="float", default=5e-6, min=1e-6, max=1e-4), + ConfigParameter(name="max_epochs", type="integer", default=2, min=1, max=5), + ConfigParameter( + name="reasoning_effort", + type="string", + default="medium", + enum=("low", "medium", "high"), + ), + ) + self.config = RecipeConfig( + model=Model.NOVA_LITE, + method=TrainingMethod.SFT_LORA, + platform=Platform.SMTJ, + parameters=self.params, + ) + + def test_frozen(self): + with self.assertRaises(AttributeError): + self.config.model = Model.NOVA_PRO + + def test_to_dict(self): + result = self.config.to_dict() + self.assertEqual(result, {"lr": 5e-6, "max_epochs": 2, "reasoning_effort": "medium"}) + + def test_repr_contains_model_and_params(self): + text = repr(self.config) + self.assertIn("NOVA_LITE", text) + self.assertIn("SFT_LORA", text) + self.assertIn("SMTJ", text) + self.assertIn("lr: float = 5e-06", text) + self.assertIn("[1e-06..0.0001]", text) + self.assertIn("max_epochs: integer = 2", text) + self.assertIn("enum=['low', 'medium', 'high']", text) + + def test_empty_parameters(self): + config = RecipeConfig( + model=Model.NOVA_LITE, + method=TrainingMethod.SFT_LORA, + platform=Platform.SMTJ, + parameters=(), + ) + self.assertEqual(config.to_dict(), {}) + self.assertIn("NOVA_LITE", repr(config)) + + def test_equality(self): + other = RecipeConfig( + model=Model.NOVA_LITE, + method=TrainingMethod.SFT_LORA, + platform=Platform.SMTJ, + parameters=self.params, + ) + self.assertEqual(self.config, other) + + def test_hashable(self): + self.assertIsInstance(hash(self.config), int) diff --git a/tests/unit/core/test_data_mixing_config.py b/tests/unit/core/test_data_mixing_config.py new file mode 100644 index 0000000..47bcfc2 --- /dev/null +++ b/tests/unit/core/test_data_mixing_config.py @@ -0,0 +1,289 @@ +# Copyright Amazon.com, Inc. or its affiliates + +# 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. +"""Unit tests for amzn_nova_forge.core.data_mixing_config (TypedDict version).""" + +import unittest +from typing import get_type_hints +from unittest.mock import patch + +from amzn_nova_forge.core.data_mixing_config import DataMixingConfig +from amzn_nova_forge.util.data_mixing import DataMixing + + +class TestDataMixingConfigTypedDict(unittest.TestCase): + """Tests that DataMixingConfig is importable and behaves as a TypedDict.""" + + def test_importable(self): + """DataMixingConfig is importable from core.data_mixing_config.""" + self.assertIsNotNone(DataMixingConfig) + + def test_plain_dict_satisfies_type(self): + """A plain dict with correct keys satisfies DataMixingConfig at runtime.""" + config: DataMixingConfig = { + "customer_data_percent": 50.0, + "nova_code_percent": 50.0, + "nova_general_percent": 25.0, + "nova_math_percent": 25.0, + } + self.assertEqual(config["customer_data_percent"], 50.0) + self.assertEqual(config["nova_code_percent"], 50.0) + + def test_subset_of_keys_works(self): + """TypedDict with total=False allows partial construction.""" + config: DataMixingConfig = {"customer_data_percent": 100.0} + self.assertEqual(config["customer_data_percent"], 100.0) + + def test_all_declared_keys_constructible(self): + """All declared keys can be set in a dict literal.""" + config: DataMixingConfig = { + "customer_data_percent": 50.0, + "nova_code_percent": 30.0, + "nova_general_percent": 40.0, + "nova_math_percent": 20.0, + "nova_image_percent": 10.0, + } + self.assertEqual(config["nova_image_percent"], 10.0) + + def test_type_hints_available(self): + """get_type_hints exposes the declared fields.""" + hints = get_type_hints(DataMixingConfig) + self.assertIn("customer_data_percent", hints) + self.assertIn("nova_code_percent", hints) + self.assertIn("nova_general_percent", hints) + self.assertIn("nova_math_percent", hints) + self.assertIn("nova_image_percent", hints) + + +class TestDataMixingConfigIntegration(unittest.TestCase): + """Integration tests: DataMixingConfig dicts flow through DataMixing.set_config().""" + + def _make_dm(self, default_fields=None): + """Helper to create a DataMixing with default_nova_fields set.""" + dm = DataMixing() + if default_fields is not None: + dm._default_nova_fields = default_fields + else: + dm._default_nova_fields = { + "nova_code_percent", + "nova_general_percent", + "nova_math_percent", + "customer_data_percent", + } + return dm + + def test_valid_config_flows_through_set_config(self): + """A valid DataMixingConfig dict is accepted by set_config.""" + dm = self._make_dm() + config: DataMixingConfig = { + "customer_data_percent": 50.0, + "nova_code_percent": 50.0, + "nova_general_percent": 25.0, + "nova_math_percent": 25.0, + } + dm.set_config(config, normalize=False) + self.assertEqual(dm.config["customer_data_percent"], 50.0) + self.assertEqual(dm.config["nova_code_percent"], 50.0) + + def test_range_violation_raises_valueerror(self): + """A field outside 0-100 raises ValueError with the field name.""" + dm = self._make_dm() + config = {"customer_data_percent": 150.0, "nova_code_percent": 100.0} + with self.assertRaises(ValueError) as ctx: + dm.set_config(config, normalize=False) + self.assertIn("customer_data_percent", str(ctx.exception)) + + def test_negative_range_raises_valueerror(self): + """A negative field raises ValueError with the field name.""" + dm = self._make_dm() + config = {"nova_code_percent": -5.0} + with self.assertRaises(ValueError) as ctx: + dm.set_config(config, normalize=False) + self.assertIn("nova_code_percent", str(ctx.exception)) + + def test_sum_violation_raises_valueerror(self): + """Nova fields not summing to 100 raises ValueError.""" + dm = self._make_dm() + config = { + "customer_data_percent": 50.0, + "nova_code_percent": 20.0, + "nova_general_percent": 20.0, + "nova_math_percent": 20.0, + } + with self.assertRaises(ValueError) as ctx: + dm.set_config(config, normalize=False) + msg = str(ctx.exception) + self.assertIn("sum", msg.lower()) + + def test_customer_100_with_nova_nonzero_raises_valueerror(self): + """customer=100 with nova fields > 0 raises ValueError.""" + dm = self._make_dm() + config = { + "customer_data_percent": 100.0, + "nova_code_percent": 50.0, + "nova_general_percent": 50.0, + } + with self.assertRaises(ValueError) as ctx: + dm.set_config(config, normalize=False) + msg = str(ctx.exception) + self.assertIn("0", msg) + + def test_all_zero_nova_with_customer_below_100_raises_valueerror(self): + """All nova fields = 0 with customer < 100 raises ValueError.""" + dm = self._make_dm() + config = { + "customer_data_percent": 50.0, + "nova_code_percent": 0.0, + "nova_general_percent": 0.0, + "nova_math_percent": 0.0, + } + with self.assertRaises(ValueError) as ctx: + dm.set_config(config, normalize=False) + msg = str(ctx.exception) + self.assertIn("100", msg) + + def test_extra_fields_accepted(self): + """Unknown nova_*_percent fields are accepted.""" + dm = self._make_dm( + default_fields={ + "nova_code_percent", + "nova_general_percent", + "nova_math_percent", + "nova_video_percent", + "customer_data_percent", + } + ) + config = { + "customer_data_percent": 50.0, + "nova_code_percent": 50.0, + "nova_general_percent": 25.0, + "nova_math_percent": 25.0, + "nova_video_percent": 0.0, + } + dm.set_config(config, normalize=False) + self.assertEqual(dm.config["nova_video_percent"], 0.0) + + def test_extra_fields_round_trip(self): + """Extra nova fields survive the set_config round-trip.""" + dm = self._make_dm( + default_fields={ + "nova_code_percent", + "nova_general_percent", + "nova_math_percent", + "nova_video_percent", + "customer_data_percent", + } + ) + config = { + "customer_data_percent": 50.0, + "nova_code_percent": 50.0, + "nova_general_percent": 25.0, + "nova_math_percent": 25.0, + "nova_video_percent": 0.0, + } + dm.set_config(config, normalize=False) + self.assertIn("nova_video_percent", dm.config) + self.assertEqual(dm.config["nova_video_percent"], 0.0) + + def test_dataset_catalog_stripped_with_warning(self): + """dataset_catalog key is stripped and a warning is logged.""" + dm = self._make_dm() + config = { + "nova_code_percent": 100.0, + "customer_data_percent": 50.0, + "dataset_catalog": "some_catalog", + } + with patch("amzn_nova_forge.util.data_mixing.logger") as mock_logger: + dm.set_config(config, normalize=False) + mock_logger.warning.assert_called_once() + self.assertNotIn("dataset_catalog", dm.config) + + def test_none_values_filtered_out(self): + """None values are filtered out of the stored config.""" + dm = self._make_dm() + config = { + "nova_code_percent": 100.0, + "nova_general_percent": None, + "customer_data_percent": 50.0, + } + dm.set_config(config, normalize=False) + self.assertNotIn("nova_general_percent", dm.config) + self.assertEqual(dm.config["nova_code_percent"], 100.0) + + def test_multiple_errors_surfaced(self): + """Multiple validation errors appear in a single ValueError message.""" + dm = self._make_dm() + config = { + "customer_data_percent": 150.0, + "nova_code_percent": -10.0, + } + with self.assertRaises(ValueError) as ctx: + dm.set_config(config, normalize=False) + msg = str(ctx.exception) + self.assertIn("customer_data_percent", msg) + self.assertIn("nova_code_percent", msg) + + def test_string_value_raises_valueerror(self): + """String value for a percent field raises ValueError with type info.""" + dm = self._make_dm() + config = {"customer_data_percent": "999"} + with self.assertRaises(ValueError) as ctx: + dm.set_config(config, normalize=False) + msg = str(ctx.exception) + self.assertIn("customer_data_percent", msg) + self.assertIn("must be a number", msg) + self.assertIn("str", msg) + + def test_float_tolerance(self): + """Sum within +/-0.01 of 100 should pass.""" + dm = self._make_dm() + config = { + "customer_data_percent": 50.0, + "nova_code_percent": 33.33, + "nova_general_percent": 33.33, + "nova_math_percent": 33.34, + } + # Should not raise + dm.set_config(config, normalize=False) + + def test_customer_100_no_nova_valid(self): + """customer_data_percent=100 with no nova fields is valid.""" + dm = self._make_dm() + config = {"customer_data_percent": 100.0} + dm.set_config(config, normalize=False) + + def test_integer_inputs_accepted(self): + """Integer values are accepted for percent fields.""" + dm = self._make_dm() + config = { + "customer_data_percent": 50, + "nova_code_percent": 50, + "nova_general_percent": 25, + "nova_math_percent": 25, + } + dm.set_config(config, normalize=False) + + def test_customer_zero_with_nova_sum_100(self): + """customer_data_percent=0 with nova summing to 100 is valid.""" + dm = self._make_dm() + config = { + "customer_data_percent": 0.0, + "nova_code_percent": 50.0, + "nova_general_percent": 25.0, + "nova_math_percent": 25.0, + } + dm.set_config(config, normalize=False) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/core/test_training_overrides.py b/tests/unit/core/test_training_overrides.py new file mode 100644 index 0000000..4c1e64a --- /dev/null +++ b/tests/unit/core/test_training_overrides.py @@ -0,0 +1,77 @@ +# Copyright Amazon.com, Inc. or its affiliates + +# 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. +"""Unit tests for amzn_nova_forge.core.training_overrides.""" + +import unittest +from typing import get_type_hints + +from amzn_nova_forge.core.training_overrides import TrainingOverrides + + +class TestTrainingOverrides(unittest.TestCase): + def test_plain_dict_compatible(self): + """A plain dict with known keys is a valid TrainingOverrides at runtime.""" + overrides: TrainingOverrides = {"lr": 0.001, "max_steps": 100} + assert overrides["lr"] == 0.001 + assert overrides["max_steps"] == 100 + + def test_subset_of_keys(self): + """Only a subset of keys is needed (total=False).""" + overrides: TrainingOverrides = {"global_batch_size": 8} + assert overrides["global_batch_size"] == 8 + + def test_all_keys_constructible(self): + """All declared keys can be set.""" + overrides: TrainingOverrides = { + "lr": 0.001, + "loraplus_lr_ratio": 4.0, + "lora_plus_lr_ratio": 4.0, + "max_steps": 1000, + "max_epochs": 3, + "save_steps": 500, + "warmup_steps": 100, + "global_batch_size": 16, + "max_length": 4096, + "alpha": 32, + "beta": 0.1, + "reasoning_enabled": True, + "reasoning_effort": "medium", + "val_check_interval": 50, + "validation_s3_path": "s3://bucket/val", + "validation_data_s3_path": "s3://bucket/val_data", + "top_logprobs": 5, + } + assert len(overrides) == 17 + + def test_unknown_keys_no_runtime_error(self): + """TypedDict does not reject unknown keys at runtime.""" + overrides: TrainingOverrides = {"lr": 0.001, "some_future_key": "value"} # type: ignore[typeddict-unknown-key] + assert "some_future_key" in overrides + + def test_empty_dict_valid(self): + """Empty dict is valid since total=False.""" + overrides: TrainingOverrides = {} + assert len(overrides) == 0 + + def test_type_hints_available(self): + """TypedDict exposes its annotations for introspection.""" + hints = get_type_hints(TrainingOverrides) + assert "lr" in hints + assert "max_steps" in hints + assert hints["lr"] is float + assert hints["max_steps"] is int + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/deployer/test_forge_deployer.py b/tests/unit/deployer/test_forge_deployer.py index 3fbcb00..398c1fc 100644 --- a/tests/unit/deployer/test_forge_deployer.py +++ b/tests/unit/deployer/test_forge_deployer.py @@ -813,11 +813,12 @@ def test_get_logs_with_endpoint_arn_only(self, mock_check_status, mock_validate_ region="us-east-1", ) - @patch(f"{PATCH_PREFIX}.logger") - def test_get_logs_no_arn_logs_info(self, mock_logger, mock_validate_region): + def test_get_logs_no_arn_raises_value_error(self, mock_validate_region): deployer = self._make_deployer() - deployer.get_logs() - mock_logger.info.assert_called_once_with("No endpoint ARN available. Call deploy() first.") + with self.assertRaises(ValueError) as ctx: + deployer.get_logs() + + self.assertIn("endpoint_arn", str(ctx.exception)) class TestDeploymentResultRegion(unittest.TestCase): diff --git a/tests/unit/evaluator/test_forge_evaluator.py b/tests/unit/evaluator/test_forge_evaluator.py index 18cb262..4418fe0 100644 --- a/tests/unit/evaluator/test_forge_evaluator.py +++ b/tests/unit/evaluator/test_forge_evaluator.py @@ -707,12 +707,27 @@ def test_get_logs_with_job_result(self, mock_monitor_cls): ) @patch("amzn_nova_forge.evaluator.forge_evaluator.CloudWatchLogMonitor") - @patch("amzn_nova_forge.evaluator.forge_evaluator.logger") - def test_get_logs_missing_params_logs_info(self, mock_logger, mock_monitor_cls): - self.evaluator.get_logs() + def test_get_logs_missing_params_raises_value_error(self, mock_monitor_cls): + with self.assertRaises(ValueError) as ctx: + self.evaluator.get_logs() + + self.assertIn("job_result", str(ctx.exception)) + mock_monitor_cls.assert_not_called() + + @patch("amzn_nova_forge.evaluator.forge_evaluator.CloudWatchLogMonitor") + def test_get_logs_job_id_only_raises_value_error(self, mock_monitor_cls): + with self.assertRaises(ValueError) as ctx: + self.evaluator.get_logs(job_id="some-job") + + self.assertIn("job_result", str(ctx.exception)) + mock_monitor_cls.assert_not_called() + + @patch("amzn_nova_forge.evaluator.forge_evaluator.CloudWatchLogMonitor") + def test_get_logs_started_time_only_raises_value_error(self, mock_monitor_cls): + with self.assertRaises(ValueError) as ctx: + self.evaluator.get_logs(started_time=datetime(2025, 1, 1, tzinfo=timezone.utc)) - mock_logger.info.assert_called_once() - self.assertIn("job_result", mock_logger.info.call_args[0][0]) + self.assertIn("job_result", str(ctx.exception)) mock_monitor_cls.assert_not_called() @patch("amzn_nova_forge.evaluator.forge_evaluator.CloudWatchLogMonitor") diff --git a/tests/unit/inference/test_forge_inference.py b/tests/unit/inference/test_forge_inference.py index 89afd9a..da65cc5 100644 --- a/tests/unit/inference/test_forge_inference.py +++ b/tests/unit/inference/test_forge_inference.py @@ -541,30 +541,42 @@ def test_get_logs_with_job_result(self, mock_monitor_cls): limit=None, start_from_head=False, end_time=None ) - @patch("amzn_nova_forge.inference.forge_inference.logger") - def test_get_logs_missing_params_logs_info(self, mock_logger): + def test_get_logs_missing_params_raises_value_error(self): inf = self._make_inference_with_platform(Platform.SMTJ) - inf.get_logs() # No job_result, no job_id, no started_time + with self.assertRaises(ValueError) as ctx: + inf.get_logs() # No job_result, no job_id, no started_time - mock_logger.info.assert_called_once_with( - "Provide either a job_result or explicit job_id and started_time." - ) + self.assertIn("job_result", str(ctx.exception)) + + def test_get_logs_job_id_only_raises_value_error(self): + inf = self._make_inference_with_platform(Platform.SMTJ) + + with self.assertRaises(ValueError) as ctx: + inf.get_logs(job_id="some-job") + + self.assertIn("job_result", str(ctx.exception)) + + def test_get_logs_started_time_only_raises_value_error(self): + inf = self._make_inference_with_platform(Platform.SMTJ) + + with self.assertRaises(ValueError) as ctx: + inf.get_logs(started_time=datetime(2024, 1, 1, tzinfo=timezone.utc)) + + self.assertIn("job_result", str(ctx.exception)) - @patch("amzn_nova_forge.inference.forge_inference.logger") - def test_get_logs_no_platform_logs_info(self, mock_logger): + def test_get_logs_no_platform_raises_value_error(self): with patch("boto3.session.Session") as mock_session: type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") inf = ForgeInference(region="us-east-1") # No infra -> no platform - inf.get_logs( - job_id="job-xyz", - started_time=datetime(2024, 1, 1, tzinfo=timezone.utc), - ) + with self.assertRaises(ValueError) as ctx: + inf.get_logs( + job_id="job-xyz", + started_time=datetime(2024, 1, 1, tzinfo=timezone.utc), + ) - mock_logger.info.assert_called_once_with( - "Cannot determine platform — provide infra in the constructor." - ) + self.assertIn("platform", str(ctx.exception)) @patch("amzn_nova_forge.inference.forge_inference.CloudWatchLogMonitor") def test_get_logs_smhp_includes_cluster_and_namespace(self, mock_monitor_cls): diff --git a/tests/unit/manager/test_smtj_dataprep_runtime_manager.py b/tests/unit/manager/test_smtj_dataprep_runtime_manager.py index 797cf01..37886dd 100644 --- a/tests/unit/manager/test_smtj_dataprep_runtime_manager.py +++ b/tests/unit/manager/test_smtj_dataprep_runtime_manager.py @@ -798,5 +798,55 @@ def test_put_object_failure_does_not_raise(self): mock_s3.put_object.assert_called_once() +class TestSmtjEntryScriptWheelSafety(unittest.TestCase): + """SMTJ entry script rejects wheels with unsafe paths.""" + + def _exec_install_deps_with_whl(self, whl_names): + """Run _install_deps() from the embedded entry script against a wheel + containing ``whl_names`` as members. Returns (deps_dir, target_dir).""" + import zipfile + + deps_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, deps_dir) + target_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, target_dir) + + whl_path = os.path.join(deps_dir, "test.whl") + with zipfile.ZipFile(whl_path, "w") as zf: + for name in whl_names: + zf.writestr(name, b"payload") + + # Redirect the hardcoded deps path and site-packages to temp dirs. + script_src = SMTJDataPrepRuntimeManager._SMTJ_ENTRY_SCRIPT + patched = script_src.replace('"/opt/ml/input/data/deps"', repr(deps_dir)) + + # __name__ != "__main__" so exec() only defines functions. + script_globals = {"__name__": "__not_main__"} + with patch("site.getsitepackages", return_value=[target_dir]): + exec(compile(patched, "", "exec"), script_globals) + script_globals["_install_deps"](None) + + return deps_dir, target_dir + + def test_rejects_path_traversal(self): + """Wheel containing '../escaped.txt' raises before extraction.""" + with self.assertRaises(RuntimeError) as ctx: + self._exec_install_deps_with_whl(["../escaped.txt"]) + self.assertIn("Unsafe path traversal", str(ctx.exception)) + + def test_rejects_absolute_path(self): + """Wheel containing an absolute path raises before extraction.""" + with self.assertRaises(RuntimeError) as ctx: + self._exec_install_deps_with_whl(["/etc/passwd"]) + self.assertIn("Unsafe absolute path", str(ctx.exception)) + + def test_accepts_normal_wheel_contents(self): + """Wheel with ordinary relative paths extracts without error.""" + _, target_dir = self._exec_install_deps_with_whl( + ["foo/__init__.py", "foo.dist-info/METADATA"] + ) + self.assertTrue(os.path.exists(os.path.join(target_dir, "foo", "__init__.py"))) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/model/test_nova_model_customizer.py b/tests/unit/model/test_nova_model_customizer.py index 4b15aba..bfa55c0 100644 --- a/tests/unit/model/test_nova_model_customizer.py +++ b/tests/unit/model/test_nova_model_customizer.py @@ -1030,6 +1030,12 @@ def test_is_multimodal_preserved_across_model_change(self, mock_client): # is_multimodal was not touched by model setter self.assertTrue(customizer.is_multimodal) + def test_get_logs_without_job_state_raises_value_error(self): + with self.assertRaises(ValueError) as ctx: + self.customizer.get_logs() + + self.assertIn(".train()", str(ctx.exception)) + class TestTrain(TestNovaModelCustomizer): @patch("amzn_nova_forge.recipe.recipe_builder.RecipeBuilder.build_and_validate") @@ -4746,5 +4752,100 @@ def test_smtj_non_serverless_still_raises(self, mock_client): self.assertIn("Data mixing is only supported for", str(ctx.exception)) +class TestNovaModelCustomizerGetConfig(unittest.TestCase): + def setUp(self): + self.model = Model.NOVA_MICRO + self.method = TrainingMethod.SFT_LORA + self.data_s3_path = "s3://test-bucket/data" + self.output_s3_path = "s3://test-bucket/output" + + self.mock_runtime_manager = create_autospec(SMTJRuntimeManager) + self.mock_runtime_manager.platform = Platform.SMTJ + self.mock_runtime_manager.instance_count = 2 + self.mock_runtime_manager.kms_key_id = None + + p_session = patch("boto3.session.Session") + mock_session = p_session.start() + self.addCleanup(p_session.stop) + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + + p_set_output = patch( + "amzn_nova_forge.model.nova_model_customizer.set_output_s3_path", + return_value="s3://sagemaker-nova-123456789012-us-east-1/output", + ) + p_set_output.start() + self.addCleanup(p_set_output.stop) + + def _make_customizer(self): + return NovaModelCustomizer( + model=self.model, + method=self.method, + infra=self.mock_runtime_manager, + data_s3_path=self.data_s3_path, + output_s3_path=self.output_s3_path, + ) + + @patch("amzn_nova_forge.trainer.forge_trainer.ForgeTrainer.get_config") + @patch("amzn_nova_forge.trainer.forge_trainer.ForgeTrainer.__init__", return_value=None) + def test_get_config_emits_deprecation_warning(self, mock_init, mock_get_config): + from amzn_nova_forge.core.types import RecipeConfig + + mock_get_config.return_value = RecipeConfig( + model=self.model, + method=self.method, + platform=Platform.SMTJ, + parameters=(), + ) + + customizer = self._make_customizer() + with self.assertWarns(DeprecationWarning): + customizer.get_config() + + @patch("amzn_nova_forge.trainer.forge_trainer.ForgeTrainer.get_config") + @patch("amzn_nova_forge.trainer.forge_trainer.ForgeTrainer.__init__", return_value=None) + def test_get_config_delegates_to_forge_trainer(self, mock_init, mock_get_config): + from amzn_nova_forge.core.types import ConfigParameter, RecipeConfig + + expected = RecipeConfig( + model=self.model, + method=self.method, + platform=Platform.SMTJ, + parameters=(ConfigParameter(name="lr", type="float", default=5e-6),), + ) + mock_get_config.return_value = expected + + customizer = self._make_customizer() + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + result = customizer.get_config() + + self.assertEqual(result, expected) + mock_get_config.assert_called_once_with(overrides=None) + + @patch("amzn_nova_forge.trainer.forge_trainer.ForgeTrainer.get_config") + @patch("amzn_nova_forge.trainer.forge_trainer.ForgeTrainer.__init__", return_value=None) + def test_get_config_passes_overrides(self, mock_init, mock_get_config): + from amzn_nova_forge.core.types import RecipeConfig + + mock_get_config.return_value = RecipeConfig( + model=self.model, + method=self.method, + platform=Platform.SMTJ, + parameters=(), + ) + + customizer = self._make_customizer() + overrides = {"lr": 1e-5} + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + customizer.get_config(overrides=overrides) + + mock_get_config.assert_called_once_with(overrides=overrides) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/recipe/test_recipe_builder.py b/tests/unit/recipe/test_recipe_builder.py index 13ffd42..82fcb0f 100644 --- a/tests/unit/recipe/test_recipe_builder.py +++ b/tests/unit/recipe/test_recipe_builder.py @@ -4151,5 +4151,360 @@ def test_mlflow_placeholders_not_raw_when_no_monitor( self.assertNotIn("{{mlflow_tracking_uri}}", raw) +class TestGetOverridableConfig(unittest.TestCase): + def setUp(self): + self.region = "us-east-1" + self.platform = Platform.SMTJ + self.method = TrainingMethod.SFT_LORA + + self.mock_model = Mock(spec=Model) + self.mock_model.name = "nova-micro" + self.mock_model.value = "nova_micro" + self.mock_model.version = Version.ONE + self.mock_model.model_type = "test-model" + self.mock_model.model_path = "models/test" + + self.mock_infra = Mock(spec=RuntimeManager) + self.mock_infra.instance_type = "ml.g5.12xlarge" + self.mock_infra.instance_count = 1 + + self.overrides_template = { + "lr": { + "type": "float", + "default": 5e-6, + "min": 1e-6, + "max": 1e-4, + "description": "Learning rate", + }, + "max_epochs": { + "type": "integer", + "default": 2, + "min": 1, + "max": 5, + "description": "Max epochs", + }, + "reasoning_effort": { + "type": "string", + "default": "medium", + "enum": ["low", "medium", "high"], + }, + "model_type": {"type": "string", "default": "test-model"}, + "model_name_or_path": {"type": "string", "default": "models/test"}, + "peft_scheme": {"type": "string", "default": "lora"}, + "name": {"type": "string", "default": "job"}, + "data_s3_path": {"type": "string", "default": "s3://bucket/data"}, + "output_s3_path": {"type": "string", "default": "s3://bucket/output"}, + "task": {"type": "string", "default": "sft"}, + } + self.recipe_template = {"run": {"lr": "{{lr}}", "max_epochs": "{{max_epochs}}"}} + + self.load_patcher = patch("amzn_nova_forge.recipe.recipe_builder.load_recipe_templates") + self.mock_load = self.load_patcher.start() + self.mock_load.return_value = ( + None, + self.recipe_template, + self.overrides_template, + "test-image-uri", + ) + self.addCleanup(self.load_patcher.stop) + + def _make_builder(self, **kwargs): + defaults = dict( + region=self.region, + job_name="__config_preview__", + platform=self.platform, + model=self.mock_model, + method=self.method, + instance_type="ml.g5.12xlarge", + instance_count=1, + infra=self.mock_infra, + output_s3_path="s3://bucket/output", + data_s3_path="s3://bucket/data", + ) + defaults.update(kwargs) + return RecipeBuilder(**defaults) + + def test_returns_recipe_config(self): + from amzn_nova_forge.core.types import RecipeConfig + + builder = self._make_builder() + config = builder.get_overridable_config() + self.assertIsInstance(config, RecipeConfig) + self.assertEqual(config.model, self.mock_model) + self.assertEqual(config.method, self.method) + self.assertEqual(config.platform, self.platform) + + def test_excludes_non_overridable_keys(self): + builder = self._make_builder() + config = builder.get_overridable_config() + param_names = {p.name for p in config.parameters} + for excluded in ( + "model_type", + "model_name_or_path", + "peft_scheme", + "name", + "data_s3_path", + "output_s3_path", + "task", + ): + self.assertNotIn(excluded, param_names) + + def test_includes_overridable_keys(self): + builder = self._make_builder() + config = builder.get_overridable_config() + param_names = {p.name for p in config.parameters} + self.assertIn("lr", param_names) + self.assertIn("max_epochs", param_names) + self.assertIn("reasoning_effort", param_names) + + def test_preserves_constraints(self): + builder = self._make_builder() + config = builder.get_overridable_config() + lr_param = next(p for p in config.parameters if p.name == "lr") + self.assertEqual(lr_param.min, 1e-6) + self.assertEqual(lr_param.max, 1e-4) + self.assertEqual(lr_param.description, "Learning rate") + + def test_preserves_enum(self): + builder = self._make_builder() + config = builder.get_overridable_config() + effort_param = next(p for p in config.parameters if p.name == "reasoning_effort") + self.assertEqual(effort_param.enum, ("low", "medium", "high")) + + def test_with_overrides_merges_into_defaults(self): + builder = self._make_builder() + config = builder.get_overridable_config(overrides={"lr": 1e-5}) + lr_param = next(p for p in config.parameters if p.name == "lr") + self.assertEqual(lr_param.default, 1e-5) + + def test_skips_non_dict_metadata(self): + self.overrides_template["stray_string"] = "not a dict" + self.overrides_template["stray_int"] = 42 + builder = self._make_builder() + config = builder.get_overridable_config() + param_names = {p.name for p in config.parameters} + self.assertNotIn("stray_string", param_names) + self.assertNotIn("stray_int", param_names) + + def test_config_is_frozen(self): + builder = self._make_builder() + config = builder.get_overridable_config() + with self.assertRaises(AttributeError): + config.model = self.mock_model + + +class TestReasoningEffortNoneOverride(unittest.TestCase): + """Tests for reasoning_effort=None override handling.""" + + def setUp(self): + self.region = "us-east-1" + self.job_name = "test-reasoning-none" + self.platform = Platform.SMTJ + self.instance_type = "ml.p5.48xlarge" + self.instance_count = 4 + self.data_s3 = "s3://bucket/rft-data.jsonl" + self.output_s3 = "s3://bucket/output" + + self.mock_model = Mock(spec=Model) + self.mock_model.name = "NOVA_LITE_2" + self.mock_model.value = "nova_lite_2" + self.mock_model.version = Version.TWO + self.mock_model.model_type = "nova_lite_2" + self.mock_model.model_path = "models/nova-lite-2" + + self.mock_infra = Mock(spec=RuntimeManager) + self.mock_infra.instance_type = self.instance_type + self.mock_infra.instance_count = self.instance_count + self.mock_infra.platform = Platform.SMTJ + + @patch("amzn_nova_forge.util.recipe._get_smhp_replicas_enum", return_value=None) + @patch("amzn_nova_forge.util.recipe.get_hub_recipe_metadata") + @patch("amzn_nova_forge.util.recipe.download_templates_from_s3") + def test_reasoning_effort_none_passes_validation( + self, mock_download, mock_metadata, _mock_replicas + ): + """reasoning_effort=None override must not raise ValueError during build_and_validate.""" + mock_metadata.return_value = {"recipe_uri": "s3://bucket/recipe"} + + recipe_template = { + "run": { + "name": "{{name}}", + "replicas": 1, + "lambda_arn": "{{reward_lambda_arn}}", + }, + "training_config": { + "max_steps": "{{max_steps}}", + "lr": "{{lr}}", + "reasoning_effort": "medium", + }, + } + + overrides_template = { + "name": {"default": "", "type": "string", "required": True}, + "reward_lambda_arn": {"default": "", "type": "string", "required": True}, + "max_steps": {"default": 100, "type": "integer", "min": 1, "max": 10000}, + "lr": {"default": 5e-6, "type": "float"}, + "reasoning_effort": { + "type": "string", + "enum": ["low", "medium", "high"], + "default": "medium", + }, + } + + mock_download.return_value = (recipe_template, overrides_template, "image_uri") + + rft_lambda = "arn:aws:lambda:us-east-1:123456789012:function:reward" + builder = RecipeBuilder( + region=self.region, + job_name=self.job_name, + platform=self.platform, + model=self.mock_model, + method=TrainingMethod.RFT_LORA, + instance_type=self.instance_type, + instance_count=self.instance_count, + infra=self.mock_infra, + output_s3_path=self.output_s3, + data_s3_path=self.data_s3, + rft_lambda_arn=rft_lambda, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + output_path = os.path.join(tmpdir, "recipe.yaml") + recipe_path, *_ = builder.build_and_validate( + output_recipe_path=output_path, + overrides={"reasoning_effort": None}, + validation_config=ValidationConfig(iam=False, infra=False, recipe=True), + ) + + with open(recipe_path, "r") as f: + recipe = yaml.safe_load(f) + + assert recipe["training_config"]["reasoning_effort"] is None + + @patch("amzn_nova_forge.util.recipe._get_smhp_replicas_enum", return_value=None) + @patch("amzn_nova_forge.util.recipe.get_hub_recipe_metadata") + @patch("amzn_nova_forge.util.recipe.download_templates_from_s3") + def test_reasoning_effort_none_placeholder_path_still_validates( + self, mock_download, mock_metadata, _mock_replicas + ): + """reasoning_effort=None on a {{placeholder}}-same-key template still raises. + + The fix only targets the literal-value branch (line 575-585). When the + template uses {{reasoning_effort}} as a placeholder, the overrides_template + retains its original required: True, so None is correctly rejected. + """ + mock_metadata.return_value = {"recipe_uri": "s3://bucket/recipe"} + + recipe_template = { + "run": { + "name": "{{name}}", + "replicas": 1, + "lambda_arn": "{{reward_lambda_arn}}", + }, + "training_config": { + "max_steps": "{{max_steps}}", + "lr": "{{lr}}", + "reasoning_effort": "{{reasoning_effort}}", + }, + } + + overrides_template = { + "name": {"default": "", "type": "string", "required": True}, + "reward_lambda_arn": {"default": "", "type": "string", "required": True}, + "max_steps": {"default": 100, "type": "integer", "min": 1, "max": 10000}, + "lr": {"default": 5e-6, "type": "float"}, + "reasoning_effort": { + "type": "string", + "required": True, + "enum": ["low", "medium", "high"], + "default": "medium", + }, + } + + mock_download.return_value = (recipe_template, overrides_template, "image_uri") + + rft_lambda = "arn:aws:lambda:us-east-1:123456789012:function:reward" + builder = RecipeBuilder( + region=self.region, + job_name=self.job_name, + platform=self.platform, + model=self.mock_model, + method=TrainingMethod.RFT_LORA, + instance_type=self.instance_type, + instance_count=self.instance_count, + infra=self.mock_infra, + output_s3_path=self.output_s3, + data_s3_path=self.data_s3, + rft_lambda_arn=rft_lambda, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + output_path = os.path.join(tmpdir, "recipe.yaml") + with self.assertRaises(ValueError) as ctx: + builder.build_and_validate( + output_recipe_path=output_path, + overrides={"reasoning_effort": None}, + validation_config=ValidationConfig(iam=False, infra=False, recipe=True), + ) + self.assertIn("expects string", str(ctx.exception)) + self.assertIn("NoneType", str(ctx.exception)) + + @patch("amzn_nova_forge.util.recipe._get_smhp_replicas_enum", return_value=None) + @patch("amzn_nova_forge.util.recipe.get_hub_recipe_metadata") + @patch("amzn_nova_forge.util.recipe.download_templates_from_s3") + def test_non_nullable_field_none_still_raises( + self, mock_download, mock_metadata, _mock_replicas + ): + """max_steps=None must still raise — only fields typed as X | None are nullable.""" + mock_metadata.return_value = {"recipe_uri": "s3://bucket/recipe"} + + recipe_template = { + "run": { + "name": "{{name}}", + "replicas": 1, + "lambda_arn": "{{reward_lambda_arn}}", + }, + "training_config": { + "max_steps": 100, + "lr": "{{lr}}", + }, + } + + overrides_template = { + "name": {"default": "", "type": "string", "required": True}, + "reward_lambda_arn": {"default": "", "type": "string", "required": True}, + "max_steps": {"default": 100, "type": "integer", "min": 1, "max": 10000}, + "lr": {"default": 5e-6, "type": "float"}, + } + + mock_download.return_value = (recipe_template, overrides_template, "image_uri") + + rft_lambda = "arn:aws:lambda:us-east-1:123456789012:function:reward" + builder = RecipeBuilder( + region=self.region, + job_name=self.job_name, + platform=self.platform, + model=self.mock_model, + method=TrainingMethod.RFT_LORA, + instance_type=self.instance_type, + instance_count=self.instance_count, + infra=self.mock_infra, + output_s3_path=self.output_s3, + data_s3_path=self.data_s3, + rft_lambda_arn=rft_lambda, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + output_path = os.path.join(tmpdir, "recipe.yaml") + with self.assertRaises(ValueError) as ctx: + builder.build_and_validate( + output_recipe_path=output_path, + overrides={"max_steps": None}, + validation_config=ValidationConfig(iam=False, infra=False, recipe=True), + ) + self.assertIn("expects int", str(ctx.exception)) + self.assertIn("NoneType", str(ctx.exception)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/trainer/test_forge_trainer.py b/tests/unit/trainer/test_forge_trainer.py index 7b59650..256c20f 100644 --- a/tests/unit/trainer/test_forge_trainer.py +++ b/tests/unit/trainer/test_forge_trainer.py @@ -685,13 +685,33 @@ def test_get_logs_with_explicit_ids(self, MockMonitor): ) @patch("amzn_nova_forge.trainer.forge_trainer.CloudWatchLogMonitor") - def test_get_logs_missing_info_returns_early(self, MockMonitor): + def test_get_logs_missing_info_raises_value_error(self, MockMonitor): trainer = self._make_trainer() - with patch("amzn_nova_forge.trainer.forge_trainer.logger") as mock_logger: + with self.assertRaises(ValueError) as ctx: trainer.get_logs() - mock_logger.info.assert_called_once() + self.assertIn("job_result", str(ctx.exception)) + MockMonitor.assert_not_called() + + @patch("amzn_nova_forge.trainer.forge_trainer.CloudWatchLogMonitor") + def test_get_logs_job_id_only_raises_value_error(self, MockMonitor): + trainer = self._make_trainer() + + with self.assertRaises(ValueError) as ctx: + trainer.get_logs(job_id="some-job") + + self.assertIn("job_result", str(ctx.exception)) + MockMonitor.assert_not_called() + + @patch("amzn_nova_forge.trainer.forge_trainer.CloudWatchLogMonitor") + def test_get_logs_started_time_only_raises_value_error(self, MockMonitor): + trainer = self._make_trainer() + + with self.assertRaises(ValueError) as ctx: + trainer.get_logs(started_time=datetime(2025, 1, 1, tzinfo=timezone.utc)) + + self.assertIn("job_result", str(ctx.exception)) MockMonitor.assert_not_called() @patch("amzn_nova_forge.trainer.forge_trainer.CloudWatchLogMonitor") @@ -1130,5 +1150,104 @@ def test_val_check_interval_none_no_hyperparameters( self.assertIsNone(job_config.trainer_config_hyperparameters) +class TestForgeTrainerGetConfig(unittest.TestCase): + """Tests for ForgeTrainer.get_config().""" + + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def _make_trainer(self, mock_session, mock_set_output, **kwargs): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + defaults = dict( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=_make_smtj_infra(), + training_data_s3_path="s3://bucket/data", + ) + defaults.update(kwargs) + return ForgeTrainer(**defaults) + + @patch("amzn_nova_forge.trainer.forge_trainer.RecipeBuilder") + def test_returns_recipe_config(self, mock_rb_cls): + from amzn_nova_forge.core.types import ConfigParameter, RecipeConfig + + expected = RecipeConfig( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + platform=Platform.SMTJ, + parameters=(ConfigParameter(name="lr", type="float", default=5e-6),), + ) + mock_rb_cls.return_value.get_overridable_config.return_value = expected + + trainer = self._make_trainer() + result = trainer.get_config() + + self.assertEqual(result, expected) + mock_rb_cls.return_value.get_overridable_config.assert_called_once_with(overrides=None) + + @patch("amzn_nova_forge.trainer.forge_trainer.RecipeBuilder") + def test_passes_overrides_through(self, mock_rb_cls): + from amzn_nova_forge.core.types import RecipeConfig + + mock_rb_cls.return_value.get_overridable_config.return_value = RecipeConfig( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + platform=Platform.SMTJ, + parameters=(), + ) + + trainer = self._make_trainer() + overrides = {"lr": 1e-5, "max_epochs": 3} + trainer.get_config(overrides=overrides) + + mock_rb_cls.return_value.get_overridable_config.assert_called_once_with(overrides=overrides) + + @patch("amzn_nova_forge.trainer.forge_trainer.RecipeBuilder") + def test_does_not_start_job(self, mock_rb_cls): + from amzn_nova_forge.core.types import RecipeConfig + + mock_rb_cls.return_value.get_overridable_config.return_value = RecipeConfig( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + platform=Platform.SMTJ, + parameters=(), + ) + + trainer = self._make_trainer() + trainer.get_config() + + trainer.infra.execute.assert_not_called() + + @patch("amzn_nova_forge.trainer.forge_trainer.RecipeBuilder") + def test_creates_recipe_builder_with_trainer_state(self, mock_rb_cls): + from amzn_nova_forge.core.types import RecipeConfig + + mock_rb_cls.return_value.get_overridable_config.return_value = RecipeConfig( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + platform=Platform.SMTJ, + parameters=(), + ) + + trainer = self._make_trainer() + trainer.get_config() + + call_kwargs = mock_rb_cls.call_args.kwargs + self.assertEqual(call_kwargs["model"], Model.NOVA_MICRO) + self.assertEqual(call_kwargs["method"], TrainingMethod.SFT_LORA) + self.assertEqual(call_kwargs["platform"], Platform.SMTJ) + self.assertEqual(call_kwargs["job_name"], "__config_preview__") + + def test_get_config_rft_multiturn_raises(self): + trainer = self._make_trainer( + model=Model.NOVA_LITE_2, method=TrainingMethod.RFT_MULTITURN_LORA + ) + with self.assertRaises(ValueError) as ctx: + trainer.get_config() + self.assertIn("rft_multiturn_infra", str(ctx.exception)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/util/test_data_mixing.py b/tests/unit/util/test_data_mixing.py index c133159..c79fa3b 100644 --- a/tests/unit/util/test_data_mixing.py +++ b/tests/unit/util/test_data_mixing.py @@ -15,10 +15,16 @@ Unit tests for the DataMixing class. """ -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest +from amzn_nova_forge.core.constants import ( + DATAMIX_CUSTOMER_DATA_FIELD, + DATAMIX_DATASET_CATALOG_FIELD, + DATAMIX_NOVA_PREFIX, + DATAMIX_PERCENT_SUFFIX, +) from amzn_nova_forge.util.data_mixing import DataMixing @@ -112,18 +118,17 @@ def test_set_config_rejects_dataset_catalog(self): with patch("amzn_nova_forge.util.data_mixing.logger") as mock_logger: data_mixing.set_config(config) mock_logger.warning.assert_called_once() - assert "dataset_catalog" in data_mixing.config + assert "dataset_catalog" not in data_mixing.config assert "nova_code_percent" in data_mixing.config - def test_set_config_invalid_nova_field(self): - """Test set_config with invalid nova field raises error.""" + def test_set_config_accepts_unknown_nova_field(self): + """Test set_config accepts unknown nova fields.""" data_mixing = DataMixing() data_mixing._default_nova_fields = {"nova_code_percent", "nova_general_percent"} - config = {"nova_invalid_percent": 50, "nova_code_percent": 50} - - with pytest.raises(ValueError, match="Invalid nova field 'nova_invalid_percent'"): - data_mixing.set_config(config) + config = {"nova_new_percent": 50, "nova_code_percent": 50} + data_mixing.set_config(config) + assert data_mixing.config["nova_new_percent"] == 50 def test_load_defaults_from_template(self): """Test loading defaults from template.""" @@ -198,9 +203,7 @@ def test_validate_valid_config(self): "nova_general_percent": 40, "customer_data_percent": 50, } - # Validation happens in set_config, not as a separate method - with patch("amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config"): - data_mixing.set_config(config) # Should not raise + data_mixing.set_config(config) def test_validate_nova_sum_not_100(self): """Test validation fails when nova fields don't sum to 100.""" @@ -211,12 +214,8 @@ def test_validate_nova_sum_not_100(self): "customer_data_percent": 50, } - with patch( - "amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config", - side_effect=ValueError("Nova data percentages must sum to 100"), - ): - with pytest.raises(ValueError, match="Nova data percentages must sum to 100"): - data_mixing.set_config(config) + with pytest.raises(ValueError, match="Nova data percentages must sum to 100"): + data_mixing.set_config(config) def test_validate_customer_data_out_of_range(self): """Test validation fails when customer_data_percent is out of range.""" @@ -226,12 +225,8 @@ def test_validate_customer_data_out_of_range(self): "customer_data_percent": 150, # Out of range } - with patch( - "amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config", - side_effect=ValueError("customer_data_percent must be between 0 and 100"), - ): - with pytest.raises(ValueError, match="customer_data_percent must be between 0 and 100"): - data_mixing.set_config(config) + with pytest.raises(ValueError, match="customer_data_percent must be between 0 and 100"): + data_mixing.set_config(config) def test_validate_nova_field_out_of_range(self): """Test validation fails when nova field is out of range.""" @@ -241,12 +236,23 @@ def test_validate_nova_field_out_of_range(self): "nova_general_percent": -50, # Also out of range } - with patch( - "amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config", - side_effect=ValueError("nova_code_percent must be between 0 and 100"), - ): - with pytest.raises(ValueError, match="nova_code_percent must be between 0 and 100"): - data_mixing.set_config(config) + with pytest.raises(ValueError, match="nova_code_percent must be between 0 and 100"): + data_mixing.set_config(config) + + def test_validate_multiple_errors_surfaced(self): + """Test that multiple validation errors are all surfaced in one message.""" + data_mixing = DataMixing() + config = { + "customer_data_percent": 150, + "nova_code_percent": -10, + } + + with pytest.raises(ValueError) as exc_info: + data_mixing.set_config(config) + + msg = str(exc_info.value) + assert "customer_data_percent" in msg + assert "nova_code_percent" in msg def test_validate_customer_100_with_nova_data(self): """Test validation fails when customer is 100% but nova data is non-zero.""" @@ -257,17 +263,8 @@ def test_validate_customer_100_with_nova_data(self): "customer_data_percent": 100, } - with patch( - "amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config", - side_effect=ValueError( - "Since customer_data_percent is 100 %, all nova data should sum to 0 %" - ), - ): - with pytest.raises( - ValueError, - match="Since customer_data_percent is 100 %, all nova data should sum to 0 %", - ): - data_mixing.set_config(config) + with pytest.raises(ValueError, match="should sum to 0"): + data_mixing.set_config(config) def test_validate_with_floating_point_sum(self): """Test validation handles floating point errors in sum.""" @@ -277,9 +274,7 @@ def test_validate_with_floating_point_sum(self): "nova_general_percent": 33.33, "nova_math_percent": 33.34, # Sum is 99.99999... } - # Validation happens in set_config - with patch("amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config"): - data_mixing.set_config(config) # Should not raise due to floating point tolerance + data_mixing.set_config(config) def test_validate_with_none_values(self): """Test validation handles None values correctly.""" @@ -289,9 +284,7 @@ def test_validate_with_none_values(self): "nova_general_percent": None, "customer_data_percent": 50, } - # Validation happens in set_config - with patch("amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config"): - data_mixing.set_config(config) # Should not raise + data_mixing.set_config(config) def test_update_default_nova_fields(self): """Test that nova fields are added to defaults when set.""" @@ -315,10 +308,10 @@ def test_update_default_nova_fields(self): def test_constants(self): """Test class constants are defined correctly.""" - assert DataMixing.NOVA_PREFIX == "nova_" - assert DataMixing.PERCENT_SUFFIX == "_percent" - assert DataMixing.CUSTOMER_DATA_FIELD == "customer_data_percent" - assert DataMixing.DATASET_CATALOG_FIELD == "dataset_catalog" + assert DATAMIX_NOVA_PREFIX == "nova_" + assert DATAMIX_PERCENT_SUFFIX == "_percent" + assert DATAMIX_CUSTOMER_DATA_FIELD == "customer_data_percent" + assert DATAMIX_DATASET_CATALOG_FIELD == "dataset_catalog" def test_set_config_adds_new_nova_fields_when_no_defaults(self): """Test that new nova fields are added to defaults when no template is loaded.""" @@ -327,8 +320,7 @@ def test_set_config_adds_new_nova_fields_when_no_defaults(self): assert data_mixing._default_nova_fields == set() config = {"nova_new_percent": 100, "customer_data_percent": 50} - with patch("amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config"): - data_mixing.set_config(config) + data_mixing.set_config(config) # Fields are not automatically added to defaults in current implementation # unless template is loaded first @@ -350,8 +342,7 @@ def test_complex_workflow(self): "nova_general_percent": 75, "customer_data_percent": 50, } - with patch("amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config"): - data_mixing.set_config(new_config) + data_mixing.set_config(new_config) # Get and verify config config = data_mixing.get_config() @@ -361,8 +352,7 @@ def test_edge_case_empty_nova_fields(self): """Test edge case with no nova fields.""" data_mixing = DataMixing() config = {"customer_data_percent": 100} - with patch("amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config"): - data_mixing.set_config(config) # Should not raise + data_mixing.set_config(config) def test_edge_case_zero_customer_data(self): """Test edge case with 0% customer data.""" @@ -372,8 +362,7 @@ def test_edge_case_zero_customer_data(self): "nova_general_percent": 30, "customer_data_percent": 0, } - with patch("amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config"): - data_mixing.set_config(config) # Should not raise + data_mixing.set_config(config) def test_load_defaults_with_non_dict_values(self): """Test loading defaults with non-dict values in template.""" @@ -418,8 +407,7 @@ def test_multiple_validations(self): # Valid config 1 config1 = {"nova_code_percent": 100, "customer_data_percent": 50} - with patch("amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config"): - data_mixing.set_config(config1) + data_mixing.set_config(config1) # Valid config 2 config2 = { @@ -427,8 +415,7 @@ def test_multiple_validations(self): "nova_general_percent": 70, "customer_data_percent": 80, } - with patch("amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config"): - data_mixing.set_config(config2) + data_mixing.set_config(config2) # Invalid config config3 = { @@ -436,12 +423,8 @@ def test_multiple_validations(self): "nova_general_percent": 50, # Sum is 80, not 100 "customer_data_percent": 80, } - with patch( - "amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config", - side_effect=ValueError("Invalid config"), - ): - with pytest.raises(ValueError): - data_mixing.set_config(config3) + with pytest.raises(ValueError): + data_mixing.set_config(config3) def test_validate_with_zero_nova_percentages(self): """Test validation with nova fields set to 0.""" @@ -452,8 +435,7 @@ def test_validate_with_zero_nova_percentages(self): "nova_math_percent": 0, "customer_data_percent": 50, } - with patch("amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config"): - data_mixing.set_config(config) # Should not raise + data_mixing.set_config(config) def test_post_init_updates_default_fields(self): """Test that initialization sets correct default fields.""" @@ -468,10 +450,10 @@ def test_post_init_updates_default_fields(self): assert data_mixing._default_nova_fields == set() def test_precedence_with_new_nova_fields(self): - """Test precedence when new nova fields are introduced.""" + """Test that new nova fields are accepted.""" data_mixing = DataMixing() - # First load a template to establish known fields + # Load a template to establish known fields template = { "nova_code_percent": {"default": 50}, "nova_general_percent": {"default": 50}, @@ -479,36 +461,19 @@ def test_precedence_with_new_nova_fields(self): } data_mixing._load_defaults_from_template(template) - # When template has been loaded, new fields can't be added dynamically - invalid_config = { - "nova_code_percent": 30, - "nova_general_percent": 30, - "nova_science_percent": 40, # Invalid new field - "customer_data_percent": 70, - } - - # Should reject the new field - with pytest.raises(ValueError, match="Invalid nova field 'nova_science_percent'"): - data_mixing.set_config(invalid_config) - - # But if no template is loaded, all nova fields can be added - data_mixing2 = DataMixing() - - # Without template, new nova fields can be added + # New nova fields are accepted even with a template loaded config_with_new_field = { "nova_code_percent": 30, "nova_general_percent": 30, "nova_science_percent": 40, "customer_data_percent": 70, } - # This should work without template - data_mixing2.set_config(config_with_new_field) + data_mixing.set_config(config_with_new_field) - # Verify fields were added - assert data_mixing2.config["nova_code_percent"] == 30 - assert data_mixing2.config["nova_general_percent"] == 30 - assert data_mixing2.config["nova_science_percent"] == 40 - assert data_mixing2.config["customer_data_percent"] == 70 + assert data_mixing.config["nova_code_percent"] == 30 + assert data_mixing.config["nova_general_percent"] == 30 + assert data_mixing.config["nova_science_percent"] == 40 + assert data_mixing.config["customer_data_percent"] == 70 def test_precedence_normalization_behavior(self): """Test how normalization affects precedence scenarios.""" @@ -569,7 +534,7 @@ def test_precedence_with_dataset_catalog(self): full_config = data_mixing.get_config() assert full_config["dataset_catalog"] == "catalog_v1" - # Try to override dataset_catalog (should be rejected) + # Try to override dataset_catalog (should be stripped by validation) override_config = { "nova_code_percent": 100, "customer_data_percent": 60, @@ -607,10 +572,8 @@ def test_precedence_with_validation_errors(self): with pytest.raises(ValueError, match="Nova data percentages must sum to 100"): data_mixing.set_config(invalid_config) - # In the current implementation, config is updated before validation - # so if validation fails, the config is left in an invalid state - assert data_mixing.config["nova_code_percent"] == 50 # Changed despite error - assert data_mixing.config["nova_general_percent"] == 50 # Changed despite error + assert data_mixing.config["nova_code_percent"] == 50 + assert data_mixing.config["nova_general_percent"] == 50 def test_precedence_with_none_values(self): """Test precedence handling with None values in configs.""" @@ -624,17 +587,15 @@ def test_precedence_with_none_values(self): } data_mixing._load_defaults_from_template(template) - # Config with None values (should be handled gracefully) config_with_none = { "nova_code_percent": 100, - "nova_general_percent": None, # None value + "nova_general_percent": None, "customer_data_percent": 50, } - with patch("amzn_nova_forge.validation.validator.Validator.validate_data_mixing_config"): - data_mixing.set_config(config_with_none, normalize=False) + data_mixing.set_config(config_with_none, normalize=False) assert data_mixing.config["nova_code_percent"] == 100 - assert data_mixing.config["nova_general_percent"] is None + assert "nova_general_percent" not in data_mixing.config assert data_mixing.config["customer_data_percent"] == 50 def test_is_data_mixing_field_nova_percent_fields(self): diff --git a/tests/unit/validation/test_validator.py b/tests/unit/validation/test_validator.py index a0c3fc8..894502a 100644 --- a/tests/unit/validation/test_validator.py +++ b/tests/unit/validation/test_validator.py @@ -2792,6 +2792,62 @@ def test_wrong_type_still_raises(self): f"Expected type error, got: {errors}", ) + def test_none_value_accepted_when_not_required(self): + """None value must pass validation when the field is not required (e.g. reasoning_effort).""" + overrides_template = { + "reasoning_effort": { + "type": "str", + "required": False, + "enum": ["low", "medium", "high"], + "default": "medium", + }, + } + recipe = { + "training_config": { + "reasoning_effort": None, + }, + } + errors = [] + Validator._validate_recipe( + recipe=recipe, + overrides_template=overrides_template, + instance_type="ml.p4d.24xlarge", + errors=errors, + method=TrainingMethod.RFT_LORA, + platform=Platform.SMTJ, + rft_lambda_arn="arn:aws:lambda:us-west-2:123456789012:function:reward-fn", + ) + self.assertEqual(errors, []) + + def test_none_value_required_str_field_without_null_enum_still_raises(self): + """None on a required str field must still error when null is NOT in the enum.""" + overrides_template = { + "reasoning_effort": { + "type": "str", + "required": True, + "enum": ["low", "medium", "high"], + "default": "medium", + }, + } + recipe = { + "training_config": { + "reasoning_effort": None, + }, + } + errors = [] + Validator._validate_recipe( + recipe=recipe, + overrides_template=overrides_template, + instance_type="ml.p4d.24xlarge", + errors=errors, + method=TrainingMethod.RFT_LORA, + platform=Platform.SMTJ, + rft_lambda_arn="arn:aws:lambda:us-west-2:123456789012:function:reward-fn", + ) + self.assertEqual(len(errors), 1) + self.assertIn("expects str", errors[0]) + self.assertIn("NoneType", errors[0]) + class TestValidationDataS3Path(unittest.TestCase): """Tests for validation_data_s3_path preflight validation."""