From 306d291f73d1ba812562ea3f1755f003be103eff Mon Sep 17 00:00:00 2001 From: Mahima Chaudhary Date: Wed, 13 May 2026 23:08:15 +0000 Subject: [PATCH] Release 2026-05-13 --- README.md | 47 +- docs/spec.md | 4092 ----------------- docs/spec/dataset.md | 621 +++ docs/spec/enums.md | 98 + docs/spec/index.md | 32 + docs/spec/job-results.md | 473 ++ docs/spec/model-customizer.md | 713 +++ docs/spec/notifications.md | 223 + docs/spec/rft-multiturn.md | 212 + docs/spec/runtime-managers.md | 500 ++ docs/spec/service-classes.md | 863 ++++ docs/spec/utilities.md | 457 ++ docs/{ => user-guides}/data_prep.md | 138 +- docs/{ => user-guides}/iam_setup.md | 5 +- docs/{ => user-guides}/instance_type_spec.md | 0 docs/{ => user-guides}/job_notifications.md | 0 docs/{ => user-guides}/rft_multiturn.md | 2 +- docs/{ => user-guides}/troubleshooting.md | 0 src/amzn_nova_forge/__init__.py | 2 + src/amzn_nova_forge/__version__.py | 2 +- .../core/result/eval_result.py | 20 +- .../core/result/inference_result.py | 18 +- src/amzn_nova_forge/core/result/job_result.py | 15 +- .../core/result/training_result.py | 16 +- src/amzn_nova_forge/core/types.py | 13 +- src/amzn_nova_forge/dataset/__init__.py | 2 + .../dataset/cloudwatch_dataset_loader.py | 135 + src/amzn_nova_forge/dataset/data_state.py | 24 +- src/amzn_nova_forge/dataset/dataset_loader.py | 10 +- .../dataset_validator/dataset_validator.py | 9 +- src/amzn_nova_forge/dataset/file_utils.py | 11 +- .../dataset/operations/filter_operation.py | 7 +- .../invalid_records_filter_operation.py | 3 +- .../language_detection_filter_operation.py | 9 +- .../dataset/operations/transform_operation.py | 11 +- .../dataset/operations/utils.py | 13 +- .../deployer/forge_deployer.py | 16 +- .../evaluator/forge_evaluator.py | 3 + src/amzn_nova_forge/iam/iam_role_creator.py | 15 +- .../inference/forge_inference.py | 2 + .../manager/runtime_manager.py | 73 +- .../model/nova_model_customizer.py | 19 +- .../model/nova_model_customizer_util.py | 4 +- src/amzn_nova_forge/monitor/log_monitor.py | 30 +- .../templates/smhp_notification_cf_stack.yaml | 4 +- src/amzn_nova_forge/recipe/recipe_builder.py | 56 +- .../rft_multiturn/base_infra.py | 4 +- .../rft_multiturn/ec2_infra.py | 2 +- .../rft_multiturn/rft_multiturn.py | 4 +- .../telemetry/telemetry_logging.py | 4 +- src/amzn_nova_forge/trainer/forge_trainer.py | 46 +- src/amzn_nova_forge/util/bedrock.py | 20 +- src/amzn_nova_forge/util/checkpoint_util.py | 3 +- src/amzn_nova_forge/util/dataset_writer.py | 8 +- src/amzn_nova_forge/util/mlflow.py | 10 +- src/amzn_nova_forge/util/recipe.py | 17 +- src/amzn_nova_forge/util/s3_utils.py | 2 +- src/amzn_nova_forge/util/sagemaker.py | 9 +- src/amzn_nova_forge/util/subprocess_utils.py | 88 + src/amzn_nova_forge/validation/validator.py | 28 +- tests/unit/conftest.py | 15 + tests/unit/dataset/test_auto_output_path.py | 341 ++ .../dataset/test_cloudwatch_dataset_loader.py | 281 ++ tests/unit/dataset/test_data_state.py | 24 + tests/unit/dataset/test_dataset_loader.py | 2 +- .../dataset/test_dataset_loader_telemetry.py | 4 +- .../test_huggingface_dataset_loader.py | 1 - ...est_language_detection_filter_operation.py | 2 +- .../unit/dataset/test_multimodal_s3_upload.py | 2 +- tests/unit/deployer/test_forge_deployer.py | 43 +- tests/unit/evaluator/test_forge_evaluator.py | 2 + tests/unit/iam/test_iam_role_creator.py | 26 + tests/unit/inference/test_forge_inference.py | 2 + tests/unit/manager/test_runtime_manager.py | 240 +- .../test_serverless_runtime_manager.py | 249 + .../test_smtj_dataprep_runtime_manager.py | 37 + tests/unit/model/result/test_eval_result.py | 18 + .../model/result/test_inference_result.py | 38 + tests/unit/model/result/test_job_result.py | 28 +- .../unit/model/result/test_training_result.py | 39 +- .../unit/model/test_nova_model_customizer.py | 157 +- tests/unit/monitor/test_log_monitor.py | 40 +- tests/unit/recipe/test_recipe_builder.py | 323 +- .../unit/telemetry/test_telemetry_logging.py | 2 +- tests/unit/trainer/test_forge_trainer.py | 379 ++ tests/unit/util/test_bedrock.py | 73 + tests/unit/util/test_dataset_writer.py | 32 + tests/unit/util/test_mlflow.py | 2 +- tests/unit/util/test_recipe.py | 28 + tests/unit/util/test_s3_utils.py | 21 + tests/unit/util/test_subprocess_utils.py | 181 + tests/unit/validation/test_validator.py | 141 + 92 files changed, 7724 insertions(+), 4312 deletions(-) delete mode 100644 docs/spec.md create mode 100644 docs/spec/dataset.md create mode 100644 docs/spec/enums.md create mode 100644 docs/spec/index.md create mode 100644 docs/spec/job-results.md create mode 100644 docs/spec/model-customizer.md create mode 100644 docs/spec/notifications.md create mode 100644 docs/spec/rft-multiturn.md create mode 100644 docs/spec/runtime-managers.md create mode 100644 docs/spec/service-classes.md create mode 100644 docs/spec/utilities.md rename docs/{ => user-guides}/data_prep.md (85%) rename docs/{ => user-guides}/iam_setup.md (97%) rename docs/{ => user-guides}/instance_type_spec.md (100%) rename docs/{ => user-guides}/job_notifications.md (100%) rename docs/{ => user-guides}/rft_multiturn.md (99%) rename docs/{ => user-guides}/troubleshooting.md (100%) create mode 100644 src/amzn_nova_forge/dataset/cloudwatch_dataset_loader.py create mode 100644 src/amzn_nova_forge/util/subprocess_utils.py create mode 100644 tests/unit/dataset/test_auto_output_path.py create mode 100644 tests/unit/dataset/test_cloudwatch_dataset_loader.py create mode 100644 tests/unit/util/test_dataset_writer.py create mode 100644 tests/unit/util/test_subprocess_utils.py diff --git a/README.md b/README.md index 87ecd5e..f805b62 100644 --- a/README.md +++ b/README.md @@ -33,7 +33,7 @@ Nova Forge SDK is tested on: * Python 3.12 ### IAM Roles/Policies -* You will need an IAM role with sufficient permissions in order to use the Nova Forge SDK. You can find a list of these permissions in the `docs/iam_setup.md` file. +* You will need an IAM role with sufficient permissions in order to use the Nova Forge SDK. You can find a list of these permissions in the `docs/user-guides/iam_setup.md` file. ### Instances @@ -41,7 +41,7 @@ Nova customization jobs also require access to enough of the right instance type - The requested instance type and count should be compatible with the requested job. The SDK will validate your instance configuration for you. - The [SageMaker account quotas](https://docs.aws.amazon.com/general/latest/gr/sagemaker.html) for using the requested instance type in training jobs (for SMTJ) or HyperPod clusters (for SMHP) should allow the requested number of instances. - (For SMHP) The selected HyperPod cluster should have a [Restricted Instance Group](https://docs.aws.amazon.com/sagemaker/latest/dg/nova-hp-cluster.html) with enough instances of the right type to run the requested job. The SDK will validate that your cluster contains a valid instance group. -- You can look in the `docs/instance_type_spec.md` file for the different instance types and combinations for specific jobs and methods. +- You can look in the `docs/user-guides/instance_type_spec.md` file for the different instance types and combinations for specific jobs and methods. ### HyperPod CLI @@ -144,7 +144,7 @@ loader.validate(method=ValidateMethod.INVALID_RECORDS, training_method=TrainingM loader.save("s3://my-bucket/prepared-data.jsonl") ``` -For the complete guide — including column mappings, dataset splitting, filtering, chaining operations, and end-to-end examples — see **[Data Preparation Guide](docs/data_prep.md)**. +For the complete guide — including column mappings, dataset splitting, filtering, chaining operations, and end-to-end examples — see **[Data Preparation Guide](docs/user-guides/data_prep.md)**. For a hands-on notebook walkthrough, see [`samples/dataprep_quickstart.ipynb`](samples/dataprep_quickstart.ipynb). @@ -161,10 +161,10 @@ The Nova Forge SDK is organized into the following modules: | **Monitor** | Job monitoring and logging | `CloudWatchLogMonitor`, `MLflowMonitor` | | **RFT Multiturn** | Reinforcement fine-tuning infrastructure | `RFTMultiturnInfrastructure` | -* For detailed API documentation: See [`docs/spec.md`](docs/spec.md) +* For detailed API documentation: See [`docs/spec/`](docs/spec/index.md) * For usage examples: See [`samples/nova_quickstart.ipynb`](samples/nova_quickstart.ipynb) * For RFT Singleturn examples: See [`samples/rft_singleturn_quickstart.ipynb`](samples/rft_singleturn_quickstart.ipynb) -* For RFT Multiturn documentation: See [`docs/rft_multiturn.md`](docs/rft_multiturn.md) +* For RFT Multiturn documentation: See [`docs/user-guides/rft_multiturn.md`](docs/user-guides/rft_multiturn.md) * For RFT Multiturn examples: See [`samples/rft_multiturn_quickstart.ipynb`](samples/rft_multiturn_quickstart.ipynb) ### Service Classes (Recommended) @@ -228,16 +228,16 @@ result.show() For the equivalent workflow using the legacy `NovaModelCustomizer`, see the [Model Module](#model-module-deprecated) section below. -For detailed API documentation on all service classes, see [`docs/spec.md`](docs/spec.md). +For detailed API documentation on all service classes, see [`docs/spec/`](docs/spec/service-classes.md). ### Dataset Module Handles data loading, transformation, validation, filtering, and persistence for training datasets. Supports JSONL, JSON, CSV, Parquet, and Arrow formats from local files or S3. -See the [Data Preparation](#data-preparation) section above for usage overview, or the full **[Data Preparation Guide](docs/data_prep.md)** for detailed documentation. +See the [Data Preparation](#data-preparation) section above for usage overview, or the full **[Data Preparation Guide](docs/user-guides/data_prep.md)** for detailed documentation. ### Manager Module Manages runtime infrastructure for executing training and evaluation jobs. -For the allowed instance types for each model/method combination, see `docs/instance_type_spec.md`. +For the allowed instance types for each model/method combination, see `docs/user-guides/instance_type_spec.md`. **Main Methods:** - `execute()` - Start a training or evaluation job @@ -274,7 +274,7 @@ result = manager.scale_cluster( target_instance_count=8 ) ``` -For more cluster scaling documentation, see [`docs/spec.md`](docs/spec.md). +For more cluster scaling documentation, see [`docs/spec/runtime-managers.md`](docs/spec/runtime-managers.md). ### Model Module (Deprecated) @@ -400,6 +400,7 @@ Data mixing allows you to blend your custom training data with Nova's high-quali **Key Features:** - Available for CPT and SFT training for Nova 1 and Nova 2 (both LoRA and Full-Rank) on SageMaker HyperPod +- Available for SFT (text-only, LoRA and Full-Rank) on Nova 2 Lite on SMTJServerless - Mix customer data (0-100%) with Nova's curated data - Nova data categories include general knowledge and code - Nova data percentages must sum to 100% @@ -407,11 +408,21 @@ Data mixing allows you to blend your custom training data with Nova's high-quali **Example Usage:** ```python -# Initialize with data mixing enabled +# Initialize with data mixing enabled (HyperPod) trainer = ForgeTrainer( model=Model.NOVA_LITE_2, method=TrainingMethod.SFT_LORA, - infra=SMHPRuntimeManager(...), # Must use HyperPod + infra=SMHPRuntimeManager(...), + training_data_s3_path="s3://bucket/data.jsonl", + data_mixing_enabled=True, + config=ForgeConfig(output_s3_path="s3://bucket/output"), +) + +# Or use SMTJServerless +trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.SFT_LORA, + infra=SMTJServerlessRuntimeManager(...), training_data_s3_path="s3://bucket/data.jsonl", data_mixing_enabled=True, config=ForgeConfig(output_s3_path="s3://bucket/output"), @@ -433,7 +444,7 @@ trainer.data_mixing.set_config({ ``` **Important Notes:** - The `dataset_catalog` field is system-managed and cannot be set by users -- Data mixing is only available on SageMaker HyperPod platform for Forge customers. +- Data mixing is available on SageMaker HyperPod and SMTJServerless for Forge customers. - Refer to the [Get Forge Subscription]('https://docs.aws.amazon.com/sagemaker/latest/dg/nova-forge.html#nova-forge-prereq-access') page to enable Nova subscription in your account to use this feature. ### Job Notifications @@ -450,7 +461,7 @@ Get email notifications when your training jobs complete, fail, or are stopped. **Platform Support:** - **SMTJ** (SageMaker Training Jobs): Minimal configuration required - **SMTJServerless** (SageMaker Serverless): No instance type needed — SageMaker manages compute automatically -- **SMHP** (SageMaker HyperPod): Requires kubectl Lambda layer + additional parameters (see [`docs/spec.md`](spec.md) for more details) +- **SMHP** (SageMaker HyperPod): Requires kubectl Lambda layer + additional parameters (see [`docs/spec/runtime-managers.md`](docs/spec/runtime-managers.md) for more details) - **Bedrock** (Amazon Bedrock): Fully managed, no infrastructure configuration required **Quick Example:** @@ -473,14 +484,14 @@ result.enable_job_notifications( **Important Notes:** - Users must confirm their email subscription by clicking the link in the AWS SNS confirmation email -- SMHP job notifications requires a kubectl Lambda layer (see [Job Notifications Guide](docs/job_notifications.md)) +- SMHP job notifications requires a kubectl Lambda layer (see [Job Notifications Guide](docs/user-guides/job_notifications.md)) - Notification infrastructure is created once per region (SMTJ) or once per cluster (SMHP) and shared across jobs. -- See [`docs/job_notifications.md`](docs/job_notifications.md) for detailed setup instructions, troubleshooting, and advanced usage -- See [`docs/spec.md`](docs/spec.md) for complete API documentation on job notifications. +- See [`docs/user-guides/job_notifications.md`](docs/user-guides/job_notifications.md) for detailed setup instructions, troubleshooting, and advanced usage +- See [`docs/spec/notifications.md`](docs/spec/notifications.md) for complete API documentation on job notifications. ### Batch Sample Tracing -Diagnose gradient spikes by identifying which training data lines were used in a specific training step. Enable with `enable_batch_sample_tracing=True` on `ForgeTrainer`, then call `trainer.trace_batch(result, step=N)` after the job completes. See [`docs/spec.md`](docs/spec.md) for full API details. +Diagnose gradient spikes by identifying which training data lines were used in a specific training step. Enable with `enable_batch_sample_tracing=True` on `ForgeTrainer`, then call `trainer.trace_batch(result, step=N)` after the job completes. See [`docs/spec/service-classes.md`](docs/spec/service-classes.md) for full API details. --- ## Telemetry @@ -497,7 +508,7 @@ This comprehensive SDK enables end-to-end customization of Amazon Nova models wi To get started customizing Nova models, please see the following files: * Notebook with "quick start" examples to start customizing at `samples/nova_quickstart.ipynb` -* Specification document with detailed information about each module at `docs/spec.md` +* Specification document with detailed information about each module at [`docs/spec/`](docs/spec/index.md) --- ## Security Best Practices for SDK Users diff --git a/docs/spec.md b/docs/spec.md deleted file mode 100644 index c888ac7..0000000 --- a/docs/spec.md +++ /dev/null @@ -1,4092 +0,0 @@ -# Nova Forge SDK - API Specification - -## Table of Contents -1. [Service Classes (Recommended)](#service-classes-recommended) -2. [NovaModelCustomizer (Deprecated)](#novamodelcustomizer) -3. [Runtime Managers](#runtime-managers) -4. [Dataset Loaders](#dataset-loaders) -5. [Job Results](#job-results) -6. [Utility Functions](#utility-functions) -7. [Monitoring](#monitoring) -8. [Enums and Configuration](#enums-and-configuration) ---- -## Service Classes (Recommended) - -The modular service classes are the recommended API for Nova model customization. Each class handles a single concern — training, evaluation, deployment, or inference — and can be used independently. - -All service classes accept an optional `ForgeConfig` dataclass for shared configuration (KMS keys, output paths, caching, etc.). - -### ForgeConfig - -Shared configuration dataclass for all service classes. - -**Signature:** -```python -@dataclass -class ForgeConfig: - kms_key_id: Optional[str] = None - output_s3_path: Optional[str] = None - generated_recipe_dir: Optional[str] = None - validation_config: Optional[ValidationConfig] = None - image_uri: Optional[str] = None - mlflow_monitor: Optional[MLflowMonitor] = None - enable_job_caching: bool = False - job_cache_dir: str = "~/.nova-forge/cache" - job_caching_config: Optional[JobCachingConfig] = None -``` - -**Parameters:** -- `kms_key_id` (Optional[str]): KMS key ID for S3 encryption -- `output_s3_path` (Optional[str]): S3 path for output artifacts. Auto-generated if not provided -- `generated_recipe_dir` (Optional[str]): Local path to save generated recipe files -- `validation_config` (Optional[ValidationConfig]): Controls pre-flight validation. Fields: `iam` (bool), `infra` (bool), `recipe` (bool) — all default to True -- `image_uri` (Optional[str]): Custom container image URI override -- `mlflow_monitor` (Optional[MLflowMonitor]): MLflow monitoring configuration (SageMaker only) -- `enable_job_caching` (bool): Enable caching of completed job results for reuse. Default: False -- `job_cache_dir` (str): Directory for cached job results. Default: `~/.nova-forge/cache` -- `job_caching_config` (Optional[JobCachingConfig]): Advanced caching configuration. Fields: `include_core` (bool), `include_recipe` (bool), `include_infra` (bool), `include_params` (List[str]), `exclude_params` (List[str]), `allowed_statuses` (List[JobStatus]) - -**Example:** -```python -from amzn_nova_forge.core import ForgeConfig, ValidationConfig -from amzn_nova_forge.monitor import MLflowMonitor - -config = ForgeConfig( - output_s3_path="s3://my-bucket/output", - kms_key_id="my-kms-key-id", - validation_config=ValidationConfig(iam=True, infra=True, recipe=True), - mlflow_monitor=MLflowMonitor( - tracking_uri="arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", - experiment_name="nova-customization" - ), - enable_job_caching=True -) -``` ---- - -### ForgeTrainer - -Handles training job configuration and execution for Nova models. - -#### Constructor - -**Signature:** -```python -def __init__( - self, - model: Model, - method: TrainingMethod, - infra: RuntimeManager, - training_data_s3_path: Optional[str] = None, - model_s3_path: Optional[str] = None, - data_mixing_enabled: bool = False, - holdout_data_s3_path: Optional[str] = None, - config: Optional[ForgeConfig] = None, - region: Optional[str] = None, - is_multimodal: Optional[bool] = None, - hub_content_version: Optional[str] = None, - enable_batch_sample_tracing: bool = False, -) -``` - -**Parameters:** -- `model` (Model): The Nova model to train (e.g., `Model.NOVA_MICRO`, `Model.NOVA_LITE_2`) -- `method` (TrainingMethod): The fine-tuning method (e.g., `TrainingMethod.SFT_LORA`, `TrainingMethod.RFT`) -- `infra` (RuntimeManager): Runtime infrastructure manager (e.g., `SMTJRuntimeManager`, `SMHPRuntimeManager`, `BedrockRuntimeManager`) -- `training_data_s3_path` (Optional[str]): S3 path to the training dataset -- `model_s3_path` (Optional[str]): S3 path for the base or previously trained model -- `data_mixing_enabled` (bool): Enable data mixing for CPT and SFT training on SMHP. Default: False -- `holdout_data_s3_path` (Optional[str]): S3 path to holdout/validation data (optional, used for CPT) -- `config` (Optional[ForgeConfig]): Shared configuration. If not provided, defaults are used -- `region` (Optional[str]): AWS region. Auto-detected if not provided -- `is_multimodal` (Optional[bool]): Explicitly set multimodal mode when `data_mixing_enabled=True`. If None, auto-detects from data -- `hub_content_version` (Optional[str]): Version of the hub content to retrieve from SageMaker Hub. If None, uses the latest version -- `enable_batch_sample_tracing` (bool): Activate per-step batch hashing during training, which enables `trace_batch()` post-training. Supported platform/method combinations are validated at construction time. Default: False - -**Raises:** -- `ValueError`: If `enable_batch_sample_tracing=True` is used with an unsupported platform or training method - -**Example:** -```python -from amzn_nova_forge.trainer import ForgeTrainer -from amzn_nova_forge.core import ForgeConfig -from amzn_nova_forge.manager import SMTJRuntimeManager -from amzn_nova_forge.model.model_enums import Model, TrainingMethod - -infra = SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=2) - -trainer = ForgeTrainer( - model=Model.NOVA_MICRO, - method=TrainingMethod.SFT_LORA, - infra=infra, - training_data_s3_path="s3://my-bucket/training-data/data.jsonl", - config=ForgeConfig(output_s3_path="s3://my-bucket/output") -) -``` ---- - -#### Methods - -##### `train()` -Generates the recipe YAML, configures the runtime, and launches a training job. - -**Signature:** -```python -def train( - self, - job_name: str, - recipe_path: Optional[str] = None, - overrides: Optional[Dict[str, Any]] = None, - rft_lambda_arn: Optional[str] = None, - dry_run: bool = False, - rft_multiturn_infra=None, -) -> Optional[TrainingResult] -``` - -**Parameters:** -- `job_name` (str): User-defined name for the training job -- `recipe_path` (Optional[str]): Path for a YAML recipe file (S3 or local) -- `overrides` (Optional[Dict[str, Any]]): Dictionary of configuration overrides (e.g., `max_epochs`, `lr`, `warmup_steps`, `global_batch_size`) -- `rft_lambda_arn` (Optional[str]): Rewards Lambda ARN (only for RFT methods). Takes priority over `rft_lambda_arn` on the RuntimeManager -- `dry_run` (bool): If True, performs validation only without starting a job. Default: False -- `rft_multiturn_infra`: Optional RFTMultiturnInfrastructure for RFT multiturn training - -**Returns:** -- `TrainingResult`: Metadata object containing `job_id`, `method`, `started_time`, `model_artifacts`, and `model_type`. Returns `None` if `dry_run=True` - -**Raises:** -- `Exception`: If job execution fails -- `ValueError`: If training method is not supported - -**Example:** -```python -result = trainer.train( - job_name="my-training-job", - overrides={ - "max_epochs": 10, - "lr": 5e-6, - "warmup_steps": 20, - "global_batch_size": 128 - } -) -print(f"Training job started: {result.job_id}") -print(f"Checkpoint path: {result.model_artifacts.checkpoint_s3_path}") -``` ---- - -##### `get_logs()` -Retrieves and displays CloudWatch logs for a training job. - -**Signature:** -```python -def get_logs( - self, - job_result=None, - job_id=None, - started_time=None, - limit=None, - start_from_head: bool = False, - end_time=None, -) -> None -``` - -**Parameters:** -- `job_result` (Optional[TrainingResult]): Job result to retrieve logs for. If not provided, uses `job_id` -- `job_id` (Optional[str]): Job identifier. Used if `job_result` is not provided -- `started_time` (Optional[datetime]): Job start time to filter logs -- `limit` (Optional[int]): Maximum number of log lines to retrieve -- `start_from_head` (bool): If True, start from the beginning of logs. Default: False -- `end_time` (Optional[int]): End time in epoch milliseconds for searching a log time range - -**Returns:** -- None (prints logs to console) - -**Example:** -```python -trainer.get_logs(job_result=result, limit=100, start_from_head=True) -``` ---- - -##### `trace_batch()` -Extracts the lines from your input training data that were used in a specific training step's batch. Useful for diagnosing gradient spikes or training anomalies — given a step number, it matches the container's batch hash logs against your source data and writes the matched lines to an output file. - -The training job must have been launched with `enable_batch_sample_tracing=True` so that batch hash logs are written during training. Supported platform/method combinations are validated at `ForgeTrainer` construction time. If you create a new `ForgeTrainer` instance to trace a previously-launched job, the flag is not strictly required on the new instance — but a warning will be emitted. - -**Signature:** -```python -def trace_batch( - self, - training_result: TrainingResult, - step: int, - output_path: str | None = None, - cache_dir: str = "~/.nova-forge/batch_trace_cache", -) -> Path | None -``` - -**Parameters:** -- `training_result` (TrainingResult): Result from a completed training job. The method extracts `training_result.model_artifacts.output_s3_path` and `training_result.job_id` to locate the batch hash logs at `{output_s3_path}/{job_id}/batch_tracing/`. -- `step` (int): Training step number to investigate (must match the step numbers in the container's batch hash logs). -- `output_path` (Optional[str]): Path for the output file containing matched lines. Default: `step__samples.jsonl` in the current working directory. -- `cache_dir` (str): Directory for caching downloaded files and fingerprint indices. Default: `~/.nova-forge/batch_trace_cache`. The cache stores downloaded S3 files (training data, hash logs) and a CSV fingerprint index. For large datasets the cache may grow to match the source data size. The cache is keyed by S3 path — if you replace the file at an existing S3 URI, delete the cache directory to avoid stale matches. - -**Returns:** -- `Path | None`: Path to the output JSONL file containing the matched lines (verbatim copies from your source data, sorted by line number). Returns `None` if either the step had no logged batch data (step out of range or job still running) or the step's batch contained no samples from your file. - -**Raises:** -- `ValueError`: If `training_data_s3_path` or `training_result.model_artifacts.output_s3_path` is not available -- `BatchTraceError`: If batch tracing encounters an unrecoverable error (e.g., missing log files, AWS auth failure) - Import: `from amzn_nova_forge.trainer.utils import BatchTraceError` - -**Example:** -```python -trainer = ForgeTrainer( - model=Model.NOVA_LITE_2, - method=TrainingMethod.CPT, - infra=infra, - training_data_s3_path="s3://my-bucket/data.jsonl", - enable_batch_sample_tracing=True, -) - -result = trainer.train(job_name="my-cpt-job") - -# After job completes, investigate step 42. -# Output writes to ./step_42_samples.jsonl by default. -matched_file = trainer.trace_batch(result, step=42) -if matched_file: - print(f"Matched lines written to: {matched_file}") - -# Explicit output path: -matched_file = trainer.trace_batch(result, step=42, output_path="/tmp/flagged.jsonl") -``` ---- - -### ForgeEvaluator - -Handles evaluation job configuration and execution for Nova models. - -#### Constructor - -**Signature:** -```python -def __init__( - self, - model: Model, - infra: RuntimeManager, - data_s3_path: Optional[str] = None, - config: Optional[ForgeConfig] = None, - region: Optional[str] = None, - hub_content_version: Optional[str] = None, -) -``` - -**Parameters:** -- `model` (Model): The Nova model to evaluate -- `infra` (RuntimeManager): Runtime infrastructure manager -- `data_s3_path` (Optional[str]): S3 path to evaluation data (required for BYOD evaluation tasks) -- `config` (Optional[ForgeConfig]): Shared configuration -- `region` (Optional[str]): AWS region. Auto-detected if not provided -- `hub_content_version` (Optional[str]): Version of the hub content to retrieve from SageMaker Hub. If None, uses the latest version - -**Example:** -```python -from amzn_nova_forge.evaluator import ForgeEvaluator -from amzn_nova_forge.manager import SMTJRuntimeManager -from amzn_nova_forge.model.model_enums import Model - -infra = SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=2) - -evaluator = ForgeEvaluator( - model=Model.NOVA_MICRO, - infra=infra, - data_s3_path="s3://my-bucket/eval-data/data.jsonl" -) -``` ---- - -#### Methods - -##### `evaluate()` -Generates the recipe YAML, configures the runtime, and launches an evaluation job. - -**Signature:** -```python -def evaluate( - self, - job_name: str, - eval_task: EvaluationTask, - model_path: Optional[str] = None, - task_config: Optional[EvalTaskConfig] = None, - recipe_path: Optional[str] = None, - overrides: Optional[Dict[str, Any]] = None, - dry_run: bool = False, - job_result: Optional[TrainingResult] = None, - rft_multiturn_infra=None, -) -> Optional[EvaluationResult] -``` - -**Parameters:** -- `job_name` (str): User-defined name for the evaluation job -- `eval_task` (EvaluationTask): The evaluation task (e.g., `EvaluationTask.MMLU`) -- `model_path` (Optional[str]): S3 path to the model to evaluate. If not provided, extracted from `job_result` -- `task_config` (Optional[EvalTaskConfig]): Task-specific configuration. Fields: `subtask`, `processor`, `rl_env`, `override_data_s3_path` -- `recipe_path` (Optional[str]): Path for a YAML recipe file (S3 or local) -- `overrides` (Optional[Dict[str, Any]]): Inference configuration overrides (e.g., `max_new_tokens`, `temperature`, `top_p`) -- `dry_run` (bool): If True, performs validation only. Default: False -- `job_result` (Optional[TrainingResult]): Training result to extract checkpoint path from -- `rft_multiturn_infra`: Optional RFTMultiturnInfrastructure for RFT evaluation - -**Returns:** -- `EvaluationResult`: Metadata object containing `job_id`, `started_time`, `eval_output_path`, and `eval_task`. Returns `None` if `dry_run=True` - -**Example:** -```python -from amzn_nova_forge.core import EvaluationTask - -eval_result = evaluator.evaluate( - job_name="my-eval-job", - eval_task=EvaluationTask.MMLU, - model_path="s3://my-bucket/checkpoints/my-model", - overrides={ - "max_new_tokens": 2048, - "temperature": 0, - "top_p": 1.0 - } -) -print(f"Evaluation job started: {eval_result.job_id}") - -# Chain from training result -eval_result = evaluator.evaluate( - job_name="my-eval-job", - eval_task=EvaluationTask.MMLU, - job_result=training_result -) -``` ---- - -##### `get_logs()` -Retrieves and displays CloudWatch logs for an evaluation job. - -**Signature:** -```python -def get_logs( - self, - job_result=None, - job_id=None, - started_time=None, - limit=None, - start_from_head: bool = False, - end_time=None, -) -> None -``` - -**Parameters:** -- `job_result` (Optional[EvaluationResult]): Job result to retrieve logs for -- `job_id` (Optional[str]): Job identifier -- `started_time` (Optional[datetime]): Job start time to filter logs -- `limit` (Optional[int]): Maximum number of log lines -- `start_from_head` (bool): If True, start from the beginning of logs. Default: False -- `end_time` (Optional[int]): End time in epoch milliseconds for searching a log time range - -**Returns:** -- None (prints logs to console) - -**Example:** -```python -evaluator.get_logs(job_result=eval_result, limit=50) -``` ---- - -### ForgeDeployer - -Handles model deployment to Amazon Bedrock and SageMaker endpoints. - -#### Constructor - -**Signature:** -```python -def __init__( - self, - region: str, - model: Model, - deployment_mode: DeploymentMode = DeploymentMode.FAIL_IF_EXISTS, - config: Optional[ForgeConfig] = None, - method: Optional[TrainingMethod] = None, -) -``` - -**Parameters:** -- `region` (str): AWS region for deployment -- `model` (Model): The Nova model being deployed -- `deployment_mode` (DeploymentMode): Behavior when endpoint already exists. Default: `FAIL_IF_EXISTS` -- `config` (Optional[ForgeConfig]): Shared configuration -- `method` (Optional[TrainingMethod]): Training method used (needed for SageMaker deployment image selection) - -**Example:** -```python -from amzn_nova_forge.deployer import ForgeDeployer -from amzn_nova_forge.model.model_enums import Model, DeploymentMode - -deployer = ForgeDeployer( - region="us-east-1", - model=Model.NOVA_MICRO, - deployment_mode=DeploymentMode.FAIL_IF_EXISTS -) -``` ---- - -#### Methods - -##### `deploy()` -Creates a custom model and deploys it to Amazon Bedrock or SageMaker in a single step. - -**Signature:** -```python -def deploy( - self, - model_artifact_path: str, - deploy_platform: DeployPlatform = DeployPlatform.BEDROCK_OD, - endpoint_name: Optional[str] = None, - unit_count: int = 1, - execution_role_name: Optional[str] = None, - sagemaker_instance_type: str = "ml.p5.48xlarge", - sagemaker_environment: Optional[SageMakerEndpointEnvironment] = None, - skip_model_reuse: bool = False, -) -> DeploymentResult -``` - -**Parameters:** -- `model_artifact_path` (str): S3 path to the trained model checkpoint -- `deploy_platform` (DeployPlatform): Platform to deploy to (`BEDROCK_OD`, `BEDROCK_PT`, or `SAGEMAKER`). Default: `BEDROCK_OD` -- `endpoint_name` (Optional[str]): Name of the endpoint (auto-generated if not provided) -- `unit_count` (int): Number of PT units (Bedrock PT) or instances (SageMaker). Default: 1 -- `execution_role_name` (Optional[str]): IAM role name. If omitted, the SDK creates a default role -- `sagemaker_instance_type` (str): Instance type for SageMaker deployment. Default: `"ml.p5.48xlarge"` -- `sagemaker_environment` (Optional[SageMakerEndpointEnvironment]): SageMaker endpoint environment config. Fields: - - `CONTEXT_LENGTH` (int, default: 4000), `MAX_CONCURRENCY` (int, default: 1) - - Optional generation defaults: `DEFAULT_TEMPERATURE` (0–2), `DEFAULT_TOP_P` (1e-10–1), `DEFAULT_TOP_K` (-1 to disable, or ≥1), `DEFAULT_MAX_NEW_TOKENS` (≥1), `DEFAULT_LOGPROBS` (1–20) - - Optional speculative decoding: `SPECULATIVE_DECODING_METHOD` (`"eagle3"` or `"suffix"`), `DISABLE_SPECULATIVE_DECODING` (`"true"` or `"false"`), `NUM_SPECULATIVE_TOKENS` (1–10), `SUFFIX_DECODING_MAX_TREE_DEPTH`, `SUFFIX_DECODING_MAX_CACHED_REQUESTS`, `SUFFIX_DECODING_MAX_SPEC_FACTOR`, `SUFFIX_DECODING_MIN_TOKEN_PROB` - - Optional memory/quantization: `KV_CACHE_DTYPE` (`"fp8"`), `QUANTIZATION_DTYPE` (`"fp8"`) -- `skip_model_reuse` (bool): If True, always create a new model. Default: False - -**Returns:** -- `DeploymentResult`: Contains `endpoint` (EndpointInfo), `platform`, `endpoint_name`, `uri`, `model_artifact_path`, and `created_at` - -**Raises:** -- `Exception`: When unable to deploy the model -- `ValueError`: If platform is not supported - -**Example:** -```python -deployment = deployer.deploy( - model_artifact_path="s3://escrow-bucket/my-model-artifacts/", - deploy_platform=DeployPlatform.BEDROCK_OD, - endpoint_name="my-custom-nova-model" -) -print(f"Model deployed: {deployment.endpoint.uri}") -``` ---- - -##### `create_custom_model()` -Creates a Bedrock custom model from S3 artifacts without deploying to an endpoint. - -**Signature:** -```python -def create_custom_model( - self, - model_artifact_path: str, - endpoint_name: Optional[str] = None, - execution_role_name: Optional[str] = None, - tags: Optional[List[Dict[str, str]]] = None, - skip_model_reuse: bool = False, -) -> ModelDeployResult -``` - -**Parameters:** -- `model_artifact_path` (str): S3 path to trained model checkpoint -- `endpoint_name` (Optional[str]): Optional name prefix for the model name -- `execution_role_name` (Optional[str]): IAM role name for Bedrock -- `tags` (Optional[List[Dict[str, str]]]): Optional list of `{"key": str, "value": str}` dicts for tracking -- `skip_model_reuse` (bool): If True, always create a new model. Default: False - -**Returns:** -- `ModelDeployResult`: Contains `model_arn`, `model_name`, `escrow_uri`, and `created_at` - -**Example:** -```python -publish_result = deployer.create_custom_model( - model_artifact_path="s3://escrow-bucket/my-model-artifacts/" -) -print(f"Model ARN: {publish_result.model_arn}") -publish_result.dump(file_path="./results/") -``` ---- - -##### `deploy_to_bedrock()` -Deploys a published Bedrock custom model to an endpoint. - -**Signature:** -```python -def deploy_to_bedrock( - self, - model_deploy_result: Optional[ModelDeployResult] = None, - model_arn: Optional[str] = None, - deploy_platform: DeployPlatform = DeployPlatform.BEDROCK_OD, - pt_units: Optional[int] = None, - endpoint_name: Optional[str] = None, -) -> DeploymentResult -``` - -**Parameters:** -- `model_deploy_result` (Optional[ModelDeployResult]): Result from `create_custom_model()`. Cannot be combined with `model_arn` -- `model_arn` (Optional[str]): Direct model ARN. Cannot be combined with `model_deploy_result` -- `deploy_platform` (DeployPlatform): `BEDROCK_OD` (default) or `BEDROCK_PT` -- `pt_units` (Optional[int]): Number of PT units (required for `BEDROCK_PT`) -- `endpoint_name` (Optional[str]): Endpoint name (auto-generated if not provided) - -**Returns:** -- `DeploymentResult`: Contains `endpoint`, `created_at`, and `model_publish` - -**Raises:** -- `ValueError`: When both `model_deploy_result` and `model_arn` are provided, or when no model ARN is available -- `RuntimeError`: When deployment creation fails - -**Example:** -```python -# Two-step deploy: create model, then deploy -publish_result = deployer.create_custom_model( - model_artifact_path="s3://escrow-bucket/my-model-artifacts/" -) -deployment = deployer.deploy_to_bedrock( - model_deploy_result=publish_result, - endpoint_name="my-endpoint" -) - -# Or deploy from an existing model ARN -deployment = deployer.deploy_to_bedrock( - model_arn="arn:aws:bedrock:us-east-1:123456789012:custom-model/my-model" -) -``` ---- - -##### `find_published_model()` -Finds an existing published model by platform and escrow path to enable model reuse. - -**Signature:** -```python -def find_published_model( - self, - platform: str, - escrow_path: str, - skip_model_reuse: bool = False, -) -> Optional[str] -``` - -**Parameters:** -- `platform` (str): Target platform (`"bedrock"` or `"sagemaker"`) -- `escrow_path` (str): S3 path of the model artifacts -- `skip_model_reuse` (bool): If True, always returns None (skips lookup). Default: False - -**Returns:** -- `Optional[str]`: Existing model ARN if found, otherwise None - -**Example:** -```python -existing_arn = deployer.find_published_model( - platform="bedrock", - escrow_path="s3://escrow-bucket/my-model-artifacts/" -) -if existing_arn: - print(f"Reusing existing model: {existing_arn}") -``` ---- - -##### `get_status()` -Gets the deployment status for a DeploymentResult. - -**Signature:** -```python -def get_status(self, result: DeploymentResult) -> JobStatus -``` - -**Parameters:** -- `result` (DeploymentResult): The deployment result to check - -**Returns:** -- `JobStatus`: Current status (`IN_PROGRESS`, `COMPLETED`, or `FAILED`) - ---- - -##### `get_status_by_arn()` -Gets the deployment status by endpoint ARN and platform. - -**Signature:** -```python -def get_status_by_arn( - self, - endpoint_arn: str, - platform: DeployPlatform, -) -> Optional[JobStatus] -``` - -**Parameters:** -- `endpoint_arn` (str): The endpoint ARN to check -- `platform` (DeployPlatform): The deployment platform - -**Returns:** -- `Optional[JobStatus]`: Current status, or None if status cannot be determined - ---- - -##### `get_logs()` -Retrieves and displays logs for a deployment. - -**Signature:** -```python -def get_logs( - self, - job_result=None, - endpoint_arn=None, - platform=None, -) -> None -``` - -**Parameters:** -- `job_result` (Optional[DeploymentResult]): Deployment result to retrieve logs for -- `endpoint_arn` (Optional[str]): Endpoint ARN (used if `job_result` is not provided) -- `platform` (Optional[DeployPlatform]): Deployment platform (used with `endpoint_arn`) - -**Returns:** -- None (prints logs to console) - ---- - -### ForgeInference - -Handles single and batch inference on trained Nova models. - -#### Constructor - -**Signature:** -```python -def __init__( - self, - region: Optional[str] = None, - model: Optional[Model] = None, - infra: Optional[RuntimeManager] = None, - config: Optional[ForgeConfig] = None, - method: Optional[TrainingMethod] = None, - hub_content_version: Optional[str] = None, -) -``` - -**Parameters:** -- `region` (Optional[str]): AWS region. Auto-detected if not provided -- `model` (Optional[Model]): The Nova model (required for batch inference) -- `infra` (Optional[RuntimeManager]): Runtime infrastructure manager (required for batch inference) -- `config` (Optional[ForgeConfig]): Shared configuration -- `method` (Optional[TrainingMethod]): Training method (used for batch inference recipe generation) -- `hub_content_version` (Optional[str]): Version of the hub content to retrieve from SageMaker Hub. If None, uses the latest version - -**Example:** -```python -from amzn_nova_forge.inference import ForgeInference - -# For single inference (minimal setup) -inference = ForgeInference(region="us-east-1") - -# For batch inference -inference = ForgeInference( - region="us-east-1", - model=Model.NOVA_MICRO, - infra=SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=1), - method=TrainingMethod.SFT_LORA -) -``` ---- - -#### Methods - -##### `invoke()` -Invokes a single inference on a deployed model endpoint. - -**Signature:** -```python -def invoke( - self, - endpoint_arn: str, - request_body: Dict[str, Any], -) -> Any -``` - -**Parameters:** -- `endpoint_arn` (str): Endpoint ARN to invoke -- `request_body` (Dict[str, Any]): Inference request body - -**Returns:** -- `Any`: Inference response - -**Example:** -```python -response = inference.invoke( - endpoint_arn="arn:aws:bedrock:us-east-1:123456789012:endpoint/my-endpoint", - request_body={ - "messages": [{"role": "user", "content": "Hello! How are you?"}], - "max_tokens": 100, - "stream": False - } -) -``` ---- - -##### `invoke_batch()` -Launches a batch inference job on a trained model. - -**Signature:** -```python -def invoke_batch( - self, - job_name: str, - input_path: str, - output_s3_path: str, - model_path: Optional[str] = None, - recipe_path: Optional[str] = None, - overrides: Optional[Dict[str, Any]] = None, - dry_run: bool = False, - job_result: Optional[TrainingResult] = None, -) -> Optional[InferenceResult] -``` - -**Parameters:** -- `job_name` (str): Name for the batch inference job -- `input_path` (str): S3 path to input data -- `output_s3_path` (str): S3 path for inference outputs -- `model_path` (Optional[str]): S3 path to the model checkpoint -- `recipe_path` (Optional[str]): Path for a YAML recipe file -- `overrides` (Optional[Dict[str, Any]]): Inference configuration overrides (e.g., `max_new_tokens`, `temperature`, `top_p`) -- `dry_run` (bool): If True, performs validation only. Default: False -- `job_result` (Optional[TrainingResult]): Training result to extract checkpoint path from - -**Returns:** -- `InferenceResult`: Metadata object containing `job_id`, `started_time`, and `inference_output_path`. Returns `None` if `dry_run=True` - -**Example:** -```python -inference_result = inference.invoke_batch( - job_name="batch-inference-job", - input_path="s3://my-bucket/inference-input", - output_s3_path="s3://my-bucket/inference-output", - model_path="s3://my-bucket/trained-model" -) -print(f"Batch inference started: {inference_result.job_id}") -``` ---- - -##### `get_logs()` -Retrieves and displays CloudWatch logs for an inference job. - -**Signature:** -```python -def get_logs( - self, - job_result=None, - job_id=None, - started_time=None, - limit=None, - start_from_head: bool = False, - end_time=None, -) -> None -``` - -**Parameters:** -- `job_result` (Optional[InferenceResult]): Job result to retrieve logs for -- `job_id` (Optional[str]): Job identifier -- `started_time` (Optional[datetime]): Job start time to filter logs -- `limit` (Optional[int]): Maximum number of log lines -- `start_from_head` (bool): If True, start from the beginning of logs. Default: False -- `end_time` (Optional[int]): End time in epoch milliseconds for searching a log time range - -**Returns:** -- None (prints logs to console) - -**Example:** -```python -inference.get_logs(job_result=inference_result, limit=100) -``` ---- - -### End-to-End Example (Service Classes) - -```python -from amzn_nova_forge.trainer import ForgeTrainer -from amzn_nova_forge.evaluator import ForgeEvaluator -from amzn_nova_forge.deployer import ForgeDeployer -from amzn_nova_forge.inference import ForgeInference -from amzn_nova_forge.core import ForgeConfig -from amzn_nova_forge.manager import SMTJRuntimeManager -from amzn_nova_forge.model.model_enums import Model, TrainingMethod, DeployPlatform -from amzn_nova_forge.core import EvaluationTask - -# Shared configuration -config = ForgeConfig( - output_s3_path="s3://my-bucket/output", - enable_job_caching=True -) -infra = SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=2) - -# 1. Train -trainer = ForgeTrainer( - model=Model.NOVA_MICRO, - method=TrainingMethod.SFT_LORA, - infra=infra, - training_data_s3_path="s3://my-bucket/data.jsonl", - config=config -) -train_result = trainer.train(job_name="my-training-job") - -# 2. Evaluate -evaluator = ForgeEvaluator(model=Model.NOVA_MICRO, infra=infra, config=config) -eval_result = evaluator.evaluate( - job_name="my-eval-job", - eval_task=EvaluationTask.MMLU, - job_result=train_result -) - -# 3. Deploy -deployer = ForgeDeployer(region="us-east-1", model=Model.NOVA_MICRO) -deployment = deployer.deploy( - model_artifact_path=train_result.model_artifacts.checkpoint_s3_path, - deploy_platform=DeployPlatform.BEDROCK_OD, - endpoint_name="my-nova-endpoint" -) - -# 4. Inference -inference_client = ForgeInference(region="us-east-1") -response = inference_client.invoke( - endpoint_arn=deployment.endpoint.uri, - request_body={ - "messages": [{"role": "user", "content": "Hello!"}], - "max_tokens": 100 - } -) -``` - ---- -## NovaModelCustomizer - -> **Deprecated**: `NovaModelCustomizer` is a legacy facade. Use `ForgeTrainer`, `ForgeEvaluator`, `ForgeDeployer`, and `ForgeInference` instead. - -The main entrypoint class for customizing and training Nova models. - -### Constructor - -#### `__init__()` -Initializes a NovaModelCustomizer instance. - -**Signature:** -```python -def __init__( - self, - model: Model, - method: TrainingMethod, - infra: RuntimeManager, - data_s3_path: Optional[str] = None, - output_s3_path: Optional[str] = None, - model_path: Optional[str] = None, - validation_config: Optional[Dict[str, bool]] = None, - generated_recipe_dir: Optional[str] = None, - mlflow_monitor: Optional[MLflowMonitor] = None, - deployment_mode: DeploymentMode = DeploymentMode.FAIL_IF_EXISTS, - data_mixing_enabled: bool = False, - enable_job_caching: bool = False, - is_multimodal: Optional[bool] = None, - hub_content_version: Optional[str] = None, -) -``` -**Parameters:** -- `model` (Model): The Nova model to be trained (e.g., `Model.NOVA_MICRO`, `Model.NOVA_LITE`, `Model.NOVA_LITE_2`, `Model.NOVA_PRO`) -- `method` (TrainingMethod): The fine-tuning method (e.g., `TrainingMethod.SFT_LORA`, `TrainingMethod.RFT`) -- `infra` (RuntimeManager): Runtime infrastructure manager (e.g., `SMTJRuntimeManager`, `SMHPRuntimeManager`, or `BedrockRuntimeManager`) -- `data_s3_path` (Optional[str]): S3 path to the training dataset -- `output_s3_path` (Optional[str]): S3 path for output artifacts. If not provided, will be auto-generated -- `model_path` (Optional[str]): S3 path for model path -- `validation_config` (Optional[Dict[str, bool]]): Optional dict to control validation. Defaults to `{'iam': True, 'infra': True, 'recipe': True}`. - - `iam` (bool): Enable IAM permission validation (default: True) - - `infra` (bool): Enable infrastructure validation (default: True) - - `recipe` (bool): Enable recipe constraint validation (default: True) -- `generated_recipe_dir` (Optional[str]): Optional local path to save the generated recipe -- `mlflow_monitor` (Optional[MLflowMonitor]): Optional MLflow monitoring configuration for experiment tracking (SageMaker only, not supported on Bedrock) -- `deployment_mode` (DeploymentMode): Behavior when deploying to existing endpoint name. Options: FAIL_IF_EXISTS (default), UPDATE_IF_EXISTS -- `data_mixing_enabled` (bool): Enable data mixing feature for CPT and SFT training on SageMaker HyperPod. Default is False - - **Note:** The `data_mixing_enabled` parameter must be set to `True` during initialization to use data mixing features. - - **Note:** Datamixing is only supported for CPT, SFT_LORA, and SFT_FULL methods on SageMaker HyperPod (SMHP). -- `is_multimodal` (Optional[bool]): Only applicable when `data_mixing_enabled=True`. Explicitly set multimodal mode. If None (default), auto-detects from data. Set to False to skip detection for performance on large text-only datasets. Ignored when `data_mixing_enabled=False` -- `enable_job_caching` (bool): Whether to enable job result caching. When enabled, completed job results are cached to `job_cache_dir` (default: `.cached-nova-jobs/`) and reused for identical job configurations. Default: False -- `hub_content_version` (Optional[str]): Version of the hub content to retrieve from SageMaker Hub. If None, uses the latest version - -**Raises:** -- `ValueError`: If region is unsupported or model is invalid - -**Example:** -```python -from amzn_nova_forge import * - -# SageMaker Training Jobs (SMTJ) -infra = SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=2) - -customizer = NovaModelCustomizer( - model=Model.NOVA_MICRO, - method=TrainingMethod.SFT_LORA, - infra=infra, - data_s3_path="s3://my-bucket/training-data/data.jsonl", - output_s3_path="s3://my-bucket/output" -) - -# Amazon Bedrock (fully managed) -bedrock_infra = BedrockRuntimeManager( - execution_role="arn:aws:iam::123456789012:role/BedrockRole", - base_model_identifier="arn:aws:bedrock:us-east-1::custom-model/amazon.nova-2-lite-v1:0:256k:abcdefghijk" -) - -bedrock_customizer = NovaModelCustomizer( - model=Model.NOVA_MICRO, - method=TrainingMethod.SFT_LORA, - infra=bedrock_infra, - data_s3_path="s3://my-bucket/training-data/data.jsonl", - output_s3_path="s3://my-bucket/output" -) - -# With MLflow monitoring (SageMaker only) -mlflow_monitor = MLflowMonitor( - tracking_uri="arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", - experiment_name="nova-customization", - run_name="sft-run-1" -) - -customizer_with_mlflow = NovaModelCustomizer( - model=Model.NOVA_MICRO, - method=TrainingMethod.SFT_LORA, - infra=infra, - data_s3_path="s3://my-bucket/training-data/data.jsonl", - output_s3_path="s3://my-bucket/output", - mlflow_monitor=mlflow_monitor -) -``` ---- -### Methods - -#### `get_data_mixing_config()` -Get the current data mixing configuration. - -**Signature:** -```python -def get_data_mixing_config( - self -) -> Dict[str, Any] -``` - -**Returns:** -- `Dict[str, Any]`: Dictionary containing the data mixing configuration - -**Example:** -```python -config = customizer.get_data_mixing_config() -print(config) -# Output: {'customer_data_percent': 50, 'nova_code_percent': 30, 'nova_general_percent': 70} -``` - ---- - -#### `set_data_mixing_config()` -Set the data mixing configuration. - -**Signature:** -```python -def set_data_mixing_config( - self, - config: Dict[str, Any] -) -> None -``` - -**Parameters:** -- `config` (Dict[str, Any]): Dictionary containing the data mixing configuration - - `customer_data_percent` (int/float): Percentage of customer data (0-100) - - `nova_code_percent` (int/float): Percentage of Nova code data (0-100) - - `nova_general_percent` (int/float): Percentage of Nova general data (0-100) - - Nova percentages must sum to 100% - -**Raises:** -- `ValueError`: If data mixing is not enabled or configuration is invalid - -**Note:** -- The SDK automatically detects whether your dataset is multimodal (contains images, videos, or documents) by scanning the data -- The appropriate Nova dataset catalog (text-only or multimodal) and nova data mixing fields are selected automatically based on this detection - -**Example:** -```python -# Text datamixing (auto-detection) -customizer = NovaModelCustomizer( - model=Model.NOVA_LITE_2, - method=TrainingMethod.SFT_LORA, - infra=SMHPRuntimeManager(...), - data_s3_path="s3://bucket/data.jsonl", - data_mixing_enabled=True, - # is_multimodal=False, # Optional: skip auto-detection for performance -) - -customizer.set_data_mixing_config({ - "customer_data_percent": 50, - "nova_code_percent": 30, - "nova_general_percent": 70 -}) - -# Multimodal datamixing (explicit) -customizer_mm = NovaModelCustomizer( - model=Model.NOVA_LITE_2, - method=TrainingMethod.SFT_LORA, - infra=SMHPRuntimeManager(...), - data_s3_path="s3://bucket/multimodal_data.jsonl", - data_mixing_enabled=True, - is_multimodal=True, -) - -customizer_mm.set_data_mixing_config({ - "customer_data_percent": 50, - "nova_general_percent": 70, - "nova_code_percent": 30, -}) -``` - ---- - -#### `train()` -Generates the recipe YAML, configures runtime, and launches a training job. - -**Signature:** -```python -def train( - self, - job_name: str, - recipe_path: Optional[str] = None, - overrides: Optional[Dict[str, Any]] = None, - rft_lambda_arn: Optional[str] = None, - validation_data_s3_path: Optional[str] = None, - dry_run: Optional[bool] = False -) -> TrainingResult -``` - -**Parameters:** -- `job_name` (str): User-defined name for the training job -- `recipe_path` (Optional[str]): Path for a YAML recipe file (both S3 and local paths are accepted) -- `overrides` (Optional[Dict[str, Any]]): Dictionary of configuration overrides. Example overrides below: - - `max_epochs` (int): Maximum number of training epochs - - `lr` (float): Learning rate - - `warmup_steps` (int): Number of warmup steps - - `loraplus_lr_ratio` (float): LoRA+ learning rate ratio - - `global_batch_size` (int): Global batch size - - `max_length` (int): Maximum sequence length - - A full list of available overrides can be found via the [Nova Customization public documentation](https://docs.aws.amazon.com/nova/latest/userguide/customize-fine-tune-sagemaker.html) or by referencing the training recipes [here](https://docs.aws.amazon.com/sagemaker/latest/dg/nova-model-recipes.html). -- `rft_lambda_arn` (Optional[str]): Rewards Lambda ARN (only used for RFT training methods). If passed, takes priority over `rft_lambda_arn` set on the `RuntimeManager`. -- `validation_data_s3_path` (Optional[str]): Validation S3 path, only applicable for CPT (but is still optional for CPT) -- `dry_run` (Optional[bool]): Actually starts a job if False, otherwise just performs validation. - -**Returns:** -- `TrainingResult`: Metadata object (either `SMTJTrainingResult`, `SMHPTrainingResult`, or `BedrockTrainingResult`) containing: - - `job_id` (str): The training job identifier - - `method` (TrainingMethod): The training method used - - `started_time` (datetime): Job start timestamp - - `model_artifacts` (ModelArtifacts): Paths to model checkpoints and outputs - - `checkpoint_s3_path` (str, Optional): Path to the model checkpoint/trained model. For `SMTJServerless`, populated after job completion via `get_model_artifacts()`. - - `output_s3_path` (str): Path to the metrics and output tar file. - - `output_model_arn` (str, Optional): Model package ARN for `SMTJServerless` jobs. Use as `model_path` for iterative training. - - `model_type` (Model): Model type of the model being trained - -**Raises:** -- `Exception`: If job execution fails -- `ValueError`: If training method is not supported - -**Example:** -```python -result = customizer.train( - job_name="my-training-job", - overrides={ - 'max_epochs': 10, - 'lr': 5e-6, - 'warmup_steps': 20, - 'global_batch_size': 128 - } -) -print(f"Training job started: {result.job_id}") -print(f"Checkpoint path: {result.model_artifacts.checkpoint_s3_path}") -``` ---- -#### `evaluate()` -Generates the recipe YAML, configures runtime, and launches an evaluation job. - -**Signature:** -```python -def evaluate( - self, - job_name: str, - eval_task: EvaluationTask, - model_path: Optional[str] = None, - subtask: Optional[str] = None, - data_s3_path: Optional[str] = None, - recipe_path: Optional[str] = None, - overrides: Optional[Dict[str, Any]] = None, - processor: Optional[Dict[str, Any]] = None, - rl_env: Optional[Dict[str, Any]] = None, - dry_run: Optional[bool] = False, - job_result: Optional[TrainingResult] = None -) -> EvaluationResult | None -``` - -**Parameters:** -- `job_name` (str): User-defined name for the evaluation job -- `eval_task` (EvaluationTask): The evaluation task to be performed (e.g., `EvaluationTask.MMLU`) - - The list of available tasks can be found here: [AWS Documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/nova-model-evaluation.html#nova-model-evaluation-benchmark) -- `model_path` (Optional[str]): S3 path for model to evaluate. If not provided, will attempt to extract from `job_result` or the customizer's most recent training job. -- `data_s3_path` (Optional[str]): S3 URI for the dataset. Only required for BYOD (Bring Your Own Data) evaluation tasks. -- `subtask` (Optional[str]): Subtask for evaluation (task-specific) - - The list of available subtasks per task can be found here: [Subtasks](https://docs.aws.amazon.com/sagemaker/latest/dg/nova-model-evaluation.html#nova-model-evaluation-subtasks) -- `recipe_path` (Optional[str]): Path for a YAML recipe file (both S3 and local paths are accepted) -- `overrides` (Optional[Dict[str, Any]]): Dictionary of inference configuration overrides - - `max_new_tokens` (int): Maximum tokens to generate - - `top_k` (int): Top-k sampling parameter - - `top_p` (float): Top-p (nucleus) sampling parameter - - `temperature` (float): Temperature for sampling -- `processor` (Optional[Dict[str, Any]]): Optional, Bring Your Own Metrics/RFT lambda Configuration -- `rl_env` (Optional[Dict[str, Any]]): Optional, Bring your own reinforcement learning environment config. For `RFT_EVAL`, if either `processor` or `rl_env` is explicitly passed, it takes priority over `rft_lambda_arn` set on the `RuntimeManager`. -- `dry_run` (Optional[bool]): Actually starts a job if False, otherwise just performs validation. -- `job_result` (Optional[TrainingResult]): Optional TrainingResult object to extract checkpoint path from. If provided and `model_path` is None, will automatically extract the checkpoint path from the training job's output and validate platform compatibility. - -**Returns:** -- `EvaluationResult(BaseJobResult)`: Metadata object (either `SMTJEvaluationResult`, `SMHPEvaluationResult`, or `BedrockEvaluationResult`) containing: - - `job_id` (str): The evaluation job identifier - - `started_time` (datetime): Job start timestamp - - `eval_output_path` (str): S3 path to evaluation results - - `eval_task` (EvaluationTask): The Evaluation task -- Returns `None` if `dry_run=True` - -**Example:** - -```python -from amzn_nova_forge.recipe import * - -# General eval task (with overrides) -eval_result = customizer.evaluate( - job_name="my-eval-job", - eval_task=EvaluationTask.MMLU, - model_path="s3://my-bucket/checkpoints/my-model", - overrides={ - 'max_new_tokens': 2048, - 'temperature': 0, - 'top_p': 1.0 - } -) -print(f"Evaluation job started: {eval_result.job_id}") - -# BYOM eval task (by providing processor config) -byom_eval_result = customizer.evaluate( - job_name='my-eval-test-byom', - eval_task=EvaluationTask.GEN_QA, - data_s3_path="s3://bucket/data", - processor={ - "lambda_arn": "arn:aws:lambda::123456789012:function:byom-lambda" - } -) -``` ---- -#### `deploy()` -Creates a custom model and deploys it to Amazon Bedrock or SageMaker. - -Deployment behavior when endpoint already exists is controlled by the `deployment_mode` -parameter set during NovaModelCustomizer initialization: -- FAIL_IF_EXISTS: Raise error (default, safest) -- UPDATE_IF_EXISTS: Try in-place update, fail if not supported (PT only) - -**Signature:** -```python -def deploy( - self, - model_artifact_path: Optional[str] = None, - deploy_platform: DeployPlatform = DeployPlatform.BEDROCK_OD, - unit_count: Optional[int] = None, - endpoint_name: Optional[str] = None, - job_result: Optional[TrainingResult] = None, - execution_role_name: Optional[str] = None, - sagemaker_instance_type: Optional[str] = "ml.p5.48xlarge", - sagemaker_environment: Optional[SageMakerEndpointEnvironment] = None, -) -> DeploymentResult -``` - -* **Note:** If DeployPlatform.BEDROCK_PT or DeployPlatform.SAGEMAKER is selected, you must include a value for unit_count. -* **Note:** If `model_artifact_path` is provided, we will NOT attempt to resolve `model_artifact_path` from `job_result` or the enclosing `NovaModelCustomizer` object. - -**Parameters:** -- `model_artifact_path` (Optional[str]): S3 path to the trained model checkpoint. If not provided, will attempt to extract from job_result or the `job_id` field of the Customizer. -- `deploy_platform` (DeployPlatform): Platform to deploy the model to - - `DeployPlatform.BEDROCK_OD`: Bedrock On-Demand - - `DeployPlatform.BEDROCK_PT`: Bedrock Provisioned Throughput - - `DeployPlatform.SAGEMAKER`: SageMaker -- `unit_count` (Optional[int]): Used in Bedrock Provisioned Throughput number of PT to purchase or SageMaker number of initial instances -- `endpoint_name` (Optional[str]): Name of the deployed model's endpoint (auto-generated if not provided) -- `job_result` (Optional[TrainingResult]): Training job result object to use for extracting checkpoint path and validating job completion. Also used to retrieve job_id if it's not provided. -- `execution_role_name` (Optional[str]): IAM role for Bedrock or SageMaker. If provided, used as-is — no policies created or attached. If omitted, the SDK creates and manages a default role with required policies. -- `sagemaker_instance_type`: Optional EC2 instance type for SageMaker deployment, defaults to ml.p5.48xlarge -- `sagemaker_environment` (Optional[SageMakerEndpointEnvironment]): SageMaker endpoint environment config. See `SageMakerEndpointEnvironment` for available fields and defaults. -**Returns:** -- `DeploymentResult`: Contains: - - `endpoint` (EndpointInfo): Endpoint information - - `platform` (DeployPlatform): Deployment platform - - `endpoint_name` (str): Endpoint name - - `uri` (str): Model ARN - - `model_artifact_path` (str): S3 path to artifacts - - `created_at` (datetime): Deployment creation timestamp - -**Raises:** -- `Exception`: When unable to successfully deploy the model -- `ValueError`: If platform is not supported - -**Example:** -```python -from amzn_nova_forge.model import * - -bedrock_deployment = customizer.deploy( - model_artifact_path="s3://escrow-bucket/my-model-artifacts/", - deploy_platform=DeployPlatform.BEDROCK_OD, - endpoint_name="my-custom-nova-model-bedrock" -) -print(f"Model deployed: {bedrock_deployment.endpoint.uri}") -print(f"Endpoint: {bedrock_deployment.endpoint.endpoint_name}") -print(f"Status: {bedrock_deployment.status}") - -sagemaker_deployment = customizer.deploy( - model_artifact_path="s3://escrow-bucket/my-model-artifacts/", - deploy_platform=DeployPlatform.SAGEMAKER, - unit_count=1, - endpoint_name="my-custom-nova-model-sagemaker", - sagemaker_environment=SageMakerEndpointEnvironment( - context_length=8000, - max_concurrency=4, - ) -) -print(f"Model deployed: {sagemaker_deployment.endpoint.uri}") -print(f"Endpoint: {sagemaker_deployment.endpoint.endpoint_name}") -print(f"Status: {sagemaker_deployment.status}") -``` - -Optionally, you can provide a Bedrock execution role name to be used in deployment. -Otherwise, a default Bedrock execution role will be created on your behalf. -You can also use the following method to create a Bedrock execution role with scoped down IAM permissions. - - -```python -from amzn_nova_forge.iam import create_bedrock_execution_role - -iam_client = boto3.client("iam") - -create_bedrock_execution_role( - iam_client=iam_client, - role_name="BedrockDeployModelExecutionRole", - bedrock_resource="your-model-name", # Optional: Name of the bedrock resources that IAM role should have restricted create and get access to - s3_resource="s3-bucket" # Optional: S3 resource that IAM role should have restricted read access to such as the training output bucket -) - -``` ---- -#### `create_custom_model()` -Creates a Bedrock custom model from S3 artifacts, decoupled from endpoint deployment. - -This method extracts the model-creation step from the deploy flow so users can -create a model independently of endpoint deployment, enabling retry of deployment -if it fails after model creation. - -**Signature:** -```python -def create_custom_model( - self, - model_artifact_path: Optional[str] = None, - job_result: Optional[TrainingResult] = None, - endpoint_name: Optional[str] = None, - execution_role_name: Optional[str] = None, - tags: Optional[List[Dict[str, str]]] = None, - skip_model_reuse: bool = False, -) -> ModelDeployResult -``` - -**Parameters:** -- `model_artifact_path` (Optional[str]): S3 path to trained model checkpoint. Takes precedence over `job_result` if both are provided. -- `job_result` (Optional[TrainingResult]): Training job result to extract checkpoint path from. -- `endpoint_name` (Optional[str]): Optional name prefix for the model name (auto-generated if not provided). -- `execution_role_name` (Optional[str]): IAM role name for Bedrock. Defaults to `BedrockDeployModelExecutionRole`. -- `tags` (Optional[List[Dict[str, str]]]): Optional list of `{"key": str, "value": str}` dicts for source tracking. -- `skip_model_reuse` (bool): If True, always create a new model even if one with the same escrow URI already exists. Default: False. - -**Returns:** -- `ModelDeployResult`: Contains: - - `model_arn` (str): The Bedrock custom model ARN - - `model_name` (str): The model name passed to CreateCustomModel - - `escrow_uri` (str): S3 artifacts path used to create the model - - `created_at` (datetime): UTC timestamp when the model was created - -**Raises:** -- `ValueError`: When neither `model_artifact_path` nor `job_result` is provided, or when checkpoint path cannot be resolved from `job_result`. -- `RuntimeError`: When IAM role creation or custom model creation fails. - -**Example:** -```python -# Create a custom model from training artifacts -publish_result = customizer.create_custom_model( - model_artifact_path="s3://escrow-bucket/my-model-artifacts/" -) -print(f"Model ARN: {publish_result.model_arn}") - -# Save for later use -publish_result.dump(file_path="./results/") - -# Or create from a training job result -publish_result = customizer.create_custom_model(job_result=training_result) -``` ---- -#### `deploy_to_bedrock()` -Deploys a published Bedrock custom model to an endpoint. - -Use after `create_custom_model()` to deploy the model, or provide a model ARN directly. -This decoupled approach allows retrying deployment without re-creating the model. - -When a `model_deploy_result` is provided, the model's status is validated before deployment: -- **Active**: Proceeds immediately. -- **Creating**: Waits for the model to become Active (30s poll interval, 15min timeout). -- **Failed**: Raises `ValueError`. - -**Signature:** -```python -def deploy_to_bedrock( - self, - model_deploy_result: Optional[ModelDeployResult] = None, - model_arn: Optional[str] = None, - deploy_platform: DeployPlatform = DeployPlatform.BEDROCK_OD, - pt_units: Optional[int] = None, - endpoint_name: Optional[str] = None, -) -> DeploymentResult -``` - -**Parameters:** -- `model_deploy_result` (Optional[ModelDeployResult]): Result from `create_custom_model()`. Cannot be combined with `model_arn`. -- `model_arn` (Optional[str]): Direct model ARN. Cannot be combined with `model_deploy_result`. -- `deploy_platform` (DeployPlatform): `BEDROCK_OD` (default) or `BEDROCK_PT`. -- `pt_units` (Optional[int]): Number of PT units (required for `BEDROCK_PT`). -- `endpoint_name` (Optional[str]): Endpoint name (auto-generated if not provided). - -**Returns:** -- `DeploymentResult`: Contains: - - `endpoint` (EndpointInfo): Endpoint information - - `created_at` (datetime): Deployment creation timestamp - - `model_publish` (Optional[ModelDeployResult]): The model deploy result, if available - - `escrow_uri` (Optional[str]): Convenience property delegating to `model_publish.escrow_uri` - -**Raises:** -- `ValueError`: When both `model_deploy_result` and `model_arn` are provided, or when no model ARN is available. -- `RuntimeError`: When deployment creation fails. - -**Example:** -```python -# Two-step deploy: create model, then deploy -publish_result = customizer.create_custom_model( - model_artifact_path="s3://escrow-bucket/my-model-artifacts/" -) -deployment = customizer.deploy_to_bedrock( - model_deploy_result=publish_result, - endpoint_name="my-endpoint" -) - -# Or deploy directly from a model ARN -deployment = customizer.deploy_to_bedrock( - model_arn="arn:aws:bedrock:us-east-1:123456789012:custom-model/my-model" -) - -# Retry deployment if it fails (model already created) -deployment = customizer.deploy_to_bedrock( - model_arn=publish_result.model_arn -) -``` ---- -#### `deploy_to_sagemaker()` -Deploys a model to a SageMaker Inference endpoint. - -Can reuse an existing SageMaker model via `model_deploy_result` (e.g., from a -previous deploy that created the model but failed on endpoint creation), -or create a new model from `model_artifact_path`. - -**Signature:** -```python -def deploy_to_sagemaker( - self, - instance_type: str, - model_deploy_result: Optional[ModelDeployResult] = None, - model_artifact_path: Optional[str] = None, - unit_count: int = 1, - endpoint_name: Optional[str] = None, - sagemaker_environment: Optional[SageMakerEndpointEnvironment] = None, - execution_role_name: Optional[str] = None, - skip_model_reuse: bool = False, -) -> DeploymentResult -``` - -**Parameters:** -- `instance_type` (str): SageMaker instance type (required). -- `model_deploy_result` (Optional[ModelDeployResult]): Result from a previous deploy containing the SM model ARN. Skips model creation. Cannot be combined with `model_artifact_path`. -- `model_artifact_path` (Optional[str]): S3 path to model artifacts. Creates a new SM model. -- `unit_count` (int): Number of instances. Default: 1. -- `endpoint_name` (Optional[str]): Endpoint name (auto-generated if not provided). -- `sagemaker_environment` (Optional[SageMakerEndpointEnvironment]): SageMaker endpoint environment config. See `SageMakerEndpointEnvironment` for available fields and defaults. -- `execution_role_name` (str): IAM execution role name. -- `skip_model_reuse` (bool): If True, always create a new model (skip tag-based discovery). Default: False. - -**Returns:** -- `DeploymentResult`: Contains endpoint info and `model_publish` (ModelDeployResult). - -**Raises:** -- `ValueError`: When both `model_deploy_result` and `model_artifact_path` are provided, when neither is provided, or when `instance_type` is missing. -- `RuntimeError`: When endpoint creation fails. The error message includes the model ARN and a retry command using `customizer.last_model_publish`. - -**Example:** -```python -# Deploy from S3 artifacts -result = customizer.deploy_to_sagemaker( - model_artifact_path="s3://escrow-bucket/checkpoint/", - instance_type="ml.g5.12xlarge", -) - -# Retry after endpoint failure (model already created) -result = customizer.deploy_to_sagemaker( - model_deploy_result=customizer.last_model_publish, - instance_type="ml.g5.12xlarge", - endpoint_name="my-endpoint", -) -``` ---- - -#### `invoke_inference()` -Invokes a single inference on a trained model. - -**Signature:** -```python -def invoke_inference( - self, - request_body: Dict[str, Any], - endpoint_arn: Optional[str] -) -> InferenceResult -``` -**Parameters:** -- `request_body` (Dict[str, Any]): Inference request body -- `endpoint_arn` (Optional[str]):Endpoint ARN to invoke inference. Optional if user wants to send request to an already deployed endpoint on customizer - -**Returns:** -- `InferenceResult`: Metadata object (`SingleInferenceResult`) containing: - - `job_id` (str): Batch inference job identifier - - `started_time` (datetime): Job start timestamp - - `inference_output_path` (str): Empty string - -**Example:** -```python -inference_result = customizer.invoke_inference( - request_body={ - "messages": [{"role": "user", "content": "Hello! How are you?"}], - "max_tokens": 100, - "stream": False, - }, - endpoint_arn="arn:aws:sagemaker:us-east-1:123456789012:endpoint/endpoint", -) -inference_result.show() -``` ---- -#### `batch_inference()` -Launches a batch inference job on a trained model. - -**Signature:** -```python -def batch_inference( - self, - job_name: str, - input_path: str, - output_s3_path: str, - model_path: Optional[str] = None, - recipe_path: Optional[str] = None, - overrides: Optional[Dict[str, Any]] = None, - dry_run: Optional[bool] = False -) -> InferenceResult -``` -**Parameters:** -- `job_name` (str): Name for the batch inference job -- `input_path` (str): S3 path to input data for inference -- `output_s3_path` (str): S3 path for inference outputs -- `model_path` (Optional[str]): S3 path to the model -- `recipe_path` (Optional[str]): Path for a YAML recipe file -- `overrides` (Optional[Dict[str, Any]]): Configuration overrides for inference - - `max_new_tokens` (int): Maximum tokens to generate - - `top_k` (int): Top-k sampling parameter - - `top_p` (float): Top-p (nucleus) sampling parameter - - `temperature` (float): Temperature for sampling - - `top_logprobs` (int): Number of top log probabilities to return -- `dry_run` (Optional[bool]): Actually starts a job if False, otherwise just performs validation. - -**Returns:** -- `InferenceResult`: Metadata object (`SMTJBatchInferenceResult`) containing: - - `job_id` (str): Batch inference job identifier - - `started_time` (datetime): Job start timestamp - - `inference_output_path` (str): S3 path to inference results - - Note: Batch inference is only supported on SageMaker platforms (SMTJ, SMHP) - -**Example:** -```python -inference_result = customizer.batch_inference( - job_name="batch-inference-job", - input_path="s3://my-bucket/inference-input", - output_s3_path="s3://my-bucket/inference-output", - model_path="s3://my-bucket/trained-model" -) -print(f"Batch inference started: {inference_result.job_id}") -``` -In a separate notebook cell, you can run the following commands to get the job status and download a formatted result file when the jobs completes. -```python -inference_result.get_job_status() # Gets the job status. -inference_result.get("s3://my-bucket/save-location/file-name.jsonl") # Uploads a formatted inference_results.jsonl file to the given s3 location. -``` ---- -#### `get_logs()` -Retrieves and displays CloudWatch logs for the current job. - -**Signature:** -```python -def get_logs( - self, - limit: Optional[int] = None, - start_from_head: bool = False, - end_time: Optional[int] = None -) -``` - -**Parameters:** -- `limit` (Optional[int]): Maximum number of log lines to retrieve -- `start_from_head` (bool): If True, start from the beginning of logs; if False, start from the end -- `end_time` (Optional[int]): End time in epoch milliseconds for searching a log time range - -**Returns:** -- None (prints logs to console) - -**Example:** -```python -# After starting a training job -customizer.train(job_name="my-job") -customizer.get_logs(limit=100, start_from_head=True) -``` ---- -## Runtime Managers -Runtime managers handle the infrastructure for executing training and evaluation jobs, -leveraging the `JobConfig` dataclass to do so: -```python -@dataclass -class JobConfig: - job_name: str - image_uri: str - recipe_path: str - output_s3_path: Optional[str] = None - data_s3_path: Optional[str] = None - input_s3_data_type: Optional[str] = None - mlflow_tracking_uri: Optional[str] = None - mlflow_experiment_name: Optional[str] = None - mlflow_run_name: Optional[str] = None -``` -* The specific instance types that can be used with the runtime managers (SMTJ, SMHP) can be found in `docs/instance_type_spec.md`. -* This file also defines which instance types can be used with a specific model and method. -* Bedrock is fully managed and does not require instance type configuration. - -### Shared RuntimeManager Methods - -The following methods are available on all `RuntimeManager` subclasses. - -#### Properties (shared) -- `rft_lambda` (Optional[str]): Lambda ARN, SageMaker hub-content ARN (SMTJServerless only), or local `.py` file path. Assigning a new value automatically updates `rft_lambda_arn` — if the value is a Lambda ARN or hub-content ARN it is resolved immediately; if it is a file path, `rft_lambda_arn` is cleared until `deploy_lambda()` is called. -- `rft_lambda_arn` (Optional[str]): Resolved Lambda ARN or hub-content ARN. Set immediately when `rft_lambda` is assigned an ARN (Lambda or hub-content), or populated by `deploy_lambda()` when `rft_lambda` is a file path. - -**Example:** -```python -# Set an ARN directly — rft_lambda_arn is updated immediately -runtime.rft_lambda = 'arn:aws:lambda:us-east-1:123456789012:function:my-reward-fn' -print(runtime.rft_lambda_arn) # 'arn:aws:lambda:us-east-1:123456789012:function:my-reward-fn' - -# Set a file path — rft_lambda_arn is cleared until deploy_lambda() is called -runtime.rft_lambda = 'reward.py' -print(runtime.rft_lambda_arn) # None -runtime.deploy_lambda(lambda_name='my-reward-fn') -print(runtime.rft_lambda_arn) -#'arn:aws:lambda:us-east-1:123456789012:function:my-reward-fn' -``` - ---- - -#### `deploy_lambda()` - -Packages a local Python file into a zip and creates or updates a Lambda function. The source file is read from `self.rft_lambda`, which must be set to a local `.py` file path before calling this method. - -**Signature:** -```python -def deploy_lambda( - self, - lambda_name: Optional[str] = None, - execution_role_arn: Optional[str] = None, -) -> str -``` - -**Parameters:** -- `lambda_name` (Optional[str]): Name for the Lambda function. Defaults to the source filename stem (underscores replaced with hyphens). -- `execution_role_arn` (Optional[str]): IAM role ARN for the Lambda. Falls back to the runtime manager's `execution_role` attribute if not provided. - -**Returns:** -- `str`: The deployed Lambda function ARN. Also sets `self.rft_lambda_arn` on the manager. - -**Raises:** -- `ValueError`: If `rft_lambda` is not set, is already an ARN (nothing to deploy), the source file is not found, or no execution role can be resolved. - -**Example:** -```python -runtime.rft_lambda = 'rft_training_reward.py' -lambda_arn = runtime.deploy_lambda(lambda_name='my-reward-fn') -# runtime.rft_lambda_arn is now set automatically -``` - ---- - -#### `validate_lambda()` - -Validates the RFT reward lambda with sample data from S3. Reads the lambda to validate from `self.rft_lambda` / `self.rft_lambda_arn`: -- If `rft_lambda` is an ARN (or `rft_lambda_arn` is set), invokes the deployed Lambda with samples from `data_s3_path`. -- If `rft_lambda` is a local `.py` path, validates by executing `lambda_handler` directly without deploying. - -**Signature:** -```python -def validate_lambda( - self, - data_s3_path: str, - validation_samples: int = 10, -) -> None -``` - -**Parameters:** -- `data_s3_path` (str): S3 path to the training dataset for pulling sample data. -- `validation_samples` (int): Number of samples to load from `data_s3_path` (default: 10). - -**Raises:** -- `ValueError`: If `rft_lambda` is not set, or if validation fails. - -**Example:** -```python -# Validate a local file without deploying -runtime.rft_lambda = 'rft_training_reward.py' -runtime.validate_lambda(data_s3_path='s3://bucket/data.jsonl') - -# Validate a deployed lambda -runtime.rft_lambda = 'arn:aws:lambda:us-east-1:123456789012:function:my-reward-fn' -runtime.validate_lambda(data_s3_path='s3://bucket/data.jsonl', validation_samples=20) -``` - ---- - -### SMTJRuntimeManager -Manages SageMaker Training Jobs. - -#### Constructor - -**Signature:** -```python -def __init__( - self, - instance_type: str, - instance_count: int, - execution_role: Optional[str] = None, - kms_key_id: Optional[str] = None, - encrypt_inter_container_traffic: bool = False, - subnets: Optional[list[str]] = None, - security_group_ids: Optional[list[str]] = None, - max_job_runtime: Optional[int] = 86400, - job_submit_poll_timeout: int = 30, - rft_lambda: Optional[str] = None, -) -``` - -**Parameters:** -- `instance_type` (str): EC2 instance type (e.g., "ml.p5.48xlarge", "ml.p4d.24xlarge") -- `instance_count` (int): Number of instances to use -- `execution_role` (Optional[str]): The execution role for the training job -- `kms_key_id` (Optional[str]): Optional KMS Key Id to use in S3 Bucket encryption, training jobs and deployments. -- `encrypt_inter_container_traffic` (bool): Boolean that determines whether to encrypt inter-container traffic. Default value is False. -- `subnets` (Optional[list[str]]): Optional list of strings representing subnets. Default value is None. -- `security_group_ids` (Optional[list[str]]): Optional list of strings representing security group IDs. Default value is None. -- `max_job_runtime` (Optional[int]): Max Job Runtime in seconds (default: 1 day) -- `job_submit_poll_timeout` (int): Maximum seconds to poll for the training job after submission (default: 30). Uses exponential backoff. -- `rft_lambda` (Optional[str]): Lambda ARN or local `.py` file path for RFT reward function. Can also be set or updated after construction. - -**Example:** -```python -from amzn_nova_forge.manager import * -infra = SMTJRuntimeManager( - instance_type="ml.p5.48xlarge", - instance_count=2 -) -``` -#### Properties -- `instance_type` (str): Returns the instance type -- `instance_count` (int): Returns the number of instances - -#### Methods - -##### `execute()` -Starts a SageMaker training job. - -**Signature:** -```python -def execute( - self, - job_config: JobConfig -) -> str -``` - -**Returns:** -- `str`: Training job name/ID - -##### `cleanup()` -Stops and cleans up a training job. - -**Signature:** -```python -def cleanup( - self, - job_name: str -) -> None -``` ---- -### SMHPRuntimeManager -Manages SageMaker HyperPod jobs. - -#### Constructor - -**Signature:** -```python -def __init__( - self, - instance_type: str, - instance_count: int, - cluster_name: str, - namespace: str, - kms_key_id: Optional[str] = None, - rft_lambda: Optional[str] = None, -) -``` - -**Parameters:** -- `instance_type` (str): EC2 instance type -- `instance_count` (int): Number of instances -- `cluster_name` (str): HyperPod cluster name -- `namespace` (str): Kubernetes namespace -- `kms_key_id` (Optional[str]): Optional KMS Key Id to use in S3 Bucket encryption -- `rft_lambda` (Optional[str]): Lambda ARN or local `.py` file path for RFT reward function. Can also be set or updated after construction. - -**Example:** -```python -from amzn_nova_forge.manager import * -infra = SMHPRuntimeManager( - instance_type="ml.p5.48xlarge", - instance_count=4, - cluster_name="my-hyperpod-cluster", - namespace="default" -) -``` -#### Properties -- `instance_type` (str): Returns the instance type -- `instance_count` (int): Returns the number of instances - -#### Methods - -##### `execute()` - -Starts a SageMaker HyperPod job. -**Signature:** -```python -def execute( - self, - job_config=JobConfig -) -> str -``` -**Returns:** -- `str`: HyperPod job ID - -##### `cleanup()` -Cancels and cleans up a HyperPod job. - -**Signature:** -```python -def cleanup( - self, - job_name: str -) -> None -``` - -##### `scale_cluster()` -Scale a HyperPod cluster instance group up or down. -The scaling operation is asynchronous - the cluster status will change to 'Updating' while scaling, and 'InService' when ready. - -**Signature:** -```python -def scale_cluster( - self, - instance_group_name: str, - target_instance_count: int, -) -> Dict[str, Any] -``` - -**Parameters:** -- `instance_group_name` (str): Name of the instance group to scale (e.g., 'worker-group') -- `target_instance_count` (int): Desired number of instances for the group (must be non-negative) - -**Returns:** -- `Dict[str, Any]`: Response containing: - - `ClusterArn` (str): ARN of the updated cluster - - `InstanceGroupName` (str): Name of the scaled instance group - - `InstanceType` (str): Instance type being scaled - - `PreviousCount` (int): Current instance count before scaling - - `TargetCount` (int): Target instance count after scaling - -**Raises:** -- `ValueError`: If target_instance_count is negative or instance group name is invalid -- `ClientError`: If scaling fails due to insufficient quota, capacity or other cluster issues. - -**Example:** -```python -from amzn_nova_forge.manager import * - -# Create a runtime manager for your cluster -manager = SMHPRuntimeManager( - instance_type="ml.p4d.24xlarge", - instance_count=4, - cluster_name="my-hyperpod-cluster", - namespace="default" -) - -# Scale up the worker group from 4 to 8 instances -result = manager.scale_cluster( - instance_group_name="worker-group", - target_instance_count=8 -) - -# Scale down to 2 instances -result = manager.scale_cluster( - instance_group_name="worker-group", - target_instance_count=2 -) -``` -**Notes:** -- This method only works with Restricted Instance Groups (RIGs) in HyperPod clusters. The cluster must be in 'InService' state before scaling can be initiated. -- This method can only scale up a SMHP cluster when there is sufficient Service Quota available. -You will need to request a quota increase **before** scaling up a RIG in your HyperPod cluster. -You can learn more [here](https://docs.aws.amazon.com/servicequotas/latest/userguide/request-quota-increase.html). - - Specifically, you will need to request a service quota increase for "INSTANCE_TYPE for cluster usage". - -##### `get_instance_groups()` -Gets the RIGs associated with the current cluster defined in the SMHPRuntimeManager. -Prints the values to the terminal and returns it as a list of dictionary entries. - -**Signature:** -```python -def get_instance_groups( - self -) -> List[Dict[str, Any]] -``` - -**Returns:** -- `List[Dict[str, Any]]`: Response containing: - - InstanceGroupName: Name of the instance group - - InstanceType: EC2 instance type (e.g., 'ml.p5.48xlarge') - - CurrentCount: Current number of instances in the group - -**Raises:** -- `ClientError`: If unable to describe the cluster - -**Example:** -```python -from amzn_nova_forge.manager import * - -# Create a runtime manager for your cluster -manager = SMHPRuntimeManager( - instance_type="ml.p4d.24xlarge", - instance_count=4, - cluster_name="my-hyperpod-cluster", - namespace="default" -) - -# Get the instance groups available on the current cluster. -instance_groups = manager.get_instance_groups() -``` - ---- -### BedrockRuntimeManager -Manages Amazon Bedrock model customization jobs. - -#### Constructor - -**Signature:** -```python -def __init__( - self, - execution_role: str, - base_model_identifier: Optional[str] = None, - kms_key_id: Optional[str] = None, - rft_lambda: Optional[str] = None, -) -``` - -**Parameters:** -- `execution_role` (str): IAM role ARN for Bedrock job execution -- `base_model_identifier` (Optional[str]): Base model ARN (e.g., "arn:aws:bedrock:us-east-1::foundation-model/amazon.nova-2-lite-v1:0:256k") -- `kms_key_id` (Optional[str]): Optional KMS Key Id for encryption -- `rft_lambda` (Optional[str]): Lambda ARN or local `.py` file path for RFT reward function. Can also be set or updated after construction. - -**Example:** -```python -from amzn_nova_forge.manager import * -infra = BedrockRuntimeManager( - execution_role="arn:aws:iam::123456789012:role/BedrockRole", - base_model_identifier="arn:aws:bedrock:us-east-1::custom-model/amazon.nova-2-lite-v1:0:256k:abcdefghijk" # optional: your custom model ARN for iterative training -) -``` - -#### Methods - -##### `execute()` - -Starts a Bedrock model customization job. -**Signature:** -```python -def execute( - self, - job_config: JobConfig -) -> str -``` -**Returns:** -- `str`: Bedrock job ARN - -##### `cleanup()` -Stops a Bedrock customization job. - -**Signature:** -```python -def cleanup( - self, - job_name: str -) -> None -``` ---- -### SMTJServerlessRuntimeManager -Manages SageMaker Training Jobs. - -> **Note:** `AWS_DEFAULT_REGION` must be set when using SageMaker Serverless training. -> The SageMaker SDK's `DataSet` API requires a region to connect to the SageMaker backend. -> Set it before running your script: -> ```bash -> export AWS_DEFAULT_REGION= -> ``` - -#### Constructor - -**Signature:** -```python -def __init__( - self, - model_package_group_name: str, - execution_role: Optional[str] = None, - kms_key_id: Optional[str] = None, - encrypt_inter_container_traffic: bool = False, - subnets: Optional[list[str]] = None, - security_group_ids: Optional[list[str]] = None, - max_job_runtime: Optional[int] = 86400, - rft_lambda: Optional[str] = None, - evaluator_name: Optional[str] = None, -) -``` - -**Parameters:** -- `model_package_group_name` (str): Model package group name to use with SageMaker Model registry (required for SMTJ Serverless) -- `execution_role` (Optional[str]): The execution role for the training job -- `kms_key_id` (Optional[str]): Optional KMS Key Id to use in S3 Bucket encryption, training jobs and deployments. -- `encrypt_inter_container_traffic` (bool): Boolean that determines whether to encrypt inter-container traffic. Default value is False. -- `subnets` (Optional[list[str]]): Optional list of strings representing subnets. Default value is None. -- `security_group_ids` (Optional[list[str]]): Optional list of strings representing security group IDs. Default value is None. -- `max_job_runtime` (Optional[int]): Max Job Runtime in seconds (default: 1 day) -- `rft_lambda` (Optional[str]): Lambda ARN, SageMaker hub-content ARN, or local `.py` file path for the RFT reward function. - - **Lambda ARN**: Automatically registered as a hub-content `JsonDoc` evaluator during `train()`. The hub-content ARN is passed as `EvaluatorArn` in `ServerlessJobConfig`. - - **Hub-content ARN**: Passed directly as `EvaluatorArn` — no registration needed. - - **Local `.py` file**: Call `deploy_lambda()` first to deploy and get a Lambda ARN. -- `evaluator_name` (Optional[str]): Optional human-readable name for the hub-content evaluator entry when auto-registering a Lambda ARN. Defaults to the Lambda function name. - -**Example:** -```python -from amzn_nova_forge.manager import * - -# With a Lambda ARN (auto-registered as hub-content during train()) -infra = SMTJServerlessRuntimeManager( - model_package_group_name="nova-rft-serverless", - rft_lambda="arn:aws:lambda:us-east-1:123456789012:function:my-reward-fn", -) - -# With a hub-content ARN (passed directly as EvaluatorArn) -infra = SMTJServerlessRuntimeManager( - model_package_group_name="nova-rft-serverless", - rft_lambda="arn:aws:sagemaker:us-east-1:123456789012:hub-content/my-hub/JsonDoc/my-evaluator/0.0.1", -) - -# With a local .py file (deploy_lambda() required before train()) -infra = SMTJServerlessRuntimeManager( - model_package_group_name="nova-rft-serverless", - rft_lambda="my_reward.py", -) -infra.deploy_lambda(lambda_name="my-reward-fn") -``` -#### Properties -- `model_package_group_name` (str): Model Package Group name - -#### Methods - -##### `execute()` -Starts a SageMaker training job. - -**Signature:** -```python -def execute( - self, - job_config: JobConfig -) -> str -``` - -**Returns:** -- `str`: Training job name/ID - -##### `cleanup()` -Stops and cleans up a training job. - -**Signature:** -```python -def cleanup( - self, - job_name: str -) -> None -``` ---- - -## Dataset Loaders -Dataset loaders handle loading, transforming, and saving datasets in various formats. - -### Base Class: DatasetLoader -Base class for all dataset loaders. Handles path resolution (relative, tilde, `..` paths normalized to absolute), directory detection (scans for matching files, loads in lexicographic order), and delegates single-file loading to subclasses. - -#### Constructor - -**Signature:** -```python -def __init__( - self, - **column_mappings -) -``` - -**Parameters:** -- `**column_mappings`: Keyword arguments mapping standard column names to dataset column names - - Example: `question="input"` where "question" is the standard name and "input" is your column name -#### Column Mappings -If you are transforming a plain JSON, JSONL, or CSV file from a generic format (e.g. 'input/output') to another format (e.g. Converse for SFT), you need to provide "column mappings" to connect your generic column/field name to the expected ones in the transformation function. - -For example, if your plain dataset has "input" and "output" columns, and you want to transform it for SFT (which requrires "question" and "answer"), you would provide the following: -```python -loader = JSONDatasetLoader( - question="input", - answer="output" -) -``` -Below is a list of accepted column mapping parameters for transformations. -* SFT: `question`, `answer` - * Optional: `system`, [image/video required options]: `image_format`/`video_format`, `s3_uri`, `bucket_owner` - * 2.0: `reasoning_text`, `tools`/`toolsConfig`* -* RFT: `question`, `reference_answer` - * Optional: `system`, `id`, `tools`* -* Eval: `query`, `response` - * Optional: `images`, `metadata` -* CPT: `text` - -Additional Notes: -* If you're providing multimodal data in a generic format, you need to provide ALL three of the following fields: - * `image_format` OR `video_format` + `s3_uri`, `bucket_owner` -* *`tools/toolsConfig` (SFT 2.0) and `tools` (RFT) parameters can *only* be provided when transforming from OpenAI Messages format to Converse or OpenAI. A generic format *cannot* be provided for this transformation to work. - ---- -### JSONLDatasetLoader -Loads datasets from JSONL (JSON Lines) files. - -#### Methods - -##### `load()` -Loads dataset from a JSONL file (local or S3). - -**Signature:** -```python -def load( - self, - path: str -) -> "DatasetLoader" -``` - -**Parameters:** -- `path` (str): Path to JSONL file (local path or S3 URI) - -**Returns:** -- `DatasetLoader`: Self (for method chaining) - -**Example:** -```python -from amzn_nova_forge.dataset import * -loader = JSONLDatasetLoader() -loader.load("s3://my-bucket/data/training.jsonl") -``` ---- -### JSONDatasetLoader -Loads datasets from JSON files. - -#### Methods - -##### `load()` -Loads dataset from a JSON file (local or S3). - -**Signature:** -```python -def load( - self, - path: str -) -> "DatasetLoader" -``` - -**Parameters:** -- `path` (str): Path to JSON file (local path or S3 URI) - -**Returns:** -- `DatasetLoader`: Self (for method chaining) - -**Example:** -```python -from amzn_nova_forge.dataset import * -loader = JSONDatasetLoader() -loader.load("data/training.json") -``` ---- -### CSVDatasetLoader -Loads datasets from CSV files. - -#### Methods - -##### `load()` -Loads dataset from a CSV file. - -**Signature:** -```python -def load( - self, - path: str -) -> "DatasetLoader" -``` - -**Parameters:** -- `path` (str): Path to CSV file (local path or S3 URI) - -**Returns:** -- `DatasetLoader`: Self (for method chaining) - -**Example:** -```python -from amzn_nova_forge.dataset import * -loader = CSVDatasetLoader(question="user_query", answer="bot_response") -loader.load("data/conversations.csv") -``` ---- -### ParquetDatasetLoader -Loads datasets from Apache Parquet files. Uses lazy batch-based iteration via PyArrow — only one row group is in memory at a time. - -#### Methods - -##### `load()` -Loads dataset from a Parquet file or directory (local or S3). - -**Signature:** -```python -def load( - self, - path: str -) -> "DatasetLoader" -``` - -**Parameters:** -- `path` (str): Path to Parquet file or directory (local path or S3 URI). Accepts `.parquet` and `.pq` extensions. - -**Returns:** -- `DatasetLoader`: Self (for method chaining) - -**Example:** -```python -from amzn_nova_forge.dataset import * -loader = ParquetDatasetLoader() -loader.load("s3://my-bucket/data/training.parquet") -``` ---- -### ArrowDatasetLoader -Loads datasets from Apache Arrow IPC and Feather files. Supports both IPC Stream and IPC File formats with automatic fallback. - -#### Methods - -##### `load()` -Loads dataset from an Arrow IPC or Feather file or directory (local or S3). - -**Signature:** -```python -def load( - self, - path: str -) -> "DatasetLoader" -``` - -**Parameters:** -- `path` (str): Path to Arrow file or directory (local path or S3 URI). Accepts `.arrow`, `.feather`, and `.ipc` extensions. - -**Returns:** -- `DatasetLoader`: Self (for method chaining) - -**Example:** -```python -from amzn_nova_forge.dataset import * -loader = ArrowDatasetLoader() -loader.load("s3://my-bucket/data/training.arrow") -``` ---- -### HuggingFaceDatasetLoader -Loads datasets from [HuggingFace Hub](https://huggingface.co/datasets) and streams them through the Nova Forge data prep pipeline. Records are streamed lazily — no network calls happen at `load()` time. The dataset is fetched when a terminal operation (`save()`, `show()`, etc.) executes the pipeline. - -Requires the optional `datasets` dependency: -```bash -pip install amzn-nova-forge[huggingface] -``` - -Authentication for private or gated datasets is handled by the `datasets` library's built-in credential resolution. Set the `HF_TOKEN` environment variable or run `huggingface-cli login` before using the loader. - -#### Methods - -##### `load()` -Loads a dataset from HuggingFace Hub. - -**Signature:** -```python -def load( - self, - path: str, - split: Optional[str] = None, - name: Optional[str] = None, - revision: Optional[str] = None, - data_files: Optional[Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]] = None, - data_dir: Optional[str] = None, -) -> "DatasetLoader" -``` - -**Parameters:** -- `path` (str): HuggingFace dataset identifier. -- `split` (Optional[str]): Dataset split, forwarded to `datasets.load_dataset(split=...)`. Accepts any value supported by the HF library, including a single split name (`"train"`, `"test"`), a slice (`"train[:100]"`, `"train[:10%]"`), or concatenated splits (`"train+test"`). When `None`, the loader probes available splits: if the dataset has exactly one split it is auto-selected, and if the dataset has multiple splits a `DataPrepError` is raised listing the available splits. -- `name` (Optional[str]): Configuration name for multi-config datasets (forwarded to `datasets.load_dataset(name=...)`). -- `revision` (Optional[str]): Git tag or commit hash to pin the dataset version (forwarded to `datasets.load_dataset(revision=...)`). -- `data_files` (Optional[Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]]): Specific data file(s) to load, forwarded to `datasets.load_dataset(data_files=...)`. When `None` (default), the kwarg is omitted so the library's own default applies. -- `data_dir` (Optional[str]): Subdirectory of the dataset repository to load from, forwarded to `datasets.load_dataset(data_dir=...)`. When `None` (default), the kwarg is omitted. - -**Returns:** -- `DatasetLoader`: Self (for method chaining) - -**Raises:** -- `ImportError`: If the `datasets` package is not installed. The error message includes the install command. -- `DataPrepError`: If the dataset cannot be loaded. Authentication failures (401/403) surface with credential guidance; all other failures include the dataset path and the original error. Raised lazily when the pipeline is executed, not at `load()` time. - -**Example — public dataset:** -```python -from amzn_nova_forge.dataset import * - -# Load a public dataset and save to S3 -loader = HuggingFaceDatasetLoader() -loader.load("foo/bar", split="train") \ - .save("s3://my-bucket/data/foo_train.jsonl") -``` - - ---- -### Common DatasetLoader Methods -These methods are available on all DatasetLoader subclasses. - -#### `show()` -Displays the first n rows of the dataset. -**Signature:** -```python -def show( - self, - n: int = 10 -) -> None -``` - -**Parameters:** -- `n` (int): Number of rows to display (default: 10) - -**Example:** -```python -loader.show(5) # Show first 5 rows -``` ---- -#### `split_data()` -Splits dataset into train, validation, and test sets. - -**Signature:** -```python -def split_data( - self, - train_ratio: float = 0.8, - val_ratio: float = 0.1, - test_ratio: float = 0.1, - seed: int = 42, -) -> Tuple["DatasetLoader", "DatasetLoader", "DatasetLoader"] -``` - -**Parameters:** -- `train_ratio` (float): Proportion of data for training (default: 0.8) -- `val_ratio` (float): Proportion of data for validation (default: 0.1) -- `test_ratio` (float): Proportion of data for testing (default: 0.1) -- `seed` (int): Random seed for reproducibility (default: 42) - -**Returns:** -- `Tuple[DatasetLoader, DatasetLoader, DatasetLoader]`: Three DatasetLoader objects (train, val, test) - -**Raises:** -- `DataPrepError`: If ratios don't sum to 1.0 or dataset is empty - -**Example:** -```python -train_loader, val_loader, test_loader = loader.split_data( - train_ratio=0.7, - val_ratio=0.2, - test_ratio=0.1 -) -``` ---- -#### `transform()` -Transforms dataset to the required format for a specific training method and model. Currently the following transformations are supported: -* Q/A-formatted CSV/JSON/JSONL to SFT 1.0, SFT 2.0 (without reasoningContent, Tools), RFT, Eval, CPT -* OpenAI Messages format to SFT 1.0 and SFT 2.0 (with Tools) - -**Signature:** -```python -def transform( - self, - method: TransformMethod = TransformMethod.SCHEMA, - model: Optional[Model] = None, - eval_task: Optional[EvaluationTask] = None, - **kwargs, -) -> "DatasetLoader" -``` -**Parameters:** -- `method` (TransformMethod): The transform method (default: `TransformMethod.SCHEMA`). Also accepts a `TrainingMethod` enum for backward compatibility (deprecated). -- `model` (Optional[Model]): The Nova model version (e.g., `Model.NOVA_LITE_2`). Can be passed positionally for backward compatibility. -- `eval_task` (Optional[EvaluationTask]): Evaluation task. Can be passed positionally for backward compatibility. -- `**kwargs`: Method-specific arguments passed to the operation. For `TransformMethod.SCHEMA`: - - `training_method` (TrainingMethod): Required. The training method (e.g., `TrainingMethod.SFT_LORA`). - - `model` (Model): Required. The target model. - - `eval_task` (EvaluationTask): Optional. Required when `training_method` is `EVALUATION`. - - `column_mappings` (dict): Optional. Maps standard column names to your dataset's column names. - - `multimodal_data_s3_path` (Optional[str]): S3 prefix where images will be uploaded during conversion (e.g., `s3://my-bucket/images/`). Required when the dataset contains `image_url` content blocks in OpenAI format. When saving the output to S3, the save path must use the same bucket. - - `multimodal_data_bucket_owner` (Optional[str]): AWS account ID that owns the S3 bucket. If not provided, auto-resolved via STS. - -**Returns:** -- `DatasetLoader`: Self (for method chaining) - -**Raises:** -- `ValueError`: If method/model combination is not supported -- `DataPrepError`: If transformation fails - -**Example:** -```python -# Text-only transform -loader.transform( - method=TransformMethod.SCHEMA, - training_method=TrainingMethod.SFT_LORA, - model=Model.NOVA_MICRO, -) - -# Multimodal transform with automatic S3 image upload -loader.transform( - method=TransformMethod.SCHEMA, - training_method=TrainingMethod.SFT_LORA, - model=Model.NOVA_LITE_2, - multimodal_data_s3_path="s3://my-bucket/images/", -) -``` ---- -#### `validate()` -Validates dataset when given the user's intended training method and model. - -**Signature:** -```python -def validate( - self, - method: ValidateMethod = ValidateMethod.INVALID_RECORDS, - model: Model (Optional), - eval_task: EvaluationTask (Optional), - platform: Platform (Optional) -) -> None -``` -**Parameters:** -- `method` (ValidateMethod): The validation method (default: `ValidateMethod.INVALID_RECORDS`). `ValidateMethod.SCHEMA` is deprecated but still supported for backward compatibility. -- `model` (Model): The Nova model version (e.g., `Model.NOVA_LITE`) -- `eval_task` (EvaluationTask): Optional. The evaluation task (e.g., `EvaluationTask.GEN_QA`) -- `platform` (Platform): Optional. The target platform (`Platform.SMTJ`, `Platform.SMHP`, or `Platform.BEDROCK`). Accepted for forward-compatibility; row-count checks are currently enforced only during `train()`, not `validate()`. - -**Returns:** -- None - -**Raises:** -- `ValueError`: If method/model combination is not supported or validation is unsuccessful. - -**Row-Count Checks:** - -The following row-count checks are enforced during `train()`. They do not currently run on the `loader.validate()` path. - -| Check | Applies To | Limit | -|---|---|---| -| Maximum dataset rows | `BEDROCK` / `SFT_LORA`, `SFT_FULL`, `DPO_LORA`, `DPO_FULL`, `CPT` / `NOVA_MICRO`, `NOVA_LITE`, `NOVA_PRO` | 20,000 rows | -| Minimum dataset rows | `BEDROCK` / `SFT_LORA`, `SFT_FULL`, `DPO_LORA`, `DPO_FULL`, `CPT` / `NOVA_MICRO`, `NOVA_LITE`, `NOVA_PRO` | 8 rows | -| Maximum dataset rows | `BEDROCK` / `SFT_LORA`, `SFT_FULL` / `NOVA_LITE_2` | 20,000 rows | -| Minimum dataset rows | `BEDROCK` / `SFT_LORA`, `SFT_FULL` / `NOVA_LITE_2` | 200 rows | -| Maximum dataset rows | `BEDROCK` / `RFT_LORA`, `RFT_FULL` / `NOVA_LITE_2` | 20,000 rows | -| Minimum dataset rows | `BEDROCK` / `RFT_LORA`, `RFT_FULL` / `NOVA_LITE_2` | 100 rows | - -> **Note:** Recipe-dependent checks (e.g. dataset size ≥ `global_batch_size`) are also automatically enforced during `train()` where the fully-built recipe is available. - -**Example:** -```python -loader.validate( - method=ValidateMethod.INVALID_RECORDS, - training_method=TrainingMethod.SFT_LORA, - model=Model.NOVA_MICRO -) -``` - -If you're validating a BYOD Evaluation dataset, you need to provide another parameter, `eval_task` to the `validate` function. For example: -``` -loader.validate( - method=ValidateMethod.INVALID_RECORDS, - training_method=TrainingMethod.EVALUATION, - model=Model.NOVA_LITE_2, - eval_task=EvaluationTask.GEN_QA -) - ->> Validation succeeded for 22 samples on an Evaluation BYOD dataset -``` ---- -#### `filter()` -Queues a data filtering operation on the loader. All filters are lazy — `filter()` records the intent and execution happens when `execute()` is called (or implicitly via `transform()`, `show()`, `save()`). Multiple `filter()` calls can be chained. - -**Note:** Different filters work on different data formats. Text filters (`DEFAULT_TEXT_FILTER`, `EXACT_DEDUP`, `FUZZY_DEDUP`) operate on raw data with a flat text field. `INVALID_RECORDS` works on schema-formatted data. Using a filter on an incompatible data format may produce unexpected results. - -**Signature:** -```python -def filter( - self, - method: FilterMethod = FilterMethod.DEFAULT_TEXT_FILTER, - **kwargs, -) -> "DatasetLoader" -``` - -**Text Filters** (below text filters work on raw data with a flat text field): - -| Value | Description | -|---|---| -| `FilterMethod.DEFAULT_TEXT_FILTER` | Removes records with excessive URL content | -| `FilterMethod.EXACT_DEDUP` | Removes exact duplicate records by content hash | -| `FilterMethod.FUZZY_DEDUP` | Removes near-duplicate records using MinHash LSH similarity | - -Text filters run on AWS Glue (Ray) and use S3 path chaining. See the [Filtering Data](data_prep.md#filtering-data) section in the data prep guide for kwargs (`output_path`, `runtime_manager`, etc.). - -**INVALID_RECORDS Filter** (works on data in the expected format for the target training method): - -| Value | Description | -|---|---| -| `FilterMethod.INVALID_RECORDS` | Removes records that don't conform to expected input for the target model and training technique (schema validation + reserved keyword detection) | - -INVALID_RECORDS runs locally and modifies the loader's dataset in place. - -**Parameters for `INVALID_RECORDS` (passed as kwargs):** -- `training_method` (TrainingMethod): Required. Determines which checks apply. -- `model` (Model): Required. The target model. -- `platform` (Platform): Required. The target platform (`Platform.SMTJ`, `Platform.SMHP`, or `Platform.BEDROCK`). -- `eval_task` (EvaluationTask): Optional. Required when `training_method` is `EVALUATION`. - -**Returns:** -- `DatasetLoader`: Self (for method chaining) - -> **Note:** After `INVALID_RECORDS` runs, use `loader.show()` to inspect the filtered dataset or `loader.save()` to persist the results. - -**Raises:** -- `ValueError`: If required kwargs (`training_method`, `model`, `platform`) are missing for `INVALID_RECORDS`. - -> **Note:** If your data is already in the correct schema format (e.g. previously transformed and saved), you can call `filter(INVALID_RECORDS)` directly after `load()` without `transform()`. - -**Example (full chain):** -```python -loader.load("data.jsonl") -loader.filter( - method=FilterMethod.DEFAULT_TEXT_FILTER, -).filter( - method=FilterMethod.EXACT_DEDUP, -) -loader.transform( - method=TransformMethod.SCHEMA, - training_method=TrainingMethod.SFT_LORA, - model=Model.NOVA_LITE_2, -) -loader.filter( - method=FilterMethod.INVALID_RECORDS, - training_method=TrainingMethod.SFT_LORA, - model=Model.NOVA_LITE_2, - platform=Platform.SMTJ, -) -loader.execute() # flush the pending filter -loader.validate( - method=ValidateMethod.SCHEMA, - training_method=TrainingMethod.SFT_LORA, - model=Model.NOVA_LITE_2, -) -``` - -**Example (already-formatted data — skip `transform()`):** -```python -loader = JSONLDatasetLoader() -loader.load("pre_transformed_data.jsonl") -loader.filter( - method=FilterMethod.INVALID_RECORDS, - training_method=TrainingMethod.SFT_LORA, - model=Model.NOVA_LITE_2, - platform=Platform.SMTJ, -) -loader.execute() -``` - ---- -#### `save_data()` -Saves the dataset to a local or S3 location. -**Signature:** -```python -def save_data( - self, - save_path: str -) -> str -``` -**Parameters:** -- `save_path` (str): Path where to save the file (local or S3, must end in .json or .jsonl) - -**Returns:** -- `str`: Path where the file was saved - -**Raises:** -- `DataPrepError`: If save fails or format is unsupported - -**Example:** -```python -# Save locally -loader.save_data("output/training_data.jsonl") -# Save to S3 -loader.save_data("s3://my-bucket/data/training_data.jsonl") -``` ---- -## Job Results -Job result classes provide methods to check status and retrieve results from training, evaluation, and inference jobs. - -### Base Classes - -#### BaseJobResult -Abstract base class for all job results. - -**Attributes:** -- `job_id` (str): Job identifier -- `started_time` (datetime): Job start timestamp - -**Methods:** -##### `get_job_status()` -Gets the current status of the job. - -**Signature:** -```python -def get_job_status( - self -) -> tuple[JobStatus, str] -``` - -**Returns:** -- `tuple[JobStatus, str]`: A tuple of (status enum, raw status string) - - `JobStatus.IN_PROGRESS`: Job is running - - `JobStatus.COMPLETED`: Job completed successfully - - `JobStatus.FAILED`: Job failed - -**Example:** -```python -status, raw_status = result.get_job_status() -if status == JobStatus.COMPLETED: - print("Job finished!") -``` - -##### `dump(file_path: Optional[str] = None, file_name: Optional[str] = None)` -Save the job result to file_path path - -**Signature:** -```python -def dump( - self, - file_path: Optional[str] = None, - file_name: Optional[str] = None -) -> Path -``` - -**Parameters:** -- `file_path` (Optional[str]): Directory path to save the result. Saves to current directory if not provided -- `file_name` (Optional[str]): The file name of the result. Default to `_.json` if not provided - -**Returns:** -- `Path`: The full result file path - -**Example:** -```python -result.dump() -# Result will be saved to ./{job_id}_{platform}.json under current dir -result.dump(file_path='/customized/path', file_name='customized_name.json') -# Result will be saved to /customized/path/customized_name.json -``` - -##### `load(file_path: str)` -Load the job result from the `file_path` path - -**Signature:** -```python -@classmethod -def load( - cls, - file_path: str -) -> "BaseJobResult": -``` - -**Returns:** -- JobResultObject. The instance of subclass of BaseJobResult such as SMTJEvaluationResult, SMHPEvaluationResult, BedrockEvaluationResult, SMTJTrainingResult, SMHPTrainingResult, or BedrockTrainingResult - -**Example:** -```python -job_result = BaseJobResult.load('./my_job_result.json') -``` - ---- - -##### `enable_job_notifications()` -Enable email notifications for when a job reaches a terminal state (Completed, Failed, or Stopped). - -**Signature:** -```python -def enable_job_notifications( - self, - emails: list[str], - output_s3_path: Optional[str] = None, - region: Optional[str] = "us-east-1", - **platform_kwargs -) -> None -``` - -**Parameters:** -- `emails` (list[str]): List of email addresses to notify -- `output_s3_path` (Optional[str]): S3 path where job outputs are stored. - - Only required if the SDK cannot automatically extract it from the job result's `model_artifacts` attribute. - - For most training jobs, this parameter is automatically populated and does not need to be provided explicitly. -- `region` (Optional[str]): AWS region for notification infrastructure (default: "us-east-1") -- `**platform_kwargs`: Platform-specific parameters: - - **For SMTJ:** - - `kms_key_id` (Optional[str]): Customer KMS key ID (not full ARN) for SNS topic encryption - - **For SMHP:** - - `namespace` (str): Kubernetes namespace where the PyTorchJob runs (e.g., "kubeflow", "default") (Required) - - `kubectl_layer_arn` (str): ARN of the lambda-kubectl layer (Required) - - `eks_cluster_arn` (Optional[str]): EKS cluster ARN (auto-detected if not provided) - - `vpc_id` (Optional[str]): VPC ID (auto-detected if not provided) - - `subnet_ids` (Optional[list[str]]): List of subnet IDs for Lambda (auto-detected if not provided) - - `security_group_id` (Optional[str]): Security group ID for Lambda (auto-detected if not provided) - - `polling_interval_minutes` (Optional[int]): How often to check job status in minutes (default: 5) - - `kms_key_id` (Optional[str]): Customer KMS key ID (not full ARN) for SNS topic encryption - -**Returns:** -- None - -**Raises:** -- `ValueError`: If required parameters are missing or invalid -- `NotificationManagerInfraError`: If infrastructure setup fails - -**How It Works:** -1. Creates AWS infrastructure (CloudFormation stack) if it doesn't exist: - - DynamoDB table to store job notification configurations - - SNS topic for email notifications - - Lambda function to monitor job status - - EventBridge rule (SMTJ) or scheduled rule (SMHP) to trigger Lambda - - (SMHP only) VPC endpoints for DynamoDB and S3 if needed -2. Stores job configuration in DynamoDB (including namespace for SMHP) -3. Subscribes email addresses to SNS topic (users must confirm subscription) -4. Monitors job status and sends email when job completes, fails, is stopped, or becomes degraded (SMHP only) - -**Email Confirmation:** -Users will receive a confirmation email from AWS SNS and must click the confirmation link before receiving job notifications. - -**Examples:** - -SMTJ (SageMaker Training Jobs): -```python -# Basic usage - output_s3_path is automatically extracted -result = customizer.train(job_name="my-job") -result.enable_job_notifications( - emails=["user@example.com", "team@example.com"] -) - -# With customer KMS encryption -result.enable_job_notifications( - emails=["user@example.com"], - kms_key_id="abc-123-def-456" # Just the key ID, not full ARN -) - -# With custom region -result.enable_job_notifications( - emails=["user@example.com"], - region="us-west-2" -) -``` - -SMHP (SageMaker HyperPod): -```python -# Basic usage (with auto-detection) -result = customizer.train(job_name="my-job") -result.enable_job_notifications( - emails=["user@example.com"], - namespace="kubeflow", # Required - kubectl_layer_arn="arn:aws:lambda::123456789012:layer:kubectl:1" # Required -) - -# With custom polling interval -result.enable_job_notifications( - emails=["user@example.com"], - namespace="kubeflow", - kubectl_layer_arn="arn:aws:lambda::123456789012:layer:kubectl:1", - polling_interval_minutes=10 # Check every 10 minutes instead of default 5 -) - -# With explicit VPC configuration of the cluster where jobs are being monitored. -result.enable_job_notifications( - emails=["user@example.com"], - namespace="kubeflow", - kubectl_layer_arn="arn:aws:lambda::123456789012:layer:kubectl:1", - eks_cluster_arn="arn:aws:eks::123456789012:cluster/my-cluster", - vpc_id="vpc-12345", - subnet_ids=["subnet-1", "subnet-2"], - security_group_id="sg-12345" -) -``` - -**Important Notes:** -- For SMHP, requires deploying a kubectl Lambda layer from AWS Serverless Application Repository -- For SMHP, the user will need to manually grant the Lambda function access to your EKS cluster (access-entry). - - Please refer to [`docs/job_notifications.md`](job_notifications.md) for the commands to run to set this up. -- See [`docs/job_notifications.md`](job_notifications.md) for detailed setup instructions, troubleshooting, and advanced usage - ---- -### EvaluationResult (ABC) -Result object for SageMaker Training Job evaluation tasks. - -**Attributes:** -- `job_id` (str): Job identifier -- `started_time` (datetime): Job start timestamp -- `eval_task` (EvaluationTask): Evaluation task performed -- `eval_output_path` (str): S3 path to evaluation results - -#### Subclasses -- SMTJEvaluationResult -- SMHPEvaluationResult -- BedrockEvaluationResult - -#### Methods - -##### `get()` -Downloads and returns evaluation results as a dictionary. - -**Signature:** -```python -def get( - self -) -> Dict -``` - -**Returns:** -- `Dict`: Evaluation results (empty dict if job not completed) - -**Example:** -```python -eval_result = customizer.evaluate(...) -# Wait for job to complete -results = eval_result.get() -print(results) -``` - ---- -##### `show()` -Prints evaluation results to console. - -**Signature:** -```python -def show( - self -) -> None -``` - -**Example:** -```python -eval_result.show() -``` - ---- -##### `upload_tensorboard_results()` -Uploads TensorBoard results to S3. - -**Signature:** -```python -def upload_tensorboard_results( - self, - tensorboard_s3_path: Optional[str] = None -) -> None -``` - -**Parameters:** -- `tensorboard_s3_path` (Optional[str]): Target S3 path (auto-generated if not provided) - -**Example:** -```python -eval_result.upload_tensorboard_results( - tensorboard_s3_path="s3://my-bucket/tensorboard/" -) -``` - ---- -##### `clean()` -Cleans up local cached results. - -**Signature:** -```python -def clean( - self -) -> None -``` - ---- -### SMTJBatchInferenceResult -Result object for batch inference jobs. -**Attributes:** -- `job_id` (str): Job identifier -- `started_time` (datetime): Job start timestamp -- `inference_output_path` (str): S3 path to inference outputs - -#### Methods - -##### `get()` -Downloads and returns inference results, optionally saving to S3. -**Signature:** -```python -def get( - self, - s3_path: Optional[str] = None -) -> Dict -``` -**Parameters:** -- `s3_path` (Optional[str]): S3 path to save formatted results -**Returns:** -- `Dict`: Dictionary containing list of inference results - - Each result has: `system`, `query`, `gold_response`, `inference_response`, `metadata` -**Example:** -```python -inference_result = customizer.batch_inference(...) -# Wait for job to complete -results = inference_result.get(s3_path="s3://my-bucket/formatted-results.jsonl") -``` ---- -##### `show()` -Prints inference results to console. -**Signature:** -```python -def show( - self -) -> None -``` ---- -##### `clean()` -Cleans up local cached results. -**Signature:** -```python -def clean( - self -) -> None -``` ---- -### IAM Role Creation SDK -This SDK provides utility functions for creating IAM roles with specific permissions for AWS Bedrock and SageMaker services. - -**Methods** - -#### `create_bedrock_execution_role()` -Creates an IAM role with permissions for Bedrock model creation and deployment. - -**Signature:** -```python -def create_bedrock_execution_role( - iam_client, - role_name: str, - bedrock_resource: str = "*", - s3_resource: str = "*" -) -> Dict -``` - -**Parameters:** -- `iam_client`: Boto3 IAM client -- `role_name` (str): Name of the IAM role to create -- `bedrock_resource` (Optional[str]): Specific Bedrock resource to restrict access. Defaults to "*" (all resources) -- `s3_resource` (Optional[str]): Specific S3 resource to restrict access. Defaults to "*" (all resources) - -**Returns:** -- `Dict`: IAM role details - -**Example:** -```python -import boto3 -from amzn_nova_forge.iam import create_bedrock_execution_role - -iam_client = boto3.client("iam") -create_bedrock_execution_role(iam_client, "role-name", "bedrock_resource", "s3_resource") -``` - -### `create_sagemaker_execution_role()` -Creates an IAM role with permissions for SageMaker model creation and deployment. - -**Signature:** -```python -def create_sagemaker_execution_role( - iam_client, - role_name: str, - s3_resource: str = "*", - kms_resource: str = "*", - ec2_condition: Optional[Dict[str, Any]] = None, - cloudwatch_metric_condition: Optional[Dict[str, Any]] = None, - cloudwatch_logstream_resource: str = "*", - cloudwatch_loggroup_resource: str = "*" -) -> Dict -``` - -**Parameters:** -- `iam_client`: Boto3 IAM client -- `role_name` (str): Name of the IAM role to create -- `s3_resource` (Optional[str]): Specific S3 resource to restrict access -- `kms_resource` (Optional[str]): Specific KMS resource to restrict access -- `ec2_condition` (Optional[Dict]): Conditional access for EC2 resources -- `cloudwatch_metric_condition` (Optional[Dict]): Conditional access for CloudWatch metrics -- `cloudwatch_logstream_resource` (Optional[str]): Specific CloudWatch log stream resource -- `cloudwatch_loggroup_resource` (Optional[str]): Specific CloudWatch log group resource - -**Returns:** -- `Dict`: IAM role details - -**Example:** -```python -import boto3 -from amzn_nova_forge.iam import create_sagemaker_execution_role - -iam_client = boto3.client("iam") -create_sagemaker_execution_role( - iam_client, - role_name="role-name", - s3_resource="example-bucket""", - kms_resource="encryption-key", - ec2_condition={ - "ArnLike": { - "ec2:Vpc": "arn:aws:ec2:*:*:vpc/example" - } - }, - cloudwatch_metric_condition={ - "StringEquals": { - "cloudwatch:namespace": ["example-namespace"] - } - }, - cloudwatch_loggroup_resource="example-loggroup", - cloudwatch_logstream_resource="example-logstream" - ) -``` ---- -#### `ModelDeployResult` -Result of creating a Bedrock or Sagemaker model. - -**Fields:** -- `model_arn` (str): The Bedrock custom model ARN or SageMaker model ARN. -- `model_name` (str): The model name. -- `escrow_uri` (str): S3 artifacts path used to create the model. -- `created_at` (datetime): UTC timestamp when the model was created. - -**Properties:** -- `platform` (Optional[str]): Returns `"bedrock"` or `"sagemaker"` based on strict ARN format validation, or `None` for unrecognized ARNs. -- `status` (ModelStatus): Queries the model's current status via live API call. Returns `CREATING`, `ACTIVE`, `FAILED`, or `UNKNOWN`. Logs a warning if no AWS client is available. - -**Class Methods:** -- `from_arn(model_arn, bedrock_client=None, sagemaker_client=None)` — Reconstruct from an existing model ARN. Detects platform from ARN format, calls the appropriate describe API, and recovers `escrow_uri` from tags. Works for both Bedrock and SageMaker ARNs. -- `load(file_path)` — Load from a JSON file saved by `dump()`. Automatically creates fresh AWS clients for status checking. - -**Instance Methods:** -- `dump(file_path=None, file_name=None)` — Save to JSON file for later use with `load()`. - -**Example:** -```python -# Create and persist -publish = customizer.create_custom_model(job_result=training_result) -publish.dump(file_path="./results/") - -# Load and check status -loaded = ModelDeployResult.load("./results/my-model_deploy_result.json") -print(loaded.platform) # "bedrock" or "sagemaker" -print(loaded.status) # ModelStatus.ACTIVE - -# Reconstruct from ARN -result = ModelDeployResult.from_arn( - "arn:aws:bedrock:us-east-1:123456789012:custom-model/my-model" -) -``` - -#### `ModelStatus` -Platform-independent model status enum. - -**Values:** -- `ModelStatus.CREATING` — Model is still being created (Bedrock only). -- `ModelStatus.ACTIVE` — Model is ready for deployment. -- `ModelStatus.FAILED` — Model creation failed. -- `ModelStatus.UNKNOWN` — Status could not be determined (no client, unrecognized ARN, or unexpected API response). ---- - -## Utility Functions - -### verify_reward_function() - -Verifies a reward function with sample data before using it in RFT training or evaluation. This utility helps you test your reward function implementation to ensure it works correctly and returns the expected format. - -**Signature:** -```python -def verify_reward_function( - reward_function: str, - sample_data: List[Dict[str, Any]], - region: str = "us-east-1", - validate_format: bool = True, - platform: Optional[Platform] = None, -) -> Dict[str, Any] -``` - -**Parameters:** -- `reward_function` (str): Either a Lambda ARN (string starting with `'arn:aws:lambda:'`) or a path to a local Python file containing the reward function. -- `sample_data` (List[Dict[str, Any]]): List of conversation samples to test. Each sample should be a dict with `'id'`, `'messages'`, and optionally `'reference_answer'` keys. -- `region` (str): AWS region for Lambda invocation (default: "us-east-1"). -- `validate_format` (bool): If True, validates that sample_data matches RFT format and output matches expected format (default: True). -- `platform` (Platform): Platform enum (Platform.SMHP or Platform.SMTJ). **Required when using Lambda ARN**. When set to Platform.SMHP, validates that Lambda ARN contains 'SageMaker' in the function name as required by SageMaker HyperPod. Optional for local files. - -**Returns:** -- `Dict[str, Any]`: Dictionary containing: - - `success` (bool): Always True if no exception raised - - `results` (list): List of individual test results - - `total_samples` (int): Total number of samples tested - - `successful_samples` (int): Number of successful tests - - `warnings` (list): List of warning messages (e.g., missing reference_answer) - -**Raises:** -- `ValueError`: If any validation errors are encountered, with a detailed error message listing all issues found. - -**Example** -```python -from amzn_nova_forge import verify_reward_function -from amzn_nova_forge.model import Platform - -# Test with Lambda ARN (platform required for Lambda ARNs) -result = verify_reward_function( - reward_function="arn:aws:lambda:us-east-1:123456789012:function:MySageMakerReward", - sample_data=[ - { - "id": "sample_1", - "reference_answer": "correct answer", - "messages": [ - {"role": "user", "content": "question"}, - {"role": "assistant", "content": "response"} - ] - } - ], - platform=Platform.SMHP # Required for Lambda ARNs -) - -print(f"Verification: {'PASSED' if result['success'] else 'FAILED'}") -print(f"Tested {result['total_samples']} samples, {result['successful_samples']} successful") - -if result.get('warnings'): - print(f"\nWarnings:") - for warning in result['warnings']: - print(f" - {warning}") - -# Test with local Python file (platform optional) -result = verify_reward_function( - reward_function="./my_reward_function.py", - sample_data=[ - { - "id": "sample_1", - "reference_answer": "correct answer", - "messages": [ - {"role": "user", "content": "question"}, - {"role": "assistant", "content": "response"} - ] - } - ] -) -``` -**Output Format Requirements from Lambda:** - -```python -{ - "id": "sample_1", # Required: string - "aggregate_reward_score": 0.75, # Required: float or int - "metrics_list": [ # Optional: validated if present - { - "name": "accuracy", # Required: string - "value": 0.85, # Required: float or int - "type": "Metric" # Required: "Metric" or "Reward" - } - ] -} -``` - -**Common Validation Errors:** -- Missing required fields in input (`messages` field is required) -- Missing required fields in output (`id` and `aggregate_reward_score` are required) -- Invalid data types (e.g., `aggregate_reward_score` must be a number) -- Missing `platform` parameter when using Lambda ARN -- SMHP Lambda ARN doesn't contain 'SageMaker' in function name -- Invalid `metrics_list` structure (must be list of dicts with `name`, `value`, `type`) -- Invalid metric `type` (must be "Metric" or "Reward") - -**Warnings:** -- Missing `reference_answer`: While optional in RFT datasets, reference answers are recommended for meaningful reward calculations. Without ground truth, your reward function cannot compare model outputs against expected answers. - - - -**Note:** The `metrics_list` field is optional. If provided, it will be validated for proper structure and logged during training/evaluation. - ---- -## Monitoring - -### CloudWatchLogMonitor - -Monitors CloudWatch logs and plots training metrics for Nova model training jobs. Supports both SageMaker Training Jobs (SMTJ) and SageMaker HyperPod (SMHP) platforms. - -#### Factory Methods - -##### `from_job_id()` - -Creates a CloudWatchLogMonitor from a job ID. - -**Signature:** -```python -@classmethod -def from_job_id( - cls, - job_id: str, - platform: Platform, - started_time: Optional[datetime] = None, - **kwargs, -) -> "CloudWatchLogMonitor" -``` - -**Parameters:** -- `job_id` (str): The training job identifier -- `platform` (Platform): Execution platform (`Platform.SMTJ` or `Platform.SMHP`) -- `started_time` (Optional[datetime]): Job start time (used to filter logs) -- `**kwargs`: Platform-specific parameters: - - SMHP requires: `cluster_name` (str), optional `namespace` (str, defaults to "kubeflow") - -**Returns:** -- `CloudWatchLogMonitor`: Monitor instance - -**Example:** -```python -from amzn_nova_forge.monitor import CloudWatchLogMonitor -from amzn_nova_forge.model import Platform - -# SMTJ -monitor = CloudWatchLogMonitor.from_job_id( - job_id="my-training-job", - platform=Platform.SMTJ, - started_time=datetime(2026, 1, 15, 12, 0, 0) -) - -# SMHP -monitor = CloudWatchLogMonitor.from_job_id( - job_id="my-hyperpod-job", - platform=Platform.SMHP, - cluster_name="my-cluster", - namespace="kubeflow" -) -``` - ---- - -##### `from_job_result()` - -Creates a CloudWatchLogMonitor from a training job result object. - -**Signature:** -```python -@classmethod -def from_job_result( - cls, - job_result: BaseJobResult, - cloudwatch_logs_client=None -) -> "CloudWatchLogMonitor" -``` - -**Parameters:** -- `job_result` (BaseJobResult): A training or evaluation result object (e.g., `TrainingResult`) -- `cloudwatch_logs_client` (Optional): Boto3 CloudWatch Logs client (auto-created if not provided) - -**Returns:** -- `CloudWatchLogMonitor`: Monitor instance - -**Example:** -```python -result = customizer.train(job_name="my-job") -monitor = CloudWatchLogMonitor.from_job_result(job_result=result) -``` - ---- - -#### Methods - -##### `get_logs()` - -Retrieves CloudWatch log events for the job. - -**Signature:** -```python -def get_logs( - self, - limit: Optional[int] = None, - start_from_head: bool = False, - end_time: Optional[int] = None, -) -> List[Dict] -``` - -**Parameters:** -- `limit` (Optional[int]): Maximum number of log events to retrieve -- `start_from_head` (bool): If True, start from the beginning of logs; if False, start from the end -- `end_time` (Optional[int]): End time in epoch milliseconds - -**Returns:** -- `List[Dict]`: List of log event dictionaries, each containing a `"message"` key - -**Example:** -```python -logs = monitor.get_logs(limit=100) -``` - ---- - -##### `show_logs()` - -Prints CloudWatch log messages to the console. - -**Signature:** -```python -def show_logs( - self, - limit: Optional[int] = None, - start_from_head: bool = False, - end_time: Optional[int] = None, -) -> None -``` - -**Parameters:** -- `limit` (Optional[int]): Maximum number of log events to display -- `start_from_head` (bool): If True, start from the beginning of logs; if False, start from the end -- `end_time` (Optional[int]): End time in epoch milliseconds - -**Example:** -```python -monitor.show_logs(limit=50, start_from_head=True) -``` - ---- - -##### `plot_metrics()` - -Parses training metrics from CloudWatch logs and displays them as matplotlib plots. Automatically fetches the latest logs if the job is still in progress or logs have not been retrieved yet. - -**Signature:** -```python -def plot_metrics( - self, - training_method: TrainingMethod, - metrics: Optional[List[str]] = None, - starting_step: Optional[int] = None, - ending_step: Optional[int] = None, -) -> None -``` - -**Parameters:** -- `training_method` (TrainingMethod): The training method used for the job (e.g., `TrainingMethod.SFT_LORA`, `TrainingMethod.CPT`, `TrainingMethod.RFT_LORA`) -- `metrics` (Optional[List[str]]): List of metric names to plot. Available metrics depend on training method: - - CPT / SFT: `"training_loss"` - - RFT: `"reward_score"` -- `starting_step` (Optional[int]): Filter to only show metrics from this global step onward -- `ending_step` (Optional[int]): Filter to only show metrics up to this global step - -**Raises:** -- `ValueError`: If `starting_step` > `ending_step`, or if no logs are found for the job -- `NotImplementedError`: If an unsupported metric is requested for the given training method/platform - -**Example:** -```python -from amzn_nova_forge.monitor import CloudWatchLogMonitor -from amzn_nova_forge.model import Platform, TrainingMethod - -# Create monitor from a training result -monitor = CloudWatchLogMonitor.from_job_result(job_result=training_result) - -# Plot training loss for an SFT job -monitor.plot_metrics( - training_method=TrainingMethod.SFT_LORA, - metrics=["training_loss"] -) - -# Plot reward score for an RFT job, filtered to steps 50-200 -monitor.plot_metrics( - training_method=TrainingMethod.RFT_LORA, - metrics=["reward_score"], - starting_step=50, - ending_step=200 -) -``` - ---- - -### MLflowMonitor - -MLflow monitoring configuration for Nova model training. This class provides experiment tracking capabilities through MLflow integration. - -**Note:** MLflow monitoring is only supported for SageMaker platforms (SMTJ, SMHP). It is not available for Bedrock platform. - -**MLflow Integration Features:** -- Automatic logging of training metrics -- Model artifact and checkpoint tracking -- Hyperparameter recording -- Support for SageMaker MLflow tracking servers -- Custom MLflow tracking server support (with proper network configuration) - - -#### Constructor - -**Signature:** -```python -def __init__( - self, - tracking_uri: Optional[str] = None, - experiment_name: Optional[str] = None, - run_name: Optional[str] = None, -) -``` - -**Parameters:** -- `tracking_uri` (Optional[str]): MLflow tracking server URI or SageMaker MLflow app ARN. If not provided, attempts to use a default SageMaker MLflow tracking server if one exists -- `experiment_name` (Optional[str]): Name of the MLflow experiment. If not provided, will use the job name -- `run_name` (Optional[str]): Name of the MLflow run. If not provided, will be auto-generated - -**Raises:** -- `ValueError`: If MLflow configuration validation fails - -**Example:** -```python -from amzn_nova_forge.monitor import * - -# With explicit tracking URI -monitor = MLflowMonitor( - tracking_uri="arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", - experiment_name="nova-customization", - run_name="sft-run-1" -) - -# With default tracking URI (if available) -monitor = MLflowMonitor( - experiment_name="nova-customization", - run_name="sft-run-1" -) - -# Use with NovaModelCustomizer -customizer = NovaModelCustomizer( - model=Model.NOVA_LITE_2, - method=TrainingMethod.SFT_LORA, - infra=runtime_manager, - data_s3_path="s3://bucket/data", - mlflow_monitor=monitor -) -``` - -#### Methods - -##### `to_dict()` - -Converts MLflow configuration to dictionary format for use in recipe overrides. - -**Signature:** -```python -def to_dict( - self -) -> dict -``` - -**Returns:** -- `dict`: Dictionary with mlflow_* keys for recipe configuration. Returns empty dict if no tracking URI is available - -**Example:** -```python -monitor = MLflowMonitor( - tracking_uri="arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", - experiment_name="nova-customization" -) - -config_dict = monitor.to_dict() -# Returns: { -# "mlflow_tracking_uri": "arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", -# "mlflow_experiment_name": "nova-customization" -# } -``` - -##### `get_presigned_url()` - -Generates a presigned URL for accessing the MLflow tracking server UI directly without navigating through the AWS Console. - -**Signature:** -```python -def get_presigned_url( - self, - session_expiration_duration_in_seconds: int = 43200, - expires_in_seconds: int = 300 -) -> str -``` - -**Parameters:** -- `session_expiration_duration_in_seconds` (int, optional): Duration in seconds for which the MLflow UI session is valid after accessing the presigned URL. Default is 43200 seconds (12 hours). Valid range: 1800-43200 seconds -- `expires_in_seconds` (int, optional): Duration in seconds for which the presigned URL itself is valid. The URL must be accessed within this time. Default is 300 seconds (5 minutes). Valid range: 5-300 seconds - -**Returns:** -- `str`: Presigned URL for accessing the MLflow tracking server UI. This URL must be used within `expires_in_seconds` - -**Raises:** -- `ValueError`: If tracking_uri is not set -- `RuntimeError`: If unable to generate presigned URL - -**Example:** -```python -monitor = MLflowMonitor( - tracking_uri="arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", - experiment_name="nova-customization" -) - -# Generate presigned URL with defaults -# URL expires in 5 minutes, but session lasts 12 hours once accessed -url = monitor.get_presigned_url() -print(f"Access MLflow UI at: {url}") - -# Generate URL with custom expiration times -url = monitor.get_presigned_url( - session_expiration_duration_in_seconds=3600, # 1 hour session - expires_in_seconds=60 # URL expires in 1 minute -) -``` - -#### MLflow Integration Notes - -When MLflow monitoring is enabled: -1. Training metrics will be automatically logged to the specified MLflow tracking server -2. Model artifacts and checkpoints will be tracked in MLflow -3. Hyperparameters and configuration will be recorded as MLflow parameters -4. You can view experiment results in the MLflow UI - -The MLflow integration supports: -- SageMaker MLflow tracking servers -- Custom MLflow tracking servers (with appropriate network configuration) -- Automatic experiment and run creation -- Metric logging during training -- Artifact tracking - ---- -## Enums and Configuration -### Model Enum -Supported Nova models with their configurations. -**Values:** -- `Model.NOVA_MICRO`: Amazon Nova Micro (Version 1) - - `model_type`: "amazon.nova-micro-v1:0:128k" - - `model_path`: "nova-micro/prod" - - `version`: Version.ONE -- `Model.NOVA_LITE`: Amazon Nova Lite (Version 1) - - `model_type`: "amazon.nova-lite-v1:0:300k" - - `model_path`: "nova-lite/prod" - - `version`: Version.ONE -- `Model.NOVA_LITE_2`: Amazon Nova Lite (Version 2) - - `model_type`: "amazon.nova-2-lite-v1:0:256k" - - `model_path`: "nova-lite-2/prod" - - `version`: Version.TWO -- `Model.NOVA_PRO`: Amazon Nova Pro (Version 1) - - `model_type`: "amazon.nova-pro-v1:0:300k" - - `model_path`: "nova-pro/prod" - - `version`: Version.ONE - -**Methods:** -##### `from_model_type()` -Gets Model enum from model type string. - -**Signature:** -```python -@classmethod -def from_model_type( - cls, - model_type: str -) -> "Model" -``` - -**Example:** -```python -model = Model.from_model_type("amazon.nova-micro-v1:0:128k") -``` ---- -### TrainingMethod Enum -Supported training methods. - -**Values:** -- `TrainingMethod.CPT`: Continued Pre-Training -- `TrainingMethod.DPO_LORA`: Direct Preference Optimization with LoRA -- `TrainingMethod.DPO_FULL`: Direct Preference Optimization (full rank) -- `TrainingMethod.SFT_LORA`: Supervised Fine-Tuning with LoRA -- `TrainingMethod.SFT_FULL`: Supervised Fine-Tuning (full rank) -- `TrainingMethod.RFT_LORA`: Reinforcement Fine-Tuning with LoRA -- `TrainingMethod.RFT_FULL`: Full reinforcement Fine-Tuning -- `TrainingMethod.EVALUATION`: Evaluation only ---- -### DeployPlatform Enum -Supported deployment platforms. - -**Values:** -- `DeployPlatform.BEDROCK_OD`: Amazon Bedrock On-Demand -- `DeployPlatform.BEDROCK_PT`: Amazon Bedrock Provisioned Throughput -- `DeployPlatform.SAGEMAKER`: Amazon SageMaker ---- -### DeploymentMode Enum -Deployment behavior when an endpoint with the same name already exists. - -**Values:** -- `DeploymentMode.FAIL_IF_EXISTS`: Raise an error if endpoint already exists (safest, default) -- `DeploymentMode.UPDATE_IF_EXISTS`: Try in-place update only, fail if not supported (PT only) - -**Note:** Only `FAIL_IF_EXISTS` and `UPDATE_IF_EXISTS` modes are currently supported. -`UPDATE_IF_EXISTS` is only applicable for Bedrock Provisioned Throughput (PT) deployments. - ---- -### EvaluationTask Enum -Supported evaluation tasks. -Common values include: -- `EvaluationTask.MMLU`: Massive Multitask Language Understanding -- `EvaluationTask.GPQA`: General Physics Question Answering -- `EvaluationTask.MATH`: Mathematical Problem Solving -- `EvaluationTask.GEN_QA`: Custom Dataset Evaluation -- The full list of available tasks can be found here: [AWS Documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/nova-model-evaluation.html#nova-model-evaluation-benchmark) ---- -### Platform Enum -Infrastructure platforms. - -**Values:** -- `Platform.SMTJ`: SageMaker Training Jobs -- `Platform.SMHP`: SageMaker HyperPod -- `Platform.BEDROCK`: Amazon Bedrock ---- -### JobStatus Enum -Job execution status. - -**Values:** -- `JobStatus.IN_PROGRESS`: Job is running -- `JobStatus.COMPLETED`: Job completed successfully -- `JobStatus.FAILED`: Job failed - ---- - -## RFT Multiturn Infrastructure - -For RFT multiturn training and evaluation, you need to set up infrastructure to run reward functions. - -### Helper Functions - -#### create_rft_execution_role - -Creates an IAM role with required permissions for RFT multiturn infrastructure. - -**Function:** -```python -def create_rft_execution_role( - region: str = "us-east-1", - role_name: Optional[str] = None, - custom_policy_path: Optional[str] = None -) -> str -``` - -**Parameters:** -- `region` (str): AWS region. Default: "us-east-1" -- `role_name` (Optional[str]): Custom role name. Default: "RFTExecutionRoleNovaSDK" -- `custom_policy_path` (Optional[str]): Path to custom policy JSON file. If not provided, uses SDK default. - -**Returns:** -- `str`: ARN of the created/existing role - -**Example:** -```python -from amzn_nova_forge import create_rft_execution_role - -# Create role with default name -role_arn = create_rft_execution_role(region="us-east-1") - -# Create role with custom name -role_arn = create_rft_execution_role(region="us-east-1", role_name="my-custom-rft-role") -``` - -#### list_rft_stacks - -Lists CloudFormation stacks in the region, optionally filtering for Nova SDK stacks. - -**Function:** -```python -def list_rft_stacks( - region: str = "us-east-1", - all_stacks: bool = False -) -> List[str] -``` - -**Parameters:** -- `region` (str): AWS region. Default: "us-east-1" -- `all_stacks` (bool): If True, list all stacks. If False, only list Nova SDK stacks (ending with "NovaForgeSDK"). Default: False - -**Returns:** -- `List[str]`: List of stack names - -**Example:** -```python -from amzn_nova_forge import list_rft_stacks - -# List only Nova SDK stacks -nova_stacks = list_rft_stacks(region="us-east-1") - -# List all CloudFormation stacks -all_stacks = list_rft_stacks(region="us-east-1", all_stacks=True) -``` - -### RFTMultiturnInfrastructure - -Manages infrastructure for RFT multiturn training (reward function workers). - -**Constructor:** -```python -def __init__( - self, - stack_name: str, - region: str = "us-east-1", - vf_env_id: Optional[VFEnvId] = None, - custom_env: Optional[CustomEnvironment] = None, - infrastructure_arn: Optional[str] = None, - python_venv_name: Optional[str] = None, - vpc_config: Optional[Dict[str, Any]] = None, - cpu: Optional[str] = None, - memory: Optional[str] = None, - rft_role_name: Optional[str] = None, -) -``` - -**Parameters:** -- `stack_name` (str): CloudFormation stack name -- `region` (str): AWS region. Default: "us-east-1" -- `vf_env_id` (Optional[VFEnvId]): Built-in environment ID (VFEnvId.WORDLE or VFEnvId.TERMINAL_BENCH) -- `custom_env` (Optional[CustomEnvironment]): Custom environment (mutually exclusive with vf_env_id) -- `infrastructure_arn` (Optional[str]): Platform ARN (EC2 instance ID, ECS cluster ARN, or None for LOCAL) -- `python_venv_name` (Optional[str]): Python virtual environment name (required for LOCAL/EC2, optional for ECS) -- `vpc_config` (Optional[Dict]): VPC configuration for ECS only. Dict with keys: - - `subnets`: List[str] - Subnet IDs - - `security_groups`: List[str] - Security group IDs -- `cpu` (Optional[str]): CPU units for ECS tasks (e.g., "2048"). Ignored for LOCAL/EC2. -- `memory` (Optional[str]): Memory in MB for ECS tasks (e.g., "4096"). Ignored for LOCAL/EC2. -- `rft_role_name` (Optional[str]): IAM role name for RFT infrastructure. If not provided, uses default role or creates one. - -**Example:** -```python -from amzn_nova_forge import RFTMultiturnInfrastructure, CustomEnvironment, VFEnvId - -# Option 1: LOCAL with built-in environment -rft_infra = RFTMultiturnInfrastructure( - stack_name="my-rft-stack", - region="us-east-1", - python_venv_name="my_rft_venv", - vf_env_id=VFEnvId.WORDLE -) - -# Option 2: ECS with custom environment and VPC config -custom_env = CustomEnvironment( - env_id="my-custom-env", - output_dir="~/custom_envs/", - env_type="single_turn" -).create(overwrite=True) - -rft_infra = RFTMultiturnInfrastructure( - stack_name="my-rft-stack", - custom_env=custom_env, - infrastructure_arn="arn:aws:ecs:us-east-1:123456789012:cluster/my-cluster", - vpc_config={ - "subnets": ["subnet-12345", "subnet-67890"], - "security_groups": ["sg-12345"] - }, - cpu="4096", - memory="8192" -) - -# Deploy infrastructure -rft_infra.setup() - -# Start training environment -rft_infra.start_training_environment() - -# Use with NovaModelCustomizer -customizer = NovaModelCustomizer( - model=Model.NOVA_LITE_2, - method=TrainingMethod.RFT_MULTITURN_LORA, - infra=runtime, - data_s3_path="s3://bucket/data.jsonl" -) - -training_result = customizer.train( - job_name="rft-training", - rft_multiturn_infra=rft_infra -) -``` - -### CustomEnvironment - -Create custom reward functions for RFT multiturn training. - -**Constructor:** -```python -def __init__( - self, - env_id: str, - local_path: str = None, - output_dir: str = "~/custom_envs", - env_type: str = "single_turn" -) -``` - -**Methods:** -- `create(overwrite: bool = False)`: Create environment structure -- `package_and_upload(bucket: Optional[str] = None)`: Upload to S3 - -**Example:** -```python -custom_env = CustomEnvironment( - env_id="my-custom-env", - output_dir="~/custom_envs/", - env_type="single_turn" -).create(overwrite=True) - -custom_env.validate() -custom_env.package_and_upload() -print(f"Uploaded to: {custom_env.s3_uri}") -``` - -### RFT Multiturn Methods - -**Infrastructure Management:** -- `setup()`: Deploy CloudFormation stack (Lambda, SQS, DynamoDB) -- `start_training_environment(vf_env_args: Dict = None)`: Start training workers -- `start_evaluation_environment(vf_env_args: Dict = None)`: Start evaluation workers -- `kill_task(env_type: EnvType)`: Stop workers -- `cleanup(delete_stack: bool = False, cleanup_environment: bool = False)`: Clean up resources - - `delete_stack`: If True, delete CloudFormation stack - - `cleanup_environment`: If True, clean up environment resources: - - LOCAL/EC2: Delete virtual environment and starter kit directories - - ECS: Deregister task definitions - -**Monitoring:** -- `get_logs(env_type: EnvType, limit: int = 100, start_from_head: bool = False, log_stream_name: Optional[str] = None, tail: bool = False)`: View worker logs - - `tail`: If True, continuously stream logs in real-time (blocks until Ctrl+C) -- `check_all_queues()`: Check SQS queue status -- `flush_all_queues()`: Clear all queues - -**Configuration:** -- `get_configuration()`: Get infrastructure config -- `get_recipe_overrides()`: Get recipe overrides for training - -**Note:** RFT multiturn only supports SageMaker HyperPod (SMHP) platform and Nova 2.0 models (NOVA_LITE_2). - ---- -## Job Notifications - -The Nova Forge SDK provides automated email notifications for training jobs when they reach terminal states (Completed, Failed, or Stopped). This feature helps you monitor long-running jobs without constantly checking their status. - -### Overview - -Job notifications are managed through platform-specific notification managers that automatically set up and manage the required AWS infrastructure: - -- **SMTJNotificationManager**: For SageMaker Training Jobs (SMTJ) -- **SMHPNotificationManager**: For SageMaker HyperPod (SMHP) - -### How It Works - -When you enable notifications for a job, the SDK automatically: - -1. **Creates AWS Infrastructure** (if it doesn't exist): - - CloudFormation stack with all required resources - - DynamoDB table to store job notification configurations - - SNS topic for email notifications - - Lambda function to handle job state changes - - EventBridge rule to monitor job status - - IAM roles and policies with appropriate permissions - -2. **Configures Job Monitoring**: - - Stores job configuration in DynamoDB - - Subscribes email addresses to SNS topic - - Monitors job status via EventBridge - -3. **Sends Notifications**: - - Detects when job reaches terminal state - - Validates output artifacts (for SMTJ, checks for manifest.json in output.tar.gz) - - Sends email notification with job details and console link - -### Using Job Notifications - -The simplest way to enable notifications is through the job result object: - -```python -from amzn_nova_forge import * - -# Start a training job -customizer = NovaModelCustomizer( - model=Model.NOVA_MICRO, - method=TrainingMethod.SFT_LORA, - infra=SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=2), - data_s3_path="s3://my-bucket/training-data/data.jsonl", - output_s3_path="s3://my-bucket/output" -) - -result = customizer.train(job_name="my-training-job") - -# Enable notifications -result.enable_job_notifications( - emails=["user@example.com", "team@example.com"], - region="us-west-2", # Optional - kms_key_id="1234abcd-12ab-34cd-56ef-1234567890ab", # Optional customer KMS key - output_s3_path="s3://my-bucket/custom-output-path" # Optional output path -) -``` -**Note:** Only provide `output_s3_path` if the 'JobResult' object doesn't have 'model_artifacts' (will be called out when you run the function). - -### Email Confirmation - -When you enable notifications: -1. Each email address receives a confirmation email from AWS SNS -2. Users must click the confirmation link in the email -3. After confirmation, they'll receive notifications for all jobs using that SNS topic -4. Confirmation is only needed once per email address per region - -### Notification Content - -Email notifications include: -- Job ID and platform (SMTJ/SMHP) -- Job status (Completed, Failed, or Stopped) -- Timestamp -- Link to AWS Console for the job -- For completed jobs: Validation status of output artifacts -- For failed jobs: Failure reason (if available) - -### Infrastructure Details (SMTJ) - -#### CloudFormation Stack - -The notification infrastructure is managed as a CloudFormation stack: -- **Stack Name**: `NovaForgeSDK-SMTJ-JobNotifications` -- **Region**: Specified when enabling notifications (default: us-east-1) -- **Resources**: DynamoDB table, SNS topic, Lambda function, EventBridge rule, IAM roles - -#### DynamoDB Table - -Stores job notification configurations: -- **Table Name**: `NovaForgeSDK-SMTJ-JobNotifications` -- **Primary Key**: `job_id` (String) -- **Attributes**: `emails` (String Set), `output_s3_path` (String), `created_at` (String), `ttl` (Number) -- **TTL**: Automatically deletes entries after 30 days - -#### SNS Topic - -Manages email subscriptions: -- **Topic Name**: `NovaForgeSDK-SMTJ-Notifications` -- **Encryption**: Optional KMS encryption -- **Subscriptions**: Email protocol with confirmation required - -#### Lambda Function - -Handles job state change events: -- **Function Name**: `NovaForgeSDK-SMTJ-NotificationHandler` -- **Runtime**: Python 3.12 -- **Timeout**: 180 seconds -- **Triggers**: EventBridge rule for SageMaker Training Job state changes - -#### EventBridge Rule - -Monitors job status: -- **Rule Name**: `NovaForgeSDK-SMTJ-Job-State-Change` -- **Event Pattern**: SageMaker Training Job State Change events -- **States Monitored**: Completed, Failed, Stopped - -### Infrastructure Details (SMHP) - -#### CloudFormation Stack - -The notification infrastructure is managed as a CloudFormation stack: -- **Stack Name**: `NovaForgeSDK-SMHP-JobNotifications-{ClusterName}` -- **Region**: Specified when enabling notifications (default: us-east-1) -- **Resources**: DynamoDB table, SNS topic, Lambda function, EventBridge rule, IAM roles, VPC endpoints - -#### DynamoDB Table - -Stores job notification configurations: -- **Table Name**: `NovaForgeSDK-SMHP-JobNotifications-{ClusterName}` -- **Primary Key**: `job_id` (String) -- **Attributes**: `output_s3_path` (String), `namespace` (String), `ttl` (Number) -- **TTL**: Automatically deletes entries after 30 days -- **Point-in-Time Recovery**: Enabled - -#### SNS Topic - -Manages email subscriptions: -- **Topic Name**: `NovaForgeSDK-SMHP-Notifications-{ClusterName}` -- **Encryption**: Optional KMS encryption -- **Subscriptions**: Email protocol with confirmation required - -#### Lambda Function - -Handles job status polling: -- **Function Name**: `NovaForgeSDK-SMHP-NotificationHandler-{ClusterName}` -- **Runtime**: Python 3.12 -- **Timeout**: 300 seconds -- **Memory**: 512 MB -- **VPC Configuration**: Deployed in VPC with access to EKS cluster -- **Layers**: kubectl layer for Kubernetes API access -- **Triggers**: EventBridge scheduled rule (default: every 5 minutes) - -#### EventBridge Rule - -Periodically checks job status: -- **Rule Name**: `NovaForgeSDK-SMHP-Job-Check-{ClusterName}` -- **Schedule**: Rate-based (default: every 5 minutes, configurable) -- **Target**: Lambda function for polling PyTorchJob status - -#### VPC Endpoints - -Enable private AWS service access for Lambda: -- **DynamoDB Gateway Endpoint**: `NovaForgeSDK-SMHP-DynamoDB-{ClusterName}` -- **SNS Interface Endpoint**: `NovaForgeSDK-SMHP-SNS-{ClusterName}` -- **S3 Gateway Endpoint**: `NovaForgeSDK-SMHP-S3-{ClusterName}` - -### Limitations and Notes - -1. **Email Confirmation**: Users must confirm their email subscription before receiving notifications. - -2. **Region-Specific**: Notification infrastructure is created per region. Jobs in different regions require separate infrastructure. - -3. **Stack Creation Restrictions**: For SMTJ, one notification stack is created per region. -For SMHP, one notification stack is created per cluster per region. - -4. **KMS Key Requirements**: If using KMS encryption: - - Provide only the key ID, not the full ARN - - The Lambda function automatically receives permissions to use the key - - The key must be in the same region as the notification infrastructure - -5. **Output Path Required**: The `output_s3_path` is required for manifest validation. The SDK will attempt to extract it from `model_artifacts` if not provided explicitly. - -6. **Hard-coded CloudFormation Stack Names**: When the CF stack is created, it will have one of the following names: `NovaForgeSDK-SMTJ-JobNotifications` or `NovaForgeSDK-SMHP-JobNotifications-{HP-Cluster}`. - -### Troubleshooting - -#### Notifications Not Received - -1. **Check email confirmation**: Ensure you clicked the confirmation link in the AWS SNS email -2. **Check spam folder**: SNS emails may be filtered as spam -3. **Verify job status**: Notifications only sent for terminal states (Completed, Failed, Stopped) -4. **Check CloudWatch Logs**: View Lambda function logs for errors - -#### Stack Creation Failures - -If CloudFormation stack creation fails: -1. Check IAM permissions for CloudFormation, DynamoDB, SNS, Lambda, EventBridge, and IAM -2. Verify no resource name conflicts exist -3. Check CloudFormation console for detailed error messages - -### API Reference - -See the [BaseJobResult.enable_job_notifications()](#enable_job_notifications) method documentation for detailed parameter information. - - ---- -## Error Handling -All SDK functions may raise exceptions. It's recommended to wrap calls in try-except blocks: -```python -try: - result = customizer.train(job_name="my-job") -except ValueError as e: - print(f"Configuration error: {e}") -except Exception as e: - print(f"Training failed: {e}") -``` -Common exceptions: -- `ValueError`: Invalid parameters or configuration -- `DataPrepError`: Dataset preparation errors -- `Exception`: General job execution or AWS API errors ---- -## Best Practices -1. **Always validate your data** using `loader.show()` before training -2. **Use overrides sparingly** - start with defaults and tune as needed -3. **Monitor logs** during training using `get_logs()` -4. **Check job status** before calling `.get()` on results -5. **Clean up resources** when done to avoid unnecessary costs -6. **Use descriptive job names** to help track and organize your experiments -7. **Save results incrementally** during long-running jobs -8. **Test with small datasets** before scaling up to full training ---- -## Additional Resources -- AWS Documentation: [Amazon Bedrock](https://docs.aws.amazon.com/bedrock/) -- AWS Documentation: [Amazon SageMaker](https://docs.aws.amazon.com/sagemaker/) -- SDK GitHub Repository: Check for updates and examples -- Support: Use AWS Support for technical assistance ---- diff --git a/docs/spec/dataset.md b/docs/spec/dataset.md new file mode 100644 index 0000000..3a359e6 --- /dev/null +++ b/docs/spec/dataset.md @@ -0,0 +1,621 @@ +# Dataset Loaders +Dataset loaders handle loading, transforming, and saving datasets in various formats. + +### Base Class: DatasetLoader +Base class for all dataset loaders. Handles path resolution (relative, tilde, `..` paths normalized to absolute), directory detection (scans for matching files, loads in lexicographic order), and delegates single-file loading to subclasses. + +#### Constructor + +**Signature:** +```python +def __init__( + self, + **column_mappings +) +``` + +**Parameters:** +- `**column_mappings`: Keyword arguments mapping standard column names to dataset column names + - Example: `question="input"` where "question" is the standard name and "input" is your column name +#### Column Mappings +If you are transforming a plain JSON, JSONL, or CSV file from a generic format (e.g. 'input/output') to another format (e.g. Converse for SFT), you need to provide "column mappings" to connect your generic column/field name to the expected ones in the transformation function. + +For example, if your plain dataset has "input" and "output" columns, and you want to transform it for SFT (which requrires "question" and "answer"), you would provide the following: +```python +loader = JSONDatasetLoader( + question="input", + answer="output" +) +``` +Below is a list of accepted column mapping parameters for transformations. +* SFT: `question`, `answer` + * Optional: `system`, [image/video required options]: `image_format`/`video_format`, `s3_uri`, `bucket_owner` + * 2.0: `reasoning_text`, `tools`/`toolsConfig`* +* RFT: `question`, `reference_answer` + * Optional: `system`, `id`, `tools`* +* Eval: `query`, `response` + * Optional: `images`, `metadata` +* CPT: `text` + +Additional Notes: +* If you're providing multimodal data in a generic format, you need to provide ALL three of the following fields: + * `image_format` OR `video_format` + `s3_uri`, `bucket_owner` +* *`tools/toolsConfig` (SFT 2.0) and `tools` (RFT) parameters can *only* be provided when transforming from OpenAI Messages format to Converse or OpenAI. A generic format *cannot* be provided for this transformation to work. + +--- +### JSONLDatasetLoader +Loads datasets from JSONL (JSON Lines) files. + +#### Methods + +##### `load()` +Loads dataset from a JSONL file (local or S3). + +**Signature:** +```python +def load( + self, + path: str +) -> "DatasetLoader" +``` + +**Parameters:** +- `path` (str): Path to JSONL file (local path or S3 URI) + +**Returns:** +- `DatasetLoader`: Self (for method chaining) + +**Example:** +```python +from amzn_nova_forge.dataset import * +loader = JSONLDatasetLoader() +loader.load("s3://my-bucket/data/training.jsonl") +``` +--- +### JSONDatasetLoader +Loads datasets from JSON files. + +#### Methods + +##### `load()` +Loads dataset from a JSON file (local or S3). + +**Signature:** +```python +def load( + self, + path: str +) -> "DatasetLoader" +``` + +**Parameters:** +- `path` (str): Path to JSON file (local path or S3 URI) + +**Returns:** +- `DatasetLoader`: Self (for method chaining) + +**Example:** +```python +from amzn_nova_forge.dataset import * +loader = JSONDatasetLoader() +loader.load("data/training.json") +``` +--- +### CSVDatasetLoader +Loads datasets from CSV files. + +#### Methods + +##### `load()` +Loads dataset from a CSV file. + +**Signature:** +```python +def load( + self, + path: str +) -> "DatasetLoader" +``` + +**Parameters:** +- `path` (str): Path to CSV file (local path or S3 URI) + +**Returns:** +- `DatasetLoader`: Self (for method chaining) + +**Example:** +```python +from amzn_nova_forge.dataset import * +loader = CSVDatasetLoader(question="user_query", answer="bot_response") +loader.load("data/conversations.csv") +``` +--- +### ParquetDatasetLoader +Loads datasets from Apache Parquet files. Uses lazy batch-based iteration via PyArrow — only one row group is in memory at a time. + +#### Methods + +##### `load()` +Loads dataset from a Parquet file or directory (local or S3). + +**Signature:** +```python +def load( + self, + path: str +) -> "DatasetLoader" +``` + +**Parameters:** +- `path` (str): Path to Parquet file or directory (local path or S3 URI). Accepts `.parquet` and `.pq` extensions. + +**Returns:** +- `DatasetLoader`: Self (for method chaining) + +**Example:** +```python +from amzn_nova_forge.dataset import * +loader = ParquetDatasetLoader() +loader.load("s3://my-bucket/data/training.parquet") +``` +--- +### ArrowDatasetLoader +Loads datasets from Apache Arrow IPC and Feather files. Supports both IPC Stream and IPC File formats with automatic fallback. + +#### Methods + +##### `load()` +Loads dataset from an Arrow IPC or Feather file or directory (local or S3). + +**Signature:** +```python +def load( + self, + path: str +) -> "DatasetLoader" +``` + +**Parameters:** +- `path` (str): Path to Arrow file or directory (local path or S3 URI). Accepts `.arrow`, `.feather`, and `.ipc` extensions. + +**Returns:** +- `DatasetLoader`: Self (for method chaining) + +**Example:** +```python +from amzn_nova_forge.dataset import * +loader = ArrowDatasetLoader() +loader.load("s3://my-bucket/data/training.arrow") +``` +--- +### HuggingFaceDatasetLoader +Loads datasets from [HuggingFace Hub](https://huggingface.co/datasets) and streams them through the Nova Forge data prep pipeline. Records are streamed lazily — no network calls happen at `load()` time. The dataset is fetched when a terminal operation (`save()`, `show()`, etc.) executes the pipeline. + +Requires the optional `datasets` dependency: +```bash +pip install amzn-nova-forge[huggingface] +``` + +Authentication for private or gated datasets is handled by the `datasets` library's built-in credential resolution. Set the `HF_TOKEN` environment variable or run `huggingface-cli login` before using the loader. + +#### Methods + +##### `load()` +Loads a dataset from HuggingFace Hub. + +**Signature:** +```python +def load( + self, + path: str, + split: Optional[str] = None, + name: Optional[str] = None, + revision: Optional[str] = None, + data_files: Optional[Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]] = None, + data_dir: Optional[str] = None, +) -> "DatasetLoader" +``` + +**Parameters:** +- `path` (str): HuggingFace dataset identifier. +- `split` (Optional[str]): Dataset split, forwarded to `datasets.load_dataset(split=...)`. Accepts any value supported by the HF library, including a single split name (`"train"`, `"test"`), a slice (`"train[:100]"`, `"train[:10%]"`), or concatenated splits (`"train+test"`). When `None`, the loader probes available splits: if the dataset has exactly one split it is auto-selected, and if the dataset has multiple splits a `DataPrepError` is raised listing the available splits. +- `name` (Optional[str]): Configuration name for multi-config datasets (forwarded to `datasets.load_dataset(name=...)`). +- `revision` (Optional[str]): Git tag or commit hash to pin the dataset version (forwarded to `datasets.load_dataset(revision=...)`). +- `data_files` (Optional[Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]]): Specific data file(s) to load, forwarded to `datasets.load_dataset(data_files=...)`. When `None` (default), the kwarg is omitted so the library's own default applies. +- `data_dir` (Optional[str]): Subdirectory of the dataset repository to load from, forwarded to `datasets.load_dataset(data_dir=...)`. When `None` (default), the kwarg is omitted. + +**Returns:** +- `DatasetLoader`: Self (for method chaining) + +**Raises:** +- `ImportError`: If the `datasets` package is not installed. The error message includes the install command. +- `DataPrepError`: If the dataset cannot be loaded. Authentication failures (401/403) surface with credential guidance; all other failures include the dataset path and the original error. Raised lazily when the pipeline is executed, not at `load()` time. + +**Example — public dataset:** +```python +from amzn_nova_forge.dataset import * + +# Load a public dataset and save to S3 +loader = HuggingFaceDatasetLoader() +loader.load("foo/bar", split="train") \ + .save("s3://my-bucket/data/foo_train.jsonl") +``` + + +--- +### CloudWatchDatasetLoader +Loads datasets from Amazon CloudWatch Logs via Insights queries and streams them through the Nova Forge data prep pipeline. The loader is lazy — `load()` stores configuration only, and the query executes when a terminal operation (`save()`, `show()`, etc.) triggers iteration. + +Uses the default AWS credential chain. Required IAM permissions: `logs:StartQuery` and `logs:GetQueryResults`. + +#### Methods + +##### `load()` +Loads data from a CloudWatch Logs Insights query. + +**Signature:** +```python +def load( + self, + log_group: str, + query: str, + start_time: datetime, + end_time: datetime, +) -> "CloudWatchDatasetLoader" +``` + +**Parameters:** +- `log_group` (str): CloudWatch log group name (e.g. `"/my-app/api-logs"`). +- `query` (str): A [CloudWatch Logs Insights query](https://docs.aws.amazon.com/AmazonCloudWatch/latest/logs/CWL_QuerySyntax.html). The query determines which fields appear as keys in the yielded dictionaries. Use `fields` to select specific fields from JSON logs, or `parse` to extract values from semi-structured text. See examples below. +- `start_time` (datetime): Query start time (inclusive). +- `end_time` (datetime): Query end time (exclusive). + +**Returns:** +- `CloudWatchDatasetLoader`: Self (for method chaining) + +**Raises:** +- `DataPrepError`: If the query fails during iteration (e.g. log group not found, invalid query syntax, insufficient permissions, query timeout). + +**Example — JSON logs:** + +If your log group contains JSON events like `{"endpoint": "/chat", "request": "What is AI?", "response": "AI is artificial intelligence."}`, use `fields` to select them: + +```python +from datetime import datetime, timezone +from amzn_nova_forge import CloudWatchDatasetLoader + +loader = CloudWatchDatasetLoader() +loader.load( + log_group="/my-app/api-logs", + query="fields request, response | filter endpoint = '/chat'", + start_time=datetime(2025, 1, 1, tzinfo=timezone.utc), + end_time=datetime(2025, 1, 8, tzinfo=timezone.utc), +) +# Yields: {"request": "What is AI?", "response": "AI is artificial intelligence."} +loader.transform( + column_mappings={"question": "request", "answer": "response"}, +) +loader.save("s3://my-bucket/data/chat_logs.jsonl") +``` + +For nested JSON objects, use dot-notation to reach into deeper fields. Array elements are accessed by index. Use `as` to alias the extracted path. For example, given logs like: + +```json +{"input": {"messages": [{"role": "user", "content": "What is AI?"}]}, "output": {"content": "AI is artificial intelligence."}} +``` + +You can extract the nested values with: + +``` +fields input.messages.0.content as question, output.content as answer +``` + +**Example — semi-structured text logs:** + +If your logs are text with a known pattern like `[INFO] endpoint=/api/chat input="What is AI?" output="AI is artificial intelligence."`, use `parse` with a regex to extract named fields: + +```python +loader = CloudWatchDatasetLoader() +loader.load( + log_group="/my-app/inference-logs", + query=''' + fields @message + | parse @message /input="(?[^"]+)" output="(?[^"]+)"/ + | filter ispresent(input) + ''', + start_time=datetime(2025, 1, 1, tzinfo=timezone.utc), + end_time=datetime(2025, 1, 8, tzinfo=timezone.utc), +) +# Yields: {"@message": "[INFO] endpoint=/api/chat ...", "input": "What is AI?", "output": "AI is artificial intelligence."} +loader.transform( + column_mappings={"question": "input", "answer": "output"}, +) +loader.save("s3://my-bucket/data/inference_logs.jsonl") +``` + + +--- +### Common DatasetLoader Methods +These methods are available on all DatasetLoader subclasses. + +#### `show()` +Displays the first n rows of the dataset. +**Signature:** +```python +def show( + self, + n: int = 10 +) -> None +``` + +**Parameters:** +- `n` (int): Number of rows to display (default: 10) + +**Example:** +```python +loader.show(5) # Show first 5 rows +``` +--- +#### `split_data()` +Splits dataset into train, validation, and test sets. + +**Signature:** +```python +def split_data( + self, + train_ratio: float = 0.8, + val_ratio: float = 0.1, + test_ratio: float = 0.1, + seed: int = 42, +) -> Tuple["DatasetLoader", "DatasetLoader", "DatasetLoader"] +``` + +**Parameters:** +- `train_ratio` (float): Proportion of data for training (default: 0.8) +- `val_ratio` (float): Proportion of data for validation (default: 0.1) +- `test_ratio` (float): Proportion of data for testing (default: 0.1) +- `seed` (int): Random seed for reproducibility (default: 42) + +**Returns:** +- `Tuple[DatasetLoader, DatasetLoader, DatasetLoader]`: Three DatasetLoader objects (train, val, test) + +**Raises:** +- `DataPrepError`: If ratios don't sum to 1.0 or dataset is empty + +**Example:** +```python +train_loader, val_loader, test_loader = loader.split_data( + train_ratio=0.7, + val_ratio=0.2, + test_ratio=0.1 +) +``` +--- +#### `transform()` +Transforms dataset to the required format for a specific training method and model. Currently the following transformations are supported: +* Q/A-formatted CSV/JSON/JSONL to SFT 1.0, SFT 2.0 (without reasoningContent, Tools), RFT, Eval, CPT +* OpenAI Messages format to SFT 1.0 and SFT 2.0 (with Tools) + +**Signature:** +```python +def transform( + self, + method: TransformMethod = TransformMethod.SCHEMA, + model: Optional[Model] = None, + eval_task: Optional[EvaluationTask] = None, + **kwargs, +) -> "DatasetLoader" +``` +**Parameters:** +- `method` (TransformMethod): The transform method (default: `TransformMethod.SCHEMA`). Also accepts a `TrainingMethod` enum for backward compatibility (deprecated). +- `model` (Optional[Model]): The Nova model version (e.g., `Model.NOVA_LITE_2`). Can be passed positionally for backward compatibility. +- `eval_task` (Optional[EvaluationTask]): Evaluation task. Can be passed positionally for backward compatibility. +- `**kwargs`: Method-specific arguments passed to the operation. For `TransformMethod.SCHEMA`: + - `training_method` (TrainingMethod): Required. The training method (e.g., `TrainingMethod.SFT_LORA`). + - `model` (Model): Required. The target model. + - `eval_task` (EvaluationTask): Optional. Required when `training_method` is `EVALUATION`. + - `column_mappings` (dict): Optional. Maps standard column names to your dataset's column names. + - `multimodal_data_s3_path` (Optional[str]): S3 prefix where images will be uploaded during conversion (e.g., `s3://my-bucket/images/`). Required when the dataset contains `image_url` content blocks in OpenAI format. When saving the output to S3, the save path must use the same bucket. + - `multimodal_data_bucket_owner` (Optional[str]): AWS account ID that owns the S3 bucket. If not provided, auto-resolved via STS. + +**Returns:** +- `DatasetLoader`: Self (for method chaining) + +**Raises:** +- `ValueError`: If method/model combination is not supported +- `DataPrepError`: If transformation fails + +**Example:** +```python +# Text-only transform +loader.transform( + method=TransformMethod.SCHEMA, + training_method=TrainingMethod.SFT_LORA, + model=Model.NOVA_MICRO, +) + +# Multimodal transform with automatic S3 image upload +loader.transform( + method=TransformMethod.SCHEMA, + training_method=TrainingMethod.SFT_LORA, + model=Model.NOVA_LITE_2, + multimodal_data_s3_path="s3://my-bucket/images/", +) +``` +--- +#### `validate()` +Validates dataset when given the user's intended training method and model. + +**Signature:** +```python +def validate( + self, + method: ValidateMethod = ValidateMethod.INVALID_RECORDS, + model: Model (Optional), + eval_task: EvaluationTask (Optional), + platform: Platform (Optional) +) -> None +``` +**Parameters:** +- `method` (ValidateMethod): The validation method (default: `ValidateMethod.INVALID_RECORDS`). `ValidateMethod.SCHEMA` is deprecated but still supported for backward compatibility. +- `model` (Model): The Nova model version (e.g., `Model.NOVA_LITE`) +- `eval_task` (EvaluationTask): Optional. The evaluation task (e.g., `EvaluationTask.GEN_QA`) +- `platform` (Platform): Optional. The target platform (`Platform.SMTJ`, `Platform.SMHP`, or `Platform.BEDROCK`). Accepted for forward-compatibility; row-count checks are currently enforced only during `train()`, not `validate()`. + +**Returns:** +- None + +**Raises:** +- `ValueError`: If method/model combination is not supported or validation is unsuccessful. + +**Row-Count Checks:** + +The following row-count checks are enforced during `train()`. They do not currently run on the `loader.validate()` path. + +| Check | Applies To | Limit | +|---|---|---| +| Maximum dataset rows | `BEDROCK` / `SFT_LORA`, `SFT_FULL`, `DPO_LORA`, `DPO_FULL`, `CPT` / `NOVA_MICRO`, `NOVA_LITE`, `NOVA_PRO` | 20,000 rows | +| Minimum dataset rows | `BEDROCK` / `SFT_LORA`, `SFT_FULL`, `DPO_LORA`, `DPO_FULL`, `CPT` / `NOVA_MICRO`, `NOVA_LITE`, `NOVA_PRO` | 8 rows | +| Maximum dataset rows | `BEDROCK` / `SFT_LORA`, `SFT_FULL` / `NOVA_LITE_2` | 20,000 rows | +| Minimum dataset rows | `BEDROCK` / `SFT_LORA`, `SFT_FULL` / `NOVA_LITE_2` | 200 rows | +| Maximum dataset rows | `BEDROCK` / `RFT_LORA`, `RFT_FULL` / `NOVA_LITE_2` | 20,000 rows | +| Minimum dataset rows | `BEDROCK` / `RFT_LORA`, `RFT_FULL` / `NOVA_LITE_2` | 100 rows | + +> **Note:** Recipe-dependent checks (e.g. dataset size ≥ `global_batch_size`) are also automatically enforced during `train()` where the fully-built recipe is available. + +**Example:** +```python +loader.validate( + method=ValidateMethod.INVALID_RECORDS, + training_method=TrainingMethod.SFT_LORA, + model=Model.NOVA_MICRO +) +``` + +If you're validating a BYOD Evaluation dataset, you need to provide another parameter, `eval_task` to the `validate` function. For example: +``` +loader.validate( + method=ValidateMethod.INVALID_RECORDS, + training_method=TrainingMethod.EVALUATION, + model=Model.NOVA_LITE_2, + eval_task=EvaluationTask.GEN_QA +) + +>> Validation succeeded for 22 samples on an Evaluation BYOD dataset +``` +--- +#### `filter()` +Queues a data filtering operation on the loader. All filters are lazy — `filter()` records the intent and execution happens when `execute()` is called (or implicitly via `transform()`, `show()`, `save()`). Multiple `filter()` calls can be chained. + +**Note:** Different filters work on different data formats. Text filters (`DEFAULT_TEXT_FILTER`, `EXACT_DEDUP`, `FUZZY_DEDUP`) operate on raw data with a flat text field. `INVALID_RECORDS` works on schema-formatted data. Using a filter on an incompatible data format may produce unexpected results. + +**Signature:** +```python +def filter( + self, + method: FilterMethod = FilterMethod.DEFAULT_TEXT_FILTER, + **kwargs, +) -> "DatasetLoader" +``` + +**Text Filters** (below text filters work on raw data with a flat text field): + +| Value | Description | +|---|---| +| `FilterMethod.DEFAULT_TEXT_FILTER` | Removes records with excessive URL content | +| `FilterMethod.EXACT_DEDUP` | Removes exact duplicate records by content hash | +| `FilterMethod.FUZZY_DEDUP` | Removes near-duplicate records using MinHash LSH similarity | + +Text filters run on AWS Glue (Ray) and use S3 path chaining. See the [Filtering Data](../user-guides/data_prep.md#filtering-data) section in the data prep guide for kwargs (`output_path`, `runtime_manager`, etc.). + +**INVALID_RECORDS Filter** (works on data in the expected format for the target training method): + +| Value | Description | +|---|---| +| `FilterMethod.INVALID_RECORDS` | Removes records that don't conform to expected input for the target model and training technique (schema validation + reserved keyword detection) | + +INVALID_RECORDS runs locally and modifies the loader's dataset in place. + +**Parameters for `INVALID_RECORDS` (passed as kwargs):** +- `training_method` (TrainingMethod): Required. Determines which checks apply. +- `model` (Model): Required. The target model. +- `platform` (Platform): Required. The target platform (`Platform.SMTJ`, `Platform.SMHP`, or `Platform.BEDROCK`). +- `eval_task` (EvaluationTask): Optional. Required when `training_method` is `EVALUATION`. + +**Returns:** +- `DatasetLoader`: Self (for method chaining) + +> **Note:** After `INVALID_RECORDS` runs, use `loader.show()` to inspect the filtered dataset or `loader.save()` to persist the results. + +**Raises:** +- `ValueError`: If required kwargs (`training_method`, `model`, `platform`) are missing for `INVALID_RECORDS`. + +> **Note:** If your data is already in the correct schema format (e.g. previously transformed and saved), you can call `filter(INVALID_RECORDS)` directly after `load()` without `transform()`. + +**Example (full chain):** +```python +loader.load("data.jsonl") +loader.filter( + method=FilterMethod.DEFAULT_TEXT_FILTER, +).filter( + method=FilterMethod.EXACT_DEDUP, +) +loader.transform( + method=TransformMethod.SCHEMA, + training_method=TrainingMethod.SFT_LORA, + model=Model.NOVA_LITE_2, +) +loader.filter( + method=FilterMethod.INVALID_RECORDS, + training_method=TrainingMethod.SFT_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMTJ, +) +loader.execute() # flush the pending filter +loader.validate( + method=ValidateMethod.SCHEMA, + training_method=TrainingMethod.SFT_LORA, + model=Model.NOVA_LITE_2, +) +``` + +**Example (already-formatted data — skip `transform()`):** +```python +loader = JSONLDatasetLoader() +loader.load("pre_transformed_data.jsonl") +loader.filter( + method=FilterMethod.INVALID_RECORDS, + training_method=TrainingMethod.SFT_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMTJ, +) +loader.execute() +``` + +--- +#### `save_data()` +Saves the dataset to a local or S3 location. +**Signature:** +```python +def save_data( + self, + save_path: str +) -> str +``` +**Parameters:** +- `save_path` (str): Path where to save the file (local or S3, must end in .json or .jsonl) + +**Returns:** +- `str`: Path where the file was saved + +**Raises:** +- `DataPrepError`: If save fails or format is unsupported + +**Example:** +```python +# Save locally +loader.save_data("output/training_data.jsonl") +# Save to S3 +loader.save_data("s3://my-bucket/data/training_data.jsonl") +``` +--- diff --git a/docs/spec/enums.md b/docs/spec/enums.md new file mode 100644 index 0000000..3c3b23e --- /dev/null +++ b/docs/spec/enums.md @@ -0,0 +1,98 @@ +# Enums and Configuration +### Model Enum +Supported Nova models with their configurations. +**Values:** +- `Model.NOVA_MICRO`: Amazon Nova Micro (Version 1) + - `model_type`: "amazon.nova-micro-v1:0:128k" + - `model_path`: "nova-micro/prod" + - `version`: Version.ONE +- `Model.NOVA_LITE`: Amazon Nova Lite (Version 1) + - `model_type`: "amazon.nova-lite-v1:0:300k" + - `model_path`: "nova-lite/prod" + - `version`: Version.ONE +- `Model.NOVA_LITE_2`: Amazon Nova Lite (Version 2) + - `model_type`: "amazon.nova-2-lite-v1:0:256k" + - `model_path`: "nova-lite-2/prod" + - `version`: Version.TWO +- `Model.NOVA_PRO`: Amazon Nova Pro (Version 1) + - `model_type`: "amazon.nova-pro-v1:0:300k" + - `model_path`: "nova-pro/prod" + - `version`: Version.ONE + +**Methods:** +##### `from_model_type()` +Gets Model enum from model type string. + +**Signature:** +```python +@classmethod +def from_model_type( + cls, + model_type: str +) -> "Model" +``` + +**Example:** +```python +model = Model.from_model_type("amazon.nova-micro-v1:0:128k") +``` +--- +### TrainingMethod Enum +Supported training methods. + +**Values:** +- `TrainingMethod.CPT`: Continued Pre-Training +- `TrainingMethod.DPO_LORA`: Direct Preference Optimization with LoRA +- `TrainingMethod.DPO_FULL`: Direct Preference Optimization (full rank) +- `TrainingMethod.SFT_LORA`: Supervised Fine-Tuning with LoRA +- `TrainingMethod.SFT_FULL`: Supervised Fine-Tuning (full rank) +- `TrainingMethod.RFT_LORA`: Reinforcement Fine-Tuning with LoRA +- `TrainingMethod.RFT_FULL`: Full reinforcement Fine-Tuning +- `TrainingMethod.EVALUATION`: Evaluation only +--- +### DeployPlatform Enum +Supported deployment platforms. + +**Values:** +- `DeployPlatform.BEDROCK_OD`: Amazon Bedrock On-Demand +- `DeployPlatform.BEDROCK_PT`: Amazon Bedrock Provisioned Throughput +- `DeployPlatform.SAGEMAKER`: Amazon SageMaker +--- +### DeploymentMode Enum +Deployment behavior when an endpoint with the same name already exists. + +**Values:** +- `DeploymentMode.FAIL_IF_EXISTS`: Raise an error if endpoint already exists (safest, default) +- `DeploymentMode.UPDATE_IF_EXISTS`: Try in-place update only, fail if not supported (PT only) + +**Note:** Only `FAIL_IF_EXISTS` and `UPDATE_IF_EXISTS` modes are currently supported. +`UPDATE_IF_EXISTS` is only applicable for Bedrock Provisioned Throughput (PT) deployments. + +--- +### EvaluationTask Enum +Supported evaluation tasks. +Common values include: +- `EvaluationTask.MMLU`: Massive Multitask Language Understanding +- `EvaluationTask.GPQA`: General Physics Question Answering +- `EvaluationTask.MATH`: Mathematical Problem Solving +- `EvaluationTask.GEN_QA`: Custom Dataset Evaluation +- The full list of available tasks can be found here: [AWS Documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/nova-model-evaluation.html#nova-model-evaluation-benchmark) +--- +### Platform Enum +Infrastructure platforms. + +**Values:** +- `Platform.SMTJ`: SageMaker Training Jobs +- `Platform.SMHP`: SageMaker HyperPod +- `Platform.BEDROCK`: Amazon Bedrock +--- +### JobStatus Enum +Job execution status. + +**Values:** +- `JobStatus.IN_PROGRESS`: Job is running +- `JobStatus.COMPLETED`: Job completed successfully +- `JobStatus.FAILED`: Job failed + +--- + diff --git a/docs/spec/index.md b/docs/spec/index.md new file mode 100644 index 0000000..9f731fc --- /dev/null +++ b/docs/spec/index.md @@ -0,0 +1,32 @@ +# Nova Forge SDK — API Specification + +## Contents + +- [Service Classes](service-classes.md) — ForgeTrainer, ForgeEvaluator, ForgeDeployer, ForgeInference, ForgeConfig +- [NovaModelCustomizer](model-customizer.md) — Deprecated unified interface +- [Runtime Managers](runtime-managers.md) — SMTJ, SMHP, Glue runtime configuration +- [Dataset Loaders](dataset.md) — DatasetLoader, transforms, filters, operations +- [Job Results](job-results.md) — TrainingResult, EvalResult, DeployResult +- [Utilities](utilities.md) — Helper functions and monitoring +- [Enums and Configuration](enums.md) — Platform, Model, TrainingMethod, etc. +- [RFT Multiturn](rft-multiturn.md) — Multi-turn RFT infrastructure +- [Notifications](notifications.md) — Job notifications and error handling + +--- + +## Best Practices +1. **Always validate your data** using `loader.show()` before training +2. **Use overrides sparingly** - start with defaults and tune as needed +3. **Monitor logs** during training using `get_logs()` +4. **Check job status** before calling `.get()` on results +5. **Clean up resources** when done to avoid unnecessary costs +6. **Use descriptive job names** to help track and organize your experiments +7. **Save results incrementally** during long-running jobs +8. **Test with small datasets** before scaling up to full training +--- +## Additional Resources +- AWS Documentation: [Amazon Bedrock](https://docs.aws.amazon.com/bedrock/) +- AWS Documentation: [Amazon SageMaker](https://docs.aws.amazon.com/sagemaker/) +- SDK GitHub Repository: Check for updates and examples +- Support: Use AWS Support for technical assistance +--- diff --git a/docs/spec/job-results.md b/docs/spec/job-results.md new file mode 100644 index 0000000..7a86e78 --- /dev/null +++ b/docs/spec/job-results.md @@ -0,0 +1,473 @@ +# Job Results +Job result classes provide methods to check status and retrieve results from training, evaluation, and inference jobs. + +### Base Classes + +#### BaseJobResult +Abstract base class for all job results. + +**Attributes:** +- `job_id` (str): Job identifier +- `started_time` (datetime): Job start timestamp + +**Methods:** +##### `get_job_status()` +Gets the current status of the job. + +**Signature:** +```python +def get_job_status( + self +) -> tuple[JobStatus, str] +``` + +**Returns:** +- `tuple[JobStatus, str]`: A tuple of (status enum, raw status string) + - `JobStatus.IN_PROGRESS`: Job is running + - `JobStatus.COMPLETED`: Job completed successfully + - `JobStatus.FAILED`: Job failed + +**Example:** +```python +status, raw_status = result.get_job_status() +if status == JobStatus.COMPLETED: + print("Job finished!") +``` + +##### `dump(file_path: Optional[str] = None, file_name: Optional[str] = None)` +Save the job result to file_path path + +**Signature:** +```python +def dump( + self, + file_path: Optional[str] = None, + file_name: Optional[str] = None +) -> Path +``` + +**Parameters:** +- `file_path` (Optional[str]): Directory path to save the result. Saves to current directory if not provided +- `file_name` (Optional[str]): The file name of the result. Default to `_.json` if not provided + +**Returns:** +- `Path`: The full result file path + +**Example:** +```python +result.dump() +# Result will be saved to ./{job_id}_{platform}.json under current dir +result.dump(file_path='/customized/path', file_name='customized_name.json') +# Result will be saved to /customized/path/customized_name.json +``` + +##### `load(file_path: str)` +Load the job result from the `file_path` path + +**Signature:** +```python +@classmethod +def load( + cls, + file_path: str +) -> "BaseJobResult": +``` + +**Returns:** +- JobResultObject. The instance of subclass of BaseJobResult such as SMTJEvaluationResult, SMHPEvaluationResult, BedrockEvaluationResult, SMTJTrainingResult, SMHPTrainingResult, or BedrockTrainingResult + +**Example:** +```python +job_result = BaseJobResult.load('./my_job_result.json') +``` + +--- + +##### `enable_job_notifications()` +Enable email notifications for when a job reaches a terminal state (Completed, Failed, or Stopped). + +**Signature:** +```python +def enable_job_notifications( + self, + emails: list[str], + output_s3_path: Optional[str] = None, + region: Optional[str] = "us-east-1", + **platform_kwargs +) -> None +``` + +**Parameters:** +- `emails` (list[str]): List of email addresses to notify +- `output_s3_path` (Optional[str]): S3 path where job outputs are stored. + - Only required if the SDK cannot automatically extract it from the job result's `model_artifacts` attribute. + - For most training jobs, this parameter is automatically populated and does not need to be provided explicitly. +- `region` (Optional[str]): AWS region for notification infrastructure (default: "us-east-1") +- `**platform_kwargs`: Platform-specific parameters: + - **For SMTJ:** + - `kms_key_id` (Optional[str]): Customer KMS key ID (not full ARN) for SNS topic encryption + - **For SMHP:** + - `namespace` (str): Kubernetes namespace where the PyTorchJob runs (e.g., "kubeflow", "default") (Required) + - `kubectl_layer_arn` (str): ARN of the lambda-kubectl layer (Required) + - `eks_cluster_arn` (Optional[str]): EKS cluster ARN (auto-detected if not provided) + - `vpc_id` (Optional[str]): VPC ID (auto-detected if not provided) + - `subnet_ids` (Optional[list[str]]): List of subnet IDs for Lambda (auto-detected if not provided) + - `security_group_id` (Optional[str]): Security group ID for Lambda (auto-detected if not provided) + - `polling_interval_minutes` (Optional[int]): How often to check job status in minutes (default: 5) + - `kms_key_id` (Optional[str]): Customer KMS key ID (not full ARN) for SNS topic encryption + +**Returns:** +- None + +**Raises:** +- `ValueError`: If required parameters are missing or invalid +- `NotificationManagerInfraError`: If infrastructure setup fails + +**How It Works:** +1. Creates AWS infrastructure (CloudFormation stack) if it doesn't exist: + - DynamoDB table to store job notification configurations + - SNS topic for email notifications + - Lambda function to monitor job status + - EventBridge rule (SMTJ) or scheduled rule (SMHP) to trigger Lambda + - (SMHP only) VPC endpoints for DynamoDB and S3 if needed +2. Stores job configuration in DynamoDB (including namespace for SMHP) +3. Subscribes email addresses to SNS topic (users must confirm subscription) +4. Monitors job status and sends email when job completes, fails, is stopped, or becomes degraded (SMHP only) + +**Email Confirmation:** +Users will receive a confirmation email from AWS SNS and must click the confirmation link before receiving job notifications. + +**Examples:** + +SMTJ (SageMaker Training Jobs): +```python +# Basic usage - output_s3_path is automatically extracted +result = customizer.train(job_name="my-job") +result.enable_job_notifications( + emails=["user@example.com", "team@example.com"] +) + +# With customer KMS encryption +result.enable_job_notifications( + emails=["user@example.com"], + kms_key_id="abc-123-def-456" # Just the key ID, not full ARN +) + +# With custom region +result.enable_job_notifications( + emails=["user@example.com"], + region="us-west-2" +) +``` + +SMHP (SageMaker HyperPod): +```python +# Basic usage (with auto-detection) +result = customizer.train(job_name="my-job") +result.enable_job_notifications( + emails=["user@example.com"], + namespace="kubeflow", # Required + kubectl_layer_arn="arn:aws:lambda::123456789012:layer:kubectl:1" # Required +) + +# With custom polling interval +result.enable_job_notifications( + emails=["user@example.com"], + namespace="kubeflow", + kubectl_layer_arn="arn:aws:lambda::123456789012:layer:kubectl:1", + polling_interval_minutes=10 # Check every 10 minutes instead of default 5 +) + +# With explicit VPC configuration of the cluster where jobs are being monitored. +result.enable_job_notifications( + emails=["user@example.com"], + namespace="kubeflow", + kubectl_layer_arn="arn:aws:lambda::123456789012:layer:kubectl:1", + eks_cluster_arn="arn:aws:eks::123456789012:cluster/my-cluster", + vpc_id="vpc-12345", + subnet_ids=["subnet-1", "subnet-2"], + security_group_id="sg-12345" +) +``` + +**Important Notes:** +- For SMHP, requires deploying a kubectl Lambda layer from AWS Serverless Application Repository +- For SMHP, the user will need to manually grant the Lambda function access to your EKS cluster (access-entry). + - Please refer to [`docs/user-guides/job_notifications.md`](../user-guides/job_notifications.md) for the commands to run to set this up. +- See [`docs/user-guides/job_notifications.md`](../user-guides/job_notifications.md) for detailed setup instructions, troubleshooting, and advanced usage + +--- +### EvaluationResult (ABC) +Result object for SageMaker Training Job evaluation tasks. + +**Attributes:** +- `job_id` (str): Job identifier +- `started_time` (datetime): Job start timestamp +- `eval_task` (EvaluationTask): Evaluation task performed +- `eval_output_path` (str): S3 path to evaluation results + +#### Subclasses +- SMTJEvaluationResult +- SMHPEvaluationResult +- BedrockEvaluationResult + +#### Methods + +##### `get()` +Downloads and returns evaluation results as a dictionary. + +**Signature:** +```python +def get( + self +) -> Dict +``` + +**Returns:** +- `Dict`: Evaluation results (empty dict if job not completed) + +**Example:** +```python +eval_result = customizer.evaluate(...) +# Wait for job to complete +results = eval_result.get() +print(results) +``` + +--- +##### `show()` +Prints evaluation results to console. + +**Signature:** +```python +def show( + self +) -> None +``` + +**Example:** +```python +eval_result.show() +``` + +--- +##### `upload_tensorboard_results()` +Uploads TensorBoard results to S3. + +**Signature:** +```python +def upload_tensorboard_results( + self, + tensorboard_s3_path: Optional[str] = None +) -> None +``` + +**Parameters:** +- `tensorboard_s3_path` (Optional[str]): Target S3 path (auto-generated if not provided) + +**Example:** +```python +eval_result.upload_tensorboard_results( + tensorboard_s3_path="s3://my-bucket/tensorboard/" +) +``` + +--- +##### `clean()` +Cleans up local cached results. + +**Signature:** +```python +def clean( + self +) -> None +``` + +--- +### SMTJBatchInferenceResult +Result object for batch inference jobs. +**Attributes:** +- `job_id` (str): Job identifier +- `started_time` (datetime): Job start timestamp +- `inference_output_path` (str): S3 path to inference outputs + +#### Methods + +##### `get()` +Downloads and returns inference results, optionally saving to S3. +**Signature:** +```python +def get( + self, + s3_path: Optional[str] = None +) -> Dict +``` +**Parameters:** +- `s3_path` (Optional[str]): S3 path to save formatted results +**Returns:** +- `Dict`: Dictionary containing list of inference results + - Each result has: `system`, `query`, `gold_response`, `inference_response`, `metadata` +**Example:** +```python +inference_result = customizer.batch_inference(...) +# Wait for job to complete +results = inference_result.get(s3_path="s3://my-bucket/formatted-results.jsonl") +``` +--- +##### `show()` +Prints inference results to console. +**Signature:** +```python +def show( + self +) -> None +``` +--- +##### `clean()` +Cleans up local cached results. +**Signature:** +```python +def clean( + self +) -> None +``` +--- +### IAM Role Creation SDK +This SDK provides utility functions for creating IAM roles with specific permissions for AWS Bedrock and SageMaker services. + +**Methods** + +#### `create_bedrock_execution_role()` +Creates an IAM role with permissions for Bedrock model creation and deployment. + +**Signature:** +```python +def create_bedrock_execution_role( + iam_client, + role_name: str, + bedrock_resource: str = "*", + s3_resource: str = "*" +) -> Dict +``` + +**Parameters:** +- `iam_client`: Boto3 IAM client +- `role_name` (str): Name of the IAM role to create +- `bedrock_resource` (Optional[str]): Specific Bedrock resource to restrict access. Defaults to "*" (all resources) +- `s3_resource` (Optional[str]): Specific S3 resource to restrict access. Defaults to "*" (all resources) + +**Returns:** +- `Dict`: IAM role details + +**Example:** +```python +import boto3 +from amzn_nova_forge.iam import create_bedrock_execution_role + +iam_client = boto3.client("iam") +create_bedrock_execution_role(iam_client, "role-name", "bedrock_resource", "s3_resource") +``` + +### `create_sagemaker_execution_role()` +Creates an IAM role with permissions for SageMaker model creation and deployment. + +**Signature:** +```python +def create_sagemaker_execution_role( + iam_client, + role_name: str, + s3_resource: str = "*", + kms_resource: str = "*", + ec2_condition: Optional[Dict[str, Any]] = None, + cloudwatch_metric_condition: Optional[Dict[str, Any]] = None, + cloudwatch_logstream_resource: str = "*", + cloudwatch_loggroup_resource: str = "*" +) -> Dict +``` + +**Parameters:** +- `iam_client`: Boto3 IAM client +- `role_name` (str): Name of the IAM role to create +- `s3_resource` (Optional[str]): Specific S3 resource to restrict access +- `kms_resource` (Optional[str]): Specific KMS resource to restrict access +- `ec2_condition` (Optional[Dict]): Conditional access for EC2 resources +- `cloudwatch_metric_condition` (Optional[Dict]): Conditional access for CloudWatch metrics +- `cloudwatch_logstream_resource` (Optional[str]): Specific CloudWatch log stream resource +- `cloudwatch_loggroup_resource` (Optional[str]): Specific CloudWatch log group resource + +**Returns:** +- `Dict`: IAM role details + +**Example:** +```python +import boto3 +from amzn_nova_forge.iam import create_sagemaker_execution_role + +iam_client = boto3.client("iam") +create_sagemaker_execution_role( + iam_client, + role_name="role-name", + s3_resource="example-bucket", + kms_resource="encryption-key", + ec2_condition={ + "ArnLike": { + "ec2:Vpc": "arn:aws:ec2:*:*:vpc/example" + } + }, + cloudwatch_metric_condition={ + "StringEquals": { + "cloudwatch:namespace": ["example-namespace"] + } + }, + cloudwatch_loggroup_resource="example-loggroup", + cloudwatch_logstream_resource="example-logstream" + ) +``` +--- +#### `ModelDeployResult` +Result of creating a Bedrock or Sagemaker model. + +**Fields:** +- `model_arn` (str): The Bedrock custom model ARN or SageMaker model ARN. +- `model_name` (str): The model name. +- `escrow_uri` (str): S3 artifacts path used to create the model. +- `created_at` (datetime): UTC timestamp when the model was created. + +**Properties:** +- `platform` (Optional[str]): Returns `"bedrock"` or `"sagemaker"` based on strict ARN format validation, or `None` for unrecognized ARNs. +- `status` (ModelStatus): Queries the model's current status via live API call. Returns `CREATING`, `ACTIVE`, `FAILED`, or `UNKNOWN`. Logs a warning if no AWS client is available. + +**Class Methods:** +- `from_arn(model_arn, bedrock_client=None, sagemaker_client=None)` — Reconstruct from an existing model ARN. Detects platform from ARN format, calls the appropriate describe API, and recovers `escrow_uri` from tags. Works for both Bedrock and SageMaker ARNs. +- `load(file_path)` — Load from a JSON file saved by `dump()`. Automatically creates fresh AWS clients for status checking. + +**Instance Methods:** +- `dump(file_path=None, file_name=None)` — Save to JSON file for later use with `load()`. + +**Example:** +```python +# Create and persist +publish = customizer.create_custom_model(job_result=training_result) +publish.dump(file_path="./results/") + +# Load and check status +loaded = ModelDeployResult.load("./results/my-model_deploy_result.json") +print(loaded.platform) # "bedrock" or "sagemaker" +print(loaded.status) # ModelStatus.ACTIVE + +# Reconstruct from ARN +result = ModelDeployResult.from_arn( + "arn:aws:bedrock:us-east-1:123456789012:custom-model/my-model" +) +``` + +#### `ModelStatus` +Platform-independent model status enum. + +**Values:** +- `ModelStatus.CREATING` — Model is still being created (Bedrock only). +- `ModelStatus.ACTIVE` — Model is ready for deployment. +- `ModelStatus.FAILED` — Model creation failed. +- `ModelStatus.UNKNOWN` — Status could not be determined (no client, unrecognized ARN, or unexpected API response). +--- + diff --git a/docs/spec/model-customizer.md b/docs/spec/model-customizer.md new file mode 100644 index 0000000..9285b23 --- /dev/null +++ b/docs/spec/model-customizer.md @@ -0,0 +1,713 @@ +# NovaModelCustomizer (Deprecated) + +> **Deprecated**: `NovaModelCustomizer` is a legacy facade. Use `ForgeTrainer`, `ForgeEvaluator`, `ForgeDeployer`, and `ForgeInference` instead. + +The main entrypoint class for customizing and training Nova models. + +### Constructor + +#### `__init__()` +Initializes a NovaModelCustomizer instance. + +**Signature:** +```python +def __init__( + self, + model: Model, + method: TrainingMethod, + infra: RuntimeManager, + data_s3_path: Optional[str] = None, + output_s3_path: Optional[str] = None, + model_path: Optional[str] = None, + validation_config: Optional[Dict[str, bool]] = None, + generated_recipe_dir: Optional[str] = None, + mlflow_monitor: Optional[MLflowMonitor] = None, + deployment_mode: DeploymentMode = DeploymentMode.FAIL_IF_EXISTS, + data_mixing_enabled: bool = False, + enable_job_caching: bool = False, + is_multimodal: Optional[bool] = None, + hub_content_version: Optional[str] = None, +) +``` +**Parameters:** +- `model` (Model): The Nova model to be trained (e.g., `Model.NOVA_MICRO`, `Model.NOVA_LITE`, `Model.NOVA_LITE_2`, `Model.NOVA_PRO`) +- `method` (TrainingMethod): The fine-tuning method (e.g., `TrainingMethod.SFT_LORA`, `TrainingMethod.RFT`) +- `infra` (RuntimeManager): Runtime infrastructure manager (e.g., `SMTJRuntimeManager`, `SMHPRuntimeManager`, or `BedrockRuntimeManager`) +- `data_s3_path` (Optional[str]): S3 path to the training dataset +- `output_s3_path` (Optional[str]): S3 path for output artifacts. If not provided, will be auto-generated +- `model_path` (Optional[str]): S3 path for model path +- `validation_config` (Optional[Dict[str, bool]]): Optional dict to control validation. Defaults to `{'iam': True, 'infra': True, 'recipe': True}`. + - `iam` (bool): Enable IAM permission validation (default: True) + - `infra` (bool): Enable infrastructure validation (default: True) + - `recipe` (bool): Enable recipe constraint validation (default: True) +- `generated_recipe_dir` (Optional[str]): Optional local path to save the generated recipe +- `mlflow_monitor` (Optional[MLflowMonitor]): Optional MLflow monitoring configuration for experiment tracking (SageMaker only, not supported on Bedrock) +- `deployment_mode` (DeploymentMode): Behavior when deploying to existing endpoint name. Options: FAIL_IF_EXISTS (default), UPDATE_IF_EXISTS +- `data_mixing_enabled` (bool): Enable data mixing feature for CPT and SFT training on SageMaker HyperPod, and SFT text-only on Nova Lite 2 on SMTJServerless. Default is False + - **Note:** The `data_mixing_enabled` parameter must be set to `True` during initialization to use data mixing features. + - **Note:** On SMHP: supported for CPT, SFT_LORA, and SFT_FULL across Nova 1 and Nova 2 models. On SMTJServerless: supported for SFT_LORA and SFT_FULL with text-only data on Nova Lite 2 only. +- `is_multimodal` (Optional[bool]): Only applicable when `data_mixing_enabled=True`. Explicitly set multimodal mode. If None (default), auto-detects from data. Set to False to skip detection for performance on large text-only datasets. Ignored when `data_mixing_enabled=False` +- `enable_job_caching` (bool): Whether to enable job result caching. When enabled, completed job results are cached to `job_cache_dir` (default: `.cached-nova-jobs/`) and reused for identical job configurations. Default: False +- `hub_content_version` (Optional[str]): Version of the hub content to retrieve from SageMaker Hub. If None, uses the latest version + +**Raises:** +- `ValueError`: If region is unsupported or model is invalid + +**Example:** +```python +from amzn_nova_forge import * + +# SageMaker Training Jobs (SMTJ) +infra = SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=2) + +customizer = NovaModelCustomizer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + data_s3_path="s3://my-bucket/training-data/data.jsonl", + output_s3_path="s3://my-bucket/output" +) + +# Amazon Bedrock (fully managed) +bedrock_infra = BedrockRuntimeManager( + execution_role="arn:aws:iam::123456789012:role/BedrockRole", + base_model_identifier="arn:aws:bedrock:us-east-1::custom-model/amazon.nova-2-lite-v1:0:256k:abcdefghijk" +) + +bedrock_customizer = NovaModelCustomizer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=bedrock_infra, + data_s3_path="s3://my-bucket/training-data/data.jsonl", + output_s3_path="s3://my-bucket/output" +) + +# With MLflow monitoring (SageMaker only) +mlflow_monitor = MLflowMonitor( + tracking_uri="arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", + experiment_name="nova-customization", + run_name="sft-run-1" +) + +customizer_with_mlflow = NovaModelCustomizer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + data_s3_path="s3://my-bucket/training-data/data.jsonl", + output_s3_path="s3://my-bucket/output", + mlflow_monitor=mlflow_monitor +) +``` +--- +### Methods + +#### `get_data_mixing_config()` +Get the current data mixing configuration. + +**Signature:** +```python +def get_data_mixing_config( + self +) -> Dict[str, Any] +``` + +**Returns:** +- `Dict[str, Any]`: Dictionary containing the data mixing configuration + +**Example:** +```python +config = customizer.get_data_mixing_config() +print(config) +# Output: {'customer_data_percent': 50, 'nova_code_percent': 30, 'nova_general_percent': 70} +``` + +--- + +#### `set_data_mixing_config()` +Set the data mixing configuration. + +**Signature:** +```python +def set_data_mixing_config( + self, + config: Dict[str, Any] +) -> None +``` + +**Parameters:** +- `config` (Dict[str, Any]): Dictionary containing the data mixing configuration + - `customer_data_percent` (int/float): Percentage of customer data (0-100) + - `nova_code_percent` (int/float): Percentage of Nova code data (0-100) + - `nova_general_percent` (int/float): Percentage of Nova general data (0-100) + - Nova percentages must sum to 100% + +**Raises:** +- `ValueError`: If data mixing is not enabled or configuration is invalid + +**Note:** +- The SDK automatically detects whether your dataset is multimodal (contains images, videos, or documents) by scanning the data +- The appropriate Nova dataset catalog (text-only or multimodal) and nova data mixing fields are selected automatically based on this detection + +**Example:** +```python +# Text datamixing (auto-detection) +customizer = NovaModelCustomizer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.SFT_LORA, + infra=SMHPRuntimeManager(...), + data_s3_path="s3://bucket/data.jsonl", + data_mixing_enabled=True, + # is_multimodal=False, # Optional: skip auto-detection for performance +) + +customizer.set_data_mixing_config({ + "customer_data_percent": 50, + "nova_code_percent": 30, + "nova_general_percent": 70 +}) + +# Multimodal datamixing (explicit) +customizer_mm = NovaModelCustomizer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.SFT_LORA, + infra=SMHPRuntimeManager(...), + data_s3_path="s3://bucket/multimodal_data.jsonl", + data_mixing_enabled=True, + is_multimodal=True, +) + +customizer_mm.set_data_mixing_config({ + "customer_data_percent": 50, + "nova_general_percent": 70, + "nova_code_percent": 30, +}) +``` + +--- + +#### `train()` +Generates the recipe YAML, configures runtime, and launches a training job. + +**Signature:** +```python +def train( + self, + job_name: str, + recipe_path: Optional[str] = None, + overrides: Optional[Dict[str, Any]] = None, + rft_lambda_arn: Optional[str] = None, + validation_data_s3_path: Optional[str] = None, + val_check_interval: Optional[int] = None, + dry_run: bool = False +) -> TrainingResult +``` + +**Parameters:** +- `job_name` (str): User-defined name for the training job +- `recipe_path` (Optional[str]): Path for a YAML recipe file (both S3 and local paths are accepted) +- `overrides` (Optional[Dict[str, Any]]): Dictionary of configuration overrides. Example overrides below: + - `max_epochs` (int): Maximum number of training epochs + - `lr` (float): Learning rate + - `warmup_steps` (int): Number of warmup steps + - `loraplus_lr_ratio` (float): LoRA+ learning rate ratio + - `global_batch_size` (int): Global batch size + - `max_length` (int): Maximum sequence length + - A full list of available overrides can be found via the [Nova Customization public documentation](https://docs.aws.amazon.com/nova/latest/userguide/customize-fine-tune-sagemaker.html) or by referencing the training recipes [here](https://docs.aws.amazon.com/sagemaker/latest/dg/nova-model-recipes.html). +- `rft_lambda_arn` (Optional[str]): Rewards Lambda ARN (only used for RFT training methods). If passed, takes priority over `rft_lambda_arn` set on the `RuntimeManager`. +- `validation_data_s3_path` (Optional[str]): Validation S3 path, applicable for CPT and SFT on SMTJ/SMTJServerless/SMHP, or any method on Bedrock (optional) +- `val_check_interval` (Optional[int]): How often (in training steps) to run validation. Defaults to 2500 if omitted. Only used when `validation_data_s3_path` is provided. +- `dry_run` (bool): Actually starts a job if False, otherwise just performs validation. + +**Returns:** +- `TrainingResult`: Metadata object (either `SMTJTrainingResult`, `SMHPTrainingResult`, or `BedrockTrainingResult`) containing: + - `job_id` (str): The training job identifier + - `method` (TrainingMethod): The training method used + - `started_time` (datetime): Job start timestamp + - `model_artifacts` (ModelArtifacts): Paths to model checkpoints and outputs + - `checkpoint_s3_path` (str, Optional): Path to the model checkpoint/trained model. For `SMTJServerless`, populated after job completion via `get_model_artifacts()`. + - `output_s3_path` (str): Path to the metrics and output tar file. + - `output_model_arn` (str, Optional): Model package ARN for `SMTJServerless` jobs. Use as `model_path` for iterative training. + - `model_type` (Model): Model type of the model being trained + +**Raises:** +- `Exception`: If job execution fails +- `ValueError`: If training method is not supported + +**Example:** +```python +result = customizer.train( + job_name="my-training-job", + overrides={ + 'max_epochs': 10, + 'lr': 5e-6, + 'warmup_steps': 20, + 'global_batch_size': 128 + } +) +print(f"Training job started: {result.job_id}") +print(f"Checkpoint path: {result.model_artifacts.checkpoint_s3_path}") +``` +--- +#### `evaluate()` +Generates the recipe YAML, configures runtime, and launches an evaluation job. + +**Signature:** +```python +def evaluate( + self, + job_name: str, + eval_task: EvaluationTask, + model_path: Optional[str] = None, + subtask: Optional[str] = None, + data_s3_path: Optional[str] = None, + recipe_path: Optional[str] = None, + overrides: Optional[Dict[str, Any]] = None, + processor: Optional[Dict[str, Any]] = None, + rl_env: Optional[Dict[str, Any]] = None, + dry_run: Optional[bool] = False, + job_result: Optional[TrainingResult] = None +) -> EvaluationResult | None +``` + +**Parameters:** +- `job_name` (str): User-defined name for the evaluation job +- `eval_task` (EvaluationTask): The evaluation task to be performed (e.g., `EvaluationTask.MMLU`) + - The list of available tasks can be found here: [AWS Documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/nova-model-evaluation.html#nova-model-evaluation-benchmark) +- `model_path` (Optional[str]): S3 path for model to evaluate. If not provided, will attempt to extract from `job_result` or the customizer's most recent training job. +- `data_s3_path` (Optional[str]): S3 URI for the dataset. Only required for BYOD (Bring Your Own Data) evaluation tasks. +- `subtask` (Optional[str]): Subtask for evaluation (task-specific) + - The list of available subtasks per task can be found here: [Subtasks](https://docs.aws.amazon.com/sagemaker/latest/dg/nova-model-evaluation.html#nova-model-evaluation-subtasks) +- `recipe_path` (Optional[str]): Path for a YAML recipe file (both S3 and local paths are accepted) +- `overrides` (Optional[Dict[str, Any]]): Dictionary of inference configuration overrides + - `max_new_tokens` (int): Maximum tokens to generate + - `top_k` (int): Top-k sampling parameter + - `top_p` (float): Top-p (nucleus) sampling parameter + - `temperature` (float): Temperature for sampling +- `processor` (Optional[Dict[str, Any]]): Optional, Bring Your Own Metrics/RFT lambda Configuration +- `rl_env` (Optional[Dict[str, Any]]): Optional, Bring your own reinforcement learning environment config. For `RFT_EVAL`, if either `processor` or `rl_env` is explicitly passed, it takes priority over `rft_lambda_arn` set on the `RuntimeManager`. +- `dry_run` (Optional[bool]): Actually starts a job if False, otherwise just performs validation. +- `job_result` (Optional[TrainingResult]): Optional TrainingResult object to extract checkpoint path from. If provided and `model_path` is None, will automatically extract the checkpoint path from the training job's output and validate platform compatibility. + +**Returns:** +- `EvaluationResult(BaseJobResult)`: Metadata object (either `SMTJEvaluationResult`, `SMHPEvaluationResult`, or `BedrockEvaluationResult`) containing: + - `job_id` (str): The evaluation job identifier + - `started_time` (datetime): Job start timestamp + - `eval_output_path` (str): S3 path to evaluation results + - `eval_task` (EvaluationTask): The Evaluation task +- Returns `None` if `dry_run=True` + +**Example:** + +```python +from amzn_nova_forge.recipe import * + +# General eval task (with overrides) +eval_result = customizer.evaluate( + job_name="my-eval-job", + eval_task=EvaluationTask.MMLU, + model_path="s3://my-bucket/checkpoints/my-model", + overrides={ + 'max_new_tokens': 2048, + 'temperature': 0, + 'top_p': 1.0 + } +) +print(f"Evaluation job started: {eval_result.job_id}") + +# BYOM eval task (by providing processor config) +byom_eval_result = customizer.evaluate( + job_name='my-eval-test-byom', + eval_task=EvaluationTask.GEN_QA, + data_s3_path="s3://bucket/data", + processor={ + "lambda_arn": "arn:aws:lambda::123456789012:function:byom-lambda" + } +) +``` +--- +#### `deploy()` +Creates a custom model and deploys it to Amazon Bedrock or SageMaker. + +Deployment behavior when endpoint already exists is controlled by the `deployment_mode` +parameter set during NovaModelCustomizer initialization: +- FAIL_IF_EXISTS: Raise error (default, safest) +- UPDATE_IF_EXISTS: Try in-place update, fail if not supported (PT only) + +**Signature:** +```python +def deploy( + self, + model_artifact_path: Optional[str] = None, + deploy_platform: DeployPlatform = DeployPlatform.BEDROCK_OD, + unit_count: Optional[int] = None, + endpoint_name: Optional[str] = None, + job_result: Optional[TrainingResult] = None, + execution_role_name: Optional[str] = None, + sagemaker_instance_type: Optional[str] = "ml.p5.48xlarge", + sagemaker_environment: Optional[SageMakerEndpointEnvironment] = None, +) -> DeploymentResult +``` + +* **Note:** If DeployPlatform.BEDROCK_PT or DeployPlatform.SAGEMAKER is selected, you must include a value for unit_count. +* **Note:** If `model_artifact_path` is provided, we will NOT attempt to resolve `model_artifact_path` from `job_result` or the enclosing `NovaModelCustomizer` object. + +**Parameters:** +- `model_artifact_path` (Optional[str]): S3 path to the trained model checkpoint. If not provided, will attempt to extract from job_result or the `job_id` field of the Customizer. +- `deploy_platform` (DeployPlatform): Platform to deploy the model to + - `DeployPlatform.BEDROCK_OD`: Bedrock On-Demand + - `DeployPlatform.BEDROCK_PT`: Bedrock Provisioned Throughput + - `DeployPlatform.SAGEMAKER`: SageMaker +- `unit_count` (Optional[int]): Used in Bedrock Provisioned Throughput number of PT to purchase or SageMaker number of initial instances +- `endpoint_name` (Optional[str]): Name of the deployed model's endpoint (auto-generated if not provided) +- `job_result` (Optional[TrainingResult]): Training job result object to use for extracting checkpoint path and validating job completion. Also used to retrieve job_id if it's not provided. +- `execution_role_name` (Optional[str]): IAM role for Bedrock or SageMaker. If provided, used as-is — no policies created or attached. If omitted, the SDK creates and manages a default role with required policies. +- `sagemaker_instance_type`: Optional EC2 instance type for SageMaker deployment, defaults to ml.p5.48xlarge +- `sagemaker_environment` (Optional[SageMakerEndpointEnvironment]): SageMaker endpoint environment config. See `SageMakerEndpointEnvironment` for available fields and defaults. +**Returns:** +- `DeploymentResult`: Contains: + - `endpoint` (EndpointInfo): Endpoint information + - `platform` (DeployPlatform): Deployment platform + - `endpoint_name` (str): Endpoint name + - `uri` (str): Model ARN + - `model_artifact_path` (str): S3 path to artifacts + - `created_at` (datetime): Deployment creation timestamp + +**Raises:** +- `Exception`: When unable to successfully deploy the model +- `ValueError`: If platform is not supported + +**Example:** +```python +from amzn_nova_forge.model import * + +bedrock_deployment = customizer.deploy( + model_artifact_path="s3://escrow-bucket/my-model-artifacts/", + deploy_platform=DeployPlatform.BEDROCK_OD, + endpoint_name="my-custom-nova-model-bedrock" +) +print(f"Model deployed: {bedrock_deployment.endpoint.uri}") +print(f"Endpoint: {bedrock_deployment.endpoint.endpoint_name}") +print(f"Status: {bedrock_deployment.status}") + +sagemaker_deployment = customizer.deploy( + model_artifact_path="s3://escrow-bucket/my-model-artifacts/", + deploy_platform=DeployPlatform.SAGEMAKER, + unit_count=1, + endpoint_name="my-custom-nova-model-sagemaker", + sagemaker_environment=SageMakerEndpointEnvironment( + context_length=8000, + max_concurrency=4, + ) +) +print(f"Model deployed: {sagemaker_deployment.endpoint.uri}") +print(f"Endpoint: {sagemaker_deployment.endpoint.endpoint_name}") +print(f"Status: {sagemaker_deployment.status}") +``` + +Optionally, you can provide a Bedrock execution role name to be used in deployment. +Otherwise, a default Bedrock execution role will be created on your behalf. +You can also use the following method to create a Bedrock execution role with scoped down IAM permissions. + + +```python +from amzn_nova_forge.iam import create_bedrock_execution_role + +iam_client = boto3.client("iam") + +create_bedrock_execution_role( + iam_client=iam_client, + role_name="BedrockDeployModelExecutionRole", + bedrock_resource="your-model-name", # Optional: Name of the bedrock resources that IAM role should have restricted create and get access to + s3_resource="s3-bucket" # Optional: S3 resource that IAM role should have restricted read access to such as the training output bucket +) + +``` +--- +#### `create_custom_model()` +Creates a Bedrock custom model from S3 artifacts, decoupled from endpoint deployment. + +This method extracts the model-creation step from the deploy flow so users can +create a model independently of endpoint deployment, enabling retry of deployment +if it fails after model creation. + +**Signature:** +```python +def create_custom_model( + self, + model_artifact_path: Optional[str] = None, + job_result: Optional[TrainingResult] = None, + endpoint_name: Optional[str] = None, + execution_role_name: Optional[str] = None, + tags: Optional[List[Dict[str, str]]] = None, + skip_model_reuse: bool = False, +) -> ModelDeployResult +``` + +**Parameters:** +- `model_artifact_path` (Optional[str]): S3 path to trained model checkpoint. Takes precedence over `job_result` if both are provided. +- `job_result` (Optional[TrainingResult]): Training job result to extract checkpoint path from. +- `endpoint_name` (Optional[str]): Optional name prefix for the model name (auto-generated if not provided). +- `execution_role_name` (Optional[str]): IAM role name for Bedrock. Defaults to `BedrockDeployModelExecutionRole`. +- `tags` (Optional[List[Dict[str, str]]]): Optional list of `{"key": str, "value": str}` dicts for source tracking. +- `skip_model_reuse` (bool): If True, always create a new model even if one with the same escrow URI already exists. Default: False. + +**Returns:** +- `ModelDeployResult`: Contains: + - `model_arn` (str): The Bedrock custom model ARN + - `model_name` (str): The model name passed to CreateCustomModel + - `escrow_uri` (str): S3 artifacts path used to create the model + - `created_at` (datetime): UTC timestamp when the model was created + +**Raises:** +- `ValueError`: When neither `model_artifact_path` nor `job_result` is provided, or when checkpoint path cannot be resolved from `job_result`. +- `RuntimeError`: When IAM role creation or custom model creation fails. + +**Example:** +```python +# Create a custom model from training artifacts +publish_result = customizer.create_custom_model( + model_artifact_path="s3://escrow-bucket/my-model-artifacts/" +) +print(f"Model ARN: {publish_result.model_arn}") + +# Save for later use +publish_result.dump(file_path="./results/") + +# Or create from a training job result +publish_result = customizer.create_custom_model(job_result=training_result) +``` +--- +#### `deploy_to_bedrock()` +Deploys a published Bedrock custom model to an endpoint. + +Use after `create_custom_model()` to deploy the model, or provide a model ARN directly. +This decoupled approach allows retrying deployment without re-creating the model. + +When a `model_deploy_result` is provided, the model's status is validated before deployment: +- **Active**: Proceeds immediately. +- **Creating**: Waits for the model to become Active (30s poll interval, 15min timeout). +- **Failed**: Raises `ValueError`. + +**Signature:** +```python +def deploy_to_bedrock( + self, + model_deploy_result: Optional[ModelDeployResult] = None, + model_arn: Optional[str] = None, + deploy_platform: DeployPlatform = DeployPlatform.BEDROCK_OD, + pt_units: Optional[int] = None, + endpoint_name: Optional[str] = None, +) -> DeploymentResult +``` + +**Parameters:** +- `model_deploy_result` (Optional[ModelDeployResult]): Result from `create_custom_model()`. Cannot be combined with `model_arn`. +- `model_arn` (Optional[str]): Direct model ARN. Cannot be combined with `model_deploy_result`. +- `deploy_platform` (DeployPlatform): `BEDROCK_OD` (default) or `BEDROCK_PT`. +- `pt_units` (Optional[int]): Number of PT units (required for `BEDROCK_PT`). +- `endpoint_name` (Optional[str]): Endpoint name (auto-generated if not provided). + +**Returns:** +- `DeploymentResult`: Contains: + - `endpoint` (EndpointInfo): Endpoint information + - `created_at` (datetime): Deployment creation timestamp + - `model_publish` (Optional[ModelDeployResult]): The model deploy result, if available + - `escrow_uri` (Optional[str]): Convenience property delegating to `model_publish.escrow_uri` + +**Raises:** +- `ValueError`: When both `model_deploy_result` and `model_arn` are provided, or when no model ARN is available. +- `RuntimeError`: When deployment creation fails. + +**Example:** +```python +# Two-step deploy: create model, then deploy +publish_result = customizer.create_custom_model( + model_artifact_path="s3://escrow-bucket/my-model-artifacts/" +) +deployment = customizer.deploy_to_bedrock( + model_deploy_result=publish_result, + endpoint_name="my-endpoint" +) + +# Or deploy directly from a model ARN +deployment = customizer.deploy_to_bedrock( + model_arn="arn:aws:bedrock:us-east-1:123456789012:custom-model/my-model" +) + +# Retry deployment if it fails (model already created) +deployment = customizer.deploy_to_bedrock( + model_arn=publish_result.model_arn +) +``` +--- +#### `deploy_to_sagemaker()` +Deploys a model to a SageMaker Inference endpoint. + +Can reuse an existing SageMaker model via `model_deploy_result` (e.g., from a +previous deploy that created the model but failed on endpoint creation), +or create a new model from `model_artifact_path`. + +**Signature:** +```python +def deploy_to_sagemaker( + self, + instance_type: str, + model_deploy_result: Optional[ModelDeployResult] = None, + model_artifact_path: Optional[str] = None, + unit_count: int = 1, + endpoint_name: Optional[str] = None, + sagemaker_environment: Optional[SageMakerEndpointEnvironment] = None, + execution_role_name: Optional[str] = None, + skip_model_reuse: bool = False, +) -> DeploymentResult +``` + +**Parameters:** +- `instance_type` (str): SageMaker instance type (required). +- `model_deploy_result` (Optional[ModelDeployResult]): Result from a previous deploy containing the SM model ARN. Skips model creation. Cannot be combined with `model_artifact_path`. +- `model_artifact_path` (Optional[str]): S3 path to model artifacts. Creates a new SM model. +- `unit_count` (int): Number of instances. Default: 1. +- `endpoint_name` (Optional[str]): Endpoint name (auto-generated if not provided). +- `sagemaker_environment` (Optional[SageMakerEndpointEnvironment]): SageMaker endpoint environment config. See `SageMakerEndpointEnvironment` for available fields and defaults. +- `execution_role_name` (str): IAM execution role name. +- `skip_model_reuse` (bool): If True, always create a new model (skip tag-based discovery). Default: False. + +**Returns:** +- `DeploymentResult`: Contains endpoint info and `model_publish` (ModelDeployResult). + +**Raises:** +- `ValueError`: When both `model_deploy_result` and `model_artifact_path` are provided, when neither is provided, or when `instance_type` is missing. +- `RuntimeError`: When endpoint creation fails. The error message includes the model ARN and a retry command using `customizer.last_model_publish`. + +**Example:** +```python +# Deploy from S3 artifacts +result = customizer.deploy_to_sagemaker( + model_artifact_path="s3://escrow-bucket/checkpoint/", + instance_type="ml.g5.12xlarge", +) + +# Retry after endpoint failure (model already created) +result = customizer.deploy_to_sagemaker( + model_deploy_result=customizer.last_model_publish, + instance_type="ml.g5.12xlarge", + endpoint_name="my-endpoint", +) +``` +--- + +#### `invoke_inference()` +Invokes a single inference on a trained model. + +**Signature:** +```python +def invoke_inference( + self, + request_body: Dict[str, Any], + endpoint_arn: Optional[str] +) -> InferenceResult +``` +**Parameters:** +- `request_body` (Dict[str, Any]): Inference request body +- `endpoint_arn` (Optional[str]):Endpoint ARN to invoke inference. Optional if user wants to send request to an already deployed endpoint on customizer + +**Returns:** +- `InferenceResult`: Metadata object (`SingleInferenceResult`) containing: + - `job_id` (str): Batch inference job identifier + - `started_time` (datetime): Job start timestamp + - `inference_output_path` (str): Empty string + +**Example:** +```python +inference_result = customizer.invoke_inference( + request_body={ + "messages": [{"role": "user", "content": "Hello! How are you?"}], + "max_tokens": 100, + "stream": False, + }, + endpoint_arn="arn:aws:sagemaker:us-east-1:123456789012:endpoint/endpoint", +) +inference_result.show() +``` +--- +#### `batch_inference()` +Launches a batch inference job on a trained model. + +**Signature:** +```python +def batch_inference( + self, + job_name: str, + input_path: str, + output_s3_path: str, + model_path: Optional[str] = None, + recipe_path: Optional[str] = None, + overrides: Optional[Dict[str, Any]] = None, + dry_run: Optional[bool] = False +) -> InferenceResult +``` +**Parameters:** +- `job_name` (str): Name for the batch inference job +- `input_path` (str): S3 path to input data for inference +- `output_s3_path` (str): S3 path for inference outputs +- `model_path` (Optional[str]): S3 path to the model +- `recipe_path` (Optional[str]): Path for a YAML recipe file +- `overrides` (Optional[Dict[str, Any]]): Configuration overrides for inference + - `max_new_tokens` (int): Maximum tokens to generate + - `top_k` (int): Top-k sampling parameter + - `top_p` (float): Top-p (nucleus) sampling parameter + - `temperature` (float): Temperature for sampling + - `top_logprobs` (int): Number of top log probabilities to return +- `dry_run` (Optional[bool]): Actually starts a job if False, otherwise just performs validation. + +**Returns:** +- `InferenceResult`: Metadata object (`SMTJBatchInferenceResult`) containing: + - `job_id` (str): Batch inference job identifier + - `started_time` (datetime): Job start timestamp + - `inference_output_path` (str): S3 path to inference results + - Note: Batch inference is only supported on SageMaker platforms (SMTJ, SMHP) + +**Example:** +```python +inference_result = customizer.batch_inference( + job_name="batch-inference-job", + input_path="s3://my-bucket/inference-input", + output_s3_path="s3://my-bucket/inference-output", + model_path="s3://my-bucket/trained-model" +) +print(f"Batch inference started: {inference_result.job_id}") +``` +In a separate notebook cell, you can run the following commands to get the job status and download a formatted result file when the jobs completes. +```python +inference_result.get_job_status() # Gets the job status. +inference_result.get("s3://my-bucket/save-location/file-name.jsonl") # Uploads a formatted inference_results.jsonl file to the given s3 location. +``` +--- +#### `get_logs()` +Retrieves and displays CloudWatch logs for the current job. + +**Signature:** +```python +def get_logs( + self, + limit: Optional[int] = None, + start_from_head: bool = False, + end_time: Optional[int] = None +) +``` + +**Parameters:** +- `limit` (Optional[int]): Maximum number of log lines to retrieve +- `start_from_head` (bool): If True, start from the beginning of logs; if False, start from the end +- `end_time` (Optional[int]): End time in epoch milliseconds for searching a log time range + +**Returns:** +- None (prints logs to console) + +**Example:** +```python +# After starting a training job +customizer.train(job_name="my-job") +customizer.get_logs(limit=100, start_from_head=True) +``` +--- diff --git a/docs/spec/notifications.md b/docs/spec/notifications.md new file mode 100644 index 0000000..f6beca5 --- /dev/null +++ b/docs/spec/notifications.md @@ -0,0 +1,223 @@ +# Job Notifications + +The Nova Forge SDK provides automated email notifications for training jobs when they reach terminal states (Completed, Failed, or Stopped). This feature helps you monitor long-running jobs without constantly checking their status. + +### Overview + +Job notifications are managed through platform-specific notification managers that automatically set up and manage the required AWS infrastructure: + +- **SMTJNotificationManager**: For SageMaker Training Jobs (SMTJ) +- **SMHPNotificationManager**: For SageMaker HyperPod (SMHP) + +### How It Works + +When you enable notifications for a job, the SDK automatically: + +1. **Creates AWS Infrastructure** (if it doesn't exist): + - CloudFormation stack with all required resources + - DynamoDB table to store job notification configurations + - SNS topic for email notifications + - Lambda function to handle job state changes + - EventBridge rule to monitor job status + - IAM roles and policies with appropriate permissions + +2. **Configures Job Monitoring**: + - Stores job configuration in DynamoDB + - Subscribes email addresses to SNS topic + - Monitors job status via EventBridge + +3. **Sends Notifications**: + - Detects when job reaches terminal state + - Validates output artifacts (for SMTJ, checks for manifest.json in output.tar.gz) + - Sends email notification with job details and console link + +### Using Job Notifications + +The simplest way to enable notifications is through the job result object: + +```python +from amzn_nova_forge import * + +# Start a training job +customizer = NovaModelCustomizer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=2), + data_s3_path="s3://my-bucket/training-data/data.jsonl", + output_s3_path="s3://my-bucket/output" +) + +result = customizer.train(job_name="my-training-job") + +# Enable notifications +result.enable_job_notifications( + emails=["user@example.com", "team@example.com"], + region="us-west-2", # Optional + kms_key_id="1234abcd-12ab-34cd-56ef-1234567890ab", # Optional customer KMS key + output_s3_path="s3://my-bucket/custom-output-path" # Optional output path +) +``` +**Note:** Only provide `output_s3_path` if the 'JobResult' object doesn't have 'model_artifacts' (will be called out when you run the function). + +### Email Confirmation + +When you enable notifications: +1. Each email address receives a confirmation email from AWS SNS +2. Users must click the confirmation link in the email +3. After confirmation, they'll receive notifications for all jobs using that SNS topic +4. Confirmation is only needed once per email address per region + +### Notification Content + +Email notifications include: +- Job ID and platform (SMTJ/SMHP) +- Job status (Completed, Failed, or Stopped) +- Timestamp +- Link to AWS Console for the job +- For completed jobs: Validation status of output artifacts +- For failed jobs: Failure reason (if available) + +### Infrastructure Details (SMTJ) + +#### CloudFormation Stack + +The notification infrastructure is managed as a CloudFormation stack: +- **Stack Name**: `NovaForgeSDK-SMTJ-JobNotifications` +- **Region**: Specified when enabling notifications (default: us-east-1) +- **Resources**: DynamoDB table, SNS topic, Lambda function, EventBridge rule, IAM roles + +#### DynamoDB Table + +Stores job notification configurations: +- **Table Name**: `NovaForgeSDK-SMTJ-JobNotifications` +- **Primary Key**: `job_id` (String) +- **Attributes**: `emails` (String Set), `output_s3_path` (String), `created_at` (String), `ttl` (Number) +- **TTL**: Automatically deletes entries after 30 days + +#### SNS Topic + +Manages email subscriptions: +- **Topic Name**: `NovaForgeSDK-SMTJ-Notifications` +- **Encryption**: Optional KMS encryption +- **Subscriptions**: Email protocol with confirmation required + +#### Lambda Function + +Handles job state change events: +- **Function Name**: `NovaForgeSDK-SMTJ-NotificationHandler` +- **Runtime**: Python 3.12 +- **Timeout**: 180 seconds +- **Triggers**: EventBridge rule for SageMaker Training Job state changes + +#### EventBridge Rule + +Monitors job status: +- **Rule Name**: `NovaForgeSDK-SMTJ-Job-State-Change` +- **Event Pattern**: SageMaker Training Job State Change events +- **States Monitored**: Completed, Failed, Stopped + +### Infrastructure Details (SMHP) + +#### CloudFormation Stack + +The notification infrastructure is managed as a CloudFormation stack: +- **Stack Name**: `NovaForgeSDK-SMHP-JobNotifications-{ClusterName}` +- **Region**: Specified when enabling notifications (default: us-east-1) +- **Resources**: DynamoDB table, SNS topic, Lambda function, EventBridge rule, IAM roles, VPC endpoints + +#### DynamoDB Table + +Stores job notification configurations: +- **Table Name**: `NovaForgeSDK-SMHP-JobNotifications-{ClusterName}` +- **Primary Key**: `job_id` (String) +- **Attributes**: `output_s3_path` (String), `namespace` (String), `ttl` (Number) +- **TTL**: Automatically deletes entries after 30 days +- **Point-in-Time Recovery**: Enabled + +#### SNS Topic + +Manages email subscriptions: +- **Topic Name**: `NovaForgeSDK-SMHP-Notifications-{ClusterName}` +- **Encryption**: Optional KMS encryption +- **Subscriptions**: Email protocol with confirmation required + +#### Lambda Function + +Handles job status polling: +- **Function Name**: `NovaForgeSDK-SMHP-NotificationHandler-{ClusterName}` +- **Runtime**: Python 3.12 +- **Timeout**: 300 seconds +- **Memory**: 512 MB +- **VPC Configuration**: Deployed in VPC with access to EKS cluster +- **Layers**: kubectl layer for Kubernetes API access +- **Triggers**: EventBridge scheduled rule (default: every 5 minutes) + +#### EventBridge Rule + +Periodically checks job status: +- **Rule Name**: `NovaForgeSDK-SMHP-Job-Check-{ClusterName}` +- **Schedule**: Rate-based (default: every 5 minutes, configurable) +- **Target**: Lambda function for polling PyTorchJob status + +#### VPC Endpoints + +Enable private AWS service access for Lambda: +- **DynamoDB Gateway Endpoint**: `NovaForgeSDK-SMHP-DynamoDB-{ClusterName}` +- **SNS Interface Endpoint**: `NovaForgeSDK-SMHP-SNS-{ClusterName}` +- **S3 Gateway Endpoint**: `NovaForgeSDK-SMHP-S3-{ClusterName}` + +### Limitations and Notes + +1. **Email Confirmation**: Users must confirm their email subscription before receiving notifications. + +2. **Region-Specific**: Notification infrastructure is created per region. Jobs in different regions require separate infrastructure. + +3. **Stack Creation Restrictions**: For SMTJ, one notification stack is created per region. +For SMHP, one notification stack is created per cluster per region. + +4. **KMS Key Requirements**: If using KMS encryption: + - Provide only the key ID, not the full ARN + - The Lambda function automatically receives permissions to use the key + - The key must be in the same region as the notification infrastructure + +5. **Output Path Required**: The `output_s3_path` is required for manifest validation. The SDK will attempt to extract it from `model_artifacts` if not provided explicitly. + +6. **Hard-coded CloudFormation Stack Names**: When the CF stack is created, it will have one of the following names: `NovaForgeSDK-SMTJ-JobNotifications` or `NovaForgeSDK-SMHP-JobNotifications-{HP-Cluster}`. + +### Troubleshooting + +#### Notifications Not Received + +1. **Check email confirmation**: Ensure you clicked the confirmation link in the AWS SNS email +2. **Check spam folder**: SNS emails may be filtered as spam +3. **Verify job status**: Notifications only sent for terminal states (Completed, Failed, Stopped) +4. **Check CloudWatch Logs**: View Lambda function logs for errors + +#### Stack Creation Failures + +If CloudFormation stack creation fails: +1. Check IAM permissions for CloudFormation, DynamoDB, SNS, Lambda, EventBridge, and IAM +2. Verify no resource name conflicts exist +3. Check CloudFormation console for detailed error messages + +### API Reference + +See the [BaseJobResult.enable_job_notifications()](#enable_job_notifications) method documentation for detailed parameter information. + + +--- +## Error Handling +All SDK functions may raise exceptions. It's recommended to wrap calls in try-except blocks: +```python +try: + result = customizer.train(job_name="my-job") +except ValueError as e: + print(f"Configuration error: {e}") +except Exception as e: + print(f"Training failed: {e}") +``` +Common exceptions: +- `ValueError`: Invalid parameters or configuration +- `DataPrepError`: Dataset preparation errors +- `Exception`: General job execution or AWS API errors +--- diff --git a/docs/spec/rft-multiturn.md b/docs/spec/rft-multiturn.md new file mode 100644 index 0000000..7c0563e --- /dev/null +++ b/docs/spec/rft-multiturn.md @@ -0,0 +1,212 @@ +# RFT Multiturn Infrastructure + +For RFT multiturn training and evaluation, you need to set up infrastructure to run reward functions. + +### Helper Functions + +#### create_rft_execution_role + +Creates an IAM role with required permissions for RFT multiturn infrastructure. + +**Function:** +```python +def create_rft_execution_role( + region: str = "us-east-1", + role_name: Optional[str] = None, + custom_policy_path: Optional[str] = None +) -> str +``` + +**Parameters:** +- `region` (str): AWS region. Default: "us-east-1" +- `role_name` (Optional[str]): Custom role name. Default: "RFTExecutionRoleNovaSDK" +- `custom_policy_path` (Optional[str]): Path to custom policy JSON file. If not provided, uses SDK default. + +**Returns:** +- `str`: ARN of the created/existing role + +**Example:** +```python +from amzn_nova_forge import create_rft_execution_role + +# Create role with default name +role_arn = create_rft_execution_role(region="us-east-1") + +# Create role with custom name +role_arn = create_rft_execution_role(region="us-east-1", role_name="my-custom-rft-role") +``` + +#### list_rft_stacks + +Lists CloudFormation stacks in the region, optionally filtering for Nova SDK stacks. + +**Function:** +```python +def list_rft_stacks( + region: str = "us-east-1", + all_stacks: bool = False +) -> List[str] +``` + +**Parameters:** +- `region` (str): AWS region. Default: "us-east-1" +- `all_stacks` (bool): If True, list all stacks. If False, only list Nova SDK stacks (ending with "NovaForgeSDK"). Default: False + +**Returns:** +- `List[str]`: List of stack names + +**Example:** +```python +from amzn_nova_forge import list_rft_stacks + +# List only Nova SDK stacks +nova_stacks = list_rft_stacks(region="us-east-1") + +# List all CloudFormation stacks +all_stacks = list_rft_stacks(region="us-east-1", all_stacks=True) +``` + +### RFTMultiturnInfrastructure + +Manages infrastructure for RFT multiturn training (reward function workers). + +**Constructor:** +```python +def __init__( + self, + stack_name: str, + region: str = "us-east-1", + vf_env_id: Optional[VFEnvId] = None, + custom_env: Optional[CustomEnvironment] = None, + infrastructure_arn: Optional[str] = None, + python_venv_name: Optional[str] = None, + vpc_config: Optional[Dict[str, Any]] = None, + cpu: Optional[str] = None, + memory: Optional[str] = None, + rft_role_name: Optional[str] = None, +) +``` + +**Parameters:** +- `stack_name` (str): CloudFormation stack name +- `region` (str): AWS region. Default: "us-east-1" +- `vf_env_id` (Optional[VFEnvId]): Built-in environment ID (VFEnvId.WORDLE or VFEnvId.TERMINAL_BENCH) +- `custom_env` (Optional[CustomEnvironment]): Custom environment (mutually exclusive with vf_env_id) +- `infrastructure_arn` (Optional[str]): Platform ARN (EC2 instance ID, ECS cluster ARN, or None for LOCAL) +- `python_venv_name` (Optional[str]): Python virtual environment name (required for LOCAL/EC2, optional for ECS) +- `vpc_config` (Optional[Dict]): VPC configuration for ECS only. Dict with keys: + - `subnets`: List[str] - Subnet IDs + - `security_groups`: List[str] - Security group IDs +- `cpu` (Optional[str]): CPU units for ECS tasks (e.g., "2048"). Ignored for LOCAL/EC2. +- `memory` (Optional[str]): Memory in MB for ECS tasks (e.g., "4096"). Ignored for LOCAL/EC2. +- `rft_role_name` (Optional[str]): IAM role name for RFT infrastructure. If not provided, uses default role or creates one. + +**Example:** +```python +from amzn_nova_forge import RFTMultiturnInfrastructure, CustomEnvironment, VFEnvId + +# Option 1: LOCAL with built-in environment +rft_infra = RFTMultiturnInfrastructure( + stack_name="my-rft-stack", + region="us-east-1", + python_venv_name="my_rft_venv", + vf_env_id=VFEnvId.WORDLE +) + +# Option 2: ECS with custom environment and VPC config +custom_env = CustomEnvironment( + env_id="my-custom-env", + output_dir="~/custom_envs/", + env_type="single_turn" +).create(overwrite=True) + +rft_infra = RFTMultiturnInfrastructure( + stack_name="my-rft-stack", + custom_env=custom_env, + infrastructure_arn="arn:aws:ecs:us-east-1:123456789012:cluster/my-cluster", + vpc_config={ + "subnets": ["subnet-12345", "subnet-67890"], + "security_groups": ["sg-12345"] + }, + cpu="4096", + memory="8192" +) + +# Deploy infrastructure +rft_infra.setup() + +# Start training environment +rft_infra.start_training_environment() + +# Use with NovaModelCustomizer +customizer = NovaModelCustomizer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=runtime, + data_s3_path="s3://bucket/data.jsonl" +) + +training_result = customizer.train( + job_name="rft-training", + rft_multiturn_infra=rft_infra +) +``` + +### CustomEnvironment + +Create custom reward functions for RFT multiturn training. + +**Constructor:** +```python +def __init__( + self, + env_id: str, + local_path: str = None, + output_dir: str = "~/custom_envs", + env_type: str = "single_turn" +) +``` + +**Methods:** +- `create(overwrite: bool = False)`: Create environment structure +- `package_and_upload(bucket: Optional[str] = None)`: Upload to S3 + +**Example:** +```python +custom_env = CustomEnvironment( + env_id="my-custom-env", + output_dir="~/custom_envs/", + env_type="single_turn" +).create(overwrite=True) + +custom_env.validate() +custom_env.package_and_upload() +print(f"Uploaded to: {custom_env.s3_uri}") +``` + +### RFT Multiturn Methods + +**Infrastructure Management:** +- `setup()`: Deploy CloudFormation stack (Lambda, SQS, DynamoDB) +- `start_training_environment(vf_env_args: Dict = None)`: Start training workers +- `start_evaluation_environment(vf_env_args: Dict = None)`: Start evaluation workers +- `kill_task(env_type: EnvType)`: Stop workers +- `cleanup(delete_stack: bool = False, cleanup_environment: bool = False)`: Clean up resources + - `delete_stack`: If True, delete CloudFormation stack + - `cleanup_environment`: If True, clean up environment resources: + - LOCAL/EC2: Delete virtual environment and starter kit directories + - ECS: Deregister task definitions + +**Monitoring:** +- `get_logs(env_type: EnvType, limit: int = 100, start_from_head: bool = False, log_stream_name: Optional[str] = None, tail: bool = False)`: View worker logs + - `tail`: If True, continuously stream logs in real-time (blocks until Ctrl+C) +- `check_all_queues()`: Check SQS queue status +- `flush_all_queues()`: Clear all queues + +**Configuration:** +- `get_configuration()`: Get infrastructure config +- `get_recipe_overrides()`: Get recipe overrides for training + +**Note:** RFT multiturn only supports SageMaker HyperPod (SMHP) platform and Nova 2.0 models (NOVA_LITE_2). + +--- diff --git a/docs/spec/runtime-managers.md b/docs/spec/runtime-managers.md new file mode 100644 index 0000000..117a6c8 --- /dev/null +++ b/docs/spec/runtime-managers.md @@ -0,0 +1,500 @@ +# Runtime Managers +Runtime managers handle the infrastructure for executing training and evaluation jobs, +leveraging the `JobConfig` dataclass to do so: +```python +@dataclass +class JobConfig: + job_name: str + image_uri: str + recipe_path: str + output_s3_path: Optional[str] = None + data_s3_path: Optional[str] = None + input_s3_data_type: Optional[str] = None + mlflow_tracking_uri: Optional[str] = None + mlflow_experiment_name: Optional[str] = None + mlflow_run_name: Optional[str] = None +``` +* The specific instance types that can be used with the runtime managers (SMTJ, SMHP) can be found in `docs/user-guides/instance_type_spec.md`. +* This file also defines which instance types can be used with a specific model and method. +* Bedrock is fully managed and does not require instance type configuration. + +### Shared RuntimeManager Methods + +The following methods are available on all `RuntimeManager` subclasses. + +#### Properties (shared) +- `rft_lambda` (Optional[str]): Lambda ARN, SageMaker hub-content ARN (SMTJServerless only), or local `.py` file path. Assigning a new value automatically updates `rft_lambda_arn` — if the value is a Lambda ARN or hub-content ARN it is resolved immediately; if it is a file path, `rft_lambda_arn` is cleared until `deploy_lambda()` is called. +- `rft_lambda_arn` (Optional[str]): Resolved Lambda ARN or hub-content ARN. Set immediately when `rft_lambda` is assigned an ARN (Lambda or hub-content), or populated by `deploy_lambda()` when `rft_lambda` is a file path. + +**Example:** +```python +# Set an ARN directly — rft_lambda_arn is updated immediately +runtime.rft_lambda = 'arn:aws:lambda:us-east-1:123456789012:function:my-reward-fn' +print(runtime.rft_lambda_arn) # 'arn:aws:lambda:us-east-1:123456789012:function:my-reward-fn' + +# Set a file path — rft_lambda_arn is cleared until deploy_lambda() is called +runtime.rft_lambda = 'reward.py' +print(runtime.rft_lambda_arn) # None +runtime.deploy_lambda(lambda_name='my-reward-fn') +print(runtime.rft_lambda_arn) +#'arn:aws:lambda:us-east-1:123456789012:function:my-reward-fn' +``` + +--- + +#### `deploy_lambda()` + +Packages a local Python file into a zip and creates or updates a Lambda function. The source file is read from `self.rft_lambda`, which must be set to a local `.py` file path before calling this method. + +**Signature:** +```python +def deploy_lambda( + self, + lambda_name: Optional[str] = None, + execution_role_arn: Optional[str] = None, +) -> str +``` + +**Parameters:** +- `lambda_name` (Optional[str]): Name for the Lambda function. Defaults to the source filename stem (underscores replaced with hyphens). +- `execution_role_arn` (Optional[str]): IAM role ARN for the Lambda. Falls back to the runtime manager's `execution_role` attribute if not provided. + +**Returns:** +- `str`: The deployed Lambda function ARN. Also sets `self.rft_lambda_arn` on the manager. + +**Raises:** +- `ValueError`: If `rft_lambda` is not set, is already an ARN (nothing to deploy), the source file is not found, or no execution role can be resolved. + +**Example:** +```python +runtime.rft_lambda = 'rft_training_reward.py' +lambda_arn = runtime.deploy_lambda(lambda_name='my-reward-fn') +# runtime.rft_lambda_arn is now set automatically +``` + +--- + +#### `validate_lambda()` + +Validates the RFT reward lambda with sample data from S3. Reads the lambda to validate from `self.rft_lambda` / `self.rft_lambda_arn`: +- If `rft_lambda` is an ARN (or `rft_lambda_arn` is set), invokes the deployed Lambda with samples from `data_s3_path`. +- If `rft_lambda` is a local `.py` path, validates by executing `lambda_handler` directly without deploying. + +**Signature:** +```python +def validate_lambda( + self, + data_s3_path: str, + validation_samples: int = 10, +) -> None +``` + +**Parameters:** +- `data_s3_path` (str): S3 path to the training dataset for pulling sample data. +- `validation_samples` (int): Number of samples to load from `data_s3_path` (default: 10). + +**Raises:** +- `ValueError`: If `rft_lambda` is not set, or if validation fails. + +**Example:** +```python +# Validate a local file without deploying +runtime.rft_lambda = 'rft_training_reward.py' +runtime.validate_lambda(data_s3_path='s3://bucket/data.jsonl') + +# Validate a deployed lambda +runtime.rft_lambda = 'arn:aws:lambda:us-east-1:123456789012:function:my-reward-fn' +runtime.validate_lambda(data_s3_path='s3://bucket/data.jsonl', validation_samples=20) +``` + +--- + +### SMTJRuntimeManager +Manages SageMaker Training Jobs. + +#### Constructor + +**Signature:** +```python +def __init__( + self, + instance_type: str, + instance_count: int, + execution_role: Optional[str] = None, + kms_key_id: Optional[str] = None, + encrypt_inter_container_traffic: bool = False, + subnets: Optional[list[str]] = None, + security_group_ids: Optional[list[str]] = None, + max_job_runtime: Optional[int] = 86400, + job_submit_poll_timeout: int = 30, + rft_lambda: Optional[str] = None, +) +``` + +**Parameters:** +- `instance_type` (str): EC2 instance type (e.g., "ml.p5.48xlarge", "ml.p4d.24xlarge") +- `instance_count` (int): Number of instances to use +- `execution_role` (Optional[str]): The execution role for the training job +- `kms_key_id` (Optional[str]): Optional KMS Key Id to use in S3 Bucket encryption, training jobs and deployments. +- `encrypt_inter_container_traffic` (bool): Boolean that determines whether to encrypt inter-container traffic. Default value is False. +- `subnets` (Optional[list[str]]): Optional list of strings representing subnets. Default value is None. +- `security_group_ids` (Optional[list[str]]): Optional list of strings representing security group IDs. Default value is None. +- `max_job_runtime` (Optional[int]): Max Job Runtime in seconds (default: 1 day) +- `job_submit_poll_timeout` (int): Maximum seconds to poll for the training job after submission (default: 30). Uses exponential backoff. +- `rft_lambda` (Optional[str]): Lambda ARN or local `.py` file path for RFT reward function. Can also be set or updated after construction. + +**Example:** +```python +from amzn_nova_forge.manager import * +infra = SMTJRuntimeManager( + instance_type="ml.p5.48xlarge", + instance_count=2 +) +``` +#### Properties +- `instance_type` (str): Returns the instance type +- `instance_count` (int): Returns the number of instances + +#### Methods + +##### `execute()` +Starts a SageMaker training job. + +**Signature:** +```python +def execute( + self, + job_config: JobConfig +) -> str +``` + +**Returns:** +- `str`: Training job name/ID + +##### `cleanup()` +Stops and cleans up a training job. + +**Signature:** +```python +def cleanup( + self, + job_name: str +) -> None +``` +--- +### SMHPRuntimeManager +Manages SageMaker HyperPod jobs. + +#### Constructor + +**Signature:** +```python +def __init__( + self, + instance_type: str, + instance_count: int, + cluster_name: str, + namespace: str, + kms_key_id: Optional[str] = None, + rft_lambda: Optional[str] = None, +) +``` + +**Parameters:** +- `instance_type` (str): EC2 instance type +- `instance_count` (int): Number of instances +- `cluster_name` (str): HyperPod cluster name +- `namespace` (str): Kubernetes namespace +- `kms_key_id` (Optional[str]): Optional KMS Key Id to use in S3 Bucket encryption +- `rft_lambda` (Optional[str]): Lambda ARN or local `.py` file path for RFT reward function. Can also be set or updated after construction. + +**Example:** +```python +from amzn_nova_forge.manager import * +infra = SMHPRuntimeManager( + instance_type="ml.p5.48xlarge", + instance_count=4, + cluster_name="my-hyperpod-cluster", + namespace="default" +) +``` +#### Properties +- `instance_type` (str): Returns the instance type +- `instance_count` (int): Returns the number of instances + +#### Methods + +##### `execute()` + +Starts a SageMaker HyperPod job. +**Signature:** +```python +def execute( + self, + job_config=JobConfig +) -> str +``` +**Returns:** +- `str`: HyperPod job ID + +##### `cleanup()` +Cancels and cleans up a HyperPod job. + +**Signature:** +```python +def cleanup( + self, + job_name: str +) -> None +``` + +##### `scale_cluster()` +Scale a HyperPod cluster instance group up or down. +The scaling operation is asynchronous - the cluster status will change to 'Updating' while scaling, and 'InService' when ready. + +**Signature:** +```python +def scale_cluster( + self, + instance_group_name: str, + target_instance_count: int, +) -> Dict[str, Any] +``` + +**Parameters:** +- `instance_group_name` (str): Name of the instance group to scale (e.g., 'worker-group') +- `target_instance_count` (int): Desired number of instances for the group (must be non-negative) + +**Returns:** +- `Dict[str, Any]`: Response containing: + - `ClusterArn` (str): ARN of the updated cluster + - `InstanceGroupName` (str): Name of the scaled instance group + - `InstanceType` (str): Instance type being scaled + - `PreviousCount` (int): Current instance count before scaling + - `TargetCount` (int): Target instance count after scaling + +**Raises:** +- `ValueError`: If target_instance_count is negative or instance group name is invalid +- `ClientError`: If scaling fails due to insufficient quota, capacity or other cluster issues. + +**Example:** +```python +from amzn_nova_forge.manager import * + +# Create a runtime manager for your cluster +manager = SMHPRuntimeManager( + instance_type="ml.p4d.24xlarge", + instance_count=4, + cluster_name="my-hyperpod-cluster", + namespace="default" +) + +# Scale up the worker group from 4 to 8 instances +result = manager.scale_cluster( + instance_group_name="worker-group", + target_instance_count=8 +) + +# Scale down to 2 instances +result = manager.scale_cluster( + instance_group_name="worker-group", + target_instance_count=2 +) +``` +**Notes:** +- This method only works with Restricted Instance Groups (RIGs) in HyperPod clusters. The cluster must be in 'InService' state before scaling can be initiated. +- This method can only scale up a SMHP cluster when there is sufficient Service Quota available. +You will need to request a quota increase **before** scaling up a RIG in your HyperPod cluster. +You can learn more [here](https://docs.aws.amazon.com/servicequotas/latest/userguide/request-quota-increase.html). + - Specifically, you will need to request a service quota increase for "INSTANCE_TYPE for cluster usage". + +##### `get_instance_groups()` +Gets the RIGs associated with the current cluster defined in the SMHPRuntimeManager. +Prints the values to the terminal and returns it as a list of dictionary entries. + +**Signature:** +```python +def get_instance_groups( + self +) -> List[Dict[str, Any]] +``` + +**Returns:** +- `List[Dict[str, Any]]`: Response containing: + - InstanceGroupName: Name of the instance group + - InstanceType: EC2 instance type (e.g., 'ml.p5.48xlarge') + - CurrentCount: Current number of instances in the group + +**Raises:** +- `ClientError`: If unable to describe the cluster + +**Example:** +```python +from amzn_nova_forge.manager import * + +# Create a runtime manager for your cluster +manager = SMHPRuntimeManager( + instance_type="ml.p4d.24xlarge", + instance_count=4, + cluster_name="my-hyperpod-cluster", + namespace="default" +) + +# Get the instance groups available on the current cluster. +instance_groups = manager.get_instance_groups() +``` + +--- +### BedrockRuntimeManager +Manages Amazon Bedrock model customization jobs. + +#### Constructor + +**Signature:** +```python +def __init__( + self, + execution_role: str, + base_model_identifier: Optional[str] = None, + kms_key_id: Optional[str] = None, + rft_lambda: Optional[str] = None, +) +``` + +**Parameters:** +- `execution_role` (str): IAM role ARN for Bedrock job execution +- `base_model_identifier` (Optional[str]): Base model ARN (e.g., "arn:aws:bedrock:us-east-1::foundation-model/amazon.nova-2-lite-v1:0:256k") +- `kms_key_id` (Optional[str]): Optional KMS Key Id for encryption +- `rft_lambda` (Optional[str]): Lambda ARN or local `.py` file path for RFT reward function. Can also be set or updated after construction. + +**Example:** +```python +from amzn_nova_forge.manager import * +infra = BedrockRuntimeManager( + execution_role="arn:aws:iam::123456789012:role/BedrockRole", + base_model_identifier="arn:aws:bedrock:us-east-1::custom-model/amazon.nova-2-lite-v1:0:256k:abcdefghijk" # optional: your custom model ARN for iterative training +) +``` + +#### Methods + +##### `execute()` + +Starts a Bedrock model customization job. +**Signature:** +```python +def execute( + self, + job_config: JobConfig +) -> str +``` +**Returns:** +- `str`: Bedrock job ARN + +##### `cleanup()` +Stops a Bedrock customization job. + +**Signature:** +```python +def cleanup( + self, + job_name: str +) -> None +``` +--- +### SMTJServerlessRuntimeManager +Manages SageMaker Training Jobs. + +> **Note:** `AWS_DEFAULT_REGION` must be set when using SageMaker Serverless training. +> The SageMaker SDK's `DataSet` API requires a region to connect to the SageMaker backend. +> Set it before running your script: +> ```bash +> export AWS_DEFAULT_REGION= +> ``` + +#### Constructor + +**Signature:** +```python +def __init__( + self, + model_package_group_name: str, + execution_role: Optional[str] = None, + kms_key_id: Optional[str] = None, + encrypt_inter_container_traffic: bool = False, + subnets: Optional[list[str]] = None, + security_group_ids: Optional[list[str]] = None, + max_job_runtime: Optional[int] = 86400, + rft_lambda: Optional[str] = None, + evaluator_name: Optional[str] = None, +) +``` + +**Parameters:** +- `model_package_group_name` (str): Model package group name to use with SageMaker Model registry (required for SMTJ Serverless) +- `execution_role` (Optional[str]): The execution role for the training job +- `kms_key_id` (Optional[str]): Optional KMS Key Id to use in S3 Bucket encryption, training jobs and deployments. +- `encrypt_inter_container_traffic` (bool): Boolean that determines whether to encrypt inter-container traffic. Default value is False. +- `subnets` (Optional[list[str]]): Optional list of strings representing subnets. Default value is None. +- `security_group_ids` (Optional[list[str]]): Optional list of strings representing security group IDs. Default value is None. +- `max_job_runtime` (Optional[int]): Max Job Runtime in seconds (default: 1 day) +- `rft_lambda` (Optional[str]): Lambda ARN, SageMaker hub-content ARN, or local `.py` file path for the RFT reward function. + - **Lambda ARN**: Automatically registered as a hub-content `JsonDoc` evaluator during `train()`. The hub-content ARN is passed as `EvaluatorArn` in `ServerlessJobConfig`. + - **Hub-content ARN**: Passed directly as `EvaluatorArn` — no registration needed. + - **Local `.py` file**: Call `deploy_lambda()` first to deploy and get a Lambda ARN. +- `evaluator_name` (Optional[str]): Optional human-readable name for the hub-content evaluator entry when auto-registering a Lambda ARN. Defaults to the Lambda function name. + +**Example:** +```python +from amzn_nova_forge.manager import * + +# With a Lambda ARN (auto-registered as hub-content during train()) +infra = SMTJServerlessRuntimeManager( + model_package_group_name="nova-rft-serverless", + rft_lambda="arn:aws:lambda:us-east-1:123456789012:function:my-reward-fn", +) + +# With a hub-content ARN (passed directly as EvaluatorArn) +infra = SMTJServerlessRuntimeManager( + model_package_group_name="nova-rft-serverless", + rft_lambda="arn:aws:sagemaker:us-east-1:123456789012:hub-content/my-hub/JsonDoc/my-evaluator/0.0.1", +) + +# With a local .py file (deploy_lambda() required before train()) +infra = SMTJServerlessRuntimeManager( + model_package_group_name="nova-rft-serverless", + rft_lambda="my_reward.py", +) +infra.deploy_lambda(lambda_name="my-reward-fn") +``` +#### Properties +- `model_package_group_name` (str): Model Package Group name + +#### Methods + +##### `execute()` +Starts a SageMaker training job. + +**Signature:** +```python +def execute( + self, + job_config: JobConfig +) -> str +``` + +**Returns:** +- `str`: Training job name/ID + +##### `cleanup()` +Stops and cleans up a training job. + +**Signature:** +```python +def cleanup( + self, + job_name: str +) -> None +``` +--- + diff --git a/docs/spec/service-classes.md b/docs/spec/service-classes.md new file mode 100644 index 0000000..5929827 --- /dev/null +++ b/docs/spec/service-classes.md @@ -0,0 +1,863 @@ +# Service Classes + +The modular service classes are the recommended API for Nova model customization. Each class handles a single concern — training, evaluation, deployment, or inference — and can be used independently. + +All service classes accept an optional `ForgeConfig` dataclass for shared configuration (KMS keys, output paths, caching, etc.). + +### ForgeConfig + +Shared configuration dataclass for all service classes. + +**Signature:** +```python +@dataclass +class ForgeConfig: + kms_key_id: Optional[str] = None + output_s3_path: Optional[str] = None + generated_recipe_dir: Optional[str] = None + validation_config: Optional[ValidationConfig] = None + image_uri: Optional[str] = None + mlflow_monitor: Optional[MLflowMonitor] = None + enable_job_caching: bool = False + job_cache_dir: str = "~/.nova-forge/cache" + job_caching_config: Optional[JobCachingConfig] = None +``` + +**Parameters:** +- `kms_key_id` (Optional[str]): KMS key ID for S3 encryption +- `output_s3_path` (Optional[str]): S3 path for output artifacts. Auto-generated if not provided +- `generated_recipe_dir` (Optional[str]): Local path to save generated recipe files +- `validation_config` (Optional[ValidationConfig]): Controls pre-flight validation. Fields: `iam` (bool), `infra` (bool), `recipe` (bool) — all default to True +- `image_uri` (Optional[str]): Custom container image URI override +- `mlflow_monitor` (Optional[MLflowMonitor]): MLflow monitoring configuration (SageMaker only) +- `enable_job_caching` (bool): Enable caching of completed job results for reuse. Default: False +- `job_cache_dir` (str): Directory for cached job results. Default: `~/.nova-forge/cache` +- `job_caching_config` (Optional[JobCachingConfig]): Advanced caching configuration. Fields: `include_core` (bool), `include_recipe` (bool), `include_infra` (bool), `include_params` (List[str]), `exclude_params` (List[str]), `allowed_statuses` (List[JobStatus]) + +**Example:** +```python +from amzn_nova_forge.core import ForgeConfig, ValidationConfig +from amzn_nova_forge.monitor import MLflowMonitor + +config = ForgeConfig( + output_s3_path="s3://my-bucket/output", + kms_key_id="my-kms-key-id", + validation_config=ValidationConfig(iam=True, infra=True, recipe=True), + mlflow_monitor=MLflowMonitor( + tracking_uri="arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", + experiment_name="nova-customization" + ), + enable_job_caching=True +) +``` +--- + +### ForgeTrainer + +Handles training job configuration and execution for Nova models. + +#### Constructor + +**Signature:** +```python +def __init__( + self, + model: Model, + method: TrainingMethod, + infra: RuntimeManager, + training_data_s3_path: Optional[str] = None, + model_s3_path: Optional[str] = None, + data_mixing_enabled: bool = False, + holdout_data_s3_path: Optional[str] = None, + val_check_interval: Optional[int] = None, + config: Optional[ForgeConfig] = None, + region: Optional[str] = None, + is_multimodal: Optional[bool] = None, + hub_content_version: Optional[str] = None, + enable_batch_sample_tracing: bool = False, +) +``` + +**Parameters:** +- `model` (Model): The Nova model to train (e.g., `Model.NOVA_MICRO`, `Model.NOVA_LITE_2`) +- `method` (TrainingMethod): The fine-tuning method (e.g., `TrainingMethod.SFT_LORA`, `TrainingMethod.RFT`) +- `infra` (RuntimeManager): Runtime infrastructure manager (e.g., `SMTJRuntimeManager`, `SMHPRuntimeManager`, `BedrockRuntimeManager`) +- `training_data_s3_path` (Optional[str]): S3 path to the training dataset +- `model_s3_path` (Optional[str]): S3 path for the base or previously trained model +- `data_mixing_enabled` (bool): Enable data mixing for CPT and SFT training on SMHP, and SFT text-only on Nova Lite 2 on SMTJServerless. Default: False +- `holdout_data_s3_path` (Optional[str]): S3 path to holdout/validation data (optional, used for CPT and SFT on SMTJ/SMTJServerless/SMHP, or any method on Bedrock) +- `val_check_interval` (Optional[int]): How often (in training steps) to run validation. Defaults to 2500 if omitted. Only used when `holdout_data_s3_path` is provided. +- `config` (Optional[ForgeConfig]): Shared configuration. If not provided, defaults are used +- `region` (Optional[str]): AWS region. Auto-detected if not provided +- `is_multimodal` (Optional[bool]): Explicitly set multimodal mode when `data_mixing_enabled=True`. If None, auto-detects from data +- `hub_content_version` (Optional[str]): Version of the hub content to retrieve from SageMaker Hub. If None, uses the latest version +- `enable_batch_sample_tracing` (bool): Activate per-step batch hashing during training, which enables `trace_batch()` post-training. Supported platform/method combinations are validated at construction time. Default: False + +**Raises:** +- `ValueError`: If `enable_batch_sample_tracing=True` is used with an unsupported platform or training method + +**Example:** +```python +from amzn_nova_forge.trainer import ForgeTrainer +from amzn_nova_forge.core import ForgeConfig +from amzn_nova_forge.manager import SMTJRuntimeManager +from amzn_nova_forge.model.model_enums import Model, TrainingMethod + +infra = SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=2) + +trainer = ForgeTrainer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + training_data_s3_path="s3://my-bucket/training-data/data.jsonl", + config=ForgeConfig(output_s3_path="s3://my-bucket/output") +) +``` +--- + +#### Methods + +##### `train()` +Generates the recipe YAML, configures the runtime, and launches a training job. + +**Signature:** +```python +def train( + self, + job_name: str, + recipe_path: Optional[str] = None, + overrides: Optional[Dict[str, Any]] = None, + rft_lambda_arn: Optional[str] = None, + dry_run: bool = False, + rft_multiturn_infra=None, +) -> Optional[TrainingResult] +``` + +**Parameters:** +- `job_name` (str): User-defined name for the training job +- `recipe_path` (Optional[str]): Path for a YAML recipe file (S3 or local) +- `overrides` (Optional[Dict[str, Any]]): Dictionary of configuration overrides (e.g., `max_epochs`, `lr`, `warmup_steps`, `global_batch_size`) +- `rft_lambda_arn` (Optional[str]): Rewards Lambda ARN (only for RFT methods). Takes priority over `rft_lambda_arn` on the RuntimeManager +- `dry_run` (bool): If True, performs validation only without starting a job. Default: False +- `rft_multiturn_infra`: Optional RFTMultiturnInfrastructure for RFT multiturn training + +**Returns:** +- `TrainingResult`: Metadata object containing `job_id`, `method`, `started_time`, `model_artifacts`, and `model_type`. Returns `None` if `dry_run=True` + +**Raises:** +- `Exception`: If job execution fails +- `ValueError`: If training method is not supported + +**Example:** +```python +result = trainer.train( + job_name="my-training-job", + overrides={ + "max_epochs": 10, + "lr": 5e-6, + "warmup_steps": 20, + "global_batch_size": 128 + } +) +print(f"Training job started: {result.job_id}") +print(f"Checkpoint path: {result.model_artifacts.checkpoint_s3_path}") +``` +--- + +##### `get_logs()` +Retrieves and displays CloudWatch logs for a training job. + +**Signature:** +```python +def get_logs( + self, + job_result=None, + job_id=None, + started_time=None, + limit=None, + start_from_head: bool = False, + end_time=None, +) -> None +``` + +**Parameters:** +- `job_result` (Optional[TrainingResult]): Job result to retrieve logs for. If not provided, uses `job_id` +- `job_id` (Optional[str]): Job identifier. Used if `job_result` is not provided +- `started_time` (Optional[datetime]): Job start time to filter logs +- `limit` (Optional[int]): Maximum number of log lines to retrieve +- `start_from_head` (bool): If True, start from the beginning of logs. Default: False +- `end_time` (Optional[int]): End time in epoch milliseconds for searching a log time range + +**Returns:** +- None (prints logs to console) + +**Example:** +```python +trainer.get_logs(job_result=result, limit=100, start_from_head=True) +``` +--- + +##### `trace_batch()` +Extracts the lines from your input training data that were used in a specific training step's batch. Useful for diagnosing gradient spikes or training anomalies — given a step number, it matches the container's batch hash logs against your source data and writes the matched lines to an output file. + +The training job must have been launched with `enable_batch_sample_tracing=True` so that batch hash logs are written during training. Supported platform/method combinations are validated at `ForgeTrainer` construction time. If you create a new `ForgeTrainer` instance to trace a previously-launched job, the flag is not strictly required on the new instance — but a warning will be emitted. + +**Signature:** +```python +def trace_batch( + self, + training_result: TrainingResult, + step: int, + output_path: str | None = None, + cache_dir: str = "~/.nova-forge/batch_trace_cache", +) -> Path | None +``` + +**Parameters:** +- `training_result` (TrainingResult): Result from a completed training job. The method extracts `training_result.model_artifacts.output_s3_path` and `training_result.job_id` to locate the batch hash logs at `{output_s3_path}/{job_id}/batch_tracing/`. +- `step` (int): Training step number to investigate (must match the step numbers in the container's batch hash logs). +- `output_path` (Optional[str]): Path for the output file containing matched lines. Default: `step__samples.jsonl` in the current working directory. +- `cache_dir` (str): Directory for caching downloaded files and fingerprint indices. Default: `~/.nova-forge/batch_trace_cache`. The cache stores downloaded S3 files (training data, hash logs) and a CSV fingerprint index. For large datasets the cache may grow to match the source data size. The cache is keyed by S3 path — if you replace the file at an existing S3 URI, delete the cache directory to avoid stale matches. + +**Returns:** +- `Path | None`: Path to the output JSONL file containing the matched lines (verbatim copies from your source data, sorted by line number). Returns `None` if either the step had no logged batch data (step out of range or job still running) or the step's batch contained no samples from your file. + +**Raises:** +- `ValueError`: If `training_data_s3_path` or `training_result.model_artifacts.output_s3_path` is not available +- `BatchTraceError`: If batch tracing encounters an unrecoverable error (e.g., missing log files, AWS auth failure) + Import: `from amzn_nova_forge.trainer.utils import BatchTraceError` + +**Example:** +```python +trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.CPT, + infra=infra, + training_data_s3_path="s3://my-bucket/data.jsonl", + enable_batch_sample_tracing=True, +) + +result = trainer.train(job_name="my-cpt-job") + +# After job completes, investigate step 42. +# Output writes to ./step_42_samples.jsonl by default. +matched_file = trainer.trace_batch(result, step=42) +if matched_file: + print(f"Matched lines written to: {matched_file}") + +# Explicit output path: +matched_file = trainer.trace_batch(result, step=42, output_path="/tmp/flagged.jsonl") +``` +--- + +### ForgeEvaluator + +Handles evaluation job configuration and execution for Nova models. + +#### Constructor + +**Signature:** +```python +def __init__( + self, + model: Model, + infra: RuntimeManager, + data_s3_path: Optional[str] = None, + config: Optional[ForgeConfig] = None, + region: Optional[str] = None, + hub_content_version: Optional[str] = None, +) +``` + +**Parameters:** +- `model` (Model): The Nova model to evaluate +- `infra` (RuntimeManager): Runtime infrastructure manager +- `data_s3_path` (Optional[str]): S3 path to evaluation data (required for BYOD evaluation tasks) +- `config` (Optional[ForgeConfig]): Shared configuration +- `region` (Optional[str]): AWS region. Auto-detected if not provided +- `hub_content_version` (Optional[str]): Version of the hub content to retrieve from SageMaker Hub. If None, uses the latest version + +**Example:** +```python +from amzn_nova_forge.evaluator import ForgeEvaluator +from amzn_nova_forge.manager import SMTJRuntimeManager +from amzn_nova_forge.model.model_enums import Model + +infra = SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=2) + +evaluator = ForgeEvaluator( + model=Model.NOVA_MICRO, + infra=infra, + data_s3_path="s3://my-bucket/eval-data/data.jsonl" +) +``` +--- + +#### Methods + +##### `evaluate()` +Generates the recipe YAML, configures the runtime, and launches an evaluation job. + +**Signature:** +```python +def evaluate( + self, + job_name: str, + eval_task: EvaluationTask, + model_path: Optional[str] = None, + task_config: Optional[EvalTaskConfig] = None, + recipe_path: Optional[str] = None, + overrides: Optional[Dict[str, Any]] = None, + dry_run: bool = False, + job_result: Optional[TrainingResult] = None, + rft_multiturn_infra=None, +) -> Optional[EvaluationResult] +``` + +**Parameters:** +- `job_name` (str): User-defined name for the evaluation job +- `eval_task` (EvaluationTask): The evaluation task (e.g., `EvaluationTask.MMLU`) +- `model_path` (Optional[str]): S3 path to the model to evaluate. If not provided, extracted from `job_result` +- `task_config` (Optional[EvalTaskConfig]): Task-specific configuration. Fields: `subtask`, `processor`, `rl_env`, `override_data_s3_path` +- `recipe_path` (Optional[str]): Path for a YAML recipe file (S3 or local) +- `overrides` (Optional[Dict[str, Any]]): Inference configuration overrides (e.g., `max_new_tokens`, `temperature`, `top_p`) +- `dry_run` (bool): If True, performs validation only. Default: False +- `job_result` (Optional[TrainingResult]): Training result to extract checkpoint path from +- `rft_multiturn_infra`: Optional RFTMultiturnInfrastructure for RFT evaluation + +**Returns:** +- `EvaluationResult`: Metadata object containing `job_id`, `started_time`, `eval_output_path`, and `eval_task`. Returns `None` if `dry_run=True` + +**Example:** +```python +from amzn_nova_forge.core import EvaluationTask + +eval_result = evaluator.evaluate( + job_name="my-eval-job", + eval_task=EvaluationTask.MMLU, + model_path="s3://my-bucket/checkpoints/my-model", + overrides={ + "max_new_tokens": 2048, + "temperature": 0, + "top_p": 1.0 + } +) +print(f"Evaluation job started: {eval_result.job_id}") + +# Chain from training result +eval_result = evaluator.evaluate( + job_name="my-eval-job", + eval_task=EvaluationTask.MMLU, + job_result=training_result +) +``` +--- + +##### `get_logs()` +Retrieves and displays CloudWatch logs for an evaluation job. + +**Signature:** +```python +def get_logs( + self, + job_result=None, + job_id=None, + started_time=None, + limit=None, + start_from_head: bool = False, + end_time=None, +) -> None +``` + +**Parameters:** +- `job_result` (Optional[EvaluationResult]): Job result to retrieve logs for +- `job_id` (Optional[str]): Job identifier +- `started_time` (Optional[datetime]): Job start time to filter logs +- `limit` (Optional[int]): Maximum number of log lines +- `start_from_head` (bool): If True, start from the beginning of logs. Default: False +- `end_time` (Optional[int]): End time in epoch milliseconds for searching a log time range + +**Returns:** +- None (prints logs to console) + +**Example:** +```python +evaluator.get_logs(job_result=eval_result, limit=50) +``` +--- + +### ForgeDeployer + +Handles model deployment to Amazon Bedrock and SageMaker endpoints. + +#### Constructor + +**Signature:** +```python +def __init__( + self, + region: str, + model: Model, + deployment_mode: DeploymentMode = DeploymentMode.FAIL_IF_EXISTS, + config: Optional[ForgeConfig] = None, + method: Optional[TrainingMethod] = None, +) +``` + +**Parameters:** +- `region` (str): AWS region for deployment +- `model` (Model): The Nova model being deployed +- `deployment_mode` (DeploymentMode): Behavior when endpoint already exists. Default: `FAIL_IF_EXISTS` +- `config` (Optional[ForgeConfig]): Shared configuration +- `method` (Optional[TrainingMethod]): Training method used (needed for SageMaker deployment image selection) + +**Example:** +```python +from amzn_nova_forge.deployer import ForgeDeployer +from amzn_nova_forge.model.model_enums import Model, DeploymentMode + +deployer = ForgeDeployer( + region="us-east-1", + model=Model.NOVA_MICRO, + deployment_mode=DeploymentMode.FAIL_IF_EXISTS +) +``` +--- + +#### Methods + +##### `deploy()` +Creates a custom model and deploys it to Amazon Bedrock or SageMaker in a single step. + +**Signature:** +```python +def deploy( + self, + model_artifact_path: str, + deploy_platform: DeployPlatform = DeployPlatform.BEDROCK_OD, + endpoint_name: Optional[str] = None, + unit_count: int = 1, + execution_role_name: Optional[str] = None, + sagemaker_instance_type: str = "ml.p5.48xlarge", + sagemaker_environment: Optional[SageMakerEndpointEnvironment] = None, + skip_model_reuse: bool = False, +) -> DeploymentResult +``` + +**Parameters:** +- `model_artifact_path` (str): S3 path to the trained model checkpoint +- `deploy_platform` (DeployPlatform): Platform to deploy to (`BEDROCK_OD`, `BEDROCK_PT`, or `SAGEMAKER`). Default: `BEDROCK_OD` +- `endpoint_name` (Optional[str]): Name of the endpoint (auto-generated if not provided) +- `unit_count` (int): Number of PT units (Bedrock PT) or instances (SageMaker). Default: 1 +- `execution_role_name` (Optional[str]): IAM role name. If omitted, the SDK creates a default role +- `sagemaker_instance_type` (str): Instance type for SageMaker deployment. Default: `"ml.p5.48xlarge"` +- `sagemaker_environment` (Optional[SageMakerEndpointEnvironment]): SageMaker endpoint environment config. Fields: + - `CONTEXT_LENGTH` (int, default: 4000), `MAX_CONCURRENCY` (int, default: 1) + - Optional generation defaults: `DEFAULT_TEMPERATURE` (0–2), `DEFAULT_TOP_P` (1e-10–1), `DEFAULT_TOP_K` (-1 to disable, or ≥1), `DEFAULT_MAX_NEW_TOKENS` (≥1), `DEFAULT_LOGPROBS` (1–20) + - Optional speculative decoding: `SPECULATIVE_DECODING_METHOD` (`"eagle3"` or `"suffix"`), `DISABLE_SPECULATIVE_DECODING` (`"true"` or `"false"`), `NUM_SPECULATIVE_TOKENS` (1–10), `SUFFIX_DECODING_MAX_TREE_DEPTH`, `SUFFIX_DECODING_MAX_CACHED_REQUESTS`, `SUFFIX_DECODING_MAX_SPEC_FACTOR`, `SUFFIX_DECODING_MIN_TOKEN_PROB` + - Optional memory/quantization: `KV_CACHE_DTYPE` (`"fp8"`), `QUANTIZATION_DTYPE` (`"fp8"`) +- `skip_model_reuse` (bool): If True, always create a new model. Default: False + +**Returns:** +- `DeploymentResult`: Contains `endpoint` (EndpointInfo), `platform`, `endpoint_name`, `uri`, `model_artifact_path`, and `created_at` + +**Raises:** +- `Exception`: When unable to deploy the model +- `ValueError`: If platform is not supported + +**Example:** +```python +deployment = deployer.deploy( + model_artifact_path="s3://escrow-bucket/my-model-artifacts/", + deploy_platform=DeployPlatform.BEDROCK_OD, + endpoint_name="my-custom-nova-model" +) +print(f"Model deployed: {deployment.endpoint.uri}") +``` +--- + +##### `create_custom_model()` +Creates a Bedrock custom model from S3 artifacts without deploying to an endpoint. + +**Signature:** +```python +def create_custom_model( + self, + model_artifact_path: str, + endpoint_name: Optional[str] = None, + execution_role_name: Optional[str] = None, + tags: Optional[List[Dict[str, str]]] = None, + skip_model_reuse: bool = False, +) -> ModelDeployResult +``` + +**Parameters:** +- `model_artifact_path` (str): S3 path to trained model checkpoint +- `endpoint_name` (Optional[str]): Optional name prefix for the model name +- `execution_role_name` (Optional[str]): IAM role name for Bedrock +- `tags` (Optional[List[Dict[str, str]]]): Optional list of `{"key": str, "value": str}` dicts for tracking +- `skip_model_reuse` (bool): If True, always create a new model. Default: False + +**Returns:** +- `ModelDeployResult`: Contains `model_arn`, `model_name`, `escrow_uri`, and `created_at` + +**Example:** +```python +publish_result = deployer.create_custom_model( + model_artifact_path="s3://escrow-bucket/my-model-artifacts/" +) +print(f"Model ARN: {publish_result.model_arn}") +publish_result.dump(file_path="./results/") +``` +--- + +##### `deploy_to_bedrock()` +Deploys a published Bedrock custom model to an endpoint. + +**Signature:** +```python +def deploy_to_bedrock( + self, + model_deploy_result: Optional[ModelDeployResult] = None, + model_arn: Optional[str] = None, + deploy_platform: DeployPlatform = DeployPlatform.BEDROCK_OD, + pt_units: Optional[int] = None, + endpoint_name: Optional[str] = None, +) -> DeploymentResult +``` + +**Parameters:** +- `model_deploy_result` (Optional[ModelDeployResult]): Result from `create_custom_model()`. Cannot be combined with `model_arn` +- `model_arn` (Optional[str]): Direct model ARN. Cannot be combined with `model_deploy_result` +- `deploy_platform` (DeployPlatform): `BEDROCK_OD` (default) or `BEDROCK_PT` +- `pt_units` (Optional[int]): Number of PT units (required for `BEDROCK_PT`) +- `endpoint_name` (Optional[str]): Endpoint name (auto-generated if not provided) + +**Returns:** +- `DeploymentResult`: Contains `endpoint`, `created_at`, and `model_publish` + +**Raises:** +- `ValueError`: When both `model_deploy_result` and `model_arn` are provided, or when no model ARN is available +- `RuntimeError`: When deployment creation fails + +**Example:** +```python +# Two-step deploy: create model, then deploy +publish_result = deployer.create_custom_model( + model_artifact_path="s3://escrow-bucket/my-model-artifacts/" +) +deployment = deployer.deploy_to_bedrock( + model_deploy_result=publish_result, + endpoint_name="my-endpoint" +) + +# Or deploy from an existing model ARN +deployment = deployer.deploy_to_bedrock( + model_arn="arn:aws:bedrock:us-east-1:123456789012:custom-model/my-model" +) +``` +--- + +##### `find_published_model()` +Finds an existing published model by platform and escrow path to enable model reuse. + +**Signature:** +```python +def find_published_model( + self, + platform: str, + escrow_path: str, + skip_model_reuse: bool = False, +) -> Optional[str] +``` + +**Parameters:** +- `platform` (str): Target platform (`"bedrock"` or `"sagemaker"`) +- `escrow_path` (str): S3 path of the model artifacts +- `skip_model_reuse` (bool): If True, always returns None (skips lookup). Default: False + +**Returns:** +- `Optional[str]`: Existing model ARN if found, otherwise None + +**Example:** +```python +existing_arn = deployer.find_published_model( + platform="bedrock", + escrow_path="s3://escrow-bucket/my-model-artifacts/" +) +if existing_arn: + print(f"Reusing existing model: {existing_arn}") +``` +--- + +##### `get_status()` +Gets the deployment status for a DeploymentResult. + +**Signature:** +```python +def get_status(self, result: DeploymentResult) -> JobStatus +``` + +**Parameters:** +- `result` (DeploymentResult): The deployment result to check + +**Returns:** +- `JobStatus`: Current status (`IN_PROGRESS`, `COMPLETED`, or `FAILED`) + +--- + +##### `get_status_by_arn()` +Gets the deployment status by endpoint ARN and platform. + +**Signature:** +```python +def get_status_by_arn( + self, + endpoint_arn: str, + platform: DeployPlatform, +) -> Optional[JobStatus] +``` + +**Parameters:** +- `endpoint_arn` (str): The endpoint ARN to check +- `platform` (DeployPlatform): The deployment platform + +**Returns:** +- `Optional[JobStatus]`: Current status, or None if status cannot be determined + +--- + +##### `get_logs()` +Retrieves and displays logs for a deployment. + +**Signature:** +```python +def get_logs( + self, + job_result=None, + endpoint_arn=None, + platform=None, +) -> None +``` + +**Parameters:** +- `job_result` (Optional[DeploymentResult]): Deployment result to retrieve logs for +- `endpoint_arn` (Optional[str]): Endpoint ARN (used if `job_result` is not provided) +- `platform` (Optional[DeployPlatform]): Deployment platform (used with `endpoint_arn`) + +**Returns:** +- None (prints logs to console) + +--- + +### ForgeInference + +Handles single and batch inference on trained Nova models. + +#### Constructor + +**Signature:** +```python +def __init__( + self, + region: Optional[str] = None, + model: Optional[Model] = None, + infra: Optional[RuntimeManager] = None, + config: Optional[ForgeConfig] = None, + method: Optional[TrainingMethod] = None, + hub_content_version: Optional[str] = None, +) +``` + +**Parameters:** +- `region` (Optional[str]): AWS region. Auto-detected if not provided +- `model` (Optional[Model]): The Nova model (required for batch inference) +- `infra` (Optional[RuntimeManager]): Runtime infrastructure manager (required for batch inference) +- `config` (Optional[ForgeConfig]): Shared configuration +- `method` (Optional[TrainingMethod]): Training method (used for batch inference recipe generation) +- `hub_content_version` (Optional[str]): Version of the hub content to retrieve from SageMaker Hub. If None, uses the latest version + +**Example:** +```python +from amzn_nova_forge.inference import ForgeInference + +# For single inference (minimal setup) +inference = ForgeInference(region="us-east-1") + +# For batch inference +inference = ForgeInference( + region="us-east-1", + model=Model.NOVA_MICRO, + infra=SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=1), + method=TrainingMethod.SFT_LORA +) +``` +--- + +#### Methods + +##### `invoke()` +Invokes a single inference on a deployed model endpoint. + +**Signature:** +```python +def invoke( + self, + endpoint_arn: str, + request_body: Dict[str, Any], +) -> Any +``` + +**Parameters:** +- `endpoint_arn` (str): Endpoint ARN to invoke +- `request_body` (Dict[str, Any]): Inference request body + +**Returns:** +- `Any`: Inference response + +**Example:** +```python +response = inference.invoke( + endpoint_arn="arn:aws:bedrock:us-east-1:123456789012:endpoint/my-endpoint", + request_body={ + "messages": [{"role": "user", "content": "Hello! How are you?"}], + "max_tokens": 100, + "stream": False + } +) +``` +--- + +##### `invoke_batch()` +Launches a batch inference job on a trained model. + +**Signature:** +```python +def invoke_batch( + self, + job_name: str, + input_path: str, + output_s3_path: str, + model_path: Optional[str] = None, + recipe_path: Optional[str] = None, + overrides: Optional[Dict[str, Any]] = None, + dry_run: bool = False, + job_result: Optional[TrainingResult] = None, +) -> Optional[InferenceResult] +``` + +**Parameters:** +- `job_name` (str): Name for the batch inference job +- `input_path` (str): S3 path to input data +- `output_s3_path` (str): S3 path for inference outputs +- `model_path` (Optional[str]): S3 path to the model checkpoint +- `recipe_path` (Optional[str]): Path for a YAML recipe file +- `overrides` (Optional[Dict[str, Any]]): Inference configuration overrides (e.g., `max_new_tokens`, `temperature`, `top_p`) +- `dry_run` (bool): If True, performs validation only. Default: False +- `job_result` (Optional[TrainingResult]): Training result to extract checkpoint path from + +**Returns:** +- `InferenceResult`: Metadata object containing `job_id`, `started_time`, and `inference_output_path`. Returns `None` if `dry_run=True` + +**Example:** +```python +inference_result = inference.invoke_batch( + job_name="batch-inference-job", + input_path="s3://my-bucket/inference-input", + output_s3_path="s3://my-bucket/inference-output", + model_path="s3://my-bucket/trained-model" +) +print(f"Batch inference started: {inference_result.job_id}") +``` +--- + +##### `get_logs()` +Retrieves and displays CloudWatch logs for an inference job. + +**Signature:** +```python +def get_logs( + self, + job_result=None, + job_id=None, + started_time=None, + limit=None, + start_from_head: bool = False, + end_time=None, +) -> None +``` + +**Parameters:** +- `job_result` (Optional[InferenceResult]): Job result to retrieve logs for +- `job_id` (Optional[str]): Job identifier +- `started_time` (Optional[datetime]): Job start time to filter logs +- `limit` (Optional[int]): Maximum number of log lines +- `start_from_head` (bool): If True, start from the beginning of logs. Default: False +- `end_time` (Optional[int]): End time in epoch milliseconds for searching a log time range + +**Returns:** +- None (prints logs to console) + +**Example:** +```python +inference.get_logs(job_result=inference_result, limit=100) +``` +--- + +### End-to-End Example (Service Classes) + +```python +from amzn_nova_forge.trainer import ForgeTrainer +from amzn_nova_forge.evaluator import ForgeEvaluator +from amzn_nova_forge.deployer import ForgeDeployer +from amzn_nova_forge.inference import ForgeInference +from amzn_nova_forge.core import ForgeConfig +from amzn_nova_forge.manager import SMTJRuntimeManager +from amzn_nova_forge.model.model_enums import Model, TrainingMethod, DeployPlatform +from amzn_nova_forge.core import EvaluationTask + +# Shared configuration +config = ForgeConfig( + output_s3_path="s3://my-bucket/output", + enable_job_caching=True +) +infra = SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=2) + +# 1. Train +trainer = ForgeTrainer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + training_data_s3_path="s3://my-bucket/data.jsonl", + config=config +) +train_result = trainer.train(job_name="my-training-job") + +# 2. Evaluate +evaluator = ForgeEvaluator(model=Model.NOVA_MICRO, infra=infra, config=config) +eval_result = evaluator.evaluate( + job_name="my-eval-job", + eval_task=EvaluationTask.MMLU, + job_result=train_result +) + +# 3. Deploy +deployer = ForgeDeployer(region="us-east-1", model=Model.NOVA_MICRO) +deployment = deployer.deploy( + model_artifact_path=train_result.model_artifacts.checkpoint_s3_path, + deploy_platform=DeployPlatform.BEDROCK_OD, + endpoint_name="my-nova-endpoint" +) + +# 4. Inference +inference_client = ForgeInference(region="us-east-1") +response = inference_client.invoke( + endpoint_arn=deployment.endpoint.uri, + request_body={ + "messages": [{"role": "user", "content": "Hello!"}], + "max_tokens": 100 + } +) +``` + +--- diff --git a/docs/spec/utilities.md b/docs/spec/utilities.md new file mode 100644 index 0000000..ae4eabe --- /dev/null +++ b/docs/spec/utilities.md @@ -0,0 +1,457 @@ +# Utility Functions + +### verify_reward_function() + +Verifies a reward function with sample data before using it in RFT training or evaluation. This utility helps you test your reward function implementation to ensure it works correctly and returns the expected format. + +**Signature:** +```python +def verify_reward_function( + reward_function: str, + sample_data: List[Dict[str, Any]], + region: str = "us-east-1", + validate_format: bool = True, + platform: Optional[Platform] = None, +) -> Dict[str, Any] +``` + +**Parameters:** +- `reward_function` (str): Either a Lambda ARN (string starting with `'arn:aws:lambda:'`) or a path to a local Python file containing the reward function. +- `sample_data` (List[Dict[str, Any]]): List of conversation samples to test. Each sample should be a dict with `'id'`, `'messages'`, and optionally `'reference_answer'` keys. +- `region` (str): AWS region for Lambda invocation (default: "us-east-1"). +- `validate_format` (bool): If True, validates that sample_data matches RFT format and output matches expected format (default: True). +- `platform` (Platform): Platform enum (Platform.SMHP or Platform.SMTJ). **Required when using Lambda ARN**. When set to Platform.SMHP, validates that Lambda ARN contains 'SageMaker' in the function name as required by SageMaker HyperPod. Optional for local files. + +**Returns:** +- `Dict[str, Any]`: Dictionary containing: + - `success` (bool): Always True if no exception raised + - `results` (list): List of individual test results + - `total_samples` (int): Total number of samples tested + - `successful_samples` (int): Number of successful tests + - `warnings` (list): List of warning messages (e.g., missing reference_answer) + +**Raises:** +- `ValueError`: If any validation errors are encountered, with a detailed error message listing all issues found. + +**Example** +```python +from amzn_nova_forge import verify_reward_function +from amzn_nova_forge.model import Platform + +# Test with Lambda ARN (platform required for Lambda ARNs) +result = verify_reward_function( + reward_function="arn:aws:lambda:us-east-1:123456789012:function:MySageMakerReward", + sample_data=[ + { + "id": "sample_1", + "reference_answer": "correct answer", + "messages": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "response"} + ] + } + ], + platform=Platform.SMHP # Required for Lambda ARNs +) + +print(f"Verification: {'PASSED' if result['success'] else 'FAILED'}") +print(f"Tested {result['total_samples']} samples, {result['successful_samples']} successful") + +if result.get('warnings'): + print(f"\nWarnings:") + for warning in result['warnings']: + print(f" - {warning}") + +# Test with local Python file (platform optional) +result = verify_reward_function( + reward_function="./my_reward_function.py", + sample_data=[ + { + "id": "sample_1", + "reference_answer": "correct answer", + "messages": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "response"} + ] + } + ] +) +``` +**Output Format Requirements from Lambda:** + +```python +{ + "id": "sample_1", # Required: string + "aggregate_reward_score": 0.75, # Required: float or int + "metrics_list": [ # Optional: validated if present + { + "name": "accuracy", # Required: string + "value": 0.85, # Required: float or int + "type": "Metric" # Required: "Metric" or "Reward" + } + ] +} +``` + +**Common Validation Errors:** +- Missing required fields in input (`messages` field is required) +- Missing required fields in output (`id` and `aggregate_reward_score` are required) +- Invalid data types (e.g., `aggregate_reward_score` must be a number) +- Missing `platform` parameter when using Lambda ARN +- SMHP Lambda ARN doesn't contain 'SageMaker' in function name +- Invalid `metrics_list` structure (must be list of dicts with `name`, `value`, `type`) +- Invalid metric `type` (must be "Metric" or "Reward") + +**Warnings:** +- Missing `reference_answer`: While optional in RFT datasets, reference answers are recommended for meaningful reward calculations. Without ground truth, your reward function cannot compare model outputs against expected answers. + + + +**Note:** The `metrics_list` field is optional. If provided, it will be validated for proper structure and logged during training/evaluation. + +--- +## Monitoring + +### CloudWatchLogMonitor + +Monitors CloudWatch logs and plots training metrics for Nova model training jobs. Supports both SageMaker Training Jobs (SMTJ) and SageMaker HyperPod (SMHP) platforms. + +#### Factory Methods + +##### `from_job_id()` + +Creates a CloudWatchLogMonitor from a job ID. + +**Signature:** +```python +@classmethod +def from_job_id( + cls, + job_id: str, + platform: Platform, + started_time: Optional[datetime] = None, + **kwargs, +) -> "CloudWatchLogMonitor" +``` + +**Parameters:** +- `job_id` (str): The training job identifier +- `platform` (Platform): Execution platform (`Platform.SMTJ` or `Platform.SMHP`) +- `started_time` (Optional[datetime]): Job start time (used to filter logs) +- `**kwargs`: Platform-specific parameters: + - SMHP requires: `cluster_name` (str), optional `namespace` (str, defaults to "kubeflow") + +**Returns:** +- `CloudWatchLogMonitor`: Monitor instance + +**Example:** +```python +from amzn_nova_forge.monitor import CloudWatchLogMonitor +from amzn_nova_forge.model import Platform + +# SMTJ +monitor = CloudWatchLogMonitor.from_job_id( + job_id="my-training-job", + platform=Platform.SMTJ, + started_time=datetime(2026, 1, 15, 12, 0, 0) +) + +# SMHP +monitor = CloudWatchLogMonitor.from_job_id( + job_id="my-hyperpod-job", + platform=Platform.SMHP, + cluster_name="my-cluster", + namespace="kubeflow" +) +``` + +--- + +##### `from_job_result()` + +Creates a CloudWatchLogMonitor from a training job result object. + +**Signature:** +```python +@classmethod +def from_job_result( + cls, + job_result: BaseJobResult, + cloudwatch_logs_client=None +) -> "CloudWatchLogMonitor" +``` + +**Parameters:** +- `job_result` (BaseJobResult): A training or evaluation result object (e.g., `TrainingResult`) +- `cloudwatch_logs_client` (Optional): Boto3 CloudWatch Logs client (auto-created if not provided) + +**Returns:** +- `CloudWatchLogMonitor`: Monitor instance + +**Example:** +```python +result = customizer.train(job_name="my-job") +monitor = CloudWatchLogMonitor.from_job_result(job_result=result) +``` + +--- + +#### Methods + +##### `get_logs()` + +Retrieves CloudWatch log events for the job. + +**Signature:** +```python +def get_logs( + self, + limit: Optional[int] = None, + start_from_head: bool = False, + end_time: Optional[int] = None, +) -> List[Dict] +``` + +**Parameters:** +- `limit` (Optional[int]): Maximum number of log events to retrieve +- `start_from_head` (bool): If True, start from the beginning of logs; if False, start from the end +- `end_time` (Optional[int]): End time in epoch milliseconds + +**Returns:** +- `List[Dict]`: List of log event dictionaries, each containing a `"message"` key + +**Example:** +```python +logs = monitor.get_logs(limit=100) +``` + +--- + +##### `show_logs()` + +Prints CloudWatch log messages to the console. + +**Signature:** +```python +def show_logs( + self, + limit: Optional[int] = None, + start_from_head: bool = False, + end_time: Optional[int] = None, +) -> None +``` + +**Parameters:** +- `limit` (Optional[int]): Maximum number of log events to display +- `start_from_head` (bool): If True, start from the beginning of logs; if False, start from the end +- `end_time` (Optional[int]): End time in epoch milliseconds + +**Example:** +```python +monitor.show_logs(limit=50, start_from_head=True) +``` + +--- + +##### `plot_metrics()` + +Parses training metrics from CloudWatch logs and displays them as matplotlib plots. Automatically fetches the latest logs if the job is still in progress or logs have not been retrieved yet. + +**Signature:** +```python +def plot_metrics( + self, + training_method: TrainingMethod, + metrics: Optional[List[str]] = None, + starting_step: Optional[int] = None, + ending_step: Optional[int] = None, +) -> None +``` + +**Parameters:** +- `training_method` (TrainingMethod): The training method used for the job (e.g., `TrainingMethod.SFT_LORA`, `TrainingMethod.CPT`, `TrainingMethod.RFT_LORA`) +- `metrics` (Optional[List[str]]): List of metric names to plot. Available metrics depend on training method: + - CPT / SFT: `"training_loss"` + - RFT: `"reward_score"` +- `starting_step` (Optional[int]): Filter to only show metrics from this global step onward +- `ending_step` (Optional[int]): Filter to only show metrics up to this global step + +**Raises:** +- `ValueError`: If `starting_step` > `ending_step`, or if no logs are found for the job +- `NotImplementedError`: If an unsupported metric is requested for the given training method/platform + +**Example:** +```python +from amzn_nova_forge.monitor import CloudWatchLogMonitor +from amzn_nova_forge.model import Platform, TrainingMethod + +# Create monitor from a training result +monitor = CloudWatchLogMonitor.from_job_result(job_result=training_result) + +# Plot training loss for an SFT job +monitor.plot_metrics( + training_method=TrainingMethod.SFT_LORA, + metrics=["training_loss"] +) + +# Plot reward score for an RFT job, filtered to steps 50-200 +monitor.plot_metrics( + training_method=TrainingMethod.RFT_LORA, + metrics=["reward_score"], + starting_step=50, + ending_step=200 +) +``` + +--- + +### MLflowMonitor + +MLflow monitoring configuration for Nova model training. This class provides experiment tracking capabilities through MLflow integration. + +**Note:** MLflow monitoring is only supported for SageMaker platforms (SMTJ, SMHP). It is not available for Bedrock platform. + +**MLflow Integration Features:** +- Automatic logging of training metrics +- Model artifact and checkpoint tracking +- Hyperparameter recording +- Support for SageMaker MLflow tracking servers +- Custom MLflow tracking server support (with proper network configuration) + + +#### Constructor + +**Signature:** +```python +def __init__( + self, + tracking_uri: Optional[str] = None, + experiment_name: Optional[str] = None, + run_name: Optional[str] = None, +) +``` + +**Parameters:** +- `tracking_uri` (Optional[str]): MLflow tracking server URI or SageMaker MLflow app ARN. If not provided, attempts to use a default SageMaker MLflow tracking server if one exists +- `experiment_name` (Optional[str]): Name of the MLflow experiment. If not provided, will use the job name +- `run_name` (Optional[str]): Name of the MLflow run. If not provided, will be auto-generated + +**Raises:** +- `ValueError`: If MLflow configuration validation fails + +**Example:** +```python +from amzn_nova_forge.monitor import * + +# With explicit tracking URI +monitor = MLflowMonitor( + tracking_uri="arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", + experiment_name="nova-customization", + run_name="sft-run-1" +) + +# With default tracking URI (if available) +monitor = MLflowMonitor( + experiment_name="nova-customization", + run_name="sft-run-1" +) + +# Use with NovaModelCustomizer +customizer = NovaModelCustomizer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.SFT_LORA, + infra=runtime_manager, + data_s3_path="s3://bucket/data", + mlflow_monitor=monitor +) +``` + +#### Methods + +##### `to_dict()` + +Converts MLflow configuration to dictionary format for use in recipe overrides. + +**Signature:** +```python +def to_dict( + self +) -> dict +``` + +**Returns:** +- `dict`: Dictionary with mlflow_* keys for recipe configuration. Returns empty dict if no tracking URI is available + +**Example:** +```python +monitor = MLflowMonitor( + tracking_uri="arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", + experiment_name="nova-customization" +) + +config_dict = monitor.to_dict() +# Returns: { +# "mlflow_tracking_uri": "arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", +# "mlflow_experiment_name": "nova-customization" +# } +``` + +##### `get_presigned_url()` + +Generates a presigned URL for accessing the MLflow tracking server UI directly without navigating through the AWS Console. + +**Signature:** +```python +def get_presigned_url( + self, + session_expiration_duration_in_seconds: int = 43200, + expires_in_seconds: int = 300 +) -> str +``` + +**Parameters:** +- `session_expiration_duration_in_seconds` (int, optional): Duration in seconds for which the MLflow UI session is valid after accessing the presigned URL. Default is 43200 seconds (12 hours). Valid range: 1800-43200 seconds +- `expires_in_seconds` (int, optional): Duration in seconds for which the presigned URL itself is valid. The URL must be accessed within this time. Default is 300 seconds (5 minutes). Valid range: 5-300 seconds + +**Returns:** +- `str`: Presigned URL for accessing the MLflow tracking server UI. This URL must be used within `expires_in_seconds` + +**Raises:** +- `ValueError`: If tracking_uri is not set +- `RuntimeError`: If unable to generate presigned URL + +**Example:** +```python +monitor = MLflowMonitor( + tracking_uri="arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", + experiment_name="nova-customization" +) + +# Generate presigned URL with defaults +# URL expires in 5 minutes, but session lasts 12 hours once accessed +url = monitor.get_presigned_url() +print(f"Access MLflow UI at: {url}") + +# Generate URL with custom expiration times +url = monitor.get_presigned_url( + session_expiration_duration_in_seconds=3600, # 1 hour session + expires_in_seconds=60 # URL expires in 1 minute +) +``` + +#### MLflow Integration Notes + +When MLflow monitoring is enabled: +1. Training metrics will be automatically logged to the specified MLflow tracking server +2. Model artifacts and checkpoints will be tracked in MLflow +3. Hyperparameters and configuration will be recorded as MLflow parameters +4. You can view experiment results in the MLflow UI + +The MLflow integration supports: +- SageMaker MLflow tracking servers +- Custom MLflow tracking servers (with appropriate network configuration) +- Automatic experiment and run creation +- Metric logging during training +- Artifact tracking + +--- diff --git a/docs/data_prep.md b/docs/user-guides/data_prep.md similarity index 85% rename from docs/data_prep.md rename to docs/user-guides/data_prep.md index 669c2e4..e62db3a 100644 --- a/docs/data_prep.md +++ b/docs/user-guides/data_prep.md @@ -11,6 +11,7 @@ This guide walks through the data preparation workflow using the Nova Forge SDK' | Theme | Operator | Purpose | Supported Runtimes | |---|---|---|---| | Load | `loader.load(path)` | Ingest data from JSONL, JSON, CSV, Parquet, or Arrow (local or S3), or from HuggingFace Hub | Local | +| Load | `loader.load(log_group, query, start_time, end_time)` | Execute a CloudWatch Logs Insights query and stream results | Local (queries CloudWatch) | | Filter | `loader.filter(method=FilterMethod.DEFAULT_TEXT_FILTER)` | Remove records with excessive URL content (web-scraped boilerplate) | SMTJ (default), AWS Glue (legacy) | | Filter | `loader.filter(method=FilterMethod.EXACT_DEDUP)` | Remove exact duplicate records by content hash | SMTJ (default), AWS Glue (legacy) | | Filter | `loader.filter(method=FilterMethod.FUZZY_DEDUP)` | Remove near-duplicates via MinHash LSH similarity | SMTJ (default), AWS Glue (legacy) | @@ -26,7 +27,7 @@ This guide walks through the data preparation workflow using the Nova Forge SDK' ## Overview -The SDK provides dataset loaders for different file formats (`JSONLDatasetLoader`, `JSONDatasetLoader`, `CSVDatasetLoader`, `ParquetDatasetLoader`, `ArrowDatasetLoader`) plus `HuggingFaceDatasetLoader` for streaming directly from HuggingFace Hub — all sharing the same chainable API. Operations are lazy — they queue up and execute when a terminal operation is called. +The SDK provides dataset loaders for different file formats (`JSONLDatasetLoader`, `JSONDatasetLoader`, `CSVDatasetLoader`, `ParquetDatasetLoader`, `ArrowDatasetLoader`) plus `HuggingFaceDatasetLoader` for streaming directly from HuggingFace Hub and `CloudWatchDatasetLoader` for querying Amazon CloudWatch Logs — all sharing the same chainable API. Operations are lazy — they queue up and execute when a terminal operation is called. ```python from amzn_nova_forge import ( @@ -284,12 +285,23 @@ loader.save("s3://my-bucket/data/train.jsonl") The in-memory terminal operations (`save()`, `show()`, `split()`, `validate()`, `transform()`, `filter(INVALID_RECORDS)`) consume the HuggingFace generator directly — no extra configuration needed. -Distributed filter operations (`DEFAULT_TEXT_FILTER`, `EXACT_DEDUP`, `FUZZY_DEDUP`, `LANGUAGE_DETECTION`) need an S3 location to land intermediate and final outputs. Because HuggingFace data isn't on S3, you must pass an explicit `output_path` to each distributed filter step — same as when the input is a local file. +Distributed filter operations (`DEFAULT_TEXT_FILTER`, `EXACT_DEDUP`, `FUZZY_DEDUP`, `LANGUAGE_DETECTION`) need an S3 location to land intermediate and final outputs. When the input is not on S3 (local files or HuggingFace), the SDK auto-derives output paths under the default data-prep bucket (`sagemaker-forge-dataprep-{account}-{region}`). You can also pass an explicit `output_path` to override this. ```python loader = HuggingFaceDatasetLoader() loader.load("foo/bar", split="train_sft") +# output_path is optional — auto-derived under the data-prep bucket +loader.filter( + method=FilterMethod.EXACT_DEDUP, + text_field="text", +) +loader.save("s3://my-bucket/data/train_dedup.jsonl") +``` + +To override the auto-derived path, pass `output_path` explicitly: + +```python loader.filter( method=FilterMethod.EXACT_DEDUP, text_field="text", @@ -300,6 +312,126 @@ loader.save("s3://my-bucket/data/train_dedup.jsonl") --- +## Loading from CloudWatch Logs + +`CloudWatchDatasetLoader` executes CloudWatch Logs Insights queries and streams the results into the Nova Forge pipeline. The loader is lazy — `load()` stores configuration only, and the query executes when a terminal operation (`save()`, `show()`, etc.) triggers iteration. + +Uses the default AWS credential chain. Required IAM permissions: `logs:StartQuery` and `logs:GetQueryResults`. + +### `load(log_group, query, start_time, end_time)` + +- `log_group` — CloudWatch log group name (e.g. `"/my-app/api-logs"`). +- `query` — A [CloudWatch Logs Insights query](https://docs.aws.amazon.com/AmazonCloudWatch/latest/logs/CWL_QuerySyntax.html). The query determines which fields appear as keys in the yielded dictionaries. See the examples below for how to write queries for different log formats. +- `start_time` — Query start time (inclusive), as a `datetime` object. +- `end_time` — Query end time (exclusive), as a `datetime` object. + +### Writing Insights queries + +The `query` parameter accepts [CloudWatch Logs Insights syntax](https://docs.aws.amazon.com/AmazonCloudWatch/latest/logs/CWL_QuerySyntax.html). The two most common patterns for training data extraction: + +**`fields` — for structured JSON logs.** If your application writes JSON events, use `fields` to select top-level or nested keys directly. CloudWatch auto-parses JSON, so dot-notation works for nested objects. + +**`parse` — for semi-structured text logs.** If your logs are plain text with a known pattern, use `parse` with a regex to extract named capture groups as fields. + +### Example 1: JSON logs + +Suppose your log group `/my-app/api-logs` contains JSON events like: + +```json +{"endpoint": "/chat", "request": "What is AI?", "response": "AI is artificial intelligence."} +``` + +Use `fields` to select the keys you need, and `filter` to narrow results: + +```python +from datetime import datetime, timezone, timedelta +from amzn_nova_forge import CloudWatchDatasetLoader, Model, TrainingMethod +from amzn_nova_forge.dataset.operations import TransformMethod, ValidateMethod + +loader = CloudWatchDatasetLoader() +loader.load( + log_group="/my-app/api-logs", + query="fields request, response | filter endpoint = '/chat'", + start_time=datetime(2025, 1, 1, tzinfo=timezone.utc), + end_time=datetime(2025, 1, 8, tzinfo=timezone.utc), +) +# Yields: {"request": "What is AI?", "response": "AI is artificial intelligence."} + +loader.transform( + method=TransformMethod.SCHEMA, + training_method=TrainingMethod.SFT_LORA, + model=Model.NOVA_LITE_2, + column_mappings={"question": "request", "answer": "response"}, +) +loader.validate( + method=ValidateMethod.INVALID_RECORDS, + training_method=TrainingMethod.SFT_LORA, + model=Model.NOVA_LITE_2, +) +loader.save("s3://my-bucket/data/chat_sft.jsonl") +``` + +For nested JSON objects, use dot-notation to reach into deeper fields. Array elements are accessed by index. Use `as` to alias the extracted path. For example, given logs like: + +```json +{"input": {"messages": [{"role": "user", "content": "What is AI?"}]}, "output": {"content": "AI is artificial intelligence."}} +``` + +You can extract the nested values with: + +``` +fields input.messages.0.content as question, output.content as answer +``` + +### Example 2: Semi-structured text logs + +Suppose your log group `/my-app/inference-logs` contains text lines like: + +``` +[INFO] endpoint=/api/chat input="Summarize this article" output="This article discusses..." +``` + +Use `parse` with a regex to extract named fields: + +```python +from datetime import datetime, timezone, timedelta +from amzn_nova_forge import CloudWatchDatasetLoader, Model, TrainingMethod +from amzn_nova_forge.dataset.operations import TransformMethod, ValidateMethod + +loader = CloudWatchDatasetLoader() +loader.load( + log_group="/my-app/inference-logs", + query=''' + fields @message + | parse @message /input="(?[^"]+)" output="(?[^"]+)"/ + | filter ispresent(input) + ''', + start_time=datetime.now(timezone.utc) - timedelta(days=7), + end_time=datetime.now(timezone.utc), +) +# Yields: {"@message": "[INFO] endpoint=/api/chat ...", "input": "Summarize this article", "output": "This article discusses..."} + +loader.transform( + method=TransformMethod.SCHEMA, + training_method=TrainingMethod.SFT_LORA, + model=Model.NOVA_LITE_2, + column_mappings={"question": "input", "answer": "output"}, +) +loader.validate( + method=ValidateMethod.INVALID_RECORDS, + training_method=TrainingMethod.SFT_LORA, + model=Model.NOVA_LITE_2, +) +loader.save("s3://my-bucket/data/inference_sft.jsonl") +``` + +### Limitations + +- CloudWatch Logs Insights returns a maximum of 10,000 results per query. For larger datasets, run multiple queries with narrower time ranges and concatenate the results. +- Queries have a 60-minute execution timeout. If your query times out, narrow the time range or simplify the query. + +--- + ## Inspecting Data `loader.show(n=10)` — Previews the current dataset state. Flushes any pending operations first. @@ -664,7 +796,7 @@ These apply to `DEFAULT_TEXT_FILTER`, `EXACT_DEDUP`, and `FUZZY_DEDUP`: | Parameter | Type | Default | Description | |---|---|---|---| -| `output_path` | `str` | Auto-derived | S3 URI for filtered output. Required when input is not on S3. | +| `output_path` | `str` | Auto-derived | S3 URI for filtered output. Auto-derived from the load path when omitted. For local or HuggingFace inputs, uses the default data-prep bucket (`s3://sagemaker-forge-dataprep-{account_id}-{region}/`). | | `input_format` | `str` | `"parquet"` | `"parquet"` or `"jsonl"` | | `output_format` | `str` | `"parquet"` | `"parquet"` or `"jsonl"` | | `text_field` | `str` | `"text"` | Column name containing the text to process. | diff --git a/docs/iam_setup.md b/docs/user-guides/iam_setup.md similarity index 97% rename from docs/iam_setup.md rename to docs/user-guides/iam_setup.md index fef3472..b545e9a 100644 --- a/docs/iam_setup.md +++ b/docs/user-guides/iam_setup.md @@ -110,7 +110,9 @@ Please refer to the "Sid" of each statement to determine which policies you need "Action": [ "logs:DescribeLogStreams", "logs:FilterLogEvents", - "logs:GetLogEvents" + "logs:GetLogEvents", + "logs:StartQuery", + "logs:GetQueryResults" ], "Resource": "arn:aws:logs:::log-group:*" }, @@ -217,6 +219,7 @@ Please refer to the "Sid" of each statement to determine which policies you need Data mixing fetches recipe templates from a cross-account S3 access point owned by the Nova Forge service. The resource is scoped to S3 access point ARNs, which allows cross-account access point calls while preventing read access to arbitrary S3 bucket objects. - [HyperPod only] If your cluster uses namespace access control, you must have access to the Kubernetes namespace +- [CloudWatch data loading only] `logs:StartQuery` and `logs:GetQueryResults` in the `AccessCloudWatchLogs` statement are required when using `CloudWatchDatasetLoader`. The other actions in that statement (`DescribeLogStreams`, `FilterLogEvents`, `GetLogEvents`) are used for job log monitoring and are not needed for data loading. ### Job Monitoring via Email Notifications diff --git a/docs/instance_type_spec.md b/docs/user-guides/instance_type_spec.md similarity index 100% rename from docs/instance_type_spec.md rename to docs/user-guides/instance_type_spec.md diff --git a/docs/job_notifications.md b/docs/user-guides/job_notifications.md similarity index 100% rename from docs/job_notifications.md rename to docs/user-guides/job_notifications.md diff --git a/docs/rft_multiturn.md b/docs/user-guides/rft_multiturn.md similarity index 99% rename from docs/rft_multiturn.md rename to docs/user-guides/rft_multiturn.md index 2fd0e23..2894837 100644 --- a/docs/rft_multiturn.md +++ b/docs/user-guides/rft_multiturn.md @@ -1307,6 +1307,6 @@ get_logs(env_type=EnvType.TRAIN, tail=True) ## Additional Resources - [Main SDK Documentation](../README.md) - Complete SDK overview and getting started guide -- [API Specification](spec.md) - Detailed API documentation for all modules +- [API Specification](../spec/index.md) - Detailed API documentation for all modules - [Quick Start Notebook](../samples/nova_quickstart.ipynb) - General Nova customization examples - [RFT Multiturn Notebook](../samples/rft_multiturn_quickstart.ipynb) - RFT multiturn specific examples diff --git a/docs/troubleshooting.md b/docs/user-guides/troubleshooting.md similarity index 100% rename from docs/troubleshooting.md rename to docs/user-guides/troubleshooting.md diff --git a/src/amzn_nova_forge/__init__.py b/src/amzn_nova_forge/__init__.py index 74af3c8..1bb9829 100644 --- a/src/amzn_nova_forge/__init__.py +++ b/src/amzn_nova_forge/__init__.py @@ -77,6 +77,7 @@ from .validation.endpoint_validator import SageMakerEndpointEnvironment _LAZY_IMPORTS = { + "CloudWatchDatasetLoader": ".dataset.cloudwatch_dataset_loader", "DefaultTextFilterOperation": ".dataset.operations.default_text_filter_operation", "ExactDedupFilterOperation": ".dataset.operations.exact_dedup_filter_operation", "FuzzyDedupFilterOperation": ".dataset.operations.fuzzy_dedup_filter_operation", @@ -98,6 +99,7 @@ def __getattr__(name: str): __all__ = [ "ArrowDatasetLoader", + "CloudWatchDatasetLoader", "CSVDatasetLoader", "HuggingFaceDatasetLoader", "JSONDatasetLoader", diff --git a/src/amzn_nova_forge/__version__.py b/src/amzn_nova_forge/__version__.py index 37c2cc2..e561852 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.5" # pragma: no cover +VERSION = "1.4.6" # pragma: no cover diff --git a/src/amzn_nova_forge/core/result/eval_result.py b/src/amzn_nova_forge/core/result/eval_result.py index 53476aa..5bdd991 100644 --- a/src/amzn_nova_forge/core/result/eval_result.py +++ b/src/amzn_nova_forge/core/result/eval_result.py @@ -50,11 +50,13 @@ def __init__( eval_task: EvaluationTask, eval_output_path: str, s3_client=None, + region: Optional[str] = None, ): self.eval_task = eval_task self.eval_output_path = eval_output_path self._cached_results_dir: Optional[str] = None - self._s3_client = s3_client or boto3.client("s3") + self._region = region + self._s3_client = s3_client or boto3.client("s3", region_name=region) super().__init__(job_id, started_time) def _download_eval_results(self) -> str: @@ -205,9 +207,12 @@ def __init__( eval_output_path: str, sagemaker_client=None, s3_client=None, + region: Optional[str] = None, ): - self._sagemaker_client = sagemaker_client or boto3.client("sagemaker") - super().__init__(job_id, started_time, eval_task, eval_output_path, s3_client) + self._sagemaker_client = sagemaker_client or boto3.client("sagemaker", region_name=region) + super().__init__( + job_id, started_time, eval_task, eval_output_path, s3_client, region=region + ) def _create_status_manager(self) -> JobStatusManager: return SMTJStatusManager(self._sagemaker_client) @@ -243,10 +248,11 @@ def __init__( eval_output_path: str, cluster_name: str, namespace: str = "kubeflow", + region: Optional[str] = None, ): self.cluster_name = cluster_name self.namespace = namespace - super().__init__(job_id, started_time, eval_task, eval_output_path) + super().__init__(job_id, started_time, eval_task, eval_output_path, region=region) def _create_status_manager(self) -> JobStatusManager: return SMHPStatusManager(self.cluster_name, self.namespace) @@ -282,11 +288,13 @@ def __init__( started_time: datetime, eval_task: EvaluationTask, eval_output_path: str, + region: Optional[str] = None, ): - super().__init__(job_id, started_time, eval_task, eval_output_path) + self._region = region + super().__init__(job_id, started_time, eval_task, eval_output_path, region=region) def _create_status_manager(self) -> JobStatusManager: - return BedrockStatusManager() + return BedrockStatusManager(region=self._region) def _to_dict(self): return { diff --git a/src/amzn_nova_forge/core/result/inference_result.py b/src/amzn_nova_forge/core/result/inference_result.py index 8c9420d..4799265 100644 --- a/src/amzn_nova_forge/core/result/inference_result.py +++ b/src/amzn_nova_forge/core/result/inference_result.py @@ -40,9 +40,16 @@ class InferenceResult(BaseJobResult, ABC): inference_output_path: str - def __init__(self, job_id: str, started_time: datetime, inference_output_path: str): + def __init__( + self, + job_id: str, + started_time: datetime, + inference_output_path: str, + region: Optional[str] = None, + ): self.inference_output_path = inference_output_path self._cached_results_dir: Optional[str] = None + self._region = region super().__init__(job_id, started_time) @staticmethod @@ -110,7 +117,7 @@ def _download_inference_results(self) -> str: bucket = parsed.netloc key = parsed.path.lstrip("/") - s3_client = boto3.client("s3") + s3_client = boto3.client("s3", region_name=self._region) # Create temp dir for caching results self._cached_results_dir = tempfile.mkdtemp(prefix=f"inference_results_{self.job_id}_") @@ -161,7 +168,7 @@ def get(self, s3_path=None) -> Dict: s3_path_stripped = s3_path[5:] # Remove 's3://' bucket, key = s3_path_stripped.split("/", 1) - s3_client = boto3.client("s3") + s3_client = boto3.client("s3", region_name=self._region) s3_client.put_object( Bucket=bucket, Key=key, @@ -233,10 +240,11 @@ def __init__( started_time: datetime, inference_output_path: str, sagemaker_client=None, + region: Optional[str] = None, ): self._cached_results_dir: Optional[str] = None - self._sagemaker_client = sagemaker_client or boto3.client("sagemaker") - super().__init__(job_id, started_time, inference_output_path) + self._sagemaker_client = sagemaker_client or boto3.client("sagemaker", region_name=region) + super().__init__(job_id, started_time, inference_output_path, region=region) def _create_status_manager(self) -> JobStatusManager: return SMTJStatusManager(self._sagemaker_client) diff --git a/src/amzn_nova_forge/core/result/job_result.py b/src/amzn_nova_forge/core/result/job_result.py index 1577632..b380fcc 100644 --- a/src/amzn_nova_forge/core/result/job_result.py +++ b/src/amzn_nova_forge/core/result/job_result.py @@ -30,6 +30,7 @@ validate_job_name, validate_namespace, ) +from amzn_nova_forge.util.subprocess_utils import _check_hyperpod_stderr logger = logging.getLogger("nova_forge_sdk") @@ -82,9 +83,9 @@ def resolve_start_time(self, job_id: str) -> datetime: class SMTJStatusManager(JobStatusManager): - def __init__(self, sagemaker_client=None): + def __init__(self, sagemaker_client=None, region: Optional[str] = None): super().__init__() - self._sagemaker_client = sagemaker_client or boto3.client("sagemaker") + self._sagemaker_client = sagemaker_client or boto3.client("sagemaker", region_name=region) def get_job_status(self, job_id: str) -> tuple[JobStatus, str]: if self._job_status == JobStatus.COMPLETED or self._job_status == JobStatus.FAILED: @@ -137,11 +138,7 @@ def _connect_cluster(self): check=True, ) - if response.stderr: - logger.error( - f"Unable to connect to HyperPod cluster {self.cluster_name}: {response.stderr}" - ) - raise RuntimeError(response.stderr) + _check_hyperpod_stderr(response.stderr) logger.info( f"Successfully connected to HyperPod cluster '{self.cluster_name}' in namespace '{self.namespace}'." @@ -251,9 +248,9 @@ def _register_bedrock_helpers( cls._get_job_details = staticmethod(get_job_details) cls._log_job_status = staticmethod(log_job_status) - def __init__(self, bedrock_client=None): + def __init__(self, bedrock_client=None, region: Optional[str] = None): super().__init__() - self._bedrock_client = bedrock_client or boto3.client("bedrock") + self._bedrock_client = bedrock_client or boto3.client("bedrock", region_name=region) def get_job_status(self, job_id: str) -> tuple[JobStatus, str]: if self._job_status == JobStatus.COMPLETED or self._job_status == JobStatus.FAILED: diff --git a/src/amzn_nova_forge/core/result/training_result.py b/src/amzn_nova_forge/core/result/training_result.py index 0ee05c9..5295f31 100644 --- a/src/amzn_nova_forge/core/result/training_result.py +++ b/src/amzn_nova_forge/core/result/training_result.py @@ -14,7 +14,7 @@ from abc import ABC from dataclasses import asdict, dataclass from datetime import datetime -from typing import Dict +from typing import Dict, Optional import boto3 @@ -70,12 +70,14 @@ def __init__( model_artifacts: ModelArtifacts, model_type: Model, sagemaker_client=None, + region: Optional[str] = None, ): - self._sagemaker_client = sagemaker_client or boto3.client("sagemaker") + self._region = region + self._sagemaker_client = sagemaker_client or boto3.client("sagemaker", region_name=region) super().__init__(job_id, started_time, method, model_artifacts, model_type) def _create_status_manager(self) -> JobStatusManager: - return SMTJStatusManager(self._sagemaker_client) + return SMTJStatusManager(self._sagemaker_client, region=self._region) def _to_dict(self): return { @@ -111,7 +113,9 @@ def __init__( cluster_name: str, model_type: Model, namespace: str = "kubeflow", + region: Optional[str] = None, ): + self._region = region self.cluster_name = cluster_name self.namespace = namespace super().__init__(job_id, started_time, method, model_artifacts, model_type) @@ -153,12 +157,14 @@ def __init__( model_artifacts: ModelArtifacts, model_type: Model, bedrock_client=None, + region: Optional[str] = None, ): - self._bedrock_client = bedrock_client or boto3.client("bedrock") + self._region = region + self._bedrock_client = bedrock_client or boto3.client("bedrock", region_name=region) super().__init__(job_id, started_time, method, model_artifacts, model_type) def _create_status_manager(self) -> JobStatusManager: - return BedrockStatusManager(self._bedrock_client) + return BedrockStatusManager(self._bedrock_client, region=self._region) def _to_dict(self): return { diff --git a/src/amzn_nova_forge/core/types.py b/src/amzn_nova_forge/core/types.py index 631a968..416a5c9 100644 --- a/src/amzn_nova_forge/core/types.py +++ b/src/amzn_nova_forge/core/types.py @@ -83,6 +83,7 @@ class EndpointInfo: endpoint_name: str uri: str model_artifact_path: str + region: Optional[str] = None @dataclass @@ -121,7 +122,9 @@ def status(self): raise RuntimeError( "Status checker not available. Ensure amzn_nova_forge.util.bedrock is imported." ) - return DeploymentResult._status_checker(self.endpoint.uri, self.endpoint.platform) + return DeploymentResult._status_checker( + self.endpoint.uri, self.endpoint.platform, self.endpoint.region + ) def validate_region(region: str) -> None: @@ -141,10 +144,16 @@ class JobConfig: output_s3_path: Optional[str] = None data_s3_path: Optional[str] = None input_s3_data_type: Optional[str] = None - validation_data_s3_path: Optional[str] = None # Validation data S3 path (for CPT and Bedrock) + validation_data_s3_path: Optional[str] = ( + None # Validation data S3 path (for CPT, SFT, and Bedrock) + ) + trainer_config_hyperparameters: Optional[Dict[str, str]] = ( + None # Extra hyperparameters passed to the training job (e.g., val_check_interval) + ) rft_lambda_arn: Optional[str] = None # RFT Lambda ARN (for RFT jobs on Bedrock) mlflow_tracking_uri: Optional[str] = None # MLflow tracking server ARN mlflow_experiment_name: Optional[str] = None mlflow_run_name: Optional[str] = None 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 diff --git a/src/amzn_nova_forge/dataset/__init__.py b/src/amzn_nova_forge/dataset/__init__.py index fa89295..d921d8b 100644 --- a/src/amzn_nova_forge/dataset/__init__.py +++ b/src/amzn_nova_forge/dataset/__init__.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. from .arrow_dataset_loader import ArrowDatasetLoader +from .cloudwatch_dataset_loader import CloudWatchDatasetLoader from .csv_dataset_loader import CSVDatasetLoader from .huggingface_dataset_loader import HuggingFaceDatasetLoader from .json_dataset_loader import JSONDatasetLoader @@ -20,6 +21,7 @@ __all__ = [ "ArrowDatasetLoader", + "CloudWatchDatasetLoader", "CSVDatasetLoader", "HuggingFaceDatasetLoader", "JSONDatasetLoader", diff --git a/src/amzn_nova_forge/dataset/cloudwatch_dataset_loader.py b/src/amzn_nova_forge/dataset/cloudwatch_dataset_loader.py new file mode 100644 index 0000000..ba94f15 --- /dev/null +++ b/src/amzn_nova_forge/dataset/cloudwatch_dataset_loader.py @@ -0,0 +1,135 @@ +# 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. +"""CloudWatch Logs Insights dataset loader.""" + +import time +from datetime import datetime, timezone + +import boto3 + +from ..telemetry import Feature, _telemetry_emitter +from ..util.logging import logger +from .dataset_loader import DatasetLoader +from .operations.base import DataPrepError + +_TERMINAL_FAILURE_STATUSES = {"Failed", "Cancelled", "Timeout", "Unknown"} + +_POLL_INTERVAL_SECONDS = 10 + + +class CloudWatchDatasetLoader(DatasetLoader): + """Load datasets from CloudWatch Logs Insights queries.""" + + _EXTENSIONS: set[str] = set() + _FORMAT: str = "cloudwatch" + + def _make_single_file_generator(self, path: str): + raise NotImplementedError( + "CloudWatch datasets are not file-based. " + "Use load(log_group=..., query=..., start_time=..., end_time=...) instead." + ) + + @_telemetry_emitter(Feature.DATA_PREP, "load") + def load( + self, + log_group: str, + query: str, + start_time: datetime, + end_time: datetime, + ) -> "CloudWatchDatasetLoader": + """Load data from a CloudWatch Logs Insights query. + + Args: + log_group: CloudWatch log group name. + query: Insights query string. + start_time: Query start time (inclusive). + end_time: Query end time (exclusive). + + Returns: + self (for method chaining) + + Raises: + DataPrepError: If the query fails during iteration. + """ + cw_log_group = log_group + cw_query = query + cw_start_time = start_time + cw_end_time = end_time + + def _generator(): + start_utc = cw_start_time.astimezone(timezone.utc) + end_utc = cw_end_time.astimezone(timezone.utc) + logger.info( + "Querying CloudWatch log group '%s' from %s to %s (UTC).\n Query: %s", + cw_log_group, + start_utc.strftime("%Y-%m-%d %H:%M:%S"), + end_utc.strftime("%Y-%m-%d %H:%M:%S"), + cw_query, + ) + + try: + client = boto3.client("logs") + response = client.start_query( + logGroupName=cw_log_group, + queryString=cw_query, + startTime=int(cw_start_time.timestamp()), + endTime=int(cw_end_time.timestamp()), + ) + except Exception as e: + raise DataPrepError(f"Failed to query CloudWatch Logs: {e}") from e + + results = self._poll_results(client, response["queryId"]) + + if not results: + logger.warning( + "CloudWatch Insights query returned 0 results for log group '%s'.", + cw_log_group, + ) + return + + for row in results: + yield {field["field"]: field["value"] for field in row} + + self.dataset = _generator + self._load_path = f"cw:///{log_group}" + self._last_state = None + self._is_materialized = False + self._session_id = None + return self + + @staticmethod + def _poll_results(client, query_id: str) -> list: + """Poll GetQueryResults until the query completes or times out.""" + elapsed = 0 + while True: + try: + result = client.get_query_results(queryId=query_id) + except Exception as e: + raise DataPrepError(f"Failed to query CloudWatch Logs: {e}") from e + + status = result["status"] + + if status == "Complete": + return result.get("results", []) + + if status in _TERMINAL_FAILURE_STATUSES: + raise DataPrepError( + f"Query did not complete successfully. Status: {status}" + ) from None + + elapsed += _POLL_INTERVAL_SECONDS + if elapsed % 30 == 0: + logger.info("Query still running... (elapsed: %ds)", elapsed) + + time.sleep(_POLL_INTERVAL_SECONDS) diff --git a/src/amzn_nova_forge/dataset/data_state.py b/src/amzn_nova_forge/dataset/data_state.py index af3aa65..a9be951 100644 --- a/src/amzn_nova_forge/dataset/data_state.py +++ b/src/amzn_nova_forge/dataset/data_state.py @@ -21,6 +21,8 @@ from enum import Enum from typing import Callable, Dict, Iterator, Optional +from amzn_nova_forge.util.s3_utils import ensure_bucket_exists, get_dataprep_bucket_name + class DataLocation(Enum): """Where the data physically resides.""" @@ -93,7 +95,8 @@ class OutputPathResolver: ///_/ Where: - - ``parent`` is the parent directory of the original load path. + - ``parent`` is the parent directory of the original load path for S3 + inputs, or the default data-prep bucket for local/HuggingFace inputs. - ``input_stem`` identifies the source file/folder (e.g. ``train`` from ``train.jsonl``). - ``session_id`` is a UTC timestamp (``YYYY-MM-DD_HH-MM-SS``) @@ -108,11 +111,24 @@ class OutputPathResolver: """ def __init__(self, load_path: str, session_id: Optional[str] = None) -> None: - base = load_path.rstrip("/") - self._parent = base.rsplit("/", 1)[0] if "/" in base else base + if load_path.startswith("s3://"): + base = load_path.rstrip("/") + self._parent: Optional[str] = base.rsplit("/", 1)[0] if "/" in base else base + else: + # Local/HuggingFace inputs: defer AWS calls until _resolve_parent() is called. + self._parent = None + self._input_stem = self._extract_stem(load_path) self._session_id = session_id or datetime.now(timezone.utc).strftime("%Y-%m-%d_%H-%M-%S") + def _resolve_parent(self) -> str: + """Return the parent path, resolving the data-prep bucket for non-S3 inputs.""" + if self._parent is None: + bucket = get_dataprep_bucket_name() + ensure_bucket_exists(bucket) + self._parent = f"s3://{bucket}" + return self._parent + def resolve_prefix(self, method, suffix: PathSuffix = PathSuffix.OUTPUT) -> str: """Return the key-safe portion: ``//_``. @@ -128,7 +144,7 @@ def resolve_path(self, method, suffix: PathSuffix = PathSuffix.OUTPUT) -> str: filter/transform operations that need a complete S3 URI or local directory path. """ - return f"{self._parent}/{self.resolve_prefix(method, suffix)}/" + return f"{self._resolve_parent()}/{self.resolve_prefix(method, suffix)}/" @staticmethod def _extract_stem(path: str) -> str: diff --git a/src/amzn_nova_forge/dataset/dataset_loader.py b/src/amzn_nova_forge/dataset/dataset_loader.py index 0d9a6c6..9020082 100644 --- a/src/amzn_nova_forge/dataset/dataset_loader.py +++ b/src/amzn_nova_forge/dataset/dataset_loader.py @@ -119,7 +119,7 @@ class DatasetLoader(ABC): ) """ - def __init__(self, **column_mappings): + def __init__(self, region: Optional[str] = None, **column_mappings): if column_mappings: logger.warning( "Passing column_mappings to the constructor is deprecated. " @@ -128,6 +128,7 @@ def __init__(self, **column_mappings): "training_method=..., model=..., " 'column_mappings={"question": "q", "answer": "a"})' ) + self._region = region self.column_mappings = column_mappings self._dataset: Callable[[], Iterator[Dict]] = lambda: iter([]) self._load_path: Optional[str] = None @@ -460,8 +461,10 @@ def filter( derived as ``////`` where ``input_stem`` is the load filename without extension and ``session`` is a UTC timestamp - (``YYYY-MM-DD_HH-MM-SS``). Required when loading - from a local file. + (``YYYY-MM-DD_HH-MM-SS``). For local or HuggingFace + inputs, the default data-prep bucket + (``sagemaker-forge-dataprep-{account}-{region}``) is + used as the parent. input_format (str): ``"parquet"`` or ``"jsonl"``. output_format (str): ``"parquet"`` or ``"jsonl"``. text_field (str): Column name containing text. @@ -546,6 +549,7 @@ def execute(self) -> "DatasetLoader": for op_type, method, kwargs in self._pending_operations: op = self._get_operation(op_type, method) kwargs["state"] = state + kwargs.setdefault("region", self._region) if "output_path" not in kwargs: kwargs["output_path"] = resolver.resolve_path(method) result = op.execute(self, **kwargs) diff --git a/src/amzn_nova_forge/dataset/dataset_validator/dataset_validator.py b/src/amzn_nova_forge/dataset/dataset_validator/dataset_validator.py index fb0fe15..4197cd0 100644 --- a/src/amzn_nova_forge/dataset/dataset_validator/dataset_validator.py +++ b/src/amzn_nova_forge/dataset/dataset_validator/dataset_validator.py @@ -112,7 +112,8 @@ def validate( "Please use the loader.transform() method to transform your data to Converse format first." ) - s3_client = boto3.client("s3") + region = kwargs.get("region") + s3_client = boto3.client("s3", region_name=region) # Validate each data entry for i, sample in enumerate(dataset): try: @@ -318,7 +319,6 @@ def _validate_content_count( def _validate_video_duration( instance: Any, validation: DatasetCheckEntry, ctx: Dict[str, Any] ) -> None: - uri = instance.source.s3Location.uri s3_client = ctx.get("s3_client") if s3_client is None: @@ -329,7 +329,7 @@ def _validate_video_duration( except ImportError: raise InfrastructureError( "pymediainfo and/or its system dependency libmediainfo not installed. " - "See 'Image/Video validation' in docs/data_prep.md for setup instructions." + "See 'Image/Video validation' in docs/user-guides/data_prep.md for setup instructions." ) try: @@ -422,7 +422,6 @@ def _validate_content_allowlist( def _validate_image_dimensions( instance: Any, validation: DatasetCheckEntry, ctx: Dict[str, Any] ) -> None: - uri = instance.source.s3Location.uri s3_client = ctx.get("s3_client") if s3_client is None: @@ -433,7 +432,7 @@ def _validate_image_dimensions( except ImportError: raise InfrastructureError( "pymediainfo and/or its system dependency libmediainfo not installed. " - "See 'Image/Video validation' in docs/data_prep.md for setup instructions." + "See 'Image/Video validation' in docs/user-guides/data_prep.md for setup instructions." ) try: bucket, key = _parse_s3_uri(uri) diff --git a/src/amzn_nova_forge/dataset/file_utils.py b/src/amzn_nova_forge/dataset/file_utils.py index fa294cb..7f336f3 100644 --- a/src/amzn_nova_forge/dataset/file_utils.py +++ b/src/amzn_nova_forge/dataset/file_utils.py @@ -15,6 +15,7 @@ import os from pathlib import Path +from typing import Optional import boto3 from botocore.exceptions import ClientError @@ -41,7 +42,7 @@ def is_directory(path: str) -> bool: return os.path.isdir(path) -def check_path_exists(path: str) -> None: +def check_path_exists(path: str, region: Optional[str] = None) -> None: """Verify that a single file path exists. Raises DataPrepError for local files that don't exist. For S3, performs a best-effort HeadObject check — logs a warning on permission errors rather than raising. @@ -56,7 +57,7 @@ def check_path_exists(path: str) -> None: s3_path = path[len("s3://") :] bucket, _, key = s3_path.partition("/") try: - client = boto3.client("s3") + client = boto3.client("s3", region_name=region) client.head_object(Bucket=bucket, Key=key) except ClientError as e: error_code = e.response.get("Error", {}).get("Code", "") @@ -150,7 +151,9 @@ def scan_local_directory(dir_path: str, expected_extensions: set[str]) -> list[s return matching -def scan_s3_directory(dir_path: str, expected_extensions: set[str]) -> list[str]: +def scan_s3_directory( + dir_path: str, expected_extensions: set[str], region: Optional[str] = None +) -> list[str]: """Scan an S3 prefix for files matching expected extensions. Uses boto3 list_objects_v2 with Delimiter='/' for non-recursive listing. @@ -163,7 +166,7 @@ def scan_s3_directory(dir_path: str, expected_extensions: set[str]) -> list[str] if not key_prefix.endswith("/"): key_prefix += "/" - client = boto3.client("s3") + client = boto3.client("s3", region_name=region) matching = [] unexpected: list[str] = [] diff --git a/src/amzn_nova_forge/dataset/operations/filter_operation.py b/src/amzn_nova_forge/dataset/operations/filter_operation.py index 1b9cb9a..c10e3ff 100644 --- a/src/amzn_nova_forge/dataset/operations/filter_operation.py +++ b/src/amzn_nova_forge/dataset/operations/filter_operation.py @@ -256,6 +256,7 @@ def _read_summary_json( output_path: str, total_key: str = "input_count", filtered_key: str = "duplicates_removed", + region: Optional[str] = None, ) -> Tuple[int, int]: """Best-effort read of ``_summary.json`` from an S3 output path. @@ -268,7 +269,7 @@ def _read_summary_json( parsed = urlparse(summary_path) bucket = parsed.netloc key = parsed.path.lstrip("/") - resp = boto3.client("s3").get_object(Bucket=bucket, Key=key) + resp = boto3.client("s3", region_name=region).get_object(Bucket=bucket, Key=key) summary = json.loads(resp["Body"].read().decode("utf-8")) return summary.get(total_key, 0), summary.get(filtered_key, 0) except Exception: @@ -276,7 +277,7 @@ def _read_summary_json( return 0, 0 -def _resolve_s3_directory_to_jsonl(s3_path: str) -> str: +def _resolve_s3_directory_to_jsonl(s3_path: str, region: Optional[str] = None) -> str: """If *s3_path* is an S3 directory, resolve it to the single .jsonl file inside. Returns the original path unchanged when it already points to a file. @@ -293,7 +294,7 @@ def _resolve_s3_directory_to_jsonl(s3_path: str) -> str: prefix_path = s3_path if s3_path.endswith("/") else s3_path + "/" bucket, prefix = prefix_path[len("s3://") :].split("/", 1) - s3 = boto3.client("s3") + s3 = boto3.client("s3", region_name=region) paginator = s3.get_paginator("list_objects_v2") jsonl_keys = [] for page in paginator.paginate(Bucket=bucket, Prefix=prefix): diff --git a/src/amzn_nova_forge/dataset/operations/invalid_records_filter_operation.py b/src/amzn_nova_forge/dataset/operations/invalid_records_filter_operation.py index df63ab4..cbb08b6 100644 --- a/src/amzn_nova_forge/dataset/operations/invalid_records_filter_operation.py +++ b/src/amzn_nova_forge/dataset/operations/invalid_records_filter_operation.py @@ -139,6 +139,7 @@ def execute(self, loader: Any, **kwargs: Any) -> FilterOperationResult: # Accept but ignore output_path — this filter operates in-memory. state = kwargs.pop("state", None) kwargs.pop("output_path", None) + region = kwargs.get("region") training_method = kwargs.get("training_method") nova_model = kwargs.get("model") @@ -202,7 +203,7 @@ def filter_generator( captured_result.total_count = 0 captured_result.filtered_count = 0 - s3_client = boto3.client("s3") + s3_client = boto3.client("s3", region_name=region) for sample in captured_dataset(): captured_result.total_count += 1 if captured_pydantic_model is not None and _sample_fails_schema( diff --git a/src/amzn_nova_forge/dataset/operations/language_detection_filter_operation.py b/src/amzn_nova_forge/dataset/operations/language_detection_filter_operation.py index 18953dd..4f3284a 100644 --- a/src/amzn_nova_forge/dataset/operations/language_detection_filter_operation.py +++ b/src/amzn_nova_forge/dataset/operations/language_detection_filter_operation.py @@ -54,7 +54,7 @@ import os import tempfile import time -from typing import Any, Dict, Tuple, Type +from typing import Any, Dict, Optional, Tuple, Type import boto3 import requests @@ -210,8 +210,9 @@ def execute(self, loader: Any, **kwargs: Any) -> FilterOperationResult: "FastText confidence is a probability in [0, 1]." ) + region = kwargs.get("region") if "model_path" not in kwargs: - kwargs["model_path"] = self._ensure_default_model() + kwargs["model_path"] = self._ensure_default_model(region=region) manager = self._resolve_runtime_manager(input_path, **kwargs) @@ -279,7 +280,7 @@ def execute(self, loader: Any, **kwargs: Any) -> FilterOperationResult: return result @staticmethod - def _ensure_default_model() -> str: + def _ensure_default_model(region: Optional[str] = None) -> str: """Return the S3 URI of the default FastText lid.176.bin model. Called only when the caller omits ``model_path``. The model is @@ -298,7 +299,7 @@ def _ensure_default_model() -> str: ensure_bucket_exists(bucket) s3_uri = f"s3://{bucket}/{_DEFAULT_MODEL_S3_KEY}" - s3 = boto3.client("s3") + s3 = boto3.client("s3", region_name=region) try: s3.head_object(Bucket=bucket, Key=_DEFAULT_MODEL_S3_KEY) logger.info("Using cached FastText lid model at %s", s3_uri) diff --git a/src/amzn_nova_forge/dataset/operations/transform_operation.py b/src/amzn_nova_forge/dataset/operations/transform_operation.py index 81f8605..6d76adc 100644 --- a/src/amzn_nova_forge/dataset/operations/transform_operation.py +++ b/src/amzn_nova_forge/dataset/operations/transform_operation.py @@ -51,6 +51,7 @@ def execute(self, loader: Any, **kwargs) -> OperationResult: training_method: Optional[TrainingMethod] = kwargs.get("training_method") model: Optional[Model] = kwargs.get("model") eval_task: Optional[EvaluationTask] = kwargs.get("eval_task") + region: Optional[str] = kwargs.get("region") if training_method is None or model is None: raise ValueError("training_method and model are required for schema transforms.") @@ -76,6 +77,7 @@ def execute(self, loader: Any, **kwargs) -> OperationResult: column_mappings, multimodal_data_s3_path, multimodal_data_bucket_owner, + region=region, ) # Transform produces in-memory data — update state to LOCAL so that # _flush_pending() will materialize the dataset and prevent the lazy @@ -135,6 +137,7 @@ def _apply_first_matching_transformer( column_mappings: dict, multimodal_data_s3_path: Optional[str], multimodal_data_bucket_owner: Optional[str], + region: Optional[str] = None, ) -> None: """Try each transformer in order; apply the first whose source schema matches.""" transformers: List[Dict[str, Any]] = transform_config.get("transformers", []) @@ -155,6 +158,7 @@ def _apply_first_matching_transformer( source_schema, multimodal_data_s3_path, multimodal_data_bucket_owner, + region=region, ) return @@ -171,6 +175,7 @@ def _wire_transform_generator( source_schema: Optional[dict], multimodal_data_s3_path: Optional[str], multimodal_data_bucket_owner: Optional[str], + region: Optional[str] = None, ) -> None: """Replace loader.dataset with a generator that applies the transformer.""" transformer_func = self._get_transformer_function(method_name) @@ -199,7 +204,7 @@ def _wire_transform_generator( bucket_owner = multimodal_data_bucket_owner if bucket_owner is None: try: - sts_client = boto3.client("sts") + sts_client = boto3.client("sts", region_name=region) account_id = sts_client.get_caller_identity()["Account"] logger.info("Auto-resolved bucket owner for %r: %s", bucket, account_id) bucket_owner = account_id @@ -222,7 +227,9 @@ def transform_generator( logger.info(" Input: in-memory (dict)") logger.info(" Output: in-memory") # Create S3 client once for the entire transform pass. - s3_client = boto3.client("s3") if transform_ctx is not None else None + s3_client = ( + boto3.client("s3", region_name=region) if transform_ctx is not None else None + ) count = 0 try: for rec in captured_dataset(): diff --git a/src/amzn_nova_forge/dataset/operations/utils.py b/src/amzn_nova_forge/dataset/operations/utils.py index 580b182..93d288b 100644 --- a/src/amzn_nova_forge/dataset/operations/utils.py +++ b/src/amzn_nova_forge/dataset/operations/utils.py @@ -18,6 +18,7 @@ import io import os import uuid +from typing import Optional import boto3 import pyarrow as pa @@ -73,7 +74,9 @@ def validate_default_bucket_access( ) -def convert_to_s3_parquet(dataset_callable, s3_base_path: str, batch_size: int = 10000) -> str: +def convert_to_s3_parquet( + dataset_callable, s3_base_path: str, batch_size: int = 10000, region: Optional[str] = None +) -> str: """Convert a dataset generator to Parquet files on S3. Consumes the generator in batches, writes each batch as a separate @@ -102,7 +105,7 @@ def convert_to_s3_parquet(dataset_callable, s3_base_path: str, batch_size: int = temp_no_scheme = temp_dir[len("s3://") :] bucket, key_prefix = temp_no_scheme.split("/", 1) - s3_client = boto3.client("s3") + s3_client = boto3.client("s3", region_name=region) part_idx = 0 batch: list[dict] = [] @@ -133,7 +136,9 @@ def _write_parquet_part( s3_client.put_object(Bucket=bucket, Key=key, Body=buf.read()) -def upload_local_file_to_s3(local_path: str, s3_base_path: str) -> str: +def upload_local_file_to_s3( + local_path: str, s3_base_path: str, region: Optional[str] = None +) -> str: """Upload a local file to a temp S3 location, preserving its format. Args: @@ -158,7 +163,7 @@ def upload_local_file_to_s3(local_path: str, s3_base_path: str) -> str: temp_no_scheme = s3_key_dir[len("s3://") :] bucket, key_prefix = temp_no_scheme.split("/", 1) - s3_client = boto3.client("s3") + s3_client = boto3.client("s3", region_name=region) key = f"{key_prefix}{filename}" with open(local_path, "rb") as f: s3_client.put_object(Bucket=bucket, Key=key, Body=f) diff --git a/src/amzn_nova_forge/deployer/forge_deployer.py b/src/amzn_nova_forge/deployer/forge_deployer.py index f76b91c..fbe8e12 100644 --- a/src/amzn_nova_forge/deployer/forge_deployer.py +++ b/src/amzn_nova_forge/deployer/forge_deployer.py @@ -202,7 +202,7 @@ def get_status(self, result: DeploymentResult) -> JobStatus: ) def get_status_by_arn(self, endpoint_arn: str, platform: DeployPlatform) -> Optional[JobStatus]: """Check deployment status by ARN.""" - status_str = check_deployment_status(endpoint_arn, platform) + status_str = check_deployment_status(endpoint_arn, platform, region=self.region) if status_str is None: return None try: @@ -239,7 +239,7 @@ def get_logs( else: platform = DeployPlatform.BEDROCK_OD - status = check_deployment_status(arn, platform) + status = check_deployment_status(arn, platform, region=self.region) logger.info(f"Deployment status for {arn}: {status}") @_telemetry_emitter( @@ -489,7 +489,9 @@ def deploy_to_bedrock( endpoint_name = name_format.replace(".", "-").replace("_", "-") # Check for existing deployment with same name - existing_deployment_arn = check_existing_deployment(endpoint_name, deploy_platform) + existing_deployment_arn = check_existing_deployment( + endpoint_name, deploy_platform, region=self.region + ) attempt_pt_update = False if existing_deployment_arn: @@ -536,7 +538,7 @@ def deploy_to_bedrock( assert existing_deployment_arn is not None try: update_provisioned_throughput_model( - existing_deployment_arn, model_arn, endpoint_name + existing_deployment_arn, model_arn, endpoint_name, region=self.region ) deployment_arn = existing_deployment_arn logger.info(f"Successfully updated existing PT deployment '{endpoint_name}'") @@ -579,6 +581,7 @@ def deploy_to_bedrock( endpoint_name=endpoint_name, uri=deployment_arn, model_artifact_path=(resolved_publish.escrow_uri if resolved_publish else model_arn), + region=self.region, ) result = DeploymentResult( @@ -640,7 +643,9 @@ def _deploy_to_bedrock( else: resolved_endpoint_name = endpoint_name - existing = check_existing_deployment(resolved_endpoint_name, deploy_platform) + existing = check_existing_deployment( + resolved_endpoint_name, deploy_platform, region=self.region + ) if existing: if self.deployment_mode == DeploymentMode.FAIL_IF_EXISTS: raise RuntimeError( @@ -773,6 +778,7 @@ def _deploy_to_sagemaker( endpoint_name=endpoint_name, uri=endpoint_arn, model_artifact_path=model_artifact_path, + region=self.region, ) return DeploymentResult( diff --git a/src/amzn_nova_forge/evaluator/forge_evaluator.py b/src/amzn_nova_forge/evaluator/forge_evaluator.py index 10b90b4..b7c83b7 100644 --- a/src/amzn_nova_forge/evaluator/forge_evaluator.py +++ b/src/amzn_nova_forge/evaluator/forge_evaluator.py @@ -305,6 +305,7 @@ def evaluate( eval_task=eval_task, started_time=start_time, eval_output_path=eval_output_s3_path, + region=self.region, ) else: cluster_name = cast(SMHPRuntimeManager, self.infra).cluster_name @@ -317,6 +318,7 @@ def evaluate( eval_output_path=eval_output_s3_path, cluster_name=cluster_name, namespace=namespace, + region=self.region, ) logger.info( @@ -361,6 +363,7 @@ def get_logs( job_id=resolved_job_id, platform=self._platform, started_time=int(resolved_started.timestamp() * 1000), + region=self.region, **kwargs, ) monitor.show_logs(limit=limit, start_from_head=start_from_head, end_time=end_time) diff --git a/src/amzn_nova_forge/iam/iam_role_creator.py b/src/amzn_nova_forge/iam/iam_role_creator.py index a32ed26..bb791d3 100644 --- a/src/amzn_nova_forge/iam/iam_role_creator.py +++ b/src/amzn_nova_forge/iam/iam_role_creator.py @@ -101,6 +101,7 @@ def create_bedrock_execution_role( role_name: str, bedrock_resource: str = "*", s3_resource: str = "*", + region: Optional[str] = None, ) -> Dict: """ Creates a new IAM Role that allows for Bedrock model creation and deployment. @@ -120,7 +121,7 @@ def create_bedrock_execution_role( Raises: Exception: If it fails at creating the new role. """ - sts_client = boto3.client("sts") + sts_client = boto3.client("sts", region_name=region) with resources.files("amzn_nova_forge.iam").joinpath("bedrock_policies.json").open() as f: policies = json.load(f) @@ -200,6 +201,7 @@ def create_sagemaker_execution_role( cloudwatch_metric_condition: Optional[Dict[str, Any]] = None, cloudwatch_logstream_resource: str = "*", cloudwatch_loggroup_resource: str = "*", + region: Optional[str] = None, ) -> Dict: """ Creates a new IAM Role that allows for SageMaker model creation and deployment. @@ -222,7 +224,7 @@ def create_sagemaker_execution_role( Raises: Exception: If it fails at creating the new role. """ - sts_client = boto3.client("sts") + sts_client = boto3.client("sts", region_name=region) with resources.files("amzn_nova_forge.iam").joinpath("sagemaker_policies.json").open() as f: policies = json.load(f) @@ -327,6 +329,7 @@ def create_smtj_dataprep_execution_role( role_name: str, s3_resource: str = "*", ecr_resource: str = "*", + region: Optional[str] = None, ) -> Dict: """Creates a new IAM Role for SageMaker Training Job data preparation pipelines. @@ -354,7 +357,7 @@ def create_smtj_dataprep_execution_role( Raises: Exception: If it fails at creating the new role. """ - sts_client = boto3.client("sts") + sts_client = boto3.client("sts", region_name=region) with resources.files("amzn_nova_forge.iam").joinpath("smtj_dataprep_policies.json").open() as f: policies = json.load(f) @@ -435,6 +438,7 @@ def create_sagemaker_invoke_role( s3_resource: str = "*", glue_job_resource: str = "*", trust_principal: Optional[Dict[str, Any]] = None, + region: Optional[str] = None, ) -> Dict: """ Creates a new IAM Role for AWS Glue data preparation jobs. @@ -463,7 +467,7 @@ def create_sagemaker_invoke_role( Raises: Exception: If it fails at creating the new role. """ - sts_client = boto3.client("sts") + sts_client = boto3.client("sts", region_name=region) with resources.files("amzn_nova_forge.iam").joinpath("glue_policies.json").open() as f: policies = json.load(f) @@ -545,6 +549,7 @@ def create_bedrock_batch_inference_execution_role( iam_client, role_name: str, s3_resource: str = "*", + region: Optional[str] = None, ) -> Dict: """ Creates a new IAM Role for Bedrock batch inference operations. @@ -562,7 +567,7 @@ def create_bedrock_batch_inference_execution_role( Raises: Exception: If it fails at creating the new role. """ - sts_client = boto3.client("sts") + sts_client = boto3.client("sts", region_name=region) with ( resources.files("amzn_nova_forge.iam") diff --git a/src/amzn_nova_forge/inference/forge_inference.py b/src/amzn_nova_forge/inference/forge_inference.py index 31ab108..d44ee9a 100644 --- a/src/amzn_nova_forge/inference/forge_inference.py +++ b/src/amzn_nova_forge/inference/forge_inference.py @@ -288,6 +288,7 @@ def invoke_batch( job_id=job_id, started_time=start_time, inference_output_path=inference_output_s3_path, + region=self.region, ) logger.info( f"Started batch inference job '{job_id}'. \n" @@ -348,6 +349,7 @@ def get_logs( job_id=resolved_job_id, platform=platform, started_time=int(resolved_started.timestamp() * 1000), + region=self.region, **kwargs, ) monitor.show_logs(limit=limit, start_from_head=start_from_head, end_time=end_time) diff --git a/src/amzn_nova_forge/manager/runtime_manager.py b/src/amzn_nova_forge/manager/runtime_manager.py index e164e0a..da8dbe5 100644 --- a/src/amzn_nova_forge/manager/runtime_manager.py +++ b/src/amzn_nova_forge/manager/runtime_manager.py @@ -61,6 +61,7 @@ extract_lambda_arn_from_hub_content, register_lambda_as_hub_content, ) +from amzn_nova_forge.util.subprocess_utils import _check_hyperpod_stderr from amzn_nova_forge.validation.endpoint_validator import is_sagemaker_arn # Maps TrainingMethod to (CustomizationTechnique, Peft|None) for ServerlessJobConfig @@ -715,7 +716,21 @@ def execute(self, job_config: JobConfig) -> str: s3_data_distribution_type="FullyReplicated", ), ) - trainer_config["input_data_config"] = [input_data] + input_data_config = [input_data] + if job_config.validation_data_s3_path: + validation_input = InputData( + channel_name="validation", + data_source=S3DataSource( + s3_uri=job_config.validation_data_s3_path, + s3_data_type=job_config.input_s3_data_type, + s3_data_distribution_type="FullyReplicated", + ), + ) + input_data_config.append(validation_input) + trainer_config["input_data_config"] = input_data_config + + if job_config.trainer_config_hyperparameters: + trainer_config["hyperparameters"] = job_config.trainer_config_hyperparameters model_trainer = ModelTrainer.from_recipe( **trainer_config @@ -1064,7 +1079,7 @@ def setup(self) -> None: self.execution_role = role_ref else: # Treat as a role name — look up or create - iam_client = boto3.client("iam") + iam_client = boto3.client("iam", region_name=self.region) try: self.execution_role = iam_client.get_role(RoleName=role_ref)["Role"]["Arn"] except iam_client.exceptions.NoSuchEntityException: @@ -1505,11 +1520,7 @@ def setup(self) -> None: check=True, ) - if response.stderr: - logger.error( - f"Unable to connect to HyperPod cluster {self.cluster_name}: {response.stderr}" - ) - raise RuntimeError(response.stderr) + _check_hyperpod_stderr(response.stderr) logger.info( f"Successfully connected to HyperPod cluster '{self.cluster_name}' in namespace '{self.namespace}'." @@ -1577,8 +1588,10 @@ def cleanup(self, job_name: str) -> None: check=True, ) - if response.stderr: - logger.error(f"Failed to cleanup HyperPod job: {response.stderr}") + try: + _check_hyperpod_stderr(response.stderr) + except RuntimeError as e: + logger.error(f"Failed to cleanup HyperPod job '{job_name}': {e}") except Exception as e: logger.error(f"Failed to cleanup HyperPod job '{job_name}': {str(e)}") @@ -1821,7 +1834,6 @@ def setup(self) -> None: raise def execute(self, job_config: JobConfig) -> str: - try: # Validate job name from amzn_nova_forge.validation.validator import Validator @@ -1991,7 +2003,6 @@ def execute(self, job_config: JobConfig) -> str: raise RuntimeError(f"Failed to create Bedrock customization job: {e}") from e def cleanup(self, job_id: str) -> None: - try: logger.info(f"Stopping Bedrock customization job: {job_id}") @@ -2077,6 +2088,7 @@ def __init__( self._execution_role = execution_role self.model_package_group_name = model_package_group_name self.evaluator_name = evaluator_name + self.hub_content_version: Optional[str] = None self.subnets = subnets self.security_group_ids = security_group_ids self.encrypt_inter_container_traffic = encrypt_inter_container_traffic @@ -2150,6 +2162,7 @@ def _resolve_base_model_arn(self, model: Model) -> str: hub_content_name=model.hub_content_name, hub_content_type="Model", region=self.region, + hub_content_version=self.hub_content_version, ) return hub_content["HubContentArn"] @@ -2191,7 +2204,10 @@ def validate_lambda( super().validate_lambda(data_s3_path, validation_samples) def _extract_hyperparameters( - self, recipe: Dict[str, Any], method: Optional[TrainingMethod] = None + self, + recipe: Dict[str, Any], + method: Optional[TrainingMethod] = None, + data_mixing_config: Optional[Dict[str, Any]] = None, ) -> Dict[str, str]: """Extract hyperparameters from recipe as string key-value pairs. @@ -2223,10 +2239,14 @@ def _extract_hyperparameters( # --- Training duration: v1 uses "max_epochs", v2 uses "max_steps" --- "max_epochs": "max_epochs", # v1 SFT/DPO "max_steps": "max_steps", # v2 SFT/RFT + "save_steps": "save_steps", # v2 SFT/RFT # --- v2 SFT only --- "reasoning_enabled": "reasoning_enabled", # --- DPO only: recipe key is "beta" (under dpo_cfg.beta) --- "beta": "adam_beta", + # --- Validation interval --- + "val_check_interval": "val_check_interval", + "fine_tuned_model": "fine_tuned_model", # v2 SFT/RFT } result: Dict[str, str] = {} @@ -2239,6 +2259,14 @@ def _extract(obj: Dict[str, Any]) -> None: _extract(v) _extract(recipe) + + # Datamix: inject percent fields directly from DataMixing.get_config() (full field names + # like customer_data_percent, nova_agents_percent). The recipe structure nests these under + # data_mixing.sources.* with short keys, so we use the config dict as the source of truth. + if data_mixing_config: + for k, v in data_mixing_config.items(): + result[k] = str(v) + logger.info(f"HyperParameters used for the job: {result}") return result @@ -2331,7 +2359,11 @@ def execute(self, job_config: JobConfig) -> str: ): evaluator_arn = self._register_lambda_as_hub_content(evaluator_arn) - hyperparams = self._extract_hyperparameters(recipe, method=job_config.method) + hyperparams = self._extract_hyperparameters( + recipe, method=job_config.method, data_mixing_config=job_config.data_mixing_config + ) + if job_config.trainer_config_hyperparameters: + hyperparams.update(job_config.trainer_config_hyperparameters) # For eval jobs with a reward lambda, the API requires lambda_arn + lambda_type # in HyperParameters. EvaluatorArn is NOT used for eval jobs. if job_config.method == TrainingMethod.EVALUATION and evaluator_arn: @@ -2408,6 +2440,21 @@ def execute(self, job_config: JobConfig) -> str: "RecordWrapperType": "None", } ] + if job_config.validation_data_s3_path: + validation_dataset_name = f"{job_config.job_name[:46]}-val-input" + validation_dataset = DataSet.create( + validation_dataset_name, job_config.validation_data_s3_path + ) + create_params["InputDataConfig"].append( + { + "ChannelName": "validation", + "DataSource": { + "DatasetSource": {"DatasetArn": validation_dataset.arn}, + }, + "CompressionType": "None", + "RecordWrapperType": "None", + } + ) if self.subnets or self.security_group_ids: create_params["VpcConfig"] = { "SecurityGroupIds": self.security_group_ids or [], diff --git a/src/amzn_nova_forge/model/nova_model_customizer.py b/src/amzn_nova_forge/model/nova_model_customizer.py index 5618401..ded7c8e 100644 --- a/src/amzn_nova_forge/model/nova_model_customizer.py +++ b/src/amzn_nova_forge/model/nova_model_customizer.py @@ -435,11 +435,14 @@ def _init_data_mixing(self, model: Model, method: TrainingMethod, platform: Plat if not self.data_mixing_enabled: return - # Data mixing is only supported on HyperPod for certain training methods - if platform != Platform.SMHP or method not in SUPPORTED_DATAMIXING_METHODS: + # Data mixing is only supported on HyperPod or SMTJServerless for certain training methods + if ( + platform not in (Platform.SMHP, Platform.SMTJServerless) + or method not in SUPPORTED_DATAMIXING_METHODS + ): raise ValueError( - f"Data mixing is only supported for {SUPPORTED_DATAMIXING_METHODS} training methods on SageMaker HyperPod. " - "Change platform to SMHP or change to a supported training method to use data mixing." + f"Data mixing is only supported for {SUPPORTED_DATAMIXING_METHODS} training methods on SageMaker HyperPod or SMTJServerless. " + "Change platform to SMHP or SMTJServerless or change to a supported training method to use data mixing." ) # Load recipe metadata and templates for non-evaluation methods @@ -522,6 +525,7 @@ def set_data_mixing_config(self, config: Dict[str, Any]) -> None: "model": self.model.value, "platform": self.platform, "dryRun": kwargs.get("dry_run", False), + "hasValidationData": kwargs.get("validation_data_s3_path") is not None, }, ) def train( @@ -532,6 +536,7 @@ def train( rft_lambda_arn: Optional[str] = None, rft_multiturn_infra: Optional[RFTMultiturnInfrastructure] = None, validation_data_s3_path: Optional[str] = None, + val_check_interval: Optional[int] = None, dry_run: bool = False, ) -> TrainingResult | None: """ @@ -552,7 +557,9 @@ def train( rft_lambda_arn: Optional Lambda ARN for RFT reward function (only used for RFT training methods). If passed, takes priority over rft_lambda_arn set on the RuntimeManager. rft_multiturn_infra: Optional RFT multiturn infrastructure, required for RFT_MULTITURN methods - validation_data_s3_path: Optional validation S3 path, only applicable for CPT (but is still optional for CPT) + validation_data_s3_path: Optional validation S3 path, applicable for CPT and SFT on SMTJ/SMTJServerless/SMHP, or any method on Bedrock (but is still optional) + val_check_interval: Optional positive integer controlling how often (in training steps) validation is run. + Defaults to 2500 if omitted. Only used when validation_data_s3_path is provided. dry_run: Actually starts a job if False, otherwise just performs validation. Default is False. Returns: @@ -595,6 +602,7 @@ def train( model_s3_path=self.model_path, data_mixing_enabled=self.data_mixing_enabled, holdout_data_s3_path=validation_data_s3_path, + val_check_interval=val_check_interval, config=self._build_forge_config(), region=self.region, is_multimodal=self.is_multimodal, @@ -1227,6 +1235,7 @@ def get_logs( job_id=self.job_id, platform=self.platform, started_time=int(self.job_started_time.timestamp() * 1000), + region=self.region, **kwargs, ) self.cloud_watch_log_monitor.show_logs( diff --git a/src/amzn_nova_forge/model/nova_model_customizer_util.py b/src/amzn_nova_forge/model/nova_model_customizer_util.py index 9c9718a..4740ae0 100644 --- a/src/amzn_nova_forge/model/nova_model_customizer_util.py +++ b/src/amzn_nova_forge/model/nova_model_customizer_util.py @@ -39,8 +39,8 @@ def set_output_s3_path( Raises: ValueError: If unable to construct the output S3 path """ - s3_client = boto3.client("s3") - sts_client = boto3.client("sts") + s3_client = boto3.client("s3", region_name=region) + sts_client = boto3.client("sts", region_name=region) account_id = sts_client.get_caller_identity()["Account"] # Strip trailing slash to avoid double-slash when path segments are concatenated later diff --git a/src/amzn_nova_forge/monitor/log_monitor.py b/src/amzn_nova_forge/monitor/log_monitor.py index 9aab227..32144d7 100644 --- a/src/amzn_nova_forge/monitor/log_monitor.py +++ b/src/amzn_nova_forge/monitor/log_monitor.py @@ -156,10 +156,12 @@ def get_metrics( class SMHPStrategy(PlatformStrategy): - def __init__(self, cluster_name: str, namespace: str, sagemaker_client=None): + def __init__( + self, cluster_name: str, namespace: str, sagemaker_client=None, region: Optional[str] = None + ): self.cluster_name = cluster_name self.namespace = namespace - self.sagemaker_client = sagemaker_client or boto3.client("sagemaker") + self.sagemaker_client = sagemaker_client or boto3.client("sagemaker", region_name=region) self._cluster_id: Optional[str] = None def get_log_group_name(self, job_id: str) -> str: @@ -240,8 +242,8 @@ def get_metrics( class BedrockStrategy(PlatformStrategy): - def __init__(self, bedrock_client=None): - self.bedrock_client = bedrock_client or boto3.client("bedrock") + def __init__(self, bedrock_client=None, region: Optional[str] = None): + self.bedrock_client = bedrock_client or boto3.client("bedrock", region_name=region) def get_log_group_name(self, job_id: str) -> str: # Bedrock customization jobs do not create CloudWatch logs @@ -301,13 +303,16 @@ def __init__( platform: Platform, started_time: Optional[int] = None, cloudwatch_logs_client=None, + region: Optional[str] = None, **kwargs, ): self.job_id = job_id self.platform = platform self.started_time = started_time - self.cloudwatch_logs_client = cloudwatch_logs_client or boto3.client("logs") - self.strategy = self._create_strategy(platform, **kwargs) + self.cloudwatch_logs_client = cloudwatch_logs_client or boto3.client( + "logs", region_name=region + ) + self.strategy = self._create_strategy(platform, region=region, **kwargs) self.log_group_name = self._get_log_group_name() self.log_stream_name = self._find_log_stream() self.job_status_manager = self._create_job_status_manager() @@ -315,6 +320,7 @@ def __init__( @staticmethod def _create_strategy(platform: Platform, **kwargs): + region = kwargs.get("region") if platform in _SMTJ_PLATFORMS: return SMTJStrategy() elif platform == Platform.SMHP: @@ -326,10 +332,10 @@ def _create_strategy(platform: Platform, **kwargs): logger.info(f"No namespace provided, using {namespace}` as default") if not cluster_name: raise ValueError("SMHP platform requires 'cluster_name' parameters") - return SMHPStrategy(cluster_name, namespace, sagemaker_client) + return SMHPStrategy(cluster_name, namespace, sagemaker_client, region=region) elif platform == Platform.BEDROCK: bedrock_client = kwargs.get("bedrock_client") - return BedrockStrategy(bedrock_client) + return BedrockStrategy(bedrock_client, region=region) else: raise NotImplementedError(f"Unsupported platform: {platform}") @@ -391,13 +397,17 @@ def _resolve_start_time_ms(job_id: str, platform: Platform, **kwargs) -> Optiona @classmethod @_telemetry_emitter(Feature.MONITOR, "from_job_result") - def from_job_result(cls, job_result: BaseJobResult, cloudwatch_logs_client=None): + def from_job_result( + cls, job_result: BaseJobResult, cloudwatch_logs_client=None, region: Optional[str] = None + ): + region = region or getattr(job_result, "_region", None) if job_result.platform in _SMTJ_PLATFORMS: return cls( job_id=job_result.job_id, platform=job_result.platform, started_time=int(job_result.started_time.timestamp() * 1000), cloudwatch_logs_client=cloudwatch_logs_client, + region=region, ) elif job_result.platform == Platform.SMHP: job_status_manager = cast(SMHPStatusManager, job_result.status_manager) @@ -408,6 +418,7 @@ def from_job_result(cls, job_result: BaseJobResult, cloudwatch_logs_client=None) cloudwatch_logs_client=cloudwatch_logs_client, cluster_name=job_status_manager.cluster_name, namespace=job_status_manager.namespace, + region=region, ) elif job_result.platform == Platform.BEDROCK: # Bedrock doesn't use CloudWatch logs, but we still create the monitor @@ -418,6 +429,7 @@ def from_job_result(cls, job_result: BaseJobResult, cloudwatch_logs_client=None) platform=job_result.platform, started_time=int(job_result.started_time.timestamp() * 1000), cloudwatch_logs_client=cloudwatch_logs_client, + region=region, ) else: raise NotImplementedError(f"Unsupported platform: {job_result.platform}") diff --git a/src/amzn_nova_forge/notifications/templates/smhp_notification_cf_stack.yaml b/src/amzn_nova_forge/notifications/templates/smhp_notification_cf_stack.yaml index 0e7373d..c07e5dc 100644 --- a/src/amzn_nova_forge/notifications/templates/smhp_notification_cf_stack.yaml +++ b/src/amzn_nova_forge/notifications/templates/smhp_notification_cf_stack.yaml @@ -23,10 +23,10 @@ Description: 'Nova Forge SDK: SMHP (SageMaker HyperPod) Notifications Infrastruc # 4. Subnet IDs (private subnets with NAT gateway) # 5. Security group ID (must allow outbound HTTPS to EKS API) # 6. Route table IDs (associated with your Lambda subnets) -# 7. kubectl Lambda layer ARN (see docs/job_notifications.md) +# 7. kubectl Lambda layer ARN (see docs/user-guides/job_notifications.md) # # Post-Deployment Steps: -# 1. Create EKS access entry for Lambda role (see docs/job_notifications.md) +# 1. Create EKS access entry for Lambda role (see docs/user-guides/job_notifications.md) # 2. Subscribe email addresses to SNS topic # 3. Use SDK to add job configurations to DynamoDB # ============================================================================ diff --git a/src/amzn_nova_forge/recipe/recipe_builder.py b/src/amzn_nova_forge/recipe/recipe_builder.py index e2728f0..719220b 100644 --- a/src/amzn_nova_forge/recipe/recipe_builder.py +++ b/src/amzn_nova_forge/recipe/recipe_builder.py @@ -71,6 +71,7 @@ def __init__( rft_lambda_arn: Optional[str] = None, rft_multiturn_infra: Optional["RFTMultiturnInfrastructure"] = None, validation_data_s3_path: Optional[str] = None, + val_check_interval: Optional[int] = None, eval_task: Optional[EvaluationTask] = None, subtask: Optional[str] = None, processor_config: Optional[Dict[str, Any]] = None, @@ -150,9 +151,12 @@ def __init__( raise ValueError("'rft_multiturn_infra' is required for RFT multiturn evaluation") # Validation data handling - # CPT: Uses validation_data_s3_path in recipe configuration - # Bedrock (SFT/RFT): Passes validation_data_s3_path directly to API via JobConfig - # SMTJ/SMHP: Not supported (only CPT uses it) + # CPT: Uses validation_data_s3_path in recipe configuration (all platforms) + # SFT: Uses validation_data_s3_path on SMTJ/SMTJServerless/SMHP + # Bedrock: Passes validation_data_s3_path directly to API via JobConfig (any method) + self.validation_data_s3_path = None + self.val_check_interval = val_check_interval + _sft_methods = (TrainingMethod.SFT_LORA, TrainingMethod.SFT_FULL) if method == TrainingMethod.CPT: self.validation_data_s3_path = validation_data_s3_path elif validation_data_s3_path is not None and platform == Platform.BEDROCK: @@ -160,10 +164,19 @@ def __init__( # It will be added to the Bedrock API request in BedrockRuntimeManager.execute() self.validation_data_s3_path = validation_data_s3_path logger.info(f"Validation data will be used for Bedrock job: {validation_data_s3_path}") + elif ( + validation_data_s3_path is not None + and method in _sft_methods + and platform in (Platform.SMTJ, Platform.SMTJServerless, Platform.SMHP) + ): + self.validation_data_s3_path = validation_data_s3_path + logger.info( + f"Validation data configured for {method.value} training: {validation_data_s3_path}" + ) elif validation_data_s3_path is not None: - # For SMTJ/SMHP non-CPT jobs, validation data is not supported + # For non-CPT/SFT methods on SMTJ/SMHP/SMTJServerless, validation data is not supported logger.info( - "'validation_data_s3_path' is only applicable for CPT on SMTJ/SMHP, or RFT/SFT method on Bedrock. Will ignore." + "'validation_data_s3_path' is only applicable for CPT, SFT on SMTJ/SMHP/SMTJServerless, or any method on Bedrock. Will ignore." ) # Eval @@ -340,6 +353,12 @@ def apply_user_provided_inputs_into_overrides_template(): overrides_template.setdefault("mlflow_run_name", {})["default"] = ( self.mlflow_run_name ) + else: + # Ensure mlflow placeholders in the recipe template are resolved to empty strings + # so they don't get passed as raw "{{...}}" values to the API + overrides_template.setdefault("mlflow_tracking_uri", {})["default"] = "" + overrides_template.setdefault("mlflow_experiment_name", {})["default"] = "" + overrides_template.setdefault("mlflow_run_name", {})["default"] = "" # RFT if self.method == TrainingMethod.RFT_LORA or self.method == TrainingMethod.RFT_FULL: @@ -347,14 +366,30 @@ def apply_user_provided_inputs_into_overrides_template(): self.rft_lambda_arn ) - # CPT - if self.method == TrainingMethod.CPT and self.validation_data_s3_path is not None: - overrides_template.setdefault("validation_s3_path", {})["default"] = ( - self.validation_data_s3_path + # CPT / SFT validation data + _methods_with_validation = ( + TrainingMethod.CPT, + TrainingMethod.SFT_LORA, + TrainingMethod.SFT_FULL, + ) + if self.method in _methods_with_validation and self.validation_data_s3_path is not None: + if self.method == TrainingMethod.CPT: + overrides_template.setdefault("validation_s3_path", {})["default"] = ( + self.validation_data_s3_path + ) + elif self.method in (TrainingMethod.SFT_LORA, TrainingMethod.SFT_FULL): + overrides_template.setdefault("validation_data_s3_path", {})["default"] = ( + self.validation_data_s3_path + ) + + # Val check interval + if self.val_check_interval is not None: + overrides_template.setdefault("val_check_interval", {})["default"] = ( + self.val_check_interval ) # Evaluation - elif self.method == TrainingMethod.EVALUATION: + if self.method == TrainingMethod.EVALUATION: overrides_template.setdefault("task", {})["default"] = ( self.eval_task.get_recipe_value() ) @@ -891,6 +926,7 @@ def build_and_validate( data_s3_path=overrides_template.get("data_s3_path", {}).get("default", None) if overrides_template else None, + validation_data_s3_path=self.validation_data_s3_path, validation_config=validation_config, rft_lambda_arn=overrides_template.get("reward_lambda_arn", {}).get("default", None) if overrides_template diff --git a/src/amzn_nova_forge/rft_multiturn/base_infra.py b/src/amzn_nova_forge/rft_multiturn/base_infra.py index 8d00d41..5be58ab 100644 --- a/src/amzn_nova_forge/rft_multiturn/base_infra.py +++ b/src/amzn_nova_forge/rft_multiturn/base_infra.py @@ -245,8 +245,8 @@ def create_rft_execution_role( Exception: If it fails at creating the new role or attaching policies. """ - iam_client = boto3.client("iam") - sts_client = boto3.client("sts") + iam_client = boto3.client("iam", region_name=region) + sts_client = boto3.client("sts", region_name=region) account_id = sts_client.get_caller_identity()["Account"] role_name = role_name or RFT_EXECUTION_ROLE_NAME diff --git a/src/amzn_nova_forge/rft_multiturn/ec2_infra.py b/src/amzn_nova_forge/rft_multiturn/ec2_infra.py index 7ee9f8d..9d6cf7d 100644 --- a/src/amzn_nova_forge/rft_multiturn/ec2_infra.py +++ b/src/amzn_nova_forge/rft_multiturn/ec2_infra.py @@ -381,7 +381,7 @@ def validate_platform(self): # Get the instance profile and ensure it has RFT permissions iam_client = boto3.client("iam", region_name=self.region) - sts_client = boto3.client("sts") + sts_client = boto3.client("sts", region_name=self.region) account_id = sts_client.get_caller_identity()["Account"] instance_profile_arn = instance["IamInstanceProfile"]["Arn"] diff --git a/src/amzn_nova_forge/rft_multiturn/rft_multiturn.py b/src/amzn_nova_forge/rft_multiturn/rft_multiturn.py index 2e8c9f0..396ec2b 100644 --- a/src/amzn_nova_forge/rft_multiturn/rft_multiturn.py +++ b/src/amzn_nova_forge/rft_multiturn/rft_multiturn.py @@ -243,7 +243,7 @@ def __init__( self.recipe_cache_dir = os.path.join(self.workspace_dir, ".nova_rft_recipes") # Initialize platform-specific infrastructure - sts_client = boto3.client("sts") + sts_client = boto3.client("sts", region_name=self.region) account_id = sts_client.get_caller_identity()["Account"] # Validate and set python_venv_name based on platform @@ -968,7 +968,7 @@ def _load_existing_stack(self, stack_name: str): response = cfn_client.describe_stacks(StackName=stack_name) outputs = response["Stacks"][0].get("Outputs", []) - sts_client = boto3.client("sts") + sts_client = boto3.client("sts", region_name=self.region) account_id = sts_client.get_caller_identity()["Account"] rollout_function_name = self._get_output(outputs, "RolloutFunctionName") diff --git a/src/amzn_nova_forge/telemetry/telemetry_logging.py b/src/amzn_nova_forge/telemetry/telemetry_logging.py index 3dafc64..123063d 100644 --- a/src/amzn_nova_forge/telemetry/telemetry_logging.py +++ b/src/amzn_nova_forge/telemetry/telemetry_logging.py @@ -216,11 +216,11 @@ def _requests_helper(url, timeout): return response -def _get_accountId(): +def _get_accountId(region: Optional[str] = None): """Return the account ID from the boto session""" try: - sts = boto3.client("sts") + sts = boto3.client("sts", region_name=region) return sts.get_caller_identity()["Account"] except Exception: return None diff --git a/src/amzn_nova_forge/trainer/forge_trainer.py b/src/amzn_nova_forge/trainer/forge_trainer.py index 2379059..43ab938 100644 --- a/src/amzn_nova_forge/trainer/forge_trainer.py +++ b/src/amzn_nova_forge/trainer/forge_trainer.py @@ -82,6 +82,7 @@ def __init__( model_s3_path: Optional[str] = None, data_mixing_enabled: bool = False, holdout_data_s3_path: Optional[str] = None, + val_check_interval: Optional[int] = None, config: Optional[ForgeConfig] = None, region: Optional[str] = None, is_multimodal: Optional[bool] = None, @@ -95,7 +96,25 @@ def __init__( self.model_s3_path = model_s3_path self.holdout_data_s3_path = holdout_data_s3_path self.hub_content_version = hub_content_version + if hub_content_version and hasattr(infra, "hub_content_version"): + infra.hub_content_version = hub_content_version self._enable_batch_sample_tracing = enable_batch_sample_tracing + self.val_check_interval: Optional[int] = None + if val_check_interval is not None: + if ( + isinstance(val_check_interval, bool) + or not isinstance(val_check_interval, int) + or val_check_interval < 1 + ): + raise ValueError( + f"val_check_interval must be a positive integer, got: {val_check_interval}" + ) + self.val_check_interval = val_check_interval + if val_check_interval is not None and holdout_data_s3_path is None: + logger.warning( + "val_check_interval is set but no holdout_data_s3_path provided. val_check_interval will have no effect without validation data." + ) + self._config = config or ForgeConfig() self.region = region or boto3.session.Session().region_name or DEFAULT_REGION @@ -148,13 +167,17 @@ def __init__( else: self._is_multimodal = False - # Data mixing setup (SMHP only, CPT/SFT methods) + # Data mixing setup (SMHP and SMTJServerless, CPT/SFT methods) self.data_mixing: Optional[DataMixing] = None if data_mixing_enabled: - if self._platform != Platform.SMHP or method not in SUPPORTED_DATAMIXING_METHODS: + _datamix_platforms = (Platform.SMHP, Platform.SMTJServerless) + if ( + self._platform not in _datamix_platforms + or method not in SUPPORTED_DATAMIXING_METHODS + ): raise ValueError( f"Data mixing is only supported for {SUPPORTED_DATAMIXING_METHODS} " - "training methods on SageMaker HyperPod." + "training methods on SageMaker HyperPod or SMTJServerless." ) self.data_mixing = DataMixing() ( @@ -196,6 +219,7 @@ def __init__( "model": self.model.value, "platform": self._platform, "dryRun": kwargs.get("dry_run", False), + "hasValidationData": self.holdout_data_s3_path is not None, }, ) def train( @@ -253,6 +277,7 @@ def train( model_path=self.model_s3_path, rft_lambda_arn=rft_lambda_arn, 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, @@ -293,9 +318,18 @@ def train( else "S3Prefix", } + hp: Dict[str, str] = {} + if self.val_check_interval is not None: + hp["val_check_interval"] = str(self.val_check_interval) + if hp: + job_config_params["trainer_config_hyperparameters"] = hp + if self._platform in (Platform.BEDROCK, Platform.SMTJServerless): job_config_params["method"] = self.method + if self._platform == Platform.SMTJServerless and self.data_mixing: + job_config_params["data_mixing_config"] = self.data_mixing.get_config() + job_id = self.infra.execute(job_config=JobConfig(**job_config_params)) training_result: TrainingResult @@ -309,7 +343,9 @@ def train( job_name=job_id, infra=self.infra, output_s3_path=resolved_output_s3_path, + region=self.region, ), + region=self.region, ) elif self._platform is Platform.BEDROCK: training_result = BedrockTrainingResult( @@ -321,6 +357,7 @@ def train( checkpoint_s3_path=None, output_s3_path=resolved_output_s3_path, ), + region=self.region, ) else: cluster_name = cast(SMHPRuntimeManager, self.infra).cluster_name @@ -334,9 +371,11 @@ def train( job_name=unique_job_name, infra=self.infra, output_s3_path=resolved_output_s3_path, + region=self.region, ), cluster_name=cluster_name, namespace=namespace, + region=self.region, ) logger.info(f"Started job '{training_result.job_id}'.") @@ -388,6 +427,7 @@ def get_logs( job_id=resolved_job_id, platform=self._platform, started_time=int(resolved_started.timestamp() * 1000), + region=self.region, **kwargs, ) monitor.show_logs(limit=limit, start_from_head=start_from_head, end_time=end_time) diff --git a/src/amzn_nova_forge/util/bedrock.py b/src/amzn_nova_forge/util/bedrock.py index 66488a2..6f6706a 100644 --- a/src/amzn_nova_forge/util/bedrock.py +++ b/src/amzn_nova_forge/util/bedrock.py @@ -172,7 +172,9 @@ def wait_for_model_ready( time.sleep(poll_interval) -def check_deployment_status(deployment_arn: str, platform: DeployPlatform) -> Optional[str]: +def check_deployment_status( + deployment_arn: str, platform: DeployPlatform, region: Optional[str] = None +) -> Optional[str]: """ Checks the current status of a Bedrock deployment. @@ -185,7 +187,7 @@ def check_deployment_status(deployment_arn: str, platform: DeployPlatform) -> Op """ status = None - bedrock_client = boto3.client("bedrock") + bedrock_client = boto3.client("bedrock", region_name=region) if platform == DeployPlatform.BEDROCK_OD: try: status = bedrock_client.get_custom_model_deployment( @@ -256,7 +258,9 @@ def get_required_bedrock_update_permissions( return [] -def check_existing_deployment(endpoint_name: str, platform: DeployPlatform) -> Optional[str]: +def check_existing_deployment( + endpoint_name: str, platform: DeployPlatform, region: Optional[str] = None +) -> Optional[str]: """ Check if a deployment with the given name exists. @@ -267,7 +271,7 @@ def check_existing_deployment(endpoint_name: str, platform: DeployPlatform) -> O Returns: Optional[str]: The ARN of the existing deployment if found, None otherwise """ - bedrock_client = boto3.client("bedrock") + bedrock_client = boto3.client("bedrock", region_name=region) try: if platform == DeployPlatform.BEDROCK_OD: @@ -290,7 +294,7 @@ def check_existing_deployment(endpoint_name: str, platform: DeployPlatform) -> O def delete_existing_deployment( - deployment_arn: str, platform: DeployPlatform, endpoint_name: str + deployment_arn: str, platform: DeployPlatform, endpoint_name: str, region: Optional[str] = None ) -> None: """ Delete an existing deployment and wait for completion. @@ -303,7 +307,7 @@ def delete_existing_deployment( Raises: Exception: If deletion fails or times out """ - bedrock_client = boto3.client("bedrock") + bedrock_client = boto3.client("bedrock", region_name=region) try: logger.info(f"Deleting existing deployment '{endpoint_name}'...") @@ -369,7 +373,7 @@ def delete_existing_deployment( def update_provisioned_throughput_model( - deployment_arn: str, new_model_arn: str, endpoint_name: str + deployment_arn: str, new_model_arn: str, endpoint_name: str, region: Optional[str] = None ) -> None: """ Update a Provisioned Throughput deployment to use a new custom model. @@ -382,7 +386,7 @@ def update_provisioned_throughput_model( Raises: Exception: If update fails """ - bedrock_client = boto3.client("bedrock") + bedrock_client = boto3.client("bedrock", region_name=region) try: logger.info(f"Updating PT deployment '{endpoint_name}' to new model...") diff --git a/src/amzn_nova_forge/util/checkpoint_util.py b/src/amzn_nova_forge/util/checkpoint_util.py index f565c30..e18bb0f 100644 --- a/src/amzn_nova_forge/util/checkpoint_util.py +++ b/src/amzn_nova_forge/util/checkpoint_util.py @@ -90,6 +90,7 @@ def extract_checkpoint_path_from_job_output( output_s3_path: Optional[str] = None, job_id: Optional[str] = None, job_result: Optional[TrainingResult] = None, + region: Optional[str] = None, ) -> str: """ Extracts the model checkpoint path from a training job's output. @@ -120,7 +121,7 @@ def extract_checkpoint_path_from_job_output( """ from amzn_nova_forge.core.result.job_result import JobStatus - s3_client = boto3.client("s3") + s3_client = boto3.client("s3", region_name=region) # If job_result is provided, check if job is completed if job_result is not None: diff --git a/src/amzn_nova_forge/util/dataset_writer.py b/src/amzn_nova_forge/util/dataset_writer.py index acdf943..5aab318 100644 --- a/src/amzn_nova_forge/util/dataset_writer.py +++ b/src/amzn_nova_forge/util/dataset_writer.py @@ -22,7 +22,7 @@ import os import tempfile from pathlib import Path -from typing import Dict, Iterator +from typing import Dict, Iterator, Optional import boto3 @@ -93,7 +93,9 @@ def save_to_local(save_path: str, dataset_iter: Iterator[Dict], is_jsonl: bool) raise DatasetWriteError(f"Failed to write to local file {save_path}: {e}") @staticmethod - def save_to_s3(save_path: str, dataset_iter: Iterator[Dict], is_jsonl: bool) -> None: + def save_to_s3( + save_path: str, dataset_iter: Iterator[Dict], is_jsonl: bool, region: Optional[str] = None + ) -> None: """ Stream dataset to S3 without loading all data into memory. Uses a temporary file and boto3's upload_file for efficient multipart upload. @@ -135,7 +137,7 @@ def save_to_s3(save_path: str, dataset_iter: Iterator[Dict], is_jsonl: bool) -> DatasetWriter.save_to_local(tmp_path, dataset_iter, is_jsonl) # Upload the temporary file to S3 - s3_client = boto3.client("s3") + s3_client = boto3.client("s3", region_name=region) s3_client.upload_file( tmp_path, bucket, key, ExtraArgs={"ContentType": "application/json"} ) diff --git a/src/amzn_nova_forge/util/mlflow.py b/src/amzn_nova_forge/util/mlflow.py index f5ff5e0..efe4309 100644 --- a/src/amzn_nova_forge/util/mlflow.py +++ b/src/amzn_nova_forge/util/mlflow.py @@ -46,10 +46,7 @@ def get_default_mlflow_tracking_uri(region_name: Optional[str] = None) -> Option try: # Create SageMaker client - if region_name: - sagemaker_client = boto3.client("sagemaker", region_name=region_name) - else: - sagemaker_client = boto3.client("sagemaker") + sagemaker_client = boto3.client("sagemaker", region_name=region_name) # Find if DefaultMLFlowApp exists try: @@ -202,10 +199,7 @@ def validate_mlflow_arn_exists( try: # Create SageMaker client - if region_name: - sagemaker_client = boto3.client("sagemaker", region_name=region_name) - else: - sagemaker_client = boto3.client("sagemaker") + sagemaker_client = boto3.client("sagemaker", region_name=region_name) # Check if the MLflow resource exists try: diff --git a/src/amzn_nova_forge/util/recipe.py b/src/amzn_nova_forge/util/recipe.py index 8957566..ac55ade 100644 --- a/src/amzn_nova_forge/util/recipe.py +++ b/src/amzn_nova_forge/util/recipe.py @@ -108,7 +108,10 @@ def _validate_extension(path: str, extension: str) -> None: def load_file_content( - file_path: str, extension: Optional[str] = None, encoding: Optional[str] = "utf-8" + file_path: str, + extension: Optional[str] = None, + encoding: Optional[str] = "utf-8", + region: Optional[str] = None, ): """ Stream file content line by line from S3 or local filesystem. @@ -134,7 +137,7 @@ def load_file_content( if s3_parts: bucket, key = s3_parts try: - s3 = boto3.client("s3") + s3 = boto3.client("s3", region_name=region) response = s3.get_object(Bucket=bucket, Key=key) # Stream from S3 using iter_lines for line in response["Body"].iter_lines(): @@ -395,7 +398,7 @@ def get_hub_recipe_metadata( raise ValueError(f"{method.name} using {instance_type} is not supported on {platform.name}") -def _get_aws_account_id() -> str: +def _get_aws_account_id(region: Optional[str] = None) -> str: """ Get the AWS account ID from current credentials. @@ -403,7 +406,7 @@ def _get_aws_account_id() -> str: AWS account ID string """ try: - sts = boto3.client("sts") + sts = boto3.client("sts", region_name=region) response = sts.get_caller_identity() return response["Account"] except Exception as e: @@ -442,9 +445,9 @@ def _download_from_s3_or_access_point(uri: str, region: Optional[str] = None) -> """ # Replace {customer_id} placeholder if present - s3 = boto3.client("s3", region_name=region) if region else boto3.client("s3") + s3 = boto3.client("s3", region_name=region) - current_account = _get_aws_account_id() + current_account = _get_aws_account_id(region=region) formatted_uri = _replace_customer_id_placeholder(uri, current_account) # Check if this is an access point ARN (with or without s3:// prefix) arn_to_check = ( @@ -468,7 +471,7 @@ def _download_from_s3_or_access_point(uri: str, region: Optional[str] = None) -> raise ValueError( f"Failed to download from S3 Access Point {e}" f"\nVerify if account {current_account} has Forge subscription. Refer: https://docs.aws.amazon.com/sagemaker/latest/dg/nova-forge.html#nova-forge-prereq-access" - f"\nAlso ensure your IAM role has the right 's3:GetObject'. Refer DataMixingForgeRecipes permission in docs/iam_setup.md" + f"\nAlso ensure your IAM role has the right 's3:GetObject'. Refer DataMixingForgeRecipes permission in docs/user-guides/iam_setup.md" f" or set data_mixing = False" ) diff --git a/src/amzn_nova_forge/util/s3_utils.py b/src/amzn_nova_forge/util/s3_utils.py index 1c46b6b..e6940ce 100644 --- a/src/amzn_nova_forge/util/s3_utils.py +++ b/src/amzn_nova_forge/util/s3_utils.py @@ -64,7 +64,7 @@ def get_dataprep_bucket_name( if region is None: region = boto3.session.Session().region_name or "us-east-1" if account_id is None: - account_id = boto3.client("sts").get_caller_identity()["Account"] + account_id = boto3.client("sts", region_name=region).get_caller_identity()["Account"] return f"{DATAPREP_BUCKET_PREFIX}-{account_id}-{region}" diff --git a/src/amzn_nova_forge/util/sagemaker.py b/src/amzn_nova_forge/util/sagemaker.py index dc9044b..c0b12a5 100644 --- a/src/amzn_nova_forge/util/sagemaker.py +++ b/src/amzn_nova_forge/util/sagemaker.py @@ -205,7 +205,7 @@ def _get_sagemaker_inference_image(region: str) -> str: def get_model_artifacts( - job_name: str, infra: RuntimeManager, output_s3_path: str + job_name: str, infra: RuntimeManager, output_s3_path: str, region: Optional[str] = None ) -> ModelArtifacts: """ Retrieve model artifacts for a job @@ -221,7 +221,7 @@ def get_model_artifacts( Raises: Exception: If unable to obtain job artifact information """ - sagemaker_client = boto3.client("sagemaker") + sagemaker_client = boto3.client("sagemaker", region_name=region) if infra.platform in (Platform.SMTJ, Platform.SMTJServerless): response = sagemaker_client.describe_training_job(TrainingJobName=job_name) @@ -294,10 +294,7 @@ def get_cluster_instance_info( Raises: Exception: If unable to describe the cluster """ - if region is None: - sagemaker_client = boto3.client("sagemaker") - else: - sagemaker_client = boto3.client("sagemaker", region_name=region) + sagemaker_client = boto3.client("sagemaker", region_name=region) try: response = sagemaker_client.describe_cluster(ClusterName=cluster_name) diff --git a/src/amzn_nova_forge/util/subprocess_utils.py b/src/amzn_nova_forge/util/subprocess_utils.py new file mode 100644 index 0000000..8b5ac7b --- /dev/null +++ b/src/amzn_nova_forge/util/subprocess_utils.py @@ -0,0 +1,88 @@ +# 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. +"""Utilities for handling subprocess stderr output from HyperPod CLI commands.""" + +import logging +import re + +logger = logging.getLogger(__name__) + +# Matches the Python warning file:line: prefix format +# e.g. "/path/to/urllib3/__init__.py:35: " or "module.py:10: " +_WARNING_FILE_PREFIX = r"(?:(?:\S+/)*\S+(?:\.\w+)?:\d*: )" + +# Benign warning class names that should not be treated as errors +_BENIGN_WARNING_NAMES = [ + r"NotOpenSSLWarning:", + r"DeprecationWarning:", + r"FutureWarning:", + r"InsecureRequestWarning:", + r"ResourceWarning:", + r"UserWarning:", + r"urllib3\.[a-zA-Z_]+Warning:", +] + +# Each pattern matches either: +# 1. file:line: WarningClass: (standard Python warning format) +# 2. WarningClass: at line start (bare warning without file prefix) +_BENIGN_STDERR_PATTERNS = [ + re.compile(rf"^(?:{_WARNING_FILE_PREFIX})?{name}", re.MULTILINE) + for name in _BENIGN_WARNING_NAMES +] + + +def _check_hyperpod_stderr(stderr: str) -> None: + """Check HyperPod CLI stderr output and raise if it contains real errors. + + The HyperPod CLI can exit 0 while writing real errors to stderr + (e.g. "Cluster with name not found"). This function distinguishes + benign warnings (urllib3, deprecation, etc.) from actual errors. + + Args: + stderr: The stderr output from a subprocess call. + + Raises: + RuntimeError: If stderr contains lines that don't match known benign patterns. + """ + if not stderr or not stderr.strip(): + return + + non_benign_lines = [] + prev_stripped = None + last_benign_stripped = None + for line in stderr.strip().splitlines(): + stripped = line.strip() + prev_line_benign = ( + last_benign_stripped is not None and last_benign_stripped == prev_stripped + ) + last_benign_stripped = None # one-shot: only protects the immediately next line + prev_stripped = stripped + + if not stripped: + continue + if any(pattern.search(stripped) for pattern in _BENIGN_STDERR_PATTERNS): + last_benign_stripped = stripped + continue + # Python warnings module indents the source line following the warning + if prev_line_benign and line.startswith((" ", "\t")): + continue + non_benign_lines.append(stripped) + + if not non_benign_lines: + logger.debug("HyperPod CLI stderr (benign warnings only): %s", stderr.strip()) + return + + error_output = "\n".join(non_benign_lines) + logger.error("HyperPod CLI stderr contains errors: %s", error_output) + raise RuntimeError(stderr.strip()) diff --git a/src/amzn_nova_forge/validation/validator.py b/src/amzn_nova_forge/validation/validator.py index 18669ad..eb7c4a0 100644 --- a/src/amzn_nova_forge/validation/validator.py +++ b/src/amzn_nova_forge/validation/validator.py @@ -131,14 +131,14 @@ def _resolve_execution_role(infra: Optional[RuntimeManager]) -> str: return execution_role @staticmethod - def _is_cross_account_role(execution_role_arn: str) -> bool: + def _is_cross_account_role(execution_role_arn: str, region: Optional[str] = None) -> bool: """Check if execution role is in a different account.""" try: # Extract account from role ARN: arn:aws:iam::ACCOUNT:role/RoleName role_account = execution_role_arn.split(":")[4] # Get current account - sts_client = boto3.client("sts") + sts_client = boto3.client("sts", region_name=region) current_account = sts_client.get_caller_identity()["Account"] return role_account != current_account @@ -160,7 +160,7 @@ def _access_cross_account_role( """ try: # Assume the cross-account role - sts_client = boto3.client("sts") + sts_client = boto3.client("sts", region_name=region_name) assumed_role = sts_client.assume_role( RoleArn=execution_role_arn, RoleSessionName="SageMakerValidation" ) @@ -502,7 +502,9 @@ def _validate_smtj_iam( role_name = execution_role.split("/")[-1] - is_cross_account_role = Validator._is_cross_account_role(execution_role) + is_cross_account_role = Validator._is_cross_account_role( + execution_role, region=region_name + ) # Check if this is a cross-account role if is_cross_account_role: @@ -615,7 +617,9 @@ def _validate_smtj_iam( errors.append(f"Failed to validate execution role permissions: {str(e)}") except Exception as e: # For cross-account roles, silently skip IAM validation on unexpected errors - if execution_role and Validator._is_cross_account_role(execution_role): + if execution_role and Validator._is_cross_account_role( + execution_role, region=region_name + ): return if not execution_role: @@ -1037,6 +1041,9 @@ def get_recipe_value(data: Dict[str, Any], key_to_find: str) -> Any: # Validate proper types are used if "type" in override_metadata: + # 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 python_type_name = TYPE_ALIASES.get(override_metadata["type"]) if python_type_name is None: continue @@ -1178,6 +1185,7 @@ def validate( overrides_template: Dict[str, Any], output_s3_path: Optional[str] = None, data_s3_path: Optional[str] = None, + validation_data_s3_path: Optional[str] = None, validation_config: Optional[ValidationConfig] = None, rft_lambda_arn: Optional[str] = None, eval_task: Optional[EvaluationTask] = None, @@ -1197,6 +1205,7 @@ def validate( overrides_template: Dict containing recipe constraints output_s3_path: Output S3 data path data_s3_path: Input S3 data path + validation_data_s3_path: Optional S3 path for validation data validation_config: Optional configuration to determine which resource validation checks to perform rft_lambda_arn: Optional Lambda ARN for RFT eval_task: Optional evaluation task @@ -1231,6 +1240,15 @@ def validate( # Recipe validation if config.recipe: + # Validate validation_data_s3_path S3 format if provided + if validation_data_s3_path: + s3_parts = _parse_s3_uri(validation_data_s3_path) + if not s3_parts: + errors.append( + f"Invalid S3 path for validation_data_s3_path: {validation_data_s3_path}. " + "Expected format: s3://bucket-name/path/to/data.jsonl" + ) + cls._validate_recipe( recipe=recipe, overrides_template=overrides_template, diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 0b6d9f0..d1ebff9 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -53,3 +53,18 @@ def _mock_telemetry(request): return with patch("amzn_nova_forge.telemetry.telemetry_logging._send_telemetry_request"): yield + + +@pytest.fixture(autouse=True) +def _mock_dataprep_bucket(): + """Prevent OutputPathResolver from making real STS/S3 calls when + rebasing non-S3 paths to the data-prep bucket during tests. + """ + with ( + patch( + "amzn_nova_forge.dataset.data_state.get_dataprep_bucket_name", + return_value="sagemaker-forge-dataprep-123456789012-us-east-1", + ), + patch("amzn_nova_forge.dataset.data_state.ensure_bucket_exists"), + ): + yield diff --git a/tests/unit/dataset/test_auto_output_path.py b/tests/unit/dataset/test_auto_output_path.py new file mode 100644 index 0000000..17ef84a --- /dev/null +++ b/tests/unit/dataset/test_auto_output_path.py @@ -0,0 +1,341 @@ +# 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 auto-derived output paths when input is not on S3. + +Verifies that filter operations auto-derive S3 output paths using the +default data-prep bucket when no explicit output_path is provided and +input is local or HuggingFace. +""" + +from __future__ import annotations + +import re +from unittest.mock import MagicMock, patch + +import pytest + +from amzn_nova_forge.core.enums import FilterMethod +from amzn_nova_forge.dataset.data_state import DataLocation, DataState +from amzn_nova_forge.dataset.operations.base import OperationResult + +FAKE_DATAPREP_BUCKET = "sagemaker-forge-dataprep-123456789012-us-east-1" +_SESSION_RE = r"\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2}" + + +def _assert_auto_derived_path(output_path, expected_parent, expected_stem, expected_method): + """Assert output_path matches ``///_output/``.""" + parent = expected_parent.rstrip("/") + pattern = ( + rf"^{re.escape(parent)}/{re.escape(expected_stem)}" + rf"/{_SESSION_RE}/{re.escape(expected_method)}_output/$" + ) + assert re.match(pattern, output_path), ( + f"Auto-derived output_path does not match expected pattern.\n" + f" Got: '{output_path}'\n" + f" Expected: '{parent}/{expected_stem}//{expected_method}_output/'" + ) + + +def _s3_state(path="s3://bucket/out/", fmt="parquet"): + """Shorthand for an S3 DataState used as a mock operation result.""" + return DataState(path=path, format=fmt, location=DataLocation.S3, generator=lambda: iter([])) + + +def _mock_op(*output_paths): + """Build a mock filter operation returning successive S3 output states.""" + op = MagicMock() + op.execute.side_effect = [ + OperationResult(status="SUCCEEDED", output_state=_s3_state(p)) for p in output_paths + ] + return op + + +class TestPrepareInputGuardsDirectCalls: + """prepare_input() without output_path raises ValueError for non-S3 inputs. + + In normal flow, execute() provides output_path via OutputPathResolver. + These tests guard against direct misuse of prepare_input(). + """ + + @pytest.fixture() + def filter_op(self): + from amzn_nova_forge.dataset.operations.filter_operation import get_filter_operation + + return get_filter_operation(FilterMethod.DEFAULT_TEXT_FILTER) + + @pytest.mark.parametrize( + "path, fmt, location", + [ + ("/home/user/data/train.jsonl", "jsonl", DataLocation.LOCAL), + ("/home/user/data/train.csv", "csv", DataLocation.LOCAL), + ("hf://fake-org/fake-dataset/train", "huggingface", DataLocation.HUGGINGFACE), + ], + ids=["local-jsonl", "local-csv", "huggingface"], + ) + def test_non_s3_input_without_output_path_raises(self, filter_op, path, fmt, location): + """Direct call to prepare_input() without output_path still raises ValueError.""" + state = DataState( + path=path, + format=fmt, + location=location, + generator=lambda: iter([{"text": "hello"}]), + ) + with pytest.raises(ValueError, match="require an S3 output path"): + filter_op.prepare_input(state) + + def test_s3_input_passes_through(self, filter_op): + """S3 input in compatible format → returned unchanged, no upload.""" + state = DataState( + path="s3://my-bucket/data/train.jsonl", + format="jsonl", + location=DataLocation.S3, + generator=lambda: iter([{"text": "hello"}]), + ) + + result = filter_op.prepare_input(state) + + assert result is state + + def test_explicit_output_path_still_honored(self, filter_op): + """Explicit output_path → used as S3 base (backward compat).""" + state = DataState( + path="/home/user/data/train.jsonl", + format="jsonl", + location=DataLocation.LOCAL, + generator=lambda: iter([{"text": "hello"}]), + ) + + with patch( + "amzn_nova_forge.dataset.operations.filter_operation.upload_local_file_to_s3", + return_value="s3://explicit-bucket/uploaded.jsonl", + ) as mock_upload: + filter_op.prepare_input(state, output_path="s3://explicit-bucket/filtered/") + + s3_base = mock_upload.call_args[0][1] + assert s3_base.startswith("s3://explicit-bucket"), ( + f"Explicit output_path should be used, got: '{s3_base}'" + ) + + +class TestExecuteAutoDerivesOutputPath: + """execute() should produce S3 output paths for non-S3 inputs. + + Auto-derived path format: s3://///_output/ + """ + + @staticmethod + def _execute_single_filter(load_path, method=FilterMethod.DEFAULT_TEXT_FILTER): + """Queue one filter on a stub loader, execute with a mock op, return the output_path.""" + from amzn_nova_forge.dataset.jsonl_dataset_loader import JSONLDatasetLoader + + op = _mock_op(f"s3://{FAKE_DATAPREP_BUCKET}/out/") + loader = JSONLDatasetLoader() + loader._load_path = load_path + loader.dataset = lambda: iter([{"text": "hello"}]) + loader.filter(method=method, text_field="text") + + with ( + patch("amzn_nova_forge.dataset.dataset_loader.get_filter_operation", return_value=op), + patch( + "amzn_nova_forge.dataset.data_state.get_dataprep_bucket_name", + return_value=FAKE_DATAPREP_BUCKET, + ), + patch("amzn_nova_forge.dataset.data_state.ensure_bucket_exists"), + ): + loader.execute() + + return op.execute.call_args.kwargs["output_path"] + + @pytest.mark.parametrize( + "load_path, expected_parent", + [ + ("/home/user/data/train.jsonl", f"s3://{FAKE_DATAPREP_BUCKET}"), + ("hf://fake-org/fake-dataset/train", f"s3://{FAKE_DATAPREP_BUCKET}"), + ], + ids=["local", "huggingface"], + ) + def test_non_s3_input(self, load_path, expected_parent): + """Non-S3 load path → output under data-prep bucket.""" + output_path = self._execute_single_filter(load_path) + _assert_auto_derived_path( + output_path, + expected_parent=expected_parent, + expected_stem="train", + expected_method="default_text_filter", + ) + + def test_s3_input_uses_load_path_parent(self): + """S3 load path → output derived from the S3 parent (existing behavior).""" + output_path = self._execute_single_filter("s3://my-bucket/data/train.jsonl") + _assert_auto_derived_path( + output_path, + expected_parent="s3://my-bucket/data", + expected_stem="train", + expected_method="default_text_filter", + ) + + +class TestThreeFilterChainOutputPathThreading: + """Verify output_path and state.path threading across three-filter chains. + + Key rules: + - output_path: auto-derived from resolver (anchored to load path), or + explicit if user-provided. Explicit paths do NOT shift the resolver base. + - state.path (input to next op): always the previous operation's output. + """ + + @staticmethod + def _execute_three_filters(load_path, filter_specs): + """Queue three filters, execute, return list of (output_path, input_state_path) per call. + + ``filter_specs`` is a list of (method, explicit_output_path_or_None, mock_result_path). + """ + from amzn_nova_forge.dataset.jsonl_dataset_loader import JSONLDatasetLoader + + mock_op = _mock_op(*[spec[2] for spec in filter_specs]) + + loader = JSONLDatasetLoader() + loader._load_path = load_path + loader.dataset = lambda: iter([{"text": "hello"}]) + + for method, explicit_path, _ in filter_specs: + kwargs = {"text_field": "text"} + if explicit_path is not None: + kwargs["output_path"] = explicit_path + loader.filter(method=method, **kwargs) + + with ( + patch( + "amzn_nova_forge.dataset.dataset_loader.get_filter_operation", return_value=mock_op + ), + patch( + "amzn_nova_forge.dataset.data_state.get_dataprep_bucket_name", + return_value=FAKE_DATAPREP_BUCKET, + ), + patch("amzn_nova_forge.dataset.data_state.ensure_bucket_exists"), + ): + loader.execute() + + return [ + (call.kwargs["output_path"], call.kwargs["state"].path) + for call in mock_op.execute.call_args_list + ] + + @pytest.mark.parametrize( + "load_path", + ["/home/user/data/train.jsonl", "hf://fake-org/fake-dataset/train"], + ids=["local", "huggingface"], + ) + def test_auto_explicit_auto(self, load_path): + """f1(auto) → f2(explicit) → f3(auto): f3 output from resolver, f3 input from f2.""" + dp = f"s3://{FAKE_DATAPREP_BUCKET}" + calls = self._execute_three_filters( + load_path, + [ + (FilterMethod.DEFAULT_TEXT_FILTER, None, f"{dp}/f1-out/"), + (FilterMethod.EXACT_DEDUP, "s3://user-bucket/custom/", "s3://user-bucket/custom/"), + (FilterMethod.FUZZY_DEDUP, None, f"{dp}/f3-out/"), + ], + ) + + f1_output, _ = calls[0] + f2_output, _ = calls[1] + f3_output, f3_input = calls[2] + + # f1: auto-derived under data-prep bucket + _assert_auto_derived_path(f1_output, dp, "train", "default_text_filter") + + # f2: explicit path honoured + assert f2_output == "s3://user-bucket/custom/" + + # f3: auto-derived under data-prep bucket (NOT relative to f2) + _assert_auto_derived_path(f3_output, dp, "train", "fuzzy_dedup") + assert "user-bucket" not in f3_output + + # f3 reads from f2's output + assert f3_input == "s3://user-bucket/custom/" + + @pytest.mark.parametrize( + "load_path", + ["/home/user/data/train.jsonl", "hf://fake-org/fake-dataset/train"], + ids=["local", "huggingface"], + ) + def test_auto_explicit_explicit(self, load_path): + """f1(auto) → f2(explicit A) → f3(explicit B): each uses its own path.""" + dp = f"s3://{FAKE_DATAPREP_BUCKET}" + calls = self._execute_three_filters( + load_path, + [ + (FilterMethod.DEFAULT_TEXT_FILTER, None, f"{dp}/f1-out/"), + (FilterMethod.EXACT_DEDUP, "s3://user-bucket/dedup/", "s3://user-bucket/dedup/"), + (FilterMethod.FUZZY_DEDUP, "s3://user-bucket/fuzzy/", "s3://user-bucket/fuzzy/"), + ], + ) + + f1_output, _ = calls[0] + f2_output, _ = calls[1] + f3_output, f3_input = calls[2] + + # f1: auto-derived + _assert_auto_derived_path(f1_output, dp, "train", "default_text_filter") + + # f2: its own explicit path + assert f2_output == "s3://user-bucket/dedup/" + + # f3: its own explicit path (NOT f2's) + assert f3_output == "s3://user-bucket/fuzzy/" + + # f3 reads from f2's output + assert f3_input == "s3://user-bucket/dedup/" + + @pytest.mark.parametrize( + "load_path", + ["/home/user/data/train.jsonl", "hf://fake-org/fake-dataset/train"], + ids=["local", "huggingface"], + ) + def test_explicit_auto(self, load_path): + """f1(explicit) → f2(auto): f2 output from resolver (data-prep bucket), f2 input from f1.""" + dp = f"s3://{FAKE_DATAPREP_BUCKET}" + calls = self._execute_three_filters( + load_path, + [ + ( + FilterMethod.DEFAULT_TEXT_FILTER, + "s3://user-bucket/first/", + "s3://user-bucket/first/", + ), + (FilterMethod.EXACT_DEDUP, None, f"{dp}/f2-out/"), + (FilterMethod.FUZZY_DEDUP, None, f"{dp}/f3-out/"), + ], + ) + + f1_output, _ = calls[0] + f2_output, f2_input = calls[1] + f3_output, f3_input = calls[2] + + # f1: explicit path honoured + assert f1_output == "s3://user-bucket/first/" + + # f2: auto-derived under data-prep bucket (NOT relative to f1's explicit path) + _assert_auto_derived_path(f2_output, dp, "train", "exact_dedup_filter") + assert "user-bucket" not in f2_output + + # f2 reads from f1's output + assert f2_input == "s3://user-bucket/first/" + + # f3: also auto-derived under data-prep bucket + _assert_auto_derived_path(f3_output, dp, "train", "fuzzy_dedup") + + # f3 reads from f2's output + assert f3_input == f"{dp}/f2-out/" diff --git a/tests/unit/dataset/test_cloudwatch_dataset_loader.py b/tests/unit/dataset/test_cloudwatch_dataset_loader.py new file mode 100644 index 0000000..e97dd09 --- /dev/null +++ b/tests/unit/dataset/test_cloudwatch_dataset_loader.py @@ -0,0 +1,281 @@ +# 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 CloudWatchDatasetLoader. + +All boto3 calls are mocked via ``unittest.mock.patch`` so no real network +calls are made. Telemetry is auto-mocked globally by +``tests/unit/conftest.py::_mock_telemetry``. +""" + +import logging +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch + +import pytest + +from amzn_nova_forge.dataset.cloudwatch_dataset_loader import ( + CloudWatchDatasetLoader, +) +from amzn_nova_forge.dataset.operations.base import DataPrepError + +LOG_GROUP = "/test/app-logs" +QUERY = "fields @timestamp, @message | limit 100" +START_TIME = datetime(2024, 6, 1, 0, 0, 0, tzinfo=timezone.utc) +END_TIME = datetime(2024, 6, 2, 0, 0, 0, tzinfo=timezone.utc) +QUERY_ID = "12345678-1234-1234-1234-123456789012" + + +def _complete_response(results): + """Build a GetQueryResults response with status Complete.""" + return {"status": "Complete", "results": results} + + +def _running_response(): + """Build a GetQueryResults response with status Running.""" + return {"status": "Running", "results": []} + + +class TestLoadIsLazy: + """load() stores config only; no boto3 calls happen at load time.""" + + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.boto3") + def test_no_boto3_calls_at_load_time(self, mock_boto3): + loader = CloudWatchDatasetLoader() + result = loader.load(LOG_GROUP, QUERY, START_TIME, END_TIME) + + mock_boto3.client.assert_not_called() + assert result is loader + assert loader._load_path == f"cw:///{LOG_GROUP}" + + +class TestStateReset: + """load() resets _last_state, _is_materialized, and _session_id.""" + + def test_state_fields_reset_on_load(self): + loader = CloudWatchDatasetLoader() + loader._last_state = "something" + loader._is_materialized = True + loader._session_id = "old-session" + + loader.load(LOG_GROUP, QUERY, START_TIME, END_TIME) + + assert loader._last_state is None + assert loader._is_materialized is False + assert loader._session_id is None + + +class TestMakeSingleFileGenerator: + """_make_single_file_generator raises NotImplementedError.""" + + def test_raises_not_implemented(self): + loader = CloudWatchDatasetLoader() + with pytest.raises(NotImplementedError) as exc_info: + loader._make_single_file_generator("anything") + + assert "not file-based" in str(exc_info.value) + + +class TestGeneratorHappyPath: + """Generator executes query and yields flat dicts on iteration.""" + + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.boto3") + def test_yields_converted_dicts(self, mock_boto3): + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + mock_client.start_query.return_value = {"queryId": QUERY_ID} + mock_client.get_query_results.return_value = _complete_response( + [ + [ + {"field": "@timestamp", "value": "2024-06-01 10:30:00.000"}, + {"field": "@message", "value": '{"action": "invoke"}'}, + {"field": "requestId", "value": "req-001"}, + ], + [ + {"field": "@timestamp", "value": "2024-06-01 10:31:00.000"}, + {"field": "@message", "value": '{"action": "complete"}'}, + {"field": "requestId", "value": "req-002"}, + ], + ] + ) + + loader = CloudWatchDatasetLoader() + loader.load(LOG_GROUP, QUERY, START_TIME, END_TIME) + results = list(loader.dataset()) + + assert results == [ + { + "@timestamp": "2024-06-01 10:30:00.000", + "@message": '{"action": "invoke"}', + "requestId": "req-001", + }, + { + "@timestamp": "2024-06-01 10:31:00.000", + "@message": '{"action": "complete"}', + "requestId": "req-002", + }, + ] + + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.boto3") + def test_start_query_receives_correct_parameters(self, mock_boto3): + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + mock_client.start_query.return_value = {"queryId": QUERY_ID} + mock_client.get_query_results.return_value = _complete_response([]) + + loader = CloudWatchDatasetLoader() + loader.load(LOG_GROUP, QUERY, START_TIME, END_TIME) + list(loader.dataset()) + + mock_client.start_query.assert_called_once_with( + logGroupName=LOG_GROUP, + queryString=QUERY, + startTime=int(START_TIME.timestamp()), + endTime=int(END_TIME.timestamp()), + ) + + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.boto3") + def test_timestamps_converted_to_epoch_seconds(self, mock_boto3): + """Datetimes are passed as int(dt.timestamp()) — epoch seconds, not millis.""" + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + mock_client.start_query.return_value = {"queryId": QUERY_ID} + mock_client.get_query_results.return_value = _complete_response([]) + + est = timezone(timedelta(hours=-5)) + start = datetime(2024, 6, 15, 7, 0, 0, tzinfo=est) + end = datetime(2024, 6, 16, 7, 0, 0, tzinfo=est) + + loader = CloudWatchDatasetLoader() + loader.load(LOG_GROUP, QUERY, start, end) + list(loader.dataset()) + + call_kwargs = mock_client.start_query.call_args.kwargs + assert call_kwargs["startTime"] == int(start.timestamp()) + assert call_kwargs["endTime"] == int(end.timestamp()) + + +class TestPollingBehavior: + """Polling calls sleep between attempts and passes correct query ID.""" + + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.time.sleep") + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.boto3") + def test_sleep_called_between_polls(self, mock_boto3, mock_sleep): + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + mock_client.start_query.return_value = {"queryId": QUERY_ID} + mock_client.get_query_results.side_effect = [ + _running_response(), + _running_response(), + _complete_response([[{"field": "@timestamp", "value": "2024-06-01"}]]), + ] + + loader = CloudWatchDatasetLoader() + loader.load(LOG_GROUP, QUERY, START_TIME, END_TIME) + list(loader.dataset()) + + assert mock_sleep.call_count == 2 + + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.time.sleep") + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.boto3") + def test_correct_query_id_passed_to_get_results(self, mock_boto3, mock_sleep): + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + mock_client.start_query.return_value = {"queryId": QUERY_ID} + mock_client.get_query_results.return_value = _complete_response([]) + + loader = CloudWatchDatasetLoader() + loader.load(LOG_GROUP, QUERY, START_TIME, END_TIME) + list(loader.dataset()) + + mock_client.get_query_results.assert_called_with(queryId=QUERY_ID) + + +class TestEmptyResults: + """Empty results yield nothing and log a warning.""" + + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.boto3") + def test_empty_results_yields_nothing_and_warns(self, mock_boto3, caplog): + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + mock_client.start_query.return_value = {"queryId": QUERY_ID} + mock_client.get_query_results.return_value = _complete_response([]) + + loader = CloudWatchDatasetLoader() + loader.load(LOG_GROUP, QUERY, START_TIME, END_TIME) + + with caplog.at_level(logging.WARNING): + results = list(loader.dataset()) + + assert results == [] + assert any("0 results" in rec.message for rec in caplog.records) + + +class TestErrorMapping: + """CloudWatch exceptions are wrapped in DataPrepError with chaining.""" + + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.boto3") + def test_start_query_error_wrapped_and_chained(self, mock_boto3): + original = RuntimeError("something went wrong") + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + mock_client.start_query.side_effect = original + + loader = CloudWatchDatasetLoader() + loader.load(LOG_GROUP, QUERY, START_TIME, END_TIME) + + with pytest.raises(DataPrepError) as exc_info: + list(loader.dataset()) + + assert "something went wrong" in str(exc_info.value) + assert exc_info.value.__cause__ is original + + @pytest.mark.parametrize( + "status", + ["Failed", "Cancelled", "Timeout", "Unknown"], + ) + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.time.sleep") + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.boto3") + def test_terminal_failure_statuses(self, mock_boto3, mock_sleep, status): + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + mock_client.start_query.return_value = {"queryId": QUERY_ID} + mock_client.get_query_results.return_value = { + "status": status, + "results": [], + } + + loader = CloudWatchDatasetLoader() + loader.load(LOG_GROUP, QUERY, START_TIME, END_TIME) + + with pytest.raises(DataPrepError) as exc_info: + list(loader.dataset()) + + assert status in str(exc_info.value) + + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.time.sleep") + @patch("amzn_nova_forge.dataset.cloudwatch_dataset_loader.boto3") + def test_poll_error_wrapped_and_chained(self, mock_boto3, mock_sleep): + original = RuntimeError("network timeout") + mock_client = MagicMock() + mock_boto3.client.return_value = mock_client + mock_client.start_query.return_value = {"queryId": QUERY_ID} + mock_client.get_query_results.side_effect = original + + loader = CloudWatchDatasetLoader() + loader.load(LOG_GROUP, QUERY, START_TIME, END_TIME) + + with pytest.raises(DataPrepError) as exc_info: + list(loader.dataset()) + + assert exc_info.value.__cause__ is original diff --git a/tests/unit/dataset/test_data_state.py b/tests/unit/dataset/test_data_state.py index a744509..f38f2c3 100644 --- a/tests/unit/dataset/test_data_state.py +++ b/tests/unit/dataset/test_data_state.py @@ -120,6 +120,30 @@ def test_session_id_fallback(self): resolver = OutputPathResolver("s3://bucket/train.jsonl") self.assertRegex(resolver._session_id, r"\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2}") + def test_local_path_rebases_to_dataprep_bucket(self): + """Local path should be rebased to the data-prep bucket.""" + resolver = OutputPathResolver("/home/user/data/train.jsonl", "2026-05-04_12-00-00") + + path = resolver.resolve_path("default_text_filter") + self.assertTrue(path.startswith("s3://"), f"Expected S3 path, got: {path}") + self.assertIn("train", path) + + def test_hf_path_rebases_to_dataprep_bucket(self): + """HuggingFace path should be rebased to the data-prep bucket.""" + resolver = OutputPathResolver("hf://fake-org/fake-dataset/train", "2026-05-04_12-00-00") + + path = resolver.resolve_path("default_text_filter") + self.assertTrue(path.startswith("s3://"), f"Expected S3 path, got: {path}") + self.assertIn("train", path) + + def test_s3_path_not_rebased(self): + """S3 path should NOT be rebased — uses the original parent.""" + resolver = OutputPathResolver("s3://my-bucket/data/train.jsonl", "2026-05-04_12-00-00") + path = resolver.resolve_path("default_text_filter") + self.assertTrue( + path.startswith("s3://my-bucket/"), f"Expected original bucket, got: {path}" + ) + @pytest.mark.parametrize( "path, expected_location", diff --git a/tests/unit/dataset/test_dataset_loader.py b/tests/unit/dataset/test_dataset_loader.py index a51c893..b38fe09 100644 --- a/tests/unit/dataset/test_dataset_loader.py +++ b/tests/unit/dataset/test_dataset_loader.py @@ -684,7 +684,7 @@ def test_save_data_s3_success(self, mock_boto_client): result_path = dataset_loader.save(save_path) self.assertEqual(result_path, save_path) - mock_boto_client.assert_called_once_with("s3") + mock_boto_client.assert_called_once_with("s3", region_name=None) # Verify upload_file was called mock_s3.upload_file.assert_called_once() diff --git a/tests/unit/dataset/test_dataset_loader_telemetry.py b/tests/unit/dataset/test_dataset_loader_telemetry.py index a2b154c..0e610c2 100644 --- a/tests/unit/dataset/test_dataset_loader_telemetry.py +++ b/tests/unit/dataset/test_dataset_loader_telemetry.py @@ -100,4 +100,6 @@ def test_eval_task_without_value_attribute(self): result = _extract_model_training_method(eval_task="custom_task") self.assertIsNotNone(result) - self.assertEqual(result["evalTask"], "custom_task") \ No newline at end of file + self.assertEqual(result["evalTask"], "custom_task") + + diff --git a/tests/unit/dataset/test_huggingface_dataset_loader.py b/tests/unit/dataset/test_huggingface_dataset_loader.py index 0d41d6c..ef49500 100644 --- a/tests/unit/dataset/test_huggingface_dataset_loader.py +++ b/tests/unit/dataset/test_huggingface_dataset_loader.py @@ -330,7 +330,6 @@ def test_data_dir_forwarded_when_provided(): @pytest.mark.parametrize("kwarg", ["data_files", "data_dir"]) def test_optional_kwargs_omitted_when_none(kwarg): - """``data_files`` and ``data_dir`` are NOT in ``load_dataset`` kwargs when caller omits them.""" fake_mod = _make_fake_datasets_module(load_dataset_return_value=[]) with patch.dict(sys.modules, {"datasets": fake_mod}): loader = HuggingFaceDatasetLoader() diff --git a/tests/unit/dataset/test_language_detection_filter_operation.py b/tests/unit/dataset/test_language_detection_filter_operation.py index c637063..22f3482 100644 --- a/tests/unit/dataset/test_language_detection_filter_operation.py +++ b/tests/unit/dataset/test_language_detection_filter_operation.py @@ -98,7 +98,7 @@ def test_missing_model_path_auto_stages(self): languages=["en"], ) - mock_ensure.assert_called_once_with() + mock_ensure.assert_called_once_with(region=None) job_config = manager.execute.call_args[0][0] self.assertEqual( job_config.extra_args["model_path"], diff --git a/tests/unit/dataset/test_multimodal_s3_upload.py b/tests/unit/dataset/test_multimodal_s3_upload.py index 709a25a..98373a3 100644 --- a/tests/unit/dataset/test_multimodal_s3_upload.py +++ b/tests/unit/dataset/test_multimodal_s3_upload.py @@ -547,7 +547,7 @@ def test_auto_resolve_bucket_owner_via_sts(self, mock_boto_client): mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} mock_s3 = MagicMock() - def side_effect(service): + def side_effect(service, **kwargs): if service == "sts": return mock_sts elif service == "s3": diff --git a/tests/unit/deployer/test_forge_deployer.py b/tests/unit/deployer/test_forge_deployer.py index e7caeda..3fbcb00 100644 --- a/tests/unit/deployer/test_forge_deployer.py +++ b/tests/unit/deployer/test_forge_deployer.py @@ -244,7 +244,12 @@ def client_side_effect(service, **kwargs): deploy_platform=DeployPlatform.BEDROCK_PT, ) - mock_update_pt.assert_called_once() + mock_update_pt.assert_called_once_with( + "arn:aws:bedrock:us-east-1:123456789012:pt/existing-pt", + "arn:aws:bedrock:us-east-1:123456789012:custom-model/my-model", + "nova-micro-us-east-1", + region="us-east-1", + ) self.assertEqual( result.endpoint.uri, "arn:aws:bedrock:us-east-1:123456789012:pt/existing-pt", @@ -751,11 +756,13 @@ def test_get_status_delegates_to_result_status(self, mock_validate_region): def test_get_status_by_arn_returns_job_status(self, mock_check_status, mock_validate_region): deployer = self._make_deployer() status = deployer.get_status_by_arn( - "arn:aws:bedrock:us-east-1:123:deployment/d", DeployPlatform.BEDROCK_OD + "arn:aws:bedrock:us-east-1:123456789012:deployment/d", DeployPlatform.BEDROCK_OD ) self.assertEqual(status, JobStatus.IN_PROGRESS) mock_check_status.assert_called_once_with( - "arn:aws:bedrock:us-east-1:123:deployment/d", DeployPlatform.BEDROCK_OD + "arn:aws:bedrock:us-east-1:123456789012:deployment/d", + DeployPlatform.BEDROCK_OD, + region="us-east-1", ) @patch(f"{PATCH_PREFIX}.check_deployment_status", return_value=None) @@ -784,22 +791,26 @@ def test_get_logs_with_job_result(self, mock_check_status, mock_validate_region) endpoint = EndpointInfo( platform=DeployPlatform.BEDROCK_OD, endpoint_name="ep", - uri="arn:aws:bedrock:us-east-1:123:deployment/d", + uri="arn:aws:bedrock:us-east-1:123456789012:deployment/d", model_artifact_path="s3://bucket/model", ) result = DeploymentResult(endpoint=endpoint, created_at=datetime.now(timezone.utc)) deployer.get_logs(job_result=result) mock_check_status.assert_called_once_with( - "arn:aws:bedrock:us-east-1:123:deployment/d", DeployPlatform.BEDROCK_OD + "arn:aws:bedrock:us-east-1:123456789012:deployment/d", + DeployPlatform.BEDROCK_OD, + region="us-east-1", ) @patch(f"{PATCH_PREFIX}.check_deployment_status", return_value="InProgress") def test_get_logs_with_endpoint_arn_only(self, mock_check_status, mock_validate_region): deployer = self._make_deployer() - deployer.get_logs(endpoint_arn="arn:aws:sagemaker:us-east-1:123:endpoint/ep") + deployer.get_logs(endpoint_arn="arn:aws:sagemaker:us-east-1:123456789012:endpoint/ep") mock_check_status.assert_called_once_with( - "arn:aws:sagemaker:us-east-1:123:endpoint/ep", DeployPlatform.SAGEMAKER + "arn:aws:sagemaker:us-east-1:123456789012:endpoint/ep", + DeployPlatform.SAGEMAKER, + region="us-east-1", ) @patch(f"{PATCH_PREFIX}.logger") @@ -809,5 +820,23 @@ def test_get_logs_no_arn_logs_info(self, mock_logger, mock_validate_region): mock_logger.info.assert_called_once_with("No endpoint ARN available. Call deploy() first.") +class TestDeploymentResultRegion(unittest.TestCase): + """Tests for region propagation through DeploymentResult.status.""" + + def test_deployment_result_status_passes_region(self): + endpoint = EndpointInfo( + platform=DeployPlatform.BEDROCK_OD, + endpoint_name="ep", + uri="arn:test", + model_artifact_path="s3://x", + region="eu-west-1", + ) + result = DeploymentResult(endpoint=endpoint, created_at=datetime.now(timezone.utc)) + with patch.object(DeploymentResult, "_status_checker") as mock_checker: + mock_checker.return_value = "Active" + _ = result.status + mock_checker.assert_called_once_with("arn:test", DeployPlatform.BEDROCK_OD, "eu-west-1") + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/evaluator/test_forge_evaluator.py b/tests/unit/evaluator/test_forge_evaluator.py index 15d837e..18cb262 100644 --- a/tests/unit/evaluator/test_forge_evaluator.py +++ b/tests/unit/evaluator/test_forge_evaluator.py @@ -700,6 +700,7 @@ def test_get_logs_with_job_result(self, mock_monitor_cls): job_id="eval-job-123", platform=Platform.SMTJ, started_time=int(started.timestamp() * 1000), + region="us-east-1", ) mock_monitor.show_logs.assert_called_once_with( limit=None, start_from_head=False, end_time=None @@ -743,6 +744,7 @@ def test_get_logs_smhp_includes_cluster_namespace(self, mock_monitor_cls): started_time=int(started.timestamp() * 1000), cluster_name="my-cluster", namespace="kubeflow", + region="us-east-1", ) diff --git a/tests/unit/iam/test_iam_role_creator.py b/tests/unit/iam/test_iam_role_creator.py index 26ae40d..48b0a29 100644 --- a/tests/unit/iam/test_iam_role_creator.py +++ b/tests/unit/iam/test_iam_role_creator.py @@ -917,5 +917,31 @@ def test_creates_exactly_two_policies(self, mock_boto_client): ) +class TestRegionPropagation(unittest.TestCase): + """Tests that the region parameter is propagated to the STS client.""" + + @patch("boto3.client") + def test_create_bedrock_execution_role_propagates_region(self, mock_boto_client): + mock_sts = MagicMock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + mock_boto_client.return_value = mock_sts + + mock_iam = MagicMock() + mock_iam.exceptions.NoSuchEntityException = type("NoSuchEntityException", (Exception,), {}) + mock_iam.exceptions.EntityAlreadyExistsException = type( + "EntityAlreadyExistsException", (Exception,), {} + ) + mock_iam.get_role.side_effect = mock_iam.exceptions.NoSuchEntityException("not found") + mock_iam.create_policy.return_value = { + "Policy": {"Arn": "arn:aws:iam::123456789012:policy/foo"} + } + + create_bedrock_execution_role( + iam_client=mock_iam, role_name="test-role", region="eu-west-1" + ) + + mock_boto_client.assert_called_once_with("sts", region_name="eu-west-1") + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/inference/test_forge_inference.py b/tests/unit/inference/test_forge_inference.py index 90299ec..89afd9a 100644 --- a/tests/unit/inference/test_forge_inference.py +++ b/tests/unit/inference/test_forge_inference.py @@ -535,6 +535,7 @@ def test_get_logs_with_job_result(self, mock_monitor_cls): job_id="job-abc", platform=Platform.SMTJ, started_time=int(datetime(2024, 1, 1, tzinfo=timezone.utc).timestamp() * 1000), + region="us-east-1", ) mock_monitor.show_logs.assert_called_once_with( limit=None, start_from_head=False, end_time=None @@ -591,6 +592,7 @@ def test_get_logs_smhp_includes_cluster_and_namespace(self, mock_monitor_cls): started_time=int(start.timestamp() * 1000), cluster_name="my-cluster", namespace="my-namespace", + region="us-east-1", ) mock_monitor.show_logs.assert_called_once() diff --git a/tests/unit/manager/test_runtime_manager.py b/tests/unit/manager/test_runtime_manager.py index e567844..d29ea03 100644 --- a/tests/unit/manager/test_runtime_manager.py +++ b/tests/unit/manager/test_runtime_manager.py @@ -14,6 +14,7 @@ import io import json import os +import subprocess import tempfile import unittest import zipfile @@ -515,6 +516,162 @@ def test_validate_lambda_local_raises_on_s3_read_failure( os.unlink(src) +class TestSMTJValidationDataset(unittest.TestCase): + """Tests for SFT validation dataset support in SMTJRuntimeManager.""" + + def setUp(self): + self.mock_role = "arn:aws:iam::123456789012:role/SageMakerExecutionRole" + self.instance_type = "ml.m5.xlarge" + self.instance_count = 1 + + @patch.object(SMTJRuntimeManager, "setup", return_value=None) + def _create_manager(self, mock_setup): + manager = SMTJRuntimeManager(self.instance_type, self.instance_count) + manager.execution_role = self.mock_role + manager.sagemaker_client = MagicMock() + manager.sagemaker_session = MagicMock() + manager.region = "us-east-1" + return manager + + @patch("amzn_nova_forge.manager.runtime_manager.InputData") + @patch("amzn_nova_forge.manager.runtime_manager.boto3.client") + @patch("amzn_nova_forge.manager.runtime_manager.ModelTrainer") + @patch.object(SMTJRuntimeManager, "setup", return_value=None) + def test_smtj_sft_with_validation_data( + self, mock_setup, mock_model_trainer_cls, mock_boto_client, mock_input_data + ): + manager = self._create_manager() + + mock_model_trainer = MagicMock() + mock_model_trainer.with_tensorboard_output_config.return_value = mock_model_trainer + mock_model_trainer_cls.from_recipe.return_value = mock_model_trainer + + manager.sagemaker_client.list_training_jobs.return_value = { + "TrainingJobSummaries": [{"TrainingJobName": "test-job-suffix"}] + } + + job_config = JobConfig( + job_name="test-job", + image_uri="123456789012.dkr.ecr.us-east-1.amazonaws.com/my-image:latest", + recipe_path="/path/to/recipe", + output_s3_path="s3://output-bucket/output", + data_s3_path="s3://input-bucket/train.jsonl", + validation_data_s3_path="s3://input-bucket/val.jsonl", + input_s3_data_type="S3Prefix", + ) + + manager.execute(job_config) + + # InputData should be called twice: once for "train", once for "validation" + self.assertEqual(mock_input_data.call_count, 2) + channel_names = [call.kwargs["channel_name"] for call in mock_input_data.call_args_list] + self.assertIn("train", channel_names) + self.assertIn("validation", channel_names) + + call_kwargs = mock_model_trainer_cls.from_recipe.call_args.kwargs + self.assertEqual(len(call_kwargs["input_data_config"]), 2) + + @patch("amzn_nova_forge.manager.runtime_manager.InputData") + @patch("amzn_nova_forge.manager.runtime_manager.boto3.client") + @patch("amzn_nova_forge.manager.runtime_manager.ModelTrainer") + @patch.object(SMTJRuntimeManager, "setup", return_value=None) + def test_smtj_sft_without_validation_data( + self, mock_setup, mock_model_trainer_cls, mock_boto_client, mock_input_data + ): + manager = self._create_manager() + + mock_model_trainer = MagicMock() + mock_model_trainer.with_tensorboard_output_config.return_value = mock_model_trainer + mock_model_trainer_cls.from_recipe.return_value = mock_model_trainer + + manager.sagemaker_client.list_training_jobs.return_value = { + "TrainingJobSummaries": [{"TrainingJobName": "test-job-suffix"}] + } + + job_config = JobConfig( + job_name="test-job", + image_uri="123456789012.dkr.ecr.us-east-1.amazonaws.com/my-image:latest", + recipe_path="/path/to/recipe", + output_s3_path="s3://output-bucket/output", + data_s3_path="s3://input-bucket/train.jsonl", + input_s3_data_type="S3Prefix", + ) + + manager.execute(job_config) + + # InputData should be called only once for "train" + self.assertEqual(mock_input_data.call_count, 1) + self.assertEqual(mock_input_data.call_args.kwargs["channel_name"], "train") + + call_kwargs = mock_model_trainer_cls.from_recipe.call_args.kwargs + self.assertEqual(len(call_kwargs["input_data_config"]), 1) + + @patch("amzn_nova_forge.manager.runtime_manager.InputData") + @patch("amzn_nova_forge.manager.runtime_manager.boto3.client") + @patch("amzn_nova_forge.manager.runtime_manager.ModelTrainer") + @patch.object(SMTJRuntimeManager, "setup", return_value=None) + def test_smtj_hyperparameters_passed_to_model_trainer( + self, mock_setup, mock_model_trainer_cls, mock_boto_client, mock_input_data + ): + manager = self._create_manager() + + mock_model_trainer = MagicMock() + mock_model_trainer.with_tensorboard_output_config.return_value = mock_model_trainer + mock_model_trainer_cls.from_recipe.return_value = mock_model_trainer + + manager.sagemaker_client.list_training_jobs.return_value = { + "TrainingJobSummaries": [{"TrainingJobName": "test-job-suffix"}] + } + + job_config = JobConfig( + job_name="test-job", + image_uri="123456789012.dkr.ecr.us-east-1.amazonaws.com/my-image:latest", + recipe_path="/path/to/recipe", + output_s3_path="s3://output-bucket/output", + data_s3_path="s3://input-bucket/train.jsonl", + input_s3_data_type="S3Prefix", + trainer_config_hyperparameters={"val_check_interval": "500"}, + ) + + manager.execute(job_config) + + call_kwargs = mock_model_trainer_cls.from_recipe.call_args.kwargs + self.assertIn("hyperparameters", call_kwargs) + self.assertEqual(call_kwargs["hyperparameters"]["val_check_interval"], "500") + + @patch("amzn_nova_forge.manager.runtime_manager.InputData") + @patch("amzn_nova_forge.manager.runtime_manager.boto3.client") + @patch("amzn_nova_forge.manager.runtime_manager.ModelTrainer") + @patch.object(SMTJRuntimeManager, "setup", return_value=None) + def test_smtj_no_hyperparameters_when_none( + self, mock_setup, mock_model_trainer_cls, mock_boto_client, mock_input_data + ): + manager = self._create_manager() + + mock_model_trainer = MagicMock() + mock_model_trainer.with_tensorboard_output_config.return_value = mock_model_trainer + mock_model_trainer_cls.from_recipe.return_value = mock_model_trainer + + manager.sagemaker_client.list_training_jobs.return_value = { + "TrainingJobSummaries": [{"TrainingJobName": "test-job-suffix"}] + } + + job_config = JobConfig( + job_name="test-job", + image_uri="123456789012.dkr.ecr.us-east-1.amazonaws.com/my-image:latest", + recipe_path="/path/to/recipe", + output_s3_path="s3://output-bucket/output", + data_s3_path="s3://input-bucket/train.jsonl", + input_s3_data_type="S3Prefix", + trainer_config_hyperparameters=None, + ) + + manager.execute(job_config) + + call_kwargs = mock_model_trainer_cls.from_recipe.call_args.kwargs + self.assertNotIn("hyperparameters", call_kwargs) + + class TestSMTJRuntimeManagerDataPrepDelegation(unittest.TestCase): """SMTJRuntimeManager.set_mode(DATA_PREP) flips execute/cleanup to the delegate.""" @@ -679,10 +836,48 @@ def test_instance_count_setter(self, mock_setup): return_value="arn:aws:iam::123456789012:role/MockRole", ) @patch("subprocess.run") - def test_initialization_fails(self, mock_run, mock_get_role): - mock_run.return_value.stderr = "Connection failed" + def test_initialization_stderr_benign(self, mock_run, mock_get_role): + mock_run.return_value.stderr = "NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+" + mock_run.return_value.stdout = "" + + # Should NOT raise — benign stderr is filtered by _check_hyperpod_stderr + manager = SMHPRuntimeManager( + self.instance_type, + self.instance_count, + self.cluster_name, + self.namespace, + ) + self.assertIsNotNone(manager) + + @patch( + "amzn_nova_forge.manager.runtime_manager.get_execution_role", + return_value="arn:aws:iam::123456789012:role/MockRole", + ) + @patch("subprocess.run") + def test_initialization_stderr_real_error(self, mock_run, mock_get_role): + mock_run.return_value.stderr = "Cluster with name not found" + mock_run.return_value.stdout = "" - with self.assertRaises(Exception): + with self.assertRaises(RuntimeError) as ctx: + SMHPRuntimeManager( + self.instance_type, + self.instance_count, + self.cluster_name, + self.namespace, + ) + self.assertIn("Cluster with name not found", str(ctx.exception)) + + @patch( + "amzn_nova_forge.manager.runtime_manager.get_execution_role", + return_value="arn:aws:iam::123456789012:role/MockRole", + ) + @patch("subprocess.run") + def test_initialization_subprocess_error(self, mock_run, mock_get_role): + mock_run.side_effect = subprocess.CalledProcessError( + 1, "hyperpod", stderr="Error: cluster not found" + ) + + with self.assertRaises(subprocess.CalledProcessError): SMHPRuntimeManager( self.instance_type, self.instance_count, @@ -841,6 +1036,45 @@ def test_cleanup_handles_error(self, mock_run, mock_get_role): self.assertEqual(str(context.exception), "Cleanup failed") + @patch( + "amzn_nova_forge.manager.runtime_manager.get_execution_role", + return_value="arn:aws:iam::123456789012:role/MockRole", + ) + @patch("subprocess.run") + def test_cleanup_stderr_benign(self, mock_run, mock_get_role): + mock_run.side_effect = [ + MagicMock(stdout="", stderr=""), + MagicMock( + stdout="", + stderr="NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+", + ), + ] + + manager = SMHPRuntimeManager( + self.instance_type, self.instance_count, self.cluster_name, self.namespace + ) + # Should not raise — benign stderr is filtered + manager.cleanup("test-job") + + @patch( + "amzn_nova_forge.manager.runtime_manager.get_execution_role", + return_value="arn:aws:iam::123456789012:role/MockRole", + ) + @patch("subprocess.run") + def test_cleanup_stderr_real_error_logged(self, mock_run, mock_get_role): + mock_run.side_effect = [ + MagicMock(stdout="", stderr=""), + MagicMock(stdout="", stderr="Cluster with name not found"), + ] + + manager = SMHPRuntimeManager( + self.instance_type, self.instance_count, self.cluster_name, self.namespace + ) + # Should NOT raise — cleanup is best-effort, real errors are logged + with self.assertLogs(level="ERROR") as log: + manager.cleanup("test-job") + self.assertTrue(any("Failed to cleanup HyperPod job" in msg for msg in log.output)) + @patch.object(SMHPRuntimeManager, "setup", return_value=None) def test_rft_lambda_arn_set_immediately_when_arn_passed(self, mock_setup): arn = "arn:aws:lambda:us-east-1:123456789012:function:SageMaker-reward" diff --git a/tests/unit/manager/test_serverless_runtime_manager.py b/tests/unit/manager/test_serverless_runtime_manager.py index ed8193a..b1619a2 100644 --- a/tests/unit/manager/test_serverless_runtime_manager.py +++ b/tests/unit/manager/test_serverless_runtime_manager.py @@ -459,6 +459,7 @@ def test_resolve_base_model_arn(self, mock_hub): hub_content_name="nova-textgeneration-lite-v2", hub_content_type="Model", region="us-east-1", + hub_content_version=None, ) # --- execute tests --- @@ -1006,6 +1007,202 @@ def test_required_calling_role_permissions(self): self.assertIn("s3:GetObject", perm_actions) +class TestSMTJServerlessValidationDataset(unittest.TestCase): + """Tests for SFT validation dataset support in SMTJServerlessRuntimeManager.""" + + def setUp(self): + self.mock_role = "arn:aws:iam::123456789012:role/SageMakerExecutionRole" + self.model_package_group_name = "test-model-package-group" + self.mock_group_arn = ( + "arn:aws:sagemaker:us-east-1:123456789012:model-package-group/test-model-package-group" + ) + self.mock_hub_content_arn = ( + "arn:aws:sagemaker:us-east-1:123456789012:" + "hub-content/SageMakerPublicHub/Model/nova-textgeneration-lite-v2/1.0.0" + ) + + @patch.object(SMTJServerlessRuntimeManager, "setup", return_value=None) + def _create_manager(self, mock_setup, **kwargs): + manager = SMTJServerlessRuntimeManager( + model_package_group_name=self.model_package_group_name, **kwargs + ) + manager.execution_role = self.mock_role + manager.sagemaker_client = MagicMock() + manager.sagemaker_session = MagicMock() + manager.region = "us-east-1" + manager.model_package_group_arn = self.mock_group_arn + return manager + + @patch("amzn_nova_forge.manager.runtime_manager.get_hub_content") + @patch("sagemaker.ai_registry.dataset.DataSet") + def test_serverless_sft_with_validation_data(self, mock_dataset_cls, mock_hub_content): + manager = self._create_manager() + mock_hub_content.return_value = {"HubContentArn": self.mock_hub_content_arn} + + mock_dataset = MagicMock() + mock_dataset.arn = "arn:aws:sagemaker:us-east-1:123456789012:dataset/train-input" + mock_val_dataset = MagicMock() + mock_val_dataset.arn = "arn:aws:sagemaker:us-east-1:123456789012:dataset/val-input" + mock_dataset_cls.create.side_effect = [mock_dataset, mock_val_dataset] + + manager.sagemaker_client.create_training_job.return_value = { + "TrainingJobArn": "arn:aws:sagemaker:us-east-1:123456789012:training-job/test-job" + } + + recipe = { + "run": {"model_type": "amazon.nova-2-lite-v1:0:256k"}, + "trainer": {"lr": 0.001}, + } + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(recipe, f) + recipe_path = f.name + + try: + job_config = JobConfig( + job_name="test-serverless-val-job", + image_uri="", + recipe_path=recipe_path, + output_s3_path="s3://output-bucket/output", + data_s3_path="s3://input-bucket/train.jsonl", + validation_data_s3_path="s3://input-bucket/val.jsonl", + method=TrainingMethod.SFT_LORA, + ) + manager.execute(job_config) + finally: + os.unlink(recipe_path) + + call_kwargs = manager.sagemaker_client.create_training_job.call_args.kwargs + self.assertIn("InputDataConfig", call_kwargs) + channel_names = [ch["ChannelName"] for ch in call_kwargs["InputDataConfig"]] + self.assertIn("train", channel_names) + self.assertIn("validation", channel_names) + self.assertEqual(len(call_kwargs["InputDataConfig"]), 2) + + @patch("amzn_nova_forge.manager.runtime_manager.get_hub_content") + @patch("sagemaker.ai_registry.dataset.DataSet") + def test_serverless_sft_without_validation_data(self, mock_dataset_cls, mock_hub_content): + manager = self._create_manager() + mock_hub_content.return_value = {"HubContentArn": self.mock_hub_content_arn} + + mock_dataset = MagicMock() + mock_dataset.arn = "arn:aws:sagemaker:us-east-1:123456789012:dataset/train-input" + mock_dataset_cls.create.return_value = mock_dataset + + manager.sagemaker_client.create_training_job.return_value = { + "TrainingJobArn": "arn:aws:sagemaker:us-east-1:123456789012:training-job/test-job" + } + + recipe = { + "run": {"model_type": "amazon.nova-2-lite-v1:0:256k"}, + "trainer": {"lr": 0.001}, + } + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(recipe, f) + recipe_path = f.name + + try: + job_config = JobConfig( + job_name="test-serverless-no-val-job", + image_uri="", + recipe_path=recipe_path, + output_s3_path="s3://output-bucket/output", + data_s3_path="s3://input-bucket/train.jsonl", + method=TrainingMethod.SFT_LORA, + ) + manager.execute(job_config) + finally: + os.unlink(recipe_path) + + call_kwargs = manager.sagemaker_client.create_training_job.call_args.kwargs + self.assertIn("InputDataConfig", call_kwargs) + channel_names = [ch["ChannelName"] for ch in call_kwargs["InputDataConfig"]] + self.assertIn("train", channel_names) + self.assertNotIn("validation", channel_names) + + @patch("amzn_nova_forge.manager.runtime_manager.get_hub_content") + @patch("sagemaker.ai_registry.dataset.DataSet") + def test_serverless_validation_dataset_name_truncation( + self, mock_dataset_cls, mock_hub_content + ): + manager = self._create_manager() + mock_hub_content.return_value = {"HubContentArn": self.mock_hub_content_arn} + + mock_dataset = MagicMock() + mock_dataset.arn = "arn:aws:sagemaker:us-east-1:123456789012:dataset/train-input" + mock_val_dataset = MagicMock() + mock_val_dataset.arn = "arn:aws:sagemaker:us-east-1:123456789012:dataset/val-input" + mock_dataset_cls.create.side_effect = [mock_dataset, mock_val_dataset] + + manager.sagemaker_client.create_training_job.return_value = { + "TrainingJobArn": "arn:aws:sagemaker:us-east-1:123456789012:training-job/test-job" + } + + long_job_name = "a" * 63 # max job name length + recipe = {"run": {"model_type": "amazon.nova-2-lite-v1:0:256k"}} + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(recipe, f) + recipe_path = f.name + + try: + job_config = JobConfig( + job_name=long_job_name, + image_uri="", + recipe_path=recipe_path, + output_s3_path="s3://output-bucket/output", + data_s3_path="s3://data-bucket/data.jsonl", + validation_data_s3_path="s3://data-bucket/val.jsonl", + method=TrainingMethod.SFT_LORA, + ) + manager.execute(job_config) + finally: + os.unlink(recipe_path) + + # Second DataSet.create call is for validation dataset + val_dataset_name = mock_dataset_cls.create.call_args_list[1][0][0] + self.assertLessEqual(len(val_dataset_name), 63) + + @patch("amzn_nova_forge.manager.runtime_manager.get_hub_content") + @patch("sagemaker.ai_registry.dataset.DataSet") + def test_serverless_val_check_interval_fallback_merge(self, mock_dataset_cls, mock_hub_content): + manager = self._create_manager() + mock_hub_content.return_value = {"HubContentArn": self.mock_hub_content_arn} + + mock_dataset = MagicMock() + mock_dataset.arn = "arn:aws:sagemaker:us-east-1:123456789012:dataset/train-input" + mock_dataset_cls.create.return_value = mock_dataset + + manager.sagemaker_client.create_training_job.return_value = { + "TrainingJobArn": "arn:aws:sagemaker:us-east-1:123456789012:training-job/test-job" + } + + # Recipe template WITHOUT val_check_interval + recipe = { + "run": {"model_type": "amazon.nova-2-lite-v1:0:256k"}, + "trainer": {"lr": 0.001}, + } + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(recipe, f) + recipe_path = f.name + + try: + job_config = JobConfig( + job_name="test-val-interval-job", + image_uri="", + recipe_path=recipe_path, + output_s3_path="s3://output-bucket/output", + data_s3_path="s3://input-bucket/train.jsonl", + method=TrainingMethod.SFT_LORA, + trainer_config_hyperparameters={"val_check_interval": "500"}, + ) + manager.execute(job_config) + finally: + os.unlink(recipe_path) + + call_kwargs = manager.sagemaker_client.create_training_job.call_args.kwargs + self.assertIn("val_check_interval", call_kwargs["HyperParameters"]) + self.assertEqual(call_kwargs["HyperParameters"]["val_check_interval"], "500") + + class TestMethodToServerlessConfig(unittest.TestCase): def test_all_expected_methods_present(self): expected = { @@ -1363,5 +1560,57 @@ def test_returns_none_when_reference_missing(self): self.assertIsNone(result) +class TestExtractHyperparameters(unittest.TestCase): + """Tests for SMTJServerlessRuntimeManager._extract_hyperparameters.""" + + @patch.object(SMTJServerlessRuntimeManager, "setup", return_value=None) + def setUp(self, mock_setup): + self.manager = SMTJServerlessRuntimeManager(model_package_group_name="test-group") + + def test_basic_recipe_extraction(self): + recipe = {"run": {"name": "my-job", "max_steps": 100, "global_batch_size": 64}} + result = self.manager._extract_hyperparameters(recipe) + self.assertEqual(result["name"], "my-job") + self.assertEqual(result["max_steps"], "100") + self.assertEqual(result["global_batch_size"], "64") + + def test_datamix_config_injected_as_hyperparameters(self): + recipe = {"run": {"name": "my-job"}} + data_mixing_config = { + "customer_data_percent": 50, + "nova_code_percent": 60, + "nova_general_percent": 40, + "dataset_catalog": "sft_1p5_text_chat", + } + result = self.manager._extract_hyperparameters( + recipe, data_mixing_config=data_mixing_config + ) + self.assertEqual(result["customer_data_percent"], "50") + self.assertEqual(result["nova_code_percent"], "60") + self.assertEqual(result["nova_general_percent"], "40") + self.assertEqual(result["dataset_catalog"], "sft_1p5_text_chat") + + def test_dataset_catalog_included_in_hyperparameters(self): + """dataset_catalog must be passed as a hyperparameter.""" + recipe = {} + data_mixing_config = {"dataset_catalog": "sft_1p5_text_chat", "customer_data_percent": 80} + result = self.manager._extract_hyperparameters( + recipe, data_mixing_config=data_mixing_config + ) + self.assertIn("dataset_catalog", result) + self.assertEqual(result["dataset_catalog"], "sft_1p5_text_chat") + + def test_no_datamix_config_does_not_error(self): + recipe = {"run": {"name": "my-job", "max_steps": 10}} + result = self.manager._extract_hyperparameters(recipe, data_mixing_config=None) + self.assertEqual(result["name"], "my-job") + self.assertNotIn("customer_data_percent", result) + + def test_empty_datamix_config_does_not_error(self): + recipe = {"run": {"name": "my-job"}} + result = self.manager._extract_hyperparameters(recipe, data_mixing_config={}) + self.assertNotIn("customer_data_percent", result) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/manager/test_smtj_dataprep_runtime_manager.py b/tests/unit/manager/test_smtj_dataprep_runtime_manager.py index cf1fb20..797cf01 100644 --- a/tests/unit/manager/test_smtj_dataprep_runtime_manager.py +++ b/tests/unit/manager/test_smtj_dataprep_runtime_manager.py @@ -596,6 +596,43 @@ def client_factory(service, **kwargs): ) mock_create.assert_called_once() + @patch("amzn_nova_forge.manager.runtime_manager.boto3") + def test_setup_passes_region_to_iam_client(self, mock_boto3): + mock_session = MagicMock() + mock_session.region_name = "eu-west-1" + mock_boto3.session.Session.return_value = mock_session + + mock_iam = MagicMock() + mock_iam.get_role.return_value = {"Role": {"Arn": "arn:aws:iam::123456789012:role/MyRole"}} + mock_s3 = MagicMock() + mock_s3.head_bucket.return_value = {} + + def client_factory(service, **kwargs): + if service == "sagemaker": + return MagicMock() + if service == "s3": + return mock_s3 + if service == "iam": + return mock_iam + return MagicMock() + + mock_boto3.client.side_effect = client_factory + + with patch.object(SMTJDataPrepRuntimeManager, "_upload_entry_script"): + with patch.object(SMTJDataPrepRuntimeManager, "_upload_bundled_whls"): + SMTJDataPrepRuntimeManager( + instance_type="ml.g5.2xlarge", + image_uri="test:latest", + execution_role_name="MyRole", + s3_artifact_bucket="test-bucket", + region="eu-west-1", + ) + + # Verify boto3.client("iam", region_name="eu-west-1") was called + iam_calls = [c for c in mock_boto3.client.call_args_list if c[0] == ("iam",)] + self.assertEqual(len(iam_calls), 1) + self.assertEqual(iam_calls[0][1]["region_name"], "eu-west-1") + def _build_manager_with_resolver(self, mock_boto3, instance_type, resolved_uri): mock_session = MagicMock() mock_session.region_name = "us-west-2" diff --git a/tests/unit/model/result/test_eval_result.py b/tests/unit/model/result/test_eval_result.py index 7fbfb7c..e7358fc 100644 --- a/tests/unit/model/result/test_eval_result.py +++ b/tests/unit/model/result/test_eval_result.py @@ -619,5 +619,23 @@ def mock_download_file(bucket, key, local_path): self.assertEqual(result_dict, test_results) +class TestRegionPropagation(unittest.TestCase): + @patch("amzn_nova_forge.core.result.eval_result.boto3") + def test_smtj_evaluation_result_passes_region_to_sagemaker(self, mock_boto3): + mock_boto3.client.return_value = Mock() + + SMTJEvaluationResult( + job_id="test-job", + started_time=datetime.now(), + eval_task=EvaluationTask.MMLU, + eval_output_path="s3://bucket/output", + sagemaker_client=None, + s3_client=Mock(), + region="eu-west-1", + ) + + mock_boto3.client.assert_called_with("sagemaker", region_name="eu-west-1") + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/model/result/test_inference_result.py b/tests/unit/model/result/test_inference_result.py index ed52abc..cf6a5c1 100644 --- a/tests/unit/model/result/test_inference_result.py +++ b/tests/unit/model/result/test_inference_result.py @@ -659,5 +659,43 @@ def test_show(self, mock_logger): mock_logger.info.assert_called() +class TestRegionPropagation(unittest.TestCase): + @patch("boto3.client") + def test_inference_result_passes_region_to_s3(self, mock_boto3_client): + mock_s3_client = Mock() + mock_boto3_client.return_value = mock_s3_client + mock_s3_client.download_file.return_value = None + + result = SMTJBatchInferenceResult( + job_id="test-job-region", + started_time=datetime.now(), + inference_output_path="s3://test-bucket/output/output.tar.gz", + sagemaker_client=Mock(), + region="eu-west-1", + ) + + # Trigger S3 client creation via _download_inference_results + try: + result._download_inference_results() + except Exception: + pass + + mock_boto3_client.assert_any_call("s3", region_name="eu-west-1") + + @patch("boto3.client") + def test_smtj_batch_inference_result_passes_region_to_sagemaker(self, mock_boto3_client): + mock_boto3_client.return_value = Mock() + + SMTJBatchInferenceResult( + job_id="test-job-region", + started_time=datetime.now(), + inference_output_path="s3://test-bucket/output/output.tar.gz", + sagemaker_client=None, + region="eu-west-1", + ) + + mock_boto3_client.assert_called_with("sagemaker", region_name="eu-west-1") + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/model/result/test_job_result.py b/tests/unit/model/result/test_job_result.py index 86ad4d5..b9a812b 100644 --- a/tests/unit/model/result/test_job_result.py +++ b/tests/unit/model/result/test_job_result.py @@ -121,7 +121,7 @@ def test_init_with_client(self): @patch("boto3.client") def test_init_without_client(self, mock_boto3_client): manager = SMTJStatusManager() - mock_boto3_client.assert_called_once_with("sagemaker") + mock_boto3_client.assert_called_once_with("sagemaker", region_name=None) def test_get_job_status_in_progress(self): self.mock_sagemaker_client.describe_training_job.return_value = { @@ -241,19 +241,33 @@ def test_connect_cluster_success(self, mock_logger, mock_run): ) @patch("subprocess.run") - @patch("amzn_nova_forge.core.result.job_result.logger") - def test_connect_cluster_error(self, mock_logger, mock_run): + def test_connect_cluster_stderr_benign(self, mock_run): + mock_result = Mock() + mock_result.stderr = "NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+" + mock_run.return_value = mock_result + + # Should NOT raise — benign stderr is filtered by _check_hyperpod_stderr + self.manager._connect_cluster() + + @patch("subprocess.run") + def test_connect_cluster_stderr_real_error(self, mock_run): mock_result = Mock() - mock_result.stderr = "Connection failed" + mock_result.stderr = "Cluster with name not found" mock_run.return_value = mock_result - with self.assertRaises(Exception): + with self.assertRaises(RuntimeError) as ctx: self.manager._connect_cluster() + self.assertIn("Cluster with name not found", str(ctx.exception)) - mock_logger.error.assert_called_once_with( - "Unable to connect to HyperPod cluster test-cluster: Connection failed" + @patch("subprocess.run") + def test_connect_cluster_subprocess_error(self, mock_run): + mock_run.side_effect = subprocess.CalledProcessError( + 1, "hyperpod", stderr="Error: cluster not found" ) + with self.assertRaises(subprocess.CalledProcessError): + self.manager._connect_cluster() + @patch("subprocess.run") def test_get_job_status_succeeded(self, mock_run): # Mock connect-cluster call diff --git a/tests/unit/model/result/test_training_result.py b/tests/unit/model/result/test_training_result.py index 20e62f8..8f433bb 100644 --- a/tests/unit/model/result/test_training_result.py +++ b/tests/unit/model/result/test_training_result.py @@ -29,6 +29,7 @@ SMTJStatusManager, ) from amzn_nova_forge.core.result.training_result import ( + BedrockTrainingResult, SMHPTrainingResult, SMTJTrainingResult, ) @@ -60,7 +61,7 @@ def test_init_with_default_client(self): self.assertEqual(result.platform, Platform.SMTJ) self.assertIsInstance(result.status_manager, SMTJStatusManager) self.assertEqual(result.model_type, Model.NOVA_MICRO) - mock_boto3.assert_called_once_with("sagemaker") + mock_boto3.assert_called_once_with("sagemaker", region_name=None) def test_init_with_custom_client(self): """Test initialization with custom SageMaker client""" @@ -401,5 +402,41 @@ def test_all_training_methods(self): self.assertEqual(result_dict["method"], method.value) +class TestRegionPropagation(unittest.TestCase): + def setUp(self): + self.model_artifacts = ModelArtifacts( + checkpoint_s3_path="s3://bucket/checkpoint/", + output_s3_path="s3://bucket/output/", + ) + + def test_smtj_training_result_passes_region(self): + """Test SMTJTrainingResult passes region to boto3 client""" + with patch("boto3.client") as mock_boto3: + SMTJTrainingResult( + job_id="test", + started_time=datetime(2024, 1, 1, 12, 0, 0), + method=TrainingMethod.SFT_LORA, + model_artifacts=self.model_artifacts, + model_type=Model.NOVA_MICRO, + sagemaker_client=None, + region="eu-west-1", + ) + mock_boto3.assert_called_once_with("sagemaker", region_name="eu-west-1") + + def test_bedrock_training_result_passes_region(self): + """Test BedrockTrainingResult passes region to boto3 client""" + with patch("boto3.client") as mock_boto3: + BedrockTrainingResult( + job_id="test", + started_time=datetime(2024, 1, 1, 12, 0, 0), + method=TrainingMethod.SFT_LORA, + model_artifacts=self.model_artifacts, + model_type=Model.NOVA_MICRO, + bedrock_client=None, + region="eu-west-1", + ) + mock_boto3.assert_called_once_with("bedrock", region_name="eu-west-1") + + 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 bb0d168..4b15aba 100644 --- a/tests/unit/model/test_nova_model_customizer.py +++ b/tests/unit/model/test_nova_model_customizer.py @@ -83,7 +83,6 @@ def close(self): class TestNovaModelCustomizer(unittest.TestCase): def setUp(self): - self.model = Model.NOVA_MICRO self.method = TrainingMethod.SFT_LORA self.data_s3_path = "s3://test-bucket/data" @@ -1156,6 +1155,14 @@ def test_train_rft_basic_success(self, mock_boto_client, mock_uuid, mock_build_a self.assertEqual(job_config.output_s3_path, self.output_s3_path) self.assertEqual(job_config.recipe_path, "mock_recipe.yaml") + @patch("amzn_nova_forge.trainer.forge_trainer.ForgeTrainer.__init__", return_value=None) + @patch("amzn_nova_forge.trainer.forge_trainer.ForgeTrainer.train") + def test_train_passes_val_check_interval_to_forge_trainer(self, mock_ft_train, mock_ft_init): + mock_ft_train.return_value = MagicMock() + self.customizer.train(job_name="test-job", val_check_interval=500) + init_kwargs = mock_ft_init.call_args.kwargs + self.assertEqual(init_kwargs["val_check_interval"], 500) + @patch("amzn_nova_forge.recipe.recipe_builder.RecipeBuilder.build_and_validate") def test_train_runtime_manager_failure(self, mock_build_and_validate): mock_build_and_validate.return_value = ( @@ -2177,7 +2184,9 @@ def test_deploy_fails_when_endpoint_exists_and_fail_if_exists_mode( self.assertIn("Change deployment_mode", error_msg) # Verify we checked for existing deployment - mock_check_existing.assert_called_once_with("test-endpoint", DeployPlatform.BEDROCK_OD) + mock_check_existing.assert_called_once_with( + "test-endpoint", DeployPlatform.BEDROCK_OD, region="us-east-1" + ) # Verify we didn't proceed with deployment mock_bedrock_role_creation.assert_not_called() @@ -2236,7 +2245,7 @@ def test_deploy_pt_update_success_uses_existing_arn( # Verify update was called mock_update_pt.assert_called_once_with( - existing_arn, "test:new-model:arn", "test-pt-endpoint" + existing_arn, "test:new-model:arn", "test-pt-endpoint", region="us-east-1" ) # Verify NO new deployment created @@ -2279,7 +2288,9 @@ def test_deploy_succeeds_when_no_existing_endpoint( # Verify we checked for existing deployment (pre-flight + deploy) self.assertEqual(mock_check_existing.call_count, 2) - mock_check_existing.assert_called_with("new-endpoint", DeployPlatform.BEDROCK_OD) + mock_check_existing.assert_called_with( + "new-endpoint", DeployPlatform.BEDROCK_OD, region="us-east-1" + ) # Verify deployment proceeded successfully mock_bedrock_role_creation.assert_called_once() @@ -3476,10 +3487,6 @@ def client_side_effect(service, **kwargs): self.assertIn("RFT lambda verification failed", error_msg) -if __name__ == "__main__": - unittest.main() - - class TestJobCachingAndPersistence(unittest.TestCase): def setUp(self): self.model = Model.NOVA_MICRO @@ -4607,3 +4614,137 @@ def test_no_arn_no_endpoint_info_returns_unknown(self): obj.model = Model.NOVA_MICRO result = _invoke_inference_extra_info(obj) assert result == {"model": Model.NOVA_MICRO.value, "platform": "UNKNOWN"} + + +class TestDataMixingServerlessPlatform(unittest.TestCase): + """Tests for data mixing support on SMTJServerless in NovaModelCustomizer.""" + + def setUp(self): + self.data_s3_path = "s3://bucket/data.jsonl" + self.output_s3_path = "s3://bucket/output" + + def _patch_clients(self, mock_client): + mock_sts = MagicMock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + mock_s3 = MagicMock() + mock_s3.head_bucket.return_value = {} + + def client_side_effect(service, **kw): + if service == "sts": + return mock_sts + elif service == "s3": + return mock_s3 + return MagicMock() + + mock_client.side_effect = client_side_effect + + @patch("amzn_nova_forge.model.nova_model_customizer.load_recipe_templates") + @patch("boto3.client") + def test_smtj_serverless_sft_lora_data_mixing_succeeds(self, mock_client, mock_load_recipes): + """SMTJServerless + SFT_LORA + data_mixing_enabled=True must not raise.""" + self._patch_clients(mock_client) + mock_load_recipes.return_value = (MagicMock(), MagicMock(), MagicMock(), MagicMock()) + infra = create_autospec(SMTJServerlessRuntimeManager) + infra.platform = Platform.SMTJServerless + + customizer = NovaModelCustomizer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + data_s3_path=self.data_s3_path, + output_s3_path=self.output_s3_path, + data_mixing_enabled=True, + ) + self.assertTrue(customizer.data_mixing_enabled) + mock_load_recipes.assert_called_once() + + @patch("amzn_nova_forge.model.nova_model_customizer.load_recipe_templates") + @patch("boto3.client") + def test_smtj_serverless_sft_full_data_mixing_succeeds(self, mock_client, mock_load_recipes): + """SMTJServerless + SFT_FULL + data_mixing_enabled=True must not raise.""" + self._patch_clients(mock_client) + mock_load_recipes.return_value = (MagicMock(), MagicMock(), MagicMock(), MagicMock()) + infra = create_autospec(SMTJServerlessRuntimeManager) + infra.platform = Platform.SMTJServerless + + customizer = NovaModelCustomizer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_FULL, + infra=infra, + data_s3_path=self.data_s3_path, + output_s3_path=self.output_s3_path, + data_mixing_enabled=True, + ) + self.assertTrue(customizer.data_mixing_enabled) + + @patch("amzn_nova_forge.model.nova_model_customizer.load_recipe_templates") + def test_init_data_mixing_accepts_smtj_serverless(self, mock_load_recipes): + """_init_data_mixing must not raise for Platform.SMTJServerless + SFT_LORA.""" + mock_load_recipes.return_value = (MagicMock(), MagicMock(), MagicMock(), MagicMock()) + infra = create_autospec(SMTJServerlessRuntimeManager) + infra.platform = Platform.SMTJServerless + + with patch("boto3.client"): + customizer = NovaModelCustomizer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + data_s3_path=self.data_s3_path, + output_s3_path=self.output_s3_path, + data_mixing_enabled=False, # don't trigger _init_data_mixing in __init__ + ) + + # Call _init_data_mixing directly — must not raise + mock_load_recipes.reset_mock() + customizer.data_mixing_enabled = True + customizer._init_data_mixing( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + platform=Platform.SMTJServerless, + ) + # Verify load_recipe_templates was called with SMTJServerless + data_mixing_enabled=True + mock_load_recipes.assert_called() + call_kwargs = mock_load_recipes.call_args.kwargs + self.assertEqual(call_kwargs["platform"], Platform.SMTJServerless) + self.assertTrue(call_kwargs["data_mixing_enabled"]) + + @patch("boto3.client") + def test_smtj_serverless_unsupported_method_raises(self, mock_client): + """SMTJServerless + RFT_LORA + data_mixing_enabled=True must raise.""" + self._patch_clients(mock_client) + infra = create_autospec(SMTJServerlessRuntimeManager) + infra.platform = Platform.SMTJServerless + + with self.assertRaises(ValueError) as ctx: + NovaModelCustomizer( + model=Model.NOVA_MICRO, + method=TrainingMethod.RFT_LORA, + infra=infra, + data_s3_path=self.data_s3_path, + output_s3_path=self.output_s3_path, + data_mixing_enabled=True, + ) + self.assertIn("Data mixing is only supported for", str(ctx.exception)) + self.assertIn("SMTJServerless", str(ctx.exception)) + + @patch("boto3.client") + def test_smtj_non_serverless_still_raises(self, mock_client): + """Regular SMTJ (non-serverless) must still raise for data mixing.""" + self._patch_clients(mock_client) + infra = create_autospec(SMTJRuntimeManager) + infra.platform = Platform.SMTJ + + with self.assertRaises(ValueError) as ctx: + NovaModelCustomizer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + data_s3_path=self.data_s3_path, + output_s3_path=self.output_s3_path, + data_mixing_enabled=True, + ) + self.assertIn("Data mixing is only supported for", str(ctx.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/monitor/test_log_monitor.py b/tests/unit/monitor/test_log_monitor.py index 0095745..19e6d34 100644 --- a/tests/unit/monitor/test_log_monitor.py +++ b/tests/unit/monitor/test_log_monitor.py @@ -21,7 +21,11 @@ SMHPStatusManager, SMTJStatusManager, ) -from amzn_nova_forge.monitor.log_monitor import CloudWatchLogMonitor +from amzn_nova_forge.monitor.log_monitor import ( + BedrockStrategy, + CloudWatchLogMonitor, + SMHPStrategy, +) class TestCloudWatchLogMonitor(unittest.TestCase): @@ -265,6 +269,7 @@ def test_from_job_result_smtj(self): platform=Platform.SMTJ, started_time=int(mock_job_result.started_time.timestamp() * 1000), cloudwatch_logs_client=mock_client, + region=None, ) def test_from_job_result_smhp(self): @@ -292,6 +297,7 @@ def test_from_job_result_smhp(self): cloudwatch_logs_client=mock_client, cluster_name="test-cluster", namespace="test-namespace", + region=None, ) @patch("boto3.client") @@ -410,7 +416,7 @@ def test_bedrock_strategy_uses_default_client_if_not_provided(self): strategy = BedrockStrategy() # Verify boto3.client was called to create bedrock client - mock_boto_client.assert_called_once_with("bedrock") + mock_boto_client.assert_called_once_with("bedrock", region_name=None) self.assertEqual(strategy.bedrock_client, mock_client) def test_bedrock_strategy_uses_provided_client(self): @@ -1030,5 +1036,35 @@ def test_plot_metrics_bedrock_get_metrics_raises_not_implemented(self): monitor.plot_metrics(TrainingMethod.SFT_FULL) +class TestRegionPropagation(unittest.TestCase): + @patch("boto3.client") + def test_cloudwatch_log_monitor_passes_region_to_logs_client(self, mock_boto_client): + mock_client = Mock() + mock_client.describe_log_streams.return_value = {"logStreams": []} + mock_boto_client.return_value = mock_client + with patch.object(CloudWatchLogMonitor, "_create_job_status_manager", return_value=Mock()): + CloudWatchLogMonitor( + job_id="test", + platform=Platform.SMTJ, + region="eu-west-1", + cloudwatch_logs_client=None, + ) + mock_boto_client.assert_any_call("logs", region_name="eu-west-1") + + @patch("boto3.client") + def test_smhp_strategy_passes_region_to_sagemaker_client(self, mock_boto_client): + mock_boto_client.return_value = Mock() + SMHPStrategy( + cluster_name="test", namespace="default", region="eu-west-1", sagemaker_client=None + ) + mock_boto_client.assert_called_once_with("sagemaker", region_name="eu-west-1") + + @patch("boto3.client") + def test_bedrock_strategy_passes_region_to_bedrock_client(self, mock_boto_client): + mock_boto_client.return_value = Mock() + BedrockStrategy(region="eu-west-1", bedrock_client=None) + mock_boto_client.assert_called_once_with("bedrock", region_name="eu-west-1") + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/recipe/test_recipe_builder.py b/tests/unit/recipe/test_recipe_builder.py index 7800eb7..13ffd42 100644 --- a/tests/unit/recipe/test_recipe_builder.py +++ b/tests/unit/recipe/test_recipe_builder.py @@ -251,7 +251,7 @@ def test_initialization_validation_data_s3_path_ignored_for_non_cpt_method(self, job_name=self.job_name, platform=self.platform, model=self.mock_model, - method=TrainingMethod.SFT_LORA, + method=TrainingMethod.DPO_LORA, instance_type=self.instance_type, instance_count=self.instance_count, infra=self.mock_infra, @@ -261,10 +261,10 @@ def test_initialization_validation_data_s3_path_ignored_for_non_cpt_method(self, ) mock_logger.info.assert_called_with( - "'validation_data_s3_path' is only applicable for CPT on SMTJ/SMHP, or RFT/SFT method on Bedrock. Will ignore." + "'validation_data_s3_path' is only applicable for CPT, SFT on SMTJ/SMHP/SMTJServerless, or any method on Bedrock. Will ignore." ) - self.assertFalse(hasattr(builder, "validation_data_s3_path")) + self.assertIsNone(builder.validation_data_s3_path) def test_initialization_dpo_lora_method(self): builder = RecipeBuilder( @@ -2628,7 +2628,7 @@ def test_build_and_validate_evaluation_uses_local_templates_for_special_tasks( ] for eval_task in special_tasks: - with self.subTest(eval_task=eval_task): + with self.subTest(eval_task=eval_task.value): mock_download_local.reset_mock() builder = RecipeBuilder( @@ -3792,6 +3792,111 @@ def test_build_and_validate_batch_sample_tracing_non_bedrock( self.assertIn("training_config", config) self.assertTrue(config["training_config"]["enable_batch_sample_tracing"]) + def test_sft_lora_smtj_stores_validation_path(self): + validation_data_s3 = "s3://bucket/validation-data" + builder = RecipeBuilder( + region=self.region, + job_name=self.job_name, + platform=Platform.SMTJ, + model=self.mock_model, + method=TrainingMethod.SFT_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, + validation_data_s3_path=validation_data_s3, + ) + + self.assertEqual(builder.validation_data_s3_path, validation_data_s3) + + def test_sft_full_smtj_stores_validation_path(self): + validation_data_s3 = "s3://bucket/validation-data" + builder = RecipeBuilder( + region=self.region, + job_name=self.job_name, + platform=Platform.SMTJ, + model=self.mock_model, + method=TrainingMethod.SFT_FULL, + 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, + validation_data_s3_path=validation_data_s3, + ) + + self.assertEqual(builder.validation_data_s3_path, validation_data_s3) + + def test_sft_lora_smtj_serverless_stores_validation_path(self): + validation_data_s3 = "s3://bucket/validation-data" + builder = RecipeBuilder( + region=self.region, + job_name=self.job_name, + platform=Platform.SMTJServerless, + model=self.mock_model, + method=TrainingMethod.SFT_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, + validation_data_s3_path=validation_data_s3, + ) + + self.assertEqual(builder.validation_data_s3_path, validation_data_s3) + + def test_sft_lora_smhp_stores_validation_path(self): + validation_data_s3 = "s3://bucket/validation-data" + builder = RecipeBuilder( + region=self.region, + job_name=self.job_name, + platform=Platform.SMHP, + model=self.mock_model, + method=TrainingMethod.SFT_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, + validation_data_s3_path=validation_data_s3, + ) + + self.assertEqual(builder.validation_data_s3_path, validation_data_s3) + + def test_sft_no_validation_data_attribute_is_none(self): + builder = RecipeBuilder( + region=self.region, + job_name=self.job_name, + platform=Platform.SMTJ, + model=self.mock_model, + method=TrainingMethod.SFT_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, + validation_data_s3_path=None, + ) + + self.assertIsNone(builder.validation_data_s3_path) + + def test_validation_data_s3_path_initialized_to_none(self): + builder = RecipeBuilder( + region=self.region, + job_name=self.job_name, + platform=Platform.SMTJ, + model=self.mock_model, + method=TrainingMethod.SFT_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, + ) + + self.assertIsNone(builder.validation_data_s3_path) + @patch("amzn_nova_forge.util.recipe.get_hub_recipe_metadata") @patch("amzn_nova_forge.util.recipe.download_templates_from_s3") @patch("amzn_nova_forge.recipe.recipe_builder.Validator") @@ -3835,6 +3940,216 @@ def test_build_and_validate_batch_sample_tracing_default_false( if "training_config" in config: self.assertNotIn("enable_batch_sample_tracing", config["training_config"]) + @patch("amzn_nova_forge.util.recipe.get_hub_recipe_metadata") + @patch("amzn_nova_forge.util.recipe.download_templates_from_s3") + @patch("amzn_nova_forge.recipe.recipe_builder.Validator") + def test_sft_lora_smhp_injects_val_check_interval( + self, mock_validator, mock_download, mock_metadata + ): + mock_metadata.return_value = {"recipe_uri": "s3://bucket/recipe"} + + recipe_template = { + "run": { + "name": "{{name}}", + "replicas": 1, + "validation_data_s3_path": "{{validation_data_s3_path}}", + "val_check_interval": "{{val_check_interval}}", + } + } + + overrides_template = { + "name": {"default": "", "type": "string"}, + "validation_data_s3_path": {"default": "", "type": "string"}, + "val_check_interval": {"default": 1000, "type": "integer"}, + } + + mock_download.return_value = (recipe_template, overrides_template, "image_uri") + + builder = RecipeBuilder( + region=self.region, + job_name=self.job_name, + platform=Platform.SMHP, + model=self.mock_model, + method=TrainingMethod.SFT_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, + validation_data_s3_path="s3://bucket/validation-data", + val_check_interval=500, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + output_path = os.path.join(tmpdir, "recipe.yaml") + recipe_path, *_ = builder.build_and_validate(output_recipe_path=output_path) + + with open(recipe_path, "r") as f: + config = yaml.safe_load(f) + + self.assertEqual(config["run"]["val_check_interval"], 500) + self.assertEqual( + config["run"]["validation_data_s3_path"], "s3://bucket/validation-data" + ) + + @patch("amzn_nova_forge.util.recipe.get_hub_recipe_metadata") + @patch("amzn_nova_forge.util.recipe.download_templates_from_s3") + @patch("amzn_nova_forge.recipe.recipe_builder.Validator") + def test_val_check_interval_not_injected_when_none( + self, mock_validator, mock_download, mock_metadata + ): + mock_metadata.return_value = {"recipe_uri": "s3://bucket/recipe"} + + recipe_template = { + "run": { + "name": "{{name}}", + "replicas": 1, + "val_check_interval": "{{val_check_interval}}", + } + } + + overrides_template = { + "name": {"default": "", "type": "string"}, + "val_check_interval": {"default": 1000, "type": "integer"}, + } + + mock_download.return_value = (recipe_template, overrides_template, "image_uri") + + builder = RecipeBuilder( + region=self.region, + job_name=self.job_name, + platform=Platform.SMHP, + model=self.mock_model, + method=TrainingMethod.SFT_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, + val_check_interval=None, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + output_path = os.path.join(tmpdir, "recipe.yaml") + recipe_path, *_ = builder.build_and_validate(output_recipe_path=output_path) + + with open(recipe_path, "r") as f: + config = yaml.safe_load(f) + + # Should use the template default (1000), not an injected value + self.assertEqual(config["run"]["val_check_interval"], 1000) + + +class TestRecipeBuilderMlflowPlaceholder(unittest.TestCase): + """Tests that mlflow placeholders are resolved even when no MLflowMonitor is provided.""" + + def setUp(self): + self.region = "us-east-1" + self.job_name = "test-job" + self.platform = Platform.SMTJ + self.method = TrainingMethod.SFT_LORA + self.instance_type = "ml.g5.12xlarge" + self.instance_count = 1 + self.data_s3 = "s3://bucket/data" + self.output_s3 = "s3://bucket/output" + + 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 = self.instance_type + self.mock_infra.instance_count = self.instance_count + + @patch("amzn_nova_forge.util.recipe.get_hub_recipe_metadata") + @patch("amzn_nova_forge.util.recipe.download_templates_from_s3") + @patch("amzn_nova_forge.recipe.recipe_builder.Validator") + def test_mlflow_placeholders_resolved_to_empty_when_no_monitor( + self, mock_validator, mock_download, mock_metadata + ): + """When no mlflow_monitor is provided, {{mlflow_*}} placeholders must resolve to ''.""" + mock_metadata.return_value = {"recipe_uri": "s3://bucket/recipe"} + + recipe_template = { + "run": { + "name": "{{name}}", + "mlflow_tracking_uri": "{{mlflow_tracking_uri}}", + "mlflow_experiment_name": "{{mlflow_experiment_name}}", + "mlflow_run_name": "{{mlflow_run_name}}", + } + } + overrides_template = {"name": {"default": "", "type": "string"}} + mock_download.return_value = (recipe_template, overrides_template, "image_uri") + + builder = RecipeBuilder( + region=self.region, + job_name=self.job_name, + platform=self.platform, + model=self.mock_model, + method=self.method, + 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, + mlflow_monitor=None, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + output_path = os.path.join(tmpdir, "recipe.yaml") + recipe_path, *_ = builder.build_and_validate(output_recipe_path=output_path) + + with open(recipe_path, "r") as f: + config = yaml.safe_load(f) + + # Placeholders must be resolved to empty strings, not left as "{{...}}" + self.assertEqual(config["run"]["mlflow_tracking_uri"], "") + self.assertEqual(config["run"]["mlflow_experiment_name"], "") + self.assertEqual(config["run"]["mlflow_run_name"], "") + + @patch("amzn_nova_forge.util.recipe.get_hub_recipe_metadata") + @patch("amzn_nova_forge.util.recipe.download_templates_from_s3") + @patch("amzn_nova_forge.recipe.recipe_builder.Validator") + def test_mlflow_placeholders_not_raw_when_no_monitor( + self, mock_validator, mock_download, mock_metadata + ): + """Ensure raw {{...}} strings are never present in the built recipe.""" + mock_metadata.return_value = {"recipe_uri": "s3://bucket/recipe"} + + recipe_template = { + "run": { + "name": "{{name}}", + "mlflow_tracking_uri": "{{mlflow_tracking_uri}}", + } + } + overrides_template = {"name": {"default": "", "type": "string"}} + mock_download.return_value = (recipe_template, overrides_template, "image_uri") + + builder = RecipeBuilder( + region=self.region, + job_name=self.job_name, + platform=self.platform, + model=self.mock_model, + method=self.method, + 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, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + output_path = os.path.join(tmpdir, "recipe.yaml") + recipe_path, *_ = builder.build_and_validate(output_recipe_path=output_path) + + with open(recipe_path, "r") as f: + raw = f.read() + + self.assertNotIn("{{mlflow_tracking_uri}}", raw) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/telemetry/test_telemetry_logging.py b/tests/unit/telemetry/test_telemetry_logging.py index fceea19..4ed8a3f 100644 --- a/tests/unit/telemetry/test_telemetry_logging.py +++ b/tests/unit/telemetry/test_telemetry_logging.py @@ -169,7 +169,7 @@ def test_returns_account_id(self, mock_boto3): result = _get_accountId() self.assertEqual(result, "123456789012") - mock_boto3.client.assert_called_once_with("sts") + mock_boto3.client.assert_called_once_with("sts", region_name=None) @patch("amzn_nova_forge.telemetry.telemetry_logging.boto3") def test_returns_none_on_exception(self, mock_boto3): diff --git a/tests/unit/trainer/test_forge_trainer.py b/tests/unit/trainer/test_forge_trainer.py index 9ace0dc..7b59650 100644 --- a/tests/unit/trainer/test_forge_trainer.py +++ b/tests/unit/trainer/test_forge_trainer.py @@ -263,6 +263,68 @@ def test_batch_sample_tracing_rejects_non_cpt(self, mock_session): ) self.assertIn("CPT", str(ctx.exception)) + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_hub_content_version_propagates_to_infra(self, mock_session, _mock_output): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + infra = _make_smtj_infra() + infra.hub_content_version = None + + ForgeTrainer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + hub_content_version="3.38.0", + ) + + self.assertEqual(infra.hub_content_version, "3.38.0") + + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_hub_content_version_none_does_not_set_infra(self, mock_session, _mock_output): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + infra = _make_smtj_infra() + infra.hub_content_version = None + + ForgeTrainer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + hub_content_version=None, + ) + + self.assertIsNone(infra.hub_content_version) + + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_hub_content_version_skips_infra_without_attribute(self, mock_session, _mock_output): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + infra = _make_smtj_infra() + # Remove the attribute so hasattr returns False + if hasattr(infra, "hub_content_version"): + del infra.hub_content_version + + ForgeTrainer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + hub_content_version="3.38.0", + ) + + self.assertFalse(hasattr(infra, "hub_content_version")) + class TestForgeTrainerTrain(unittest.TestCase): """Tests for ForgeTrainer.train().""" @@ -337,6 +399,30 @@ def test_dry_run_returns_none(self, MockRecipeBuilder): self.assertIsNone(result) infra.execute.assert_not_called() + @patch("amzn_nova_forge.trainer.forge_trainer.RecipeBuilder") + def test_dry_run_with_validation_data(self, MockRecipeBuilder): + mock_builder = MockRecipeBuilder.return_value + mock_builder.build_and_validate.return_value = ( + "/tmp/recipe.yaml", + "s3://bucket/output", + "s3://bucket/data", + "image:latest", + ) + + infra = _make_smtj_infra() + trainer = self._make_trainer( + infra=infra, + holdout_data_s3_path="s3://bucket/validation.jsonl", + ) + result = trainer.train(job_name="my-job", dry_run=True) + + self.assertIsNone(result) + infra.execute.assert_not_called() + MockRecipeBuilder.assert_called_once() + call_kwargs = MockRecipeBuilder.call_args[1] + self.assertEqual(call_kwargs["validation_data_s3_path"], "s3://bucket/validation.jsonl") + mock_builder.build_and_validate.assert_called_once() + @patch("amzn_nova_forge.trainer.forge_trainer.validate_rft_lambda_name") @patch("amzn_nova_forge.trainer.forge_trainer.get_model_artifacts") @patch("amzn_nova_forge.trainer.forge_trainer.RecipeBuilder") @@ -580,6 +666,7 @@ def test_get_logs_with_job_result(self, MockMonitor): job_id="job-abc", platform=Platform.SMTJ, started_time=int(datetime(2025, 1, 1, tzinfo=timezone.utc).timestamp() * 1000), + region="us-east-1", ) MockMonitor.return_value.show_logs.assert_called_once() @@ -594,6 +681,7 @@ def test_get_logs_with_explicit_ids(self, MockMonitor): job_id="explicit-job", platform=Platform.SMTJ, started_time=int(started.timestamp() * 1000), + region="us-east-1", ) @patch("amzn_nova_forge.trainer.forge_trainer.CloudWatchLogMonitor") @@ -620,6 +708,7 @@ def test_get_logs_smhp_includes_cluster_kwargs(self, MockMonitor): started_time=int(started.timestamp() * 1000), cluster_name="my-cluster", namespace="kubeflow", + region="us-east-1", ) @@ -751,5 +840,295 @@ def test_warns_when_tracing_not_enabled(self, mock_run): self.assertTrue(any("not enabled" in msg for msg in cm.output)) +class TestForgeTrainerDataMixingServerless(unittest.TestCase): + """Tests for data mixing support on SMTJServerless.""" + + def _make_serverless_infra(self): + from amzn_nova_forge.manager.runtime_manager import SMTJServerlessRuntimeManager + + infra = create_autospec(SMTJServerlessRuntimeManager) + infra.instance_type = None + infra.instance_count = None + infra.kms_key_id = None + infra.platform = Platform.SMTJServerless + return infra + + @patch("amzn_nova_forge.trainer.forge_trainer.load_recipe_templates") + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_data_mixing_smtj_serverless_sft_lora_works( + self, mock_session, _mock_output, mock_load_templates + ): + """SMTJServerless + SFT_LORA + data_mixing_enabled=True should succeed.""" + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + mock_load_templates.return_value = ({}, {}, None, "image:latest") + infra = self._make_serverless_infra() + + trainer = ForgeTrainer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + data_mixing_enabled=True, + ) + self.assertIsNotNone(trainer.data_mixing) + + @patch("amzn_nova_forge.trainer.forge_trainer.load_recipe_templates") + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_data_mixing_smtj_serverless_sft_full_works( + self, mock_session, _mock_output, mock_load_templates + ): + """SMTJServerless + SFT_FULL + data_mixing_enabled=True should succeed.""" + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + mock_load_templates.return_value = ({}, {}, None, "image:latest") + infra = self._make_serverless_infra() + + trainer = ForgeTrainer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_FULL, + infra=infra, + training_data_s3_path="s3://bucket/data", + data_mixing_enabled=True, + ) + self.assertIsNotNone(trainer.data_mixing) + + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_data_mixing_smtj_serverless_unsupported_method_raises( + self, mock_session, _mock_output + ): + """SMTJServerless + unsupported method (RFT_LORA) should raise.""" + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + infra = self._make_serverless_infra() + + with self.assertRaises(ValueError) as ctx: + ForgeTrainer( + model=Model.NOVA_MICRO, + method=TrainingMethod.RFT_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + data_mixing_enabled=True, + ) + self.assertIn("Data mixing is only supported", str(ctx.exception)) + self.assertIn("SMTJServerless", str(ctx.exception)) + + @patch("amzn_nova_forge.trainer.forge_trainer.load_recipe_templates") + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_data_mixing_config_in_job_config_params_for_serverless( + self, mock_session, _mock_output, mock_load_templates + ): + """data_mixing_config must be passed in JobConfig to infra.execute() for SMTJServerless.""" + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + mock_load_templates.return_value = ({}, {}, None, "image:latest") + infra = self._make_serverless_infra() + infra.execute.return_value = "serverless-job-id" + + trainer = ForgeTrainer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + data_mixing_enabled=True, + config=ForgeConfig(output_s3_path=FIXED_OUTPUT_PATH), + ) + trainer.data_mixing.set_config({"customer_data_percent": 50, "nova_code_percent": 100}) + + with ( + patch("amzn_nova_forge.trainer.forge_trainer.RecipeBuilder") as mock_rb_cls, + patch("amzn_nova_forge.trainer.forge_trainer.load_existing_result", return_value=None), + patch("amzn_nova_forge.trainer.forge_trainer.get_model_artifacts") as mock_gma, + patch("boto3.client"), + ): + mock_rb = mock_rb_cls.return_value + mock_rb.build_and_validate.return_value = ( + "/tmp/recipe.yaml", + FIXED_OUTPUT_PATH, + "s3://bucket/data", + "image:latest", + ) + mock_gma.return_value = MagicMock() + trainer.train(job_name="test-datamix-job") + + infra.execute.assert_called_once() + job_config = infra.execute.call_args.kwargs["job_config"] + self.assertEqual( + job_config.data_mixing_config, + {"customer_data_percent": 50, "nova_code_percent": 100}, + ) + + @patch("amzn_nova_forge.trainer.forge_trainer.load_recipe_templates") + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_data_mixing_not_set_for_smhp(self, mock_session, _mock_output, mock_load_templates): + """For SMHP, data_mixing_config must NOT be passed in JobConfig to infra.execute().""" + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + mock_load_templates.return_value = ({}, {}, None, "image:latest") + infra = _make_smhp_infra() + infra.execute.return_value = "smhp-job-id" + + trainer = ForgeTrainer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + data_mixing_enabled=True, + config=ForgeConfig(output_s3_path=FIXED_OUTPUT_PATH), + ) + trainer.data_mixing.set_config({"customer_data_percent": 50, "nova_code_percent": 100}) + + with ( + patch("amzn_nova_forge.trainer.forge_trainer.RecipeBuilder") as mock_rb_cls, + patch("amzn_nova_forge.trainer.forge_trainer.load_existing_result", return_value=None), + patch("amzn_nova_forge.trainer.forge_trainer.get_model_artifacts") as mock_gma, + ): + mock_rb = mock_rb_cls.return_value + mock_rb.build_and_validate.return_value = ( + "/tmp/recipe.yaml", + FIXED_OUTPUT_PATH, + "s3://bucket/data", + "image:latest", + ) + mock_gma.return_value = MagicMock() + trainer.train(job_name="test-datamix-smhp-job") + + infra.execute.assert_called_once() + job_config = infra.execute.call_args.kwargs["job_config"] + self.assertIsNone(job_config.data_mixing_config) + + +class TestForgeTrainerValCheckInterval(unittest.TestCase): + """Tests for val_check_interval in ForgeTrainer.""" + + def _make_trainer(self, **kwargs): + infra = _make_smtj_infra() + defaults = dict( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + ) + defaults.update(kwargs) + with ( + patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ), + patch("boto3.session.Session") as mock_session, + ): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + return ForgeTrainer(**defaults) + + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_val_check_interval_valid(self, mock_session, _mock_output): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + trainer = ForgeTrainer( + model=Model.NOVA_MICRO, + method=TrainingMethod.SFT_LORA, + infra=_make_smtj_infra(), + training_data_s3_path="s3://bucket/data", + holdout_data_s3_path="s3://bucket/val", + val_check_interval=500, + ) + self.assertEqual(trainer.val_check_interval, 500) + + def test_val_check_interval_zero_raises(self): + with self.assertRaises(ValueError): + self._make_trainer(val_check_interval=0) + + def test_val_check_interval_negative_raises(self): + with self.assertRaises(ValueError): + self._make_trainer(val_check_interval=-1) + + def test_val_check_interval_float_raises(self): + with self.assertRaises((ValueError, TypeError)): + self._make_trainer(val_check_interval=5.0) + + def test_val_check_interval_without_validation_data_warns(self): + with patch("amzn_nova_forge.trainer.forge_trainer.logger") as mock_logger: + self._make_trainer(val_check_interval=500, holdout_data_s3_path=None) + mock_logger.warning.assert_called_once() + self.assertIn( + "val_check_interval", + mock_logger.warning.call_args[0][0], + ) + + @patch("amzn_nova_forge.trainer.forge_trainer.get_model_artifacts") + @patch("amzn_nova_forge.trainer.forge_trainer.RecipeBuilder") + def test_val_check_interval_threads_to_job_config(self, MockRecipeBuilder, mock_get_artifacts): + mock_builder = MockRecipeBuilder.return_value + mock_builder.build_and_validate.return_value = ( + "/tmp/recipe.yaml", + "s3://bucket/output", + "s3://bucket/data", + "image:latest", + ) + infra = _make_smtj_infra() + infra.execute.return_value = "job-123" + mock_get_artifacts.return_value = ModelArtifacts( + checkpoint_s3_path="s3://bucket/checkpoint", + output_s3_path="s3://bucket/output", + ) + + with patch("boto3.client"): + trainer = self._make_trainer( + infra=infra, + val_check_interval=500, + holdout_data_s3_path="s3://bucket/val", + ) + trainer.train(job_name="my-job") + + call_args = infra.execute.call_args + job_config = call_args.kwargs["job_config"] + self.assertEqual(job_config.trainer_config_hyperparameters, {"val_check_interval": "500"}) + + @patch("amzn_nova_forge.trainer.forge_trainer.get_model_artifacts") + @patch("amzn_nova_forge.trainer.forge_trainer.RecipeBuilder") + def test_val_check_interval_none_no_hyperparameters( + self, MockRecipeBuilder, mock_get_artifacts + ): + mock_builder = MockRecipeBuilder.return_value + mock_builder.build_and_validate.return_value = ( + "/tmp/recipe.yaml", + "s3://bucket/output", + "s3://bucket/data", + "image:latest", + ) + infra = _make_smtj_infra() + infra.execute.return_value = "job-123" + mock_get_artifacts.return_value = ModelArtifacts( + checkpoint_s3_path="s3://bucket/checkpoint", + output_s3_path="s3://bucket/output", + ) + + with patch("boto3.client"): + trainer = self._make_trainer(infra=infra) + trainer.train(job_name="my-job") + + call_args = infra.execute.call_args + job_config = call_args.kwargs["job_config"] + self.assertIsNone(job_config.trainer_config_hyperparameters) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/util/test_bedrock.py b/tests/unit/util/test_bedrock.py index 1ab5a76..e3bfa54 100644 --- a/tests/unit/util/test_bedrock.py +++ b/tests/unit/util/test_bedrock.py @@ -92,5 +92,78 @@ def test_timeout(self, mock_time, mock_sleep): wait_for_model_ready(client, self.MODEL_ARN, timeout=10) +class TestRegionPropagation(unittest.TestCase): + """Verify that region is passed to boto3.client for bedrock utility functions.""" + + REGION = "eu-west-1" + + @patch("amzn_nova_forge.util.bedrock.boto3.client") + def test_check_deployment_status_passes_region(self, mock_boto_client): + from amzn_nova_forge.core.enums import DeployPlatform + from amzn_nova_forge.util.bedrock import check_deployment_status + + mock_client = MagicMock() + mock_client.get_custom_model_deployment.return_value = {"status": "Active"} + mock_boto_client.return_value = mock_client + + check_deployment_status( + "arn:aws:bedrock:eu-west-1:123456789012:deployment/test", + DeployPlatform.BEDROCK_OD, + region=self.REGION, + ) + + mock_boto_client.assert_called_once_with("bedrock", region_name=self.REGION) + + @patch("amzn_nova_forge.util.bedrock.boto3.client") + def test_check_existing_deployment_passes_region(self, mock_boto_client): + from amzn_nova_forge.core.enums import DeployPlatform + from amzn_nova_forge.util.bedrock import check_existing_deployment + + mock_client = MagicMock() + mock_client.list_custom_model_deployments.return_value = {"modelDeploymentSummaries": []} + mock_boto_client.return_value = mock_client + + check_existing_deployment("my-endpoint", DeployPlatform.BEDROCK_OD, region=self.REGION) + + mock_boto_client.assert_called_once_with("bedrock", region_name=self.REGION) + + @patch("amzn_nova_forge.util.bedrock.time.sleep") + @patch("amzn_nova_forge.util.bedrock.boto3.client") + def test_delete_existing_deployment_passes_region(self, mock_boto_client, mock_sleep): + from amzn_nova_forge.core.enums import DeployPlatform + from amzn_nova_forge.util.bedrock import delete_existing_deployment + + mock_client = MagicMock() + mock_client.delete_custom_model_deployment.return_value = {} + mock_client.get_custom_model_deployment.return_value = {"status": "DELETED"} + mock_boto_client.return_value = mock_client + + delete_existing_deployment( + "arn:aws:bedrock:eu-west-1:123456789012:deployment/test", + DeployPlatform.BEDROCK_OD, + "my-endpoint", + region=self.REGION, + ) + + mock_boto_client.assert_called_once_with("bedrock", region_name=self.REGION) + + @patch("amzn_nova_forge.util.bedrock.boto3.client") + def test_update_provisioned_throughput_model_passes_region(self, mock_boto_client): + from amzn_nova_forge.util.bedrock import update_provisioned_throughput_model + + mock_client = MagicMock() + mock_client.update_provisioned_model_throughput.return_value = {} + mock_boto_client.return_value = mock_client + + update_provisioned_throughput_model( + "arn:aws:bedrock:eu-west-1:123456789012:pt/test", + "arn:aws:bedrock:eu-west-1:123456789012:model/new", + "my-endpoint", + region=self.REGION, + ) + + mock_boto_client.assert_called_once_with("bedrock", region_name=self.REGION) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/util/test_dataset_writer.py b/tests/unit/util/test_dataset_writer.py new file mode 100644 index 0000000..2b7b744 --- /dev/null +++ b/tests/unit/util/test_dataset_writer.py @@ -0,0 +1,32 @@ +# 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.util.dataset_writer.""" + +import unittest +from unittest.mock import MagicMock, patch + +from amzn_nova_forge.util.dataset_writer import DatasetWriter + + +class TestDatasetWriterSaveToS3(unittest.TestCase): + @patch("amzn_nova_forge.util.dataset_writer.boto3") + def test_save_to_s3_propagates_region(self, mock_boto3): + mock_s3 = MagicMock() + mock_boto3.client.return_value = mock_s3 + + DatasetWriter.save_to_s3( + "s3://bucket/key.jsonl", iter([{"a": 1}]), is_jsonl=True, region="eu-west-1" + ) + + mock_boto3.client.assert_called_once_with("s3", region_name="eu-west-1") diff --git a/tests/unit/util/test_mlflow.py b/tests/unit/util/test_mlflow.py index 1d8f6ba..6abdc25 100644 --- a/tests/unit/util/test_mlflow.py +++ b/tests/unit/util/test_mlflow.py @@ -75,7 +75,7 @@ def test_get_default_mlflow_tracking_uri_no_region(self, mock_client_func): # Verify self.assertEqual(result, "arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-XYZ789") - mock_client_func.assert_called_once_with("sagemaker") + mock_client_func.assert_called_once_with("sagemaker", region_name=None) @patch("boto3.client") def test_get_default_mlflow_tracking_uri_updating_status(self, mock_client_func): diff --git a/tests/unit/util/test_recipe.py b/tests/unit/util/test_recipe.py index 538472c..b5b0563 100644 --- a/tests/unit/util/test_recipe.py +++ b/tests/unit/util/test_recipe.py @@ -30,6 +30,7 @@ from amzn_nova_forge.util.recipe import ( FileLoadError, RecipePath, + _get_aws_account_id, _get_smhp_replicas_enum, _parse_s3_uri, _validate_extension, @@ -2290,5 +2291,32 @@ def test_evaluation_method_skips_smhp_enum_lookup( mock_smhp_enum.assert_not_called() +class TestRegionPropagation(unittest.TestCase): + """Tests that region is correctly propagated to boto3.client calls.""" + + @patch("amzn_nova_forge.util.recipe.boto3.client") + def test_load_file_content_passes_region_to_s3_client(self, mock_boto_client): + mock_s3 = MagicMock() + mock_body = MagicMock() + mock_body.iter_lines.return_value = iter([b"content"]) + mock_s3.get_object.return_value = {"Body": mock_body} + mock_boto_client.return_value = mock_s3 + + list(load_file_content("s3://bucket/key", region="eu-west-1")) + + mock_boto_client.assert_called_once_with("s3", region_name="eu-west-1") + + @patch("amzn_nova_forge.util.recipe.boto3.client") + def test_get_aws_account_id_passes_region_to_sts_client(self, mock_boto_client): + mock_sts = MagicMock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + mock_boto_client.return_value = mock_sts + + result = _get_aws_account_id(region="eu-west-1") + + mock_boto_client.assert_called_once_with("sts", region_name="eu-west-1") + self.assertEqual(result, "123456789012") + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/util/test_s3_utils.py b/tests/unit/util/test_s3_utils.py index b1c158e..c98016c 100644 --- a/tests/unit/util/test_s3_utils.py +++ b/tests/unit/util/test_s3_utils.py @@ -98,5 +98,26 @@ def test_permission_error_on_forbidden(self, mock_boto3): ensure_bucket_exists("someone-elses-bucket", region="us-east-1") +class TestRegionPropagation(unittest.TestCase): + @patch("amzn_nova_forge.util.s3_utils.boto3") + def test_get_dataprep_bucket_name_passes_region_to_sts_client(self, mock_boto3): + mock_sts = MagicMock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + mock_boto3.client.return_value = mock_sts + + get_dataprep_bucket_name(region="eu-west-1") + + mock_boto3.client.assert_called_once_with("sts", region_name="eu-west-1") + + @patch("amzn_nova_forge.util.s3_utils.boto3") + def test_ensure_bucket_exists_passes_region_to_s3_client(self, mock_boto3): + mock_s3 = MagicMock() + mock_boto3.client.return_value = mock_s3 + + ensure_bucket_exists("my-bucket", region="eu-west-1") + + mock_boto3.client.assert_called_once_with("s3", region_name="eu-west-1") + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/util/test_subprocess_utils.py b/tests/unit/util/test_subprocess_utils.py new file mode 100644 index 0000000..ecd18fa --- /dev/null +++ b/tests/unit/util/test_subprocess_utils.py @@ -0,0 +1,181 @@ +# 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. +import logging +import unittest + +from amzn_nova_forge.util.subprocess_utils import _check_hyperpod_stderr + + +class TestCheckHyperpodStderr(unittest.TestCase): + """Tests for _check_hyperpod_stderr utility function.""" + + def test_empty_string(self): + _check_hyperpod_stderr("") + + def test_none(self): + _check_hyperpod_stderr(None) + + def test_whitespace_only(self): + _check_hyperpod_stderr(" \n\n ") + + def test_benign_not_openssl_warning_with_prefix(self): + _check_hyperpod_stderr( + "/opt/homebrew/lib/python3.12/site-packages/urllib3/__init__.py:35: " + "NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+" + ) + + def test_benign_insecure_request_with_prefix(self): + _check_hyperpod_stderr( + "urllib3/connectionpool.py:1045: InsecureRequestWarning: Unverified HTTPS request" + ) + + def test_benign_deprecation_with_prefix(self): + _check_hyperpod_stderr( + "/path/to/kubernetes/client.py:10: DeprecationWarning: some old API is deprecated" + ) + + def test_benign_future_warning_with_prefix(self): + _check_hyperpod_stderr( + "/path/to/module.py:42: FutureWarning: this will change in a future version" + ) + + def test_benign_resource_warning_with_prefix(self): + _check_hyperpod_stderr("/path/to/subprocess.py:99: ResourceWarning: unclosed file") + + def test_benign_user_warning_with_prefix(self): + _check_hyperpod_stderr("/path/to/lib.py:5: UserWarning: some user warning") + + def test_benign_urllib3_subclass_warning_with_prefix(self): + _check_hyperpod_stderr("/path/to/urllib3/__init__.py:35: urllib3.NotOpenSSLWarning: test") + + def test_benign_multiline_all_with_prefix(self): + stderr = ( + "/path/to/urllib3/__init__.py:35: NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+\n" + "/path/to/module.py:10: DeprecationWarning: old API\n" + "/path/to/other.py:5: FutureWarning: will change\n" + ) + _check_hyperpod_stderr(stderr) + + def test_benign_bare_not_openssl_warning(self): + _check_hyperpod_stderr("NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+") + + def test_benign_bare_deprecation_warning(self): + _check_hyperpod_stderr("DeprecationWarning: some old API is deprecated") + + def test_benign_bare_future_warning(self): + _check_hyperpod_stderr("FutureWarning: this will change in a future version") + + def test_benign_bare_insecure_request(self): + _check_hyperpod_stderr("InsecureRequestWarning: Unverified HTTPS request") + + def test_benign_bare_resource_warning(self): + _check_hyperpod_stderr("ResourceWarning: unclosed file") + + def test_benign_bare_user_warning(self): + _check_hyperpod_stderr("UserWarning: some user warning") + + def test_real_error_raises(self): + with self.assertRaises(RuntimeError) as ctx: + _check_hyperpod_stderr("Cluster with name not found") + self.assertIn("Cluster with name not found", str(ctx.exception)) + + def test_real_error_connection_failed(self): + with self.assertRaises(RuntimeError): + _check_hyperpod_stderr("Connection failed") + + def test_urllib3_error_not_suppressed(self): + """urllib3 MaxRetryError should NOT be treated as benign.""" + with self.assertRaises(RuntimeError): + _check_hyperpod_stderr( + "urllib3.exceptions.MaxRetryError: HTTPSConnectionPool(host='eks.us-east-1.amazonaws.com'): " + "Max retries exceeded" + ) + + def test_warning_in_prose_not_suppressed(self): + """A line mentioning 'DeprecationWarning' in prose should NOT be benign.""" + with self.assertRaises(RuntimeError): + _check_hyperpod_stderr("Error: DeprecationWarning handling failed in module X") + + def test_two_line_warning_format(self): + """Python warnings module outputs warning + indented source line.""" + stderr = ( + "/path/to/urllib3/__init__.py:35: NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+\n" + " import urllib3\n" + ) + _check_hyperpod_stderr(stderr) + + def test_two_line_warning_mixed_with_error(self): + """Two-line warning followed by a real error should raise.""" + stderr = ( + "/path/to/urllib3/__init__.py:35: NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+\n" + " import urllib3\n" + "Error: cluster not found\n" + ) + with self.assertRaises(RuntimeError) as ctx: + _check_hyperpod_stderr(stderr) + self.assertIn("Error: cluster not found", str(ctx.exception)) + + def test_warning_then_blank_then_indented_error_not_suppressed(self): + """An indented line after a blank should NOT be suppressed.""" + stderr = ( + "/path/to/urllib3/__init__.py:35: NotOpenSSLWarning: msg\n" + " import urllib3\n" + "\n" + " real indented error\n" + ) + with self.assertRaises(RuntimeError) as ctx: + _check_hyperpod_stderr(stderr) + self.assertIn("real indented error", str(ctx.exception)) + + def test_mixed_benign_and_error_raises(self): + stderr = ( + "/path/to/urllib3/__init__.py:35: NotOpenSSLWarning: urllib3 v2 only supports OpenSSL 1.1.1+\n" + "Error: cluster not found\n" + ) + with self.assertRaises(RuntimeError) as ctx: + _check_hyperpod_stderr(stderr) + self.assertIn("Error: cluster not found", str(ctx.exception)) + # Full stderr is returned for debugging (benign lines included) + self.assertIn("NotOpenSSLWarning", str(ctx.exception)) + + def test_mixed_non_warning_urllib3_line_raises(self): + """A bare 'urllib3 warning line' without proper format is non-benign.""" + stderr = ( + "/path/to/module.py:10: DeprecationWarning: old API\n" + "/path/to/other.py:5: FutureWarning: will change\n" + "fatal: something went wrong\n" + "urllib3 warning line\n" + ) + with self.assertRaises(RuntimeError) as ctx: + _check_hyperpod_stderr(stderr) + error_msg = str(ctx.exception) + self.assertIn("fatal: something went wrong", error_msg) + self.assertIn("urllib3 warning line", error_msg) + + def test_benign_logs_debug(self): + with self.assertLogs("amzn_nova_forge.util.subprocess_utils", level=logging.DEBUG) as cm: + _check_hyperpod_stderr( + "/path/to/urllib3/__init__.py:35: NotOpenSSLWarning: test warning" + ) + self.assertTrue(any("benign" in msg.lower() for msg in cm.output)) + + def test_error_logs_error(self): + with self.assertLogs("amzn_nova_forge.util.subprocess_utils", level=logging.ERROR) as cm: + with self.assertRaises(RuntimeError): + _check_hyperpod_stderr("real error message") + self.assertTrue(any("real error message" in msg for msg in cm.output)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/validation/test_validator.py b/tests/unit/validation/test_validator.py index 5732976..a0c3fc8 100644 --- a/tests/unit/validation/test_validator.py +++ b/tests/unit/validation/test_validator.py @@ -2563,6 +2563,21 @@ def test_validate_cluster_name_does_not_raise_exception(self): Validator.validate_cluster_name("good_cluster-name") +class TestCrossAccountRegionPropagation(unittest.TestCase): + """Test that _is_cross_account_role propagates region to boto3 clients.""" + + @patch("amzn_nova_forge.validation.validator.boto3.client") + def test_is_cross_account_role_passes_region_to_sts_client(self, mock_boto3_client): + """Verify _is_cross_account_role passes region to boto3.client('sts', region_name=...).""" + mock_sts = Mock() + mock_sts.get_caller_identity.return_value = {"Account": "12345679012"} + mock_boto3_client.return_value = mock_sts + + Validator._is_cross_account_role("arn:aws:iam::123456789012:role/test", region="eu-west-1") + + mock_boto3_client.assert_called_once_with("sts", region_name="eu-west-1") + + class TestBedrockRegionValidation(unittest.TestCase): """Test cases for Bedrock region validation""" @@ -2695,5 +2710,131 @@ def test_smtj_any_name_passes(self): validate_rft_lambda_name("SageMaker-reward", Platform.SMTJ) +class TestValidatorNoneTypeHandling(unittest.TestCase): + """Tests for the None value handling fix in _validate_recipe.""" + + def setUp(self): + self.mock_infra = Mock(spec=SMHPRuntimeManager) + self.mock_infra.cluster_name = "test-cluster" + self.mock_infra.instance_type = "ml.p5.48xlarge" + self.mock_infra.instance_count = 4 + self.mock_infra.region = "us-east-1" + + def test_none_value_optional_field_does_not_raise(self): + """Optional fields with None value (e.g. new hub content fields) must not raise.""" + recipe = { + "model_package_group": None, + "val_check_interval": None, + "max_steps": 100, + } + overrides_template = { + "model_package_group": {"type": "string"}, # optional, no required=True + "val_check_interval": {"type": "integer"}, # optional, no required=True + "max_steps": {"type": "integer", "default": 100}, + } + errors = [] + Validator._validate_recipe( + recipe=recipe, + overrides_template=overrides_template, + instance_type="ml.p5.48xlarge", + errors=errors, + method=TrainingMethod.SFT_LORA, + ) + self.assertEqual(errors, []) + + def test_none_value_required_field_still_raises(self): + """Required fields with None value must still produce a type error.""" + recipe = {"max_steps": None} + overrides_template = { + "max_steps": {"type": "integer", "required": True}, + } + errors = [] + Validator._validate_recipe( + recipe=recipe, + overrides_template=overrides_template, + instance_type="ml.p5.48xlarge", + errors=errors, + method=TrainingMethod.SFT_LORA, + ) + self.assertTrue( + any("expects integer" in e and "NoneType" in e for e in errors), + f"Expected type error for required None field, got: {errors}", + ) + + def test_valid_string_field_passes(self): + """A properly set string field must not produce errors.""" + recipe = {"model_package_group": "my-group"} + overrides_template = {"model_package_group": {"type": "string"}} + errors = [] + Validator._validate_recipe( + recipe=recipe, + overrides_template=overrides_template, + instance_type="ml.p5.48xlarge", + errors=errors, + method=TrainingMethod.SFT_LORA, + ) + self.assertEqual(errors, []) + + def test_wrong_type_still_raises(self): + """A field with wrong type (not None) must still produce a type error.""" + recipe = {"max_steps": "not-an-int"} + overrides_template = {"max_steps": {"type": "integer"}} + errors = [] + Validator._validate_recipe( + recipe=recipe, + overrides_template=overrides_template, + instance_type="ml.p5.48xlarge", + errors=errors, + method=TrainingMethod.SFT_LORA, + ) + self.assertTrue( + any("expects integer" in e for e in errors), + f"Expected type error, got: {errors}", + ) + + +class TestValidationDataS3Path(unittest.TestCase): + """Tests for validation_data_s3_path preflight validation.""" + + @patch("amzn_nova_forge.validation.validator.boto3.client") + def test_preflight_validates_validation_data_s3_path(self, mock_boto3_client): + """Assert validation runs for validation_data_s3_path when set.""" + mock_infra = Mock(spec=SMTJRuntimeManager) + mock_infra.instance_type = "ml.p5.48xlarge" + mock_infra.region = "us-east-1" + + with self.assertRaises(ValueError) as context: + Validator.validate( + platform=Platform.SMTJ, + method=TrainingMethod.SFT_LORA, + infra=mock_infra, + recipe={}, + overrides_template={}, + validation_data_s3_path="not-an-s3-path", + validation_config=ValidationConfig(iam=False, infra=False), + ) + self.assertIn("Invalid S3 path for validation_data_s3_path", str(context.exception)) + + @patch("amzn_nova_forge.validation.validator.boto3.client") + def test_preflight_skips_validation_data_when_none(self, mock_boto3_client): + """Assert no validation error when validation_data_s3_path is None.""" + mock_infra = Mock(spec=SMTJRuntimeManager) + mock_infra.instance_type = "ml.p5.48xlarge" + mock_infra.region = "us-east-1" + + try: + Validator.validate( + platform=Platform.SMTJ, + method=TrainingMethod.SFT_LORA, + infra=mock_infra, + recipe={}, + overrides_template={}, + validation_data_s3_path=None, + validation_config=ValidationConfig(iam=False, infra=False), + ) + except ValueError as e: + self.assertNotIn("validation_data_s3_path", str(e)) + + if __name__ == "__main__": unittest.main()