From ca198879f8b9b0a0ce5c102f7e9e05ac0fcf54ce Mon Sep 17 00:00:00 2001 From: Lisa Ni Date: Wed, 19 Aug 2026 17:55:40 +0000 Subject: [PATCH 1/2] Deprecate SDK, scope IAM permissions to account/region ARNs - Mark amzn-nova-forge deprecated (pyproject classifier + import-time DeprecationWarning) pointing to sagemaker>=3.19.0; bump to 1.4.142 - Scope SMTJ/Bedrock/Glue and iam:GetRole/PassRole permissions to account- and region-specific ARNs instead of wildcard '*' - Extract get_caller_account_id into util/aws_utils.py and reuse across runtime managers - Validator raises ValueError with a skip-IAM hint when a resource lambda cannot be evaluated (infra is None) - Add unit tests for aws_utils and lambda resource specs Update README and docs/spec + docs/user-guides accordingly. --- README.md | 876 +++++++----------- pyproject.toml | 5 +- src/amzn_nova_forge/__init__.py | 9 + src/amzn_nova_forge/__version__.py | 2 +- .../manager/glue_runtime_manager.py | 46 +- .../manager/runtime_manager.py | 120 ++- src/amzn_nova_forge/util/aws_utils.py | 46 + src/amzn_nova_forge/validation/validator.py | 7 +- tests/unit/util/test_aws_utils.py | 86 ++ tests/unit/validation/test_validator.py | 76 +- 10 files changed, 690 insertions(+), 583 deletions(-) create mode 100644 src/amzn_nova_forge/util/aws_utils.py create mode 100644 tests/unit/util/test_aws_utils.py diff --git a/README.md b/README.md index 7b5fa33..17797dc 100644 --- a/README.md +++ b/README.md @@ -1,201 +1,134 @@ +> **DEPRECATED** — This package (`amzn-nova-forge`) is deprecated and will no longer receive feature updates. +> Please use the **SageMaker Python SDK V3** (`pip install "sagemaker>=3.19.0"`) for Amazon Nova model customization. +> 📓 SageMaker SDK sample notebook: [Nova Serverless End-to-End Example on GitHub](https://github.com/aws/sagemaker-python-sdk/blob/master/v3-examples/model-customization-examples/serverless/serverless_e2e_example.ipynb) + # Amazon Nova Forge SDK A comprehensive Python SDK for fine-tuning and customizing Amazon Nova models. This SDK provides a unified interface for training, evaluation, deployment, and monitoring of Nova models across both SageMaker Training Jobs and SageMaker HyperPod. -## Table of Contents +--- + +# Migrating from Nova Forge SDK to SageMaker Python SDK V3 -- [Installation](#installation) -- [Setup](#setup) -- [Supported Models and Training Methods](#supported-models-and-training-methods) -- [Data Preparation](#data-preparation) -- [Core Modules Overview](#core-modules-overview) -- [Additional Features](#additional-features) -- [Telemetry](#telemetry) -- [Getting Started](#getting-started) -- [Security Best Practices for SDK Users](#security-best-practices-for-sdk-users) +## Why Migrate + +The `amzn-nova-forge` package is deprecated. Amazon Nova model customization functionality is available in the SageMaker Python SDK V3. ## Installation ```bash -pip install amzn-nova-forge +pip install "sagemaker>=3.19.0" ``` -* The SDK requires [sagemaker](https://pypi.org/project/sagemaker/), which is automatically set by pip. - - -## Setup - -In most cases, the SDK will inform you if the environment lacks the required setup to run a Nova customization job. - -Below are some common requirements which you can set up in advance before trying to run a job. - -### Supported Python Versions -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/user-guides/iam_setup.md` file. - -### Instances - -Nova customization jobs also require access to enough of the right instance type to run: -- 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/user-guides/instance_type_spec.md` file for the different instance types and combinations for specific jobs and methods. -### HyperPod CLI +Requires Python 3.10 or later. -For HyperPod-based customization jobs, the SDK uses the [SageMaker HyperPod CLI](https://github.com/aws/sagemaker-hyperpod-cli/) to connect to HyperPod Clusters and start jobs. +## Concept Mapping -#### Prerequisites (required for both Forge and Non-Forge customers) +| Forge SDK Concept | SageMaker SDK V3 Equivalent | +|---|---| +| `ForgeTrainer` (SFT) | `sagemaker.train.sft_trainer.SFTTrainer` | +| `ForgeTrainer` (CPT) | `sagemaker.train.cpt_trainer.CPTTrainer` | +| `ForgeTrainer` (DPO) | `sagemaker.train.dpo_trainer.DPOTrainer` | +| `ForgeTrainer` (RFT) | `sagemaker.train.rlvr_trainer.RLVRTrainer` | +| `ForgeTrainer` (MTRL) | `sagemaker.train.multi_turn_rl_trainer.MultiTurnRLTrainer` | +| `ForgeEvaluator` | `BenchMarkEvaluator`, `LLMAsJudgeEvaluator`, `InspectAIEvaluator`, `CustomScorerEvaluator`, `MultiTurnRLEvaluator` | +| `SMHPRuntimeManager` | `sagemaker.core.training.configs.HyperPodCompute` | +| `SMTJRuntimeManager` | `TrainingJobCompute` for serverful or omit for serverless | +| `data_mixing_enabled` | `sagemaker.train.data_mixing_config.DataMixingConfig` | +| `NovaModelCustomizer` | Individual trainer classes above | +| `ForgeDeployer` | `BedrockModelBuilder` or `ModelBuilder` | +| `ForgeInference` | SageMaker SDK `Predictor` / Bedrock `InvokeModel` | -1. **Install Helm 3.** Verify with `helm version`. If not installed: - ```bash - curl -fsSL -o get_helm.sh https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3 - chmod 700 get_helm.sh - ./get_helm.sh - rm -f ./get_helm.sh - ``` +## Full Quickstart Migration (Step-by-Step) -2. If you are using a Python virtual environment, activate it before installing the CLI: - ```bash - source /bin/activate - ``` +### Step 1: Import Modules -#### For Non-Forge Customers +**Before (Forge SDK):** -1. Clone the `release_v2` branch of the HyperPod CLI: - ```bash - git clone -b release_v2 https://github.com/aws/sagemaker-hyperpod-cli.git - ``` -2. Install the CLI: - ```bash - cd sagemaker-hyperpod-cli - pip install . - ``` -3. Verify the installation: - ```bash - hyperpod --help - ``` - -#### For Forge Customers - -1. Download the latest HyperPod CLI repo with Forge feature support from S3: - ```bash - aws s3 cp s3://nova-forge-c7363-206080352451-us-east-1/v1/ ./ --recursive - mkdir -p src/hyperpod_cli/sagemaker_hyperpod_recipes/launcher/nemo - git clone https://github.com/NVIDIA/NeMo-Framework-Launcher.git src/hyperpod_cli/sagemaker_hyperpod_recipes/launcher/nemo/nemo_framework_launcher --recursive - pip install -e . - ``` -2. Verify the installation: - ```bash - hyperpod --help - ``` ---- -## Supported Models and Training Methods - -### Models - -| Model | Version | Model Type | Context Length | -| ------------- | ------- | ------------------------------ | -------------- | -| `NOVA_MICRO` | 1.0 | `amazon.nova-micro-v1:0:128k` | 128k tokens | -| `NOVA_LITE` | 1.0 | `amazon.nova-lite-v1:0:300k` | 300k tokens | -| `NOVA_LITE_2` | 2.0 | `amazon.nova-2-lite-v1:0:256k` | 256k tokens | -| `NOVA_PRO` | 1.0 | `amazon.nova-pro-v1:0:300k` | 300k tokens | - -### Training Methods +```python +from amzn_nova_forge import ( + ForgeTrainer, + ForgeEvaluator, + ForgeDeployer, + ForgeInference, + ForgeConfig, + Model, + TrainingMethod, + DeployPlatform, + SMTJRuntimeManager, + SMHPRuntimeManager, + SMTJServerlessRuntimeManager, + BedrockRuntimeManager, + CloudWatchLogMonitor, + MLflowMonitor, + JSONLDatasetLoader, + TransformMethod, + ValidateMethod, + DataMixingConfig, + EvalTaskConfig, + EvaluationTask, +) +``` -| Method | Description | Supported Models | -|--------------|------------------------------------------|------------------------| -| `CPT` | Continued Pre-Training | All models (SMHP only) | -| `DPO_LORA` | Direct Preference Optimization with LoRA | Nova 1.0 models | -| `DPO_FULL` | Full-rank Direct Preference Optimization | Nova 1.0 models | -| `SFT_LORA` | Supervised Fine-tuning with LoRA | All models | -| `SFT_FULL` | Full-rank Supervised Fine-tuning | All models | -| `RFT_LORA` | Reinforcement Fine-tuning with LoRA | Nova 2.0 models | -| `RFT_FULL` | Full Reinforcement Fine-tuning | Nova 2.0 models | -| `RFT_MULTITURN_LORA` | RFT Multiturn with LoRA | Nova 2.0 models | -| `RFT_MULTITURN_FULL` | Full RFT Multiturn | Nova 2.0 models | -| `EVALUATION` | Model evaluation | All models | +**After (SageMaker SDK V3):** -### Platform Support +```python +from sagemaker.train import SFTTrainer, CPTTrainer, DPOTrainer +from sagemaker.train.evaluate import BenchMarkEvaluator, get_benchmarks +from sagemaker.train.data_mixing_config import DataMixingConfig +from sagemaker.core.training.configs import HyperPodCompute, TrainingJobCompute +``` -| Platform | Description | Models Supported | Supported Methods | -| ----------------- | ------------------------------------------ | ---------------- | ---------------------------------------- | -| `SMTJ` | SageMaker Training Jobs | All models | All methods | -| `SMTJServerless` | SageMaker Serverless Training (no infra) | All models | SFT, DPO, RFT_LORA, EVALUATION | -| `SMHP` | SageMaker HyperPod | All models | All methods | -| `BEDROCK` | Amazon Bedrock (Managed Service) | All models | SFT, DPO, RFT | +### Step 2: Configure Compute -## Data Preparation +**Before (Forge SDK) — SMTJ:** -Before launching a training job, your data needs to be in the right format. The SDK's dataset module handles loading, transforming, validating, filtering, and saving training data — supporting JSONL, JSON, CSV, Parquet, and Arrow formats from local files or S3. +```python +runtime = SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=4) +``` -A typical data preparation workflow: +**After (SageMaker SDK V3) — SMTJ:** ```python -from amzn_nova_forge.dataset import JSONLDatasetLoader -from amzn_nova_forge.model import TrainingMethod, Model, TransformMethod, ValidateMethod - -loader = JSONLDatasetLoader() -loader.load("s3://my-bucket/raw-data.jsonl") -loader.transform(method=TransformMethod.SCHEMA, training_method=TrainingMethod.SFT_LORA, model=Model.NOVA_LITE_2) -loader.validate(method=ValidateMethod.INVALID_RECORDS, training_method=TrainingMethod.SFT_LORA, model=Model.NOVA_LITE_2) -loader.save("s3://my-bucket/prepared-data.jsonl") +compute = TrainingJobCompute(instance_type="ml.p5.48xlarge", instance_count=4) ``` -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). +**Before (Forge SDK) — SMHP:** ---- -## Core Modules Overview +```python +runtime = SMHPRuntimeManager( + instance_type="ml.p5.48xlarge", + instance_count=4, + cluster_name="my-cluster", + namespace="default", +) +``` -The Nova Forge SDK is organized into the following modules: +**After (SageMaker SDK V3) — SMHP:** -| Module | Purpose | Key Components | -| ------------------ | --------------------------------------------- | ---------------------------------------------------------------- | -| **Dataset** | Data loading, transformation, filtering, and preparation | `JSONLDatasetLoader`, `JSONDatasetLoader`, `CSVDatasetLoader` | -| **Manager** | Runtime infrastructure management | `SMTJRuntimeManager`, `SMTJServerlessRuntimeManager`, `SMHPRuntimeManager`, `BedrockRuntimeManager` | -| **Model** | Main SDK entrypoint and orchestration | `NovaModelCustomizer` | -| **Monitor** | Job monitoring and logging | `CloudWatchLogMonitor`, `MLflowMonitor` | -| **RFT Multiturn** | Reinforcement fine-tuning infrastructure | `RFTMultiturnInfrastructure` | +```python +compute = HyperPodCompute( + cluster_name="my-cluster", + instance_type="ml.p5.48xlarge", + node_count=4, +) +``` -* 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/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) +**Before (Forge SDK) — Serverless:** -### Service Classes (Recommended) +```python +runtime = SMTJServerlessRuntimeManager(model_package_group_name="test-package") +``` -As of v1.4.0, the SDK provides modular service classes that replace `NovaModelCustomizer`. We recommend using these for all new work: +**After (SageMaker SDK V3) — Serverless:** -| Class | Purpose | -|-------|---------| -| `ForgeTrainer` | Training jobs (SFT, CPT, DPO, RFT) | -| `ForgeEvaluator` | Evaluation jobs (MMLU, GEN_QA, LLM_JUDGE, etc.) | -| `ForgeDeployer` | Deploy to SageMaker or Bedrock | -| `ForgeInference` | Real-time and batch inference | +Omit the `compute` parameter entirely. The trainer runs serverless by default. -Each class takes its own configuration in the constructor. Shared settings like `output_s3_path` and `mlflow_monitor` are passed via a `ForgeConfig` object. +### Step 3: Training (SFT) -**Example — Train and Deploy with Service Classes:** +**Before (Forge SDK):** ```python -from amzn_nova_forge import ( - ForgeTrainer, ForgeEvaluator, ForgeDeployer, ForgeInference, - ForgeConfig, Model, TrainingMethod, DeployPlatform, - SMTJRuntimeManager, -) - -# Configure infrastructure -runtime = SMTJRuntimeManager( - instance_type="ml.p5.48xlarge", - instance_count=4, -) - -# Train trainer = ForgeTrainer( model=Model.NOVA_LITE_2, method=TrainingMethod.SFT_LORA, @@ -203,477 +136,366 @@ trainer = ForgeTrainer( training_data_s3_path="s3://bucket/train.jsonl", config=ForgeConfig(output_s3_path="s3://bucket/output"), ) -training_result = trainer.train(job_name="my-sft-job") +result = trainer.train(job_name="my-sft-job", overrides={"lr": 5e-6, "warmup_steps": 100}) +``` -# Deploy -deployer = ForgeDeployer(model=Model.NOVA_LITE_2) -deployment_result = deployer.deploy( - model_artifact_path=training_result.model_artifacts.checkpoint_s3_path, - deploy_platform=DeployPlatform.SAGEMAKER, - unit_count=1, - endpoint_name="my-nova-endpoint", -) +**After (SageMaker SDK V3):** -# Invoke -inference = ForgeInference() -result = inference.invoke( - endpoint_arn=deployment_result.endpoint.endpoint_arn, - request_body={ - "messages": [{"role": "user", "content": "Hello!"}], - "max_tokens": 100, +```python +trainer = SFTTrainer( + model="nova-textgeneration-lite-v2", + compute=compute, + training_dataset="s3://bucket/train.jsonl", + s3_output_path="s3://bucket/output/", + overrides={ + "recipes.training_config.trainer.lr": 5e-6, + "recipes.training_config.trainer.warmup_steps": 100, }, ) -result.show() +job_name = trainer.train(wait=False) ``` -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/`](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/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/user-guides/instance_type_spec.md`. - -**Main Methods:** -- `execute()` - Start a training or evaluation job -- `cleanup()` - Stop and clean up a running job -- `scale_cluster()` - (SMHP only) Scale HyperPod cluster instance groups up or down -- `get_instance_groups()` - (SMHP only) View instance groups and current instance counts +### Step 4: Data Mixing (Optional) -**Key Classes:** -- `SMTJRuntimeManager` - For SageMaker Training Jobs -- `SMHPRuntimeManager` - For SageMaker HyperPod clusters -- `BedrockRuntimeManager` - For Amazon Bedrock managed service - -**Cluster Scaling (SMHP):** - -The `SMHPRuntimeManager` provides a `scale_cluster()` method to dynamically adjust the number of instances in a HyperPod cluster instance group: +**Before (Forge SDK):** ```python -from amzn_nova_forge.manager import SMHPRuntimeManager - -# Create a runtime manager for your cluster -manager = SMHPRuntimeManager( - instance_type="ml.p4d.24xlarge", - instance_count=4, - cluster_name="my-hyperpod-cluster", - namespace="default" -) - -# View the available instance groups to update -available_instance_groups = manager.get_instance_groups() - -# Scale up the worker group from 4 to 8 instances -result = manager.scale_cluster( - instance_group_name="worker-group", - target_instance_count=8 +trainer = ForgeTrainer(..., data_mixing_enabled=True) +trainer.data_mixing.set_config( + { + "customer_data_percent": 50, + "nova_code_percent": 30, + "nova_general_percent": 70, + } ) ``` -For more cluster scaling documentation, see [`docs/spec/runtime-managers.md`](docs/spec/runtime-managers.md). -### Model Module (Deprecated) +**After (SageMaker SDK V3):** -> ⚠️ `NovaModelCustomizer` is deprecated as of v1.4.0 and will be removed in a future version. It remains fully functional, but we recommend using the modular [Service Classes](#service-classes-recommended) (`ForgeTrainer`, `ForgeEvaluator`, `ForgeDeployer`, `ForgeInference`) for all new work. +```python +from sagemaker.train.data_mixing_config import DataMixingConfig -Provides the main SDK entrypoint for orchestrating model customization workflows. +data_mixing = DataMixingConfig( + customer_data_percent=50.0, + nova_data_percentages={"code": 30.0, "reasoning": 70.0}, +) +trainer = SFTTrainer(..., data_mixing_config=data_mixing) +``` -**Main Methods:** -- `train()` - Launch a training job -- `evaluate()` - Launch an evaluation job -- `deploy()` - Deploy trained model to Amazon SageMaker or Bedrock -- `batch_inference()` - Run batch inference on trained model -- `get_logs()` - Retrieve CloudWatch logs for current job -- `get_data_mixing_config()` - Get data mixing configuration -- `set_data_mixing_config()` - Set data mixing configuration +### Step 5: Monitor, Notifications & Dry Run -**Key Class:** -- `NovaModelCustomizer` - Main orchestration class +#### Log Streaming -### Monitor Module -Provides job monitoring and experiment tracking capabilities. +**Before (Forge SDK):** -**Main Methods:** -- `show_logs()` - Display CloudWatch logs -- `get_logs()` - Retrieve logs as list -- `from_job_result()` - Create monitor from job result -- `from_job_id()` - Create monitor from job ID +```python +trainer.get_logs(job_result=result, limit=50) +monitor = CloudWatchLogMonitor.from_job_id(job_id=result.job_id, platform=platform) +monitor.show_logs(limit=100) +``` -**Key Classes:** -- `CloudWatchLogMonitor` - For viewing job logs -- `MLflowMonitor` - For experiment tracking with presigned URL generation +**After (SageMaker SDK V3):** ---- +```python +# Stream logs (works on both trainer and evaluator) +trainer.stream_logs() +trainer.stream_logs(tail_logs=50) # last 50 log entries -### RFT Multiturn Module -Manages infrastructure for reinforcement fine-tuning with multi-turn conversational tasks. +evaluator.stream_logs() +``` -**Main Methods:** -- `setup()` - Deploy SAM stack and validate platform -- `start_training_environment()` - Start training environment -- `start_evaluation_environment()` - Start evaluation environment -- `get_logs()` - Retrieve environment logs -- `kill_task()` - Stop running task -- `cleanup()` - Clean up infrastructure resources -- `check_all_queues()` - Check message counts in all queues -- `flush_all_queues()` - Purge all messages from queues +#### Metrics Visualization -**Key Classes:** -- `RFTMultiturnInfrastructure` - Main infrastructure management class -- `CustomEnvironment` - For creating custom reward environments +**Before (Forge SDK):** -**Supported Platforms:** -- `LOCAL` - Local development environment -- `EC2` - Amazon EC2 instances -- `ECS` - Amazon ECS clusters +```python +monitor.plot_metrics(training_method=TrainingMethod.SFT_LORA) +``` -**Built-in Environments:** -- `VFEnvId.WORDLE` - Wordle game environment -- `VFEnvId.TERMINAL_BENCH` - Terminal benchmark environment +**After (SageMaker SDK V3):** ---- -### Iterative Training +```python +trainer.show_metrics() +``` -The Nova Forge SDK supports iterative fine-tuning of Nova models. +#### Job Notifications (SMTJ only) -This is done by progressively running fine-tuning jobs on the output checkpoint from the previous job: +**Before (Forge SDK):** -``` python -# Stage 1: Initial training on base model -stage1_trainer = ForgeTrainer( - model=Model.NOVA_LITE, - method=TrainingMethod.SFT_LORA, - infra=infra, - training_data_s3_path="s3://bucket/stage1-data.jsonl", - config=ForgeConfig(output_s3_path="s3://bucket/stage1-output"), -) +```python +result = trainer.train(job_name="my-job") +result.enable_job_notifications(emails=["user@example.com"]) +``` -stage1_result = stage1_trainer.train(job_name="stage1-training") -# Wait for completion... -stage1_checkpoint = stage1_result.model_artifacts.checkpoint_s3_path +**After (SageMaker SDK V3):** -# Stage 2: Continue training from Stage 1 checkpoint -stage2_trainer = ForgeTrainer( - model=Model.NOVA_LITE, - method=TrainingMethod.SFT_LORA, - infra=infra, - training_data_s3_path="s3://bucket/stage2-data.jsonl", - model_s3_path=stage1_checkpoint, # Use previous checkpoint - config=ForgeConfig(output_s3_path="s3://bucket/stage2-output"), +```python +trainer = SFTTrainer( + model="amazon.nova-lite-v2", + training_dataset="s3://bucket/train.jsonl", + notifications={ + "sns_topic_arn": "arn:aws:sns:us-east-1:123456789012:my-topic", + "event_bus_arn": "arn:aws:events:us-east-1:123456789012:event-bus/my-bus", + "events": ["Completed", "Failed"], + "job_name_prefix": "my-team-", + }, ) - -stage2_result = stage2_trainer.train(job_name="stage2-training") +trainer.train() ``` -**Note:** Iterative fine-tuning requires using the same model and training method (LoRA vs Full-Rank) across all stages. +Requires a pre-created SNS topic. Notifications fire on job state changes (Completed, Failed, Stopped). -### Dry Run +#### Dry Run Mode -The Nova Forge SDK supports `dry_run` mode for the following functions: `train()`, `evaluate()`, and `batch_inference()`. +**Before (Forge SDK):** -When calling any of the above functions, you can set the `dry_run` parameter to `True`. -The SDK will still generate your recipe and validate your input, but it won't begin a job. -This feature is useful whenever you want to test or validate inputs and still have a recipe generated, without starting a job. +```python +trainer.train(job_name="my-job", dry_run=True) +``` -``` python -# Training dry run -trainer.train( - job_name="train_dry_run", - dry_run=True, - ... -) +**After (SageMaker SDK V3):** -# Evaluation dry run -evaluator.evaluate( - job_name="evaluate_dry_run", - dry_run=True, - ... -) +```python +trainer.train(dry_run=True) ``` -### Data Mixing -Data mixing allows you to blend your custom training data with Nova's high-quality curated datasets, helping maintain the model's broad capabilities while adding your domain-specific knowledge. +Runs all validations (IAM, compute, dataset) without submitting a job. -**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% +### Step 6: Evaluate -**Example Usage:** +**Before (Forge SDK):** ```python -# Initialize with data mixing enabled (HyperPod) -trainer = ForgeTrainer( - model=Model.NOVA_LITE_2, - method=TrainingMethod.SFT_LORA, - 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( +evaluator = ForgeEvaluator( 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"), -) - -# Configure data mixing percentages -trainer.data_mixing.set_config({ - "customer_data_percent": 50, # 50% your data - "nova_code_percent": 30, # 30% Nova code data (30% of Nova's 50%) - "nova_general_percent": 70 # 70% Nova general data (70% of Nova's 50%) -}) - -# Or use 100% customer data (no Nova mixing) -trainer.data_mixing.set_config({ - "customer_data_percent": 100, - "nova_code_percent": 0, - "nova_general_percent": 0 -}) -``` -**Important Notes:** -- The `dataset_catalog` field is system-managed and cannot be set by users -- 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 - -Get email notifications when your training jobs complete, fail, or are stopped. The SDK automatically sets up the required AWS infrastructure (CloudFormation, DynamoDB, SNS, Lambda, EventBridge) to monitor job status and send notifications. - -**Features:** -- Automatic AWS infrastructure setup and management -- Email notifications for terminal job states (Completed, Failed, Stopped) -- Email notifications for SMHP master pods that are stuck in a crash loop -- Output artifact validation for successful jobs (manifest.json) -- Optional customer key KMS encryption for SNS topics - -**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/runtime-managers.md`](docs/spec/runtime-managers.md) for more details) -- **Bedrock** (Amazon Bedrock): Fully managed, no infrastructure configuration required - -**Quick Example:** -```python -# Start a training job -result = customizer.train(job_name="my-job") - -# Enable notifications (SMTJ) -result.enable_job_notifications( - emails=["user@example.com"] + infra=eval_infra, + data_s3_path="s3://bucket/eval-data.jsonl", + config=ForgeConfig(output_s3_path="s3://bucket/eval-output"), ) - -# Enable notifications (SMHP) -result.enable_job_notifications( - emails=["user@example.com"], - namespace="kubeflow", # Required for SMHP - kubectl_layer_arn="arn:aws:lambda::123456789012:layer:kubectl:1" # Required for SMHP +mmlu_result = evaluator.evaluate(job_name="eval-mmlu", eval_task=EvaluationTask.MMLU) +byod_result = evaluator.evaluate( + job_name="eval-byod", + eval_task=EvaluationTask.GEN_QA, + task_config=EvalTaskConfig(override_data_s3_path="s3://bucket/custom-eval.jsonl"), ) ``` -**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/user-guides/job_notifications.md)) -- Notification infrastructure is created once per region (SMTJ) or once per cluster (SMHP) and shared across jobs. -- 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/service-classes.md`](docs/spec/service-classes.md) for full API details. - - ---- -## Telemetry +**After (SageMaker SDK V3) — Benchmark (MMLU):** -The Nova Forge SDK has telemetry enabled to help us better understand user needs, diagnose issues, and deliver new features. This telemetry tracks the usage of various SDK functions. If you prefer to opt out of telemetry, you can do so by setting the `TELEMETRY_OPT_OUT` environment variable to `true`: +```python +from sagemaker.train.evaluate import BenchMarkEvaluator, get_benchmarks -```bash -export TELEMETRY_OPT_OUT=true +Benchmark = get_benchmarks() +evaluator = BenchMarkEvaluator( + benchmark=Benchmark.MMLU, + model="nova-textgeneration-lite-v2", + s3_output_path="s3://bucket/eval-output/", +) +execution = evaluator.evaluate(checkpoint_path="s3://bucket/output/checkpoint/") ``` ---- -## Getting Started -This comprehensive SDK enables end-to-end customization of Amazon Nova models with support for multiple training methods, deployment platforms, and monitoring capabilities. Each module is designed to work together seamlessly while providing flexibility for advanced use cases. - -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/`](docs/spec/index.md) - ---- -## Security Best Practices for SDK Users +**After (SageMaker SDK V3) — Custom Evaluator:** -### 1. IAM and Access Management +```python +from sagemaker.train.evaluate import CustomScorerEvaluator -**Execution Roles** +evaluator = CustomScorerEvaluator( + model="nova-textgeneration-lite-v2", + eval_dataset="s3://bucket/custom-eval.jsonl", + s3_output_path="s3://bucket/eval-output/", +) +execution = evaluator.evaluate(checkpoint_path="s3://bucket/output/checkpoint/") +``` -- Use dedicated execution roles for SageMaker training jobs with minimal required permissions -- Avoid using admin roles - follow the principle of least privilege -- Regularly audit role permissions and remove unused policies +**After (SageMaker SDK V3) — InspectAI Evaluator:** ```python -# Good: Explicit execution role -runtime = SMTJRuntimeManager( - instance_type="ml.p5.48xlarge", - instance_count=2, - execution_role="arn:aws:iam::123456789012:role/SageMakerNovaTrainingRole" +from sagemaker.train.evaluate import InspectAIEvaluator + +evaluator = InspectAIEvaluator( + model="nova-textgeneration-lite", + bedrock_model_id="us.amazon.nova-lite-v1:0", + benchmarks_path="s3://bucket/benchmarks/boolq/", + tasks=[{"name": "boolq_pt", "limit": 10}], + s3_output_path="s3://bucket/inspectai-eval-output/", + instance_type="ml.m5.large", ) - -# Avoid: Using default role without validation +execution = evaluator.evaluate() +execution.wait(target_status="Succeeded") +execution.show_results() ``` -**Required Permissions** +### Step 7: Deploy & Inference -The SDK requires specific IAM permissions. Review the [IAM](#iam-rolespolicies) section and: +**Before (Forge SDK):** -* Grant only the minimum permissions needed for your use case -* Use condition statements to restrict resource access -* Regularly review and rotate access keys - -**Verifying Your IAM Setup** +```python +# Deploy +deployer = ForgeDeployer(model=Model.NOVA_LITE_2) +result = deployer.deploy( + model_artifact_path=training_result.model_artifacts.checkpoint_s3_path, + deploy_platform=DeployPlatform.SAGEMAKER, + unit_count=1, + endpoint_name="my-endpoint", +) -The SDK automatically validates your IAM permissions before every `train()`, `evaluate()`, and `batch_inference()` call. If your role is missing required permissions, the SDK will report the specific missing permissions before attempting to start a job. +# Inference +inference = ForgeInference() +result = inference.invoke( + endpoint_arn=deployment_result.endpoint.endpoint_arn, + request_body={"messages": [{"role": "user", "content": "Hello!"}], "max_tokens": 100}, +) +result.show() +``` -If you want to verify your setup without starting a job, use `dry_run=True`. This runs the full validation (IAM permissions, recipe, infrastructure) and reports any issues, but does not submit a job: +**After (SageMaker SDK V3) — SageMaker Endpoint:** ```python -trainer = ForgeTrainer( - model=Model.NOVA_LITE_2, - method=TrainingMethod.SFT_LORA, - infra=SMTJRuntimeManager(instance_type="ml.p5.48xlarge", instance_count=4), - training_data_s3_path="s3://bucket/train.jsonl", - config=ForgeConfig(output_s3_path="s3://bucket/output"), +import json +from sagemaker.serve import ModelBuilder + +# Deploy +builder = ModelBuilder( + model=trainer, + role_arn="arn:aws:iam::123456789012:role/SageMakerRole", + instance_type="ml.p4d.24xlarge", +) +builder.accept_eula = True +builder.build(region="us-east-1") +endpoint = builder.deploy( + endpoint_name="my-endpoint", + instance_type="ml.p4d.24xlarge", ) -# Validate everything without starting a job -trainer.train(job_name="verify-setup", dry_run=True) +# Inference +response = endpoint.invoke( + body=json.dumps( + {"messages": [{"role": "user", "content": [{"type": "text", "text": "Hello!"}]}]} + ), + content_type="application/json", + accept="application/json", +) +body = json.loads(response.body.read()) ``` -The validation checks include: -- **Calling role permissions** — verifies your current role has the IAM actions needed for the chosen platform (e.g., `sagemaker:CreateTrainingJob` for SMTJ, `bedrock:CreateModelCustomizationJob` for Bedrock) using `iam:SimulatePrincipalPolicy` and policy inspection. -- **Execution role** (SMTJ/SMTJServerless) — confirms the execution role exists, trusts `sagemaker.amazonaws.com`, and has the required S3 permissions (`s3:GetObject`, `s3:PutObject`, `s3:ListBucket`). -- **Cluster access** (SMHP) — verifies the HyperPod cluster exists and your role can describe it. - -The IAM validation requires `iam:SimulatePrincipalPolicy` and policy-read permissions (e.g., `iam:GetRole`, `iam:ListRolePolicies`, `iam:ListAttachedRolePolicies`) on your calling role. If your role lacks these actions, the validation will fail with an error rather than silently skipping the checks. - -You can disable the IAM validation check and rely on the AWS service to report permission errors at job submission time: +**After (SageMaker SDK V3) — Bedrock:** ```python -from amzn_nova_forge.core import ValidationConfig +import json +import boto3 +from sagemaker.serve.bedrock_model_builder import BedrockModelBuilder -trainer = ForgeTrainer( - ... - config=ForgeConfig( - validation_config=ValidationConfig(iam=False), +# Deploy +builder = BedrockModelBuilder(model="s3://bucket/output/checkpoint/") +result = builder.deploy( + custom_model_name="my-custom-model", + role_arn="arn:aws:iam::123456789012:role/SageMakerRole", +) +model_arn = result["modelArn"] + +# Inference +bedrock_runtime = boto3.client("bedrock-runtime", region_name="us-east-1") +response = bedrock_runtime.invoke_model( + modelId=model_arn, + contentType="application/json", + accept="application/json", + body=json.dumps( + {"messages": [{"role": "user", "content": [{"type": "text", "text": "Hello!"}]}]} ), ) +body = json.loads(response["body"].read()) ``` -### 2. Credential Management +### Additional Training Methods -**AWS Credentials** +#### CPT (Continued Pre-Training) -* Never hardcode credentials in code or configuration files -* Use IAM roles instead of access keys when possible -* Rotate credentials regularly -* Use AWS Secrets Manager for application secrets -* Enable credential monitoring through AWS Config +**Before:** -**MLflow Integration** +```python +ForgeTrainer(model=Model.NOVA_LITE_2, method=TrainingMethod.CPT, infra=smhp_runtime, ...) +``` -* Secure MLflow tracking URIs with proper authentication -* Use encrypted connections to MLflow servers -* Implement access controls on experiment data -* Regularly audit MLflow access logs +**After:** -### 3. Data Security and Privacy +```python +CPTTrainer(model="nova-textgeneration-lite-v2", compute=HyperPodCompute(...), ...) +``` -**Training Data Protection** +#### DPO (Direct Preference Optimization) -- Encrypt data at rest in S3 using KMS keys -- Use S3 bucket policies to restrict access -- Validate data sources before processing +**Before:** ```python -# Ensure your S3 buckets have proper encryption and access controls -customizer = NovaModelCustomizer( - model=Model.NOVA_LITE_2, - method=TrainingMethod.SFT_LORA, - infra=runtime, - data_s3_path="s3://secure-training-bucket/encrypted-data/data.jsonl", - output_s3_path="s3://secure-output-bucket/results" -) +ForgeTrainer(model=Model.NOVA_MICRO, method=TrainingMethod.DPO_LORA, infra=runtime, ...) ``` -### 4. Network Security +**After:** -**VPC Configuration** - -* Deploy in private subnets when possible -* Use VPC endpoints for AWS service access -* Implement security groups with minimal required ports -* Enable VPC Flow Logs for network monitoring - -### 5. Secure Communication +```python +DPOTrainer(model="nova-textgeneration-micro", compute=compute, ...) +``` -- Always use HTTPS endpoints -- Never disable SSL certificate verification -- Keep TLS libraries updated +#### RLVR (Reinforcement Learning with Verifiable Rewards) -### 6. Input Validation +**Before:** -- Always validate user inputs before passing to SDK -- Sanitize data that will be stored or processed -- Check resource quotas before job submission -- Sanitize job names and resource identifiers +```python +ForgeTrainer(model=Model.NOVA_LITE_2, method=TrainingMethod.RFT_LORA, infra=runtime, ...) +``` +**After:** ```python -# The SDK includes built-in validation -loader = JSONLDatasetLoader(question="input", answer="output") -loader.load("s3://your-bucket/training-data.jsonl") -# Always validate your data format -loader.validate(method=TrainingMethod.SFT_LORA, model=Model.NOVA_LITE_2) +from sagemaker.train import RLVRTrainer + +trainer = RLVRTrainer( + model="nova-textgeneration-lite-v2", + compute=compute, + training_dataset="s3://bucket/rlvr-data.jsonl", + custom_reward_function="arn:aws:lambda:us-east-1:123456789012:function:my-reward", + s3_output_path="s3://bucket/output/", +) +trainer.train() ``` -### 7. Monitoring & Logging - -- Enable CloudTrail for API audit logs -- Use CloudWatch for operational monitoring -- Never log sensitive data (tokens, credentials, PII) -- Monitor job logs through CloudWatch -- Set up alerts for suspicious activities +#### Iterative Training (Resume from Checkpoint) -**Security Monitoring** -- Monitor failed authentication attempts -- Track unusual resource access patterns -- Log all model deployment activities +**Before:** -### 8. Deployment Security +```python +trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.SFT_LORA, + infra=runtime, + training_data_s3_path="s3://bucket/stage2-data.jsonl", + model_s3_path="s3://bucket/stage1-output/checkpoint/", + config=ForgeConfig(output_s3_path="s3://bucket/stage2-output"), +) +``` -**Bedrock Deployment** +**After:** -- Use least privilege policies for Bedrock access -- Implement endpoint access controls -- Monitor model inference patterns -- Enable request/response logging when appropriate +```python +trainer = SFTTrainer( + model="s3://bucket/stage1-output/checkpoint/", + compute=compute, + training_dataset="s3://bucket/stage2-data.jsonl", + s3_output_path="s3://bucket/stage2-output/", +) +trainer.train() +``` -### 9. Validation +## What's Different (Summary) -The SDK includes built-in validation: +- Compute is a config object (`HyperPodCompute`, `TrainingJobCompute`), not a runtime manager +- Model is a string identifier (e.g. `"nova-textgeneration-lite-v2"`), not an enum; also accepts S3 checkpoint paths for iterative training +- Deployment uses `ModelBuilder`/`BedrockModelBuilder` pattern instead of `ForgeDeployer` +- Overrides use full recipe paths (e.g. `"recipes.training_config.trainer.lr"`); use `trainer.get_resolved_recipe()` to inspect the final merged recipe +- No `ForgeConfig` object — shared settings are passed directly to trainer constructors +- Job notifications currently support SMTJ only — pass a `notifications` dict with SNS topic and EventBridge event bus ARNs -- IAM permission validation before job execution -- Input sanitization for user-provided parameters +## Support -Validation is enabled by default. \ No newline at end of file +- SageMaker Python SDK docs: https://sagemaker.readthedocs.io/en/stable/ +- SageMaker Python SDK GitHub: https://github.com/aws/sagemaker-python-sdk diff --git a/pyproject.toml b/pyproject.toml index 3c6b63a..1d9b0ac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,10 @@ build-backend = "setuptools.build_meta" [project] name = "amzn-nova-forge" -description = "A Python SDK for customizing Amazon Nova models." +description = "[DEPRECATED] Use sagemaker>=3.19.0 instead. A Python SDK for customizing Amazon Nova models." +classifiers = [ + "Development Status :: 7 - Inactive", +] dynamic = ["version"] requires-python = ">=3.12" authors = [ diff --git a/src/amzn_nova_forge/__init__.py b/src/amzn_nova_forge/__init__.py index 99ae5ec..c7f0117 100644 --- a/src/amzn_nova_forge/__init__.py +++ b/src/amzn_nova_forge/__init__.py @@ -11,6 +11,15 @@ # 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 warnings + +warnings.warn( + "amzn-nova-forge is deprecated. Please migrate to the SageMaker Python SDK V3 " + "(pip install 'sagemaker>=3.19.0').", + DeprecationWarning, + stacklevel=2, +) + from .core.data_mixing_config import DataMixingConfig from .core.enums import ( DeploymentMode, diff --git a/src/amzn_nova_forge/__version__.py b/src/amzn_nova_forge/__version__.py index d23c4ba..93ee9ed 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.9" # pragma: no cover +VERSION = "1.4.10" # pragma: no cover diff --git a/src/amzn_nova_forge/manager/glue_runtime_manager.py b/src/amzn_nova_forge/manager/glue_runtime_manager.py index 845fff0..0fc2b41 100644 --- a/src/amzn_nova_forge/manager/glue_runtime_manager.py +++ b/src/amzn_nova_forge/manager/glue_runtime_manager.py @@ -35,6 +35,7 @@ JobConfig, RuntimeManager, ) +from amzn_nova_forge.util.aws_utils import get_caller_account_id from amzn_nova_forge.util.logging import logger from amzn_nova_forge.util.s3_utils import ( GLUE_ARTIFACT_PREFIX, @@ -515,13 +516,44 @@ def required_calling_role_permissions(cls, data_s3_path=None, output_s3_path=Non permissions.extend( [ - ("glue:CreateJob", "*"), - ("glue:UpdateJob", "*"), - ("glue:StartJobRun", "*"), - ("glue:GetJobRun", "*"), - ("glue:BatchStopJobRun", "*"), - ("iam:GetRole", "*"), - ("iam:PassRole", "*"), + ( + "glue:CreateJob", + lambda infra: ( + f"arn:aws:glue:{infra.region}:{get_caller_account_id(infra.region)}:job/*" + ), + ), + ( + "glue:UpdateJob", + lambda infra: ( + f"arn:aws:glue:{infra.region}:{get_caller_account_id(infra.region)}:job/*" + ), + ), + ( + "glue:StartJobRun", + lambda infra: ( + f"arn:aws:glue:{infra.region}:{get_caller_account_id(infra.region)}:job/*" + ), + ), + ( + "glue:GetJobRun", + lambda infra: ( + f"arn:aws:glue:{infra.region}:{get_caller_account_id(infra.region)}:job/*" + ), + ), + ( + "glue:BatchStopJobRun", + lambda infra: ( + f"arn:aws:glue:{infra.region}:{get_caller_account_id(infra.region)}:job/*" + ), + ), + ( + "iam:GetRole", + lambda infra: f"arn:aws:iam::{get_caller_account_id(infra.region)}:role/*", + ), + ( + "iam:PassRole", + lambda infra: f"arn:aws:iam::{get_caller_account_id(infra.region)}:role/*", + ), ] ) diff --git a/src/amzn_nova_forge/manager/runtime_manager.py b/src/amzn_nova_forge/manager/runtime_manager.py index 17286fb..7bbbeae 100644 --- a/src/amzn_nova_forge/manager/runtime_manager.py +++ b/src/amzn_nova_forge/manager/runtime_manager.py @@ -54,6 +54,7 @@ from amzn_nova_forge.core.validation_patterns import MODEL_PACKAGE_ARN_REGEX from amzn_nova_forge.manager.mtrl_manager import MTRLOperations from amzn_nova_forge.telemetry import Feature, _telemetry_emitter +from amzn_nova_forge.util.aws_utils import get_caller_account_id from amzn_nova_forge.util.bedrock import ( get_customization_type, parse_bedrock_recipe_config, @@ -159,32 +160,6 @@ class DataPrepJobConfig(JobConfig): extra_pip_packages: List[str] = field(default_factory=list) -_account_id_cache: Optional[str] = None - - -def _get_caller_account_id(region: str = "us-east-1") -> str: - """Return the AWS account ID of the caller, cached to avoid redundant STS calls. - - Only caches successful results — transient STS failures return "*" without poisoning - the cache, so subsequent calls will retry. - """ - global _account_id_cache - if _account_id_cache is None: - try: - _account_id_cache = boto3.client("sts", region_name=region).get_caller_identity()[ - "Account" - ] - except Exception: - logger.warning( - "Failed to retrieve caller account ID via STS in region %s; " - "falling back to wildcard '*'", - region, - exc_info=True, - ) - return "*" - return _account_id_cache - - def _poll_for_training_job(sagemaker_client, job_name: str, timeout: int) -> str: """Poll ``list_training_jobs`` until the submitted job appears. @@ -628,9 +603,24 @@ def required_calling_role_permissions(cls, data_s3_path=None, output_s3_path=Non # Add SMTJ-specific permissions permissions.extend( [ - ("sagemaker:CreateTrainingJob", "*"), - ("sagemaker:DescribeTrainingJob", "*"), - ("sagemaker:StopTrainingJob", "*"), + ( + "sagemaker:CreateTrainingJob", + lambda infra: ( + f"arn:aws:sagemaker:{infra.region}:{get_caller_account_id(infra.region)}:training-job/*" + ), + ), + ( + "sagemaker:DescribeTrainingJob", + lambda infra: ( + f"arn:aws:sagemaker:{infra.region}:{get_caller_account_id(infra.region)}:training-job/*" + ), + ), + ( + "sagemaker:StopTrainingJob", + lambda infra: ( + f"arn:aws:sagemaker:{infra.region}:{get_caller_account_id(infra.region)}:training-job/*" + ), + ), "iam:GetRole", "iam:PassRole", "iam:GetPolicy", @@ -1518,11 +1508,32 @@ def required_calling_role_permissions(cls, data_s3_path=None, output_s3_path=Non permissions.extend( [ - ("sagemaker:CreateTrainingJob", "*"), - ("sagemaker:DescribeTrainingJob", "*"), - ("sagemaker:StopTrainingJob", "*"), - ("iam:GetRole", "*"), - ("iam:PassRole", "*"), + ( + "sagemaker:CreateTrainingJob", + lambda infra: ( + f"arn:aws:sagemaker:{infra.region}:{get_caller_account_id(infra.region)}:training-job/*" + ), + ), + ( + "sagemaker:DescribeTrainingJob", + lambda infra: ( + f"arn:aws:sagemaker:{infra.region}:{get_caller_account_id(infra.region)}:training-job/*" + ), + ), + ( + "sagemaker:StopTrainingJob", + lambda infra: ( + f"arn:aws:sagemaker:{infra.region}:{get_caller_account_id(infra.region)}:training-job/*" + ), + ), + ( + "iam:GetRole", + lambda infra: f"arn:aws:iam::{get_caller_account_id(infra.region)}:role/*", + ), + ( + "iam:PassRole", + lambda infra: f"arn:aws:iam::{get_caller_account_id(infra.region)}:role/*", + ), # Artifact bucket: auto-create, check existence, upload script + .whl ("s3:CreateBucket", "*"), ("s3:HeadBucket", "*"), @@ -1567,19 +1578,19 @@ def required_calling_role_permissions(cls, data_s3_path=None, output_s3_path=Non ( "sagemaker:DescribeCluster", lambda infra: ( - f"arn:aws:sagemaker:{infra.region}:{_get_caller_account_id(infra.region)}:cluster/{infra.cluster_name}" + f"arn:aws:sagemaker:{infra.region}:{get_caller_account_id(infra.region)}:cluster/{infra.cluster_name}" ), ), ( "eks:DescribeCluster", lambda infra: ( - f"arn:aws:eks:{infra.region}:{_get_caller_account_id(infra.region)}:cluster/*" + f"arn:aws:eks:{infra.region}:{get_caller_account_id(infra.region)}:cluster/*" ), ), ( "eks:ListAddons", lambda infra: ( - f"arn:aws:eks:{infra.region}:{_get_caller_account_id(infra.region)}:cluster/{infra.cluster_name}" + f"arn:aws:eks:{infra.region}:{get_caller_account_id(infra.region)}:cluster/{infra.cluster_name}" ), ), ("sagemaker:ListClusters", "*"), @@ -2155,9 +2166,24 @@ def required_calling_role_permissions(cls, data_s3_path=None, output_s3_path=Non # Add Bedrock-specific permissions permissions.extend( [ - ("bedrock:CreateModelCustomizationJob", "*"), - ("bedrock:StopModelCustomizationJob", "*"), - ("bedrock:GetModelCustomizationJob", "*"), + ( + "bedrock:CreateModelCustomizationJob", + lambda infra: ( + f"arn:aws:bedrock:{infra.region}:{get_caller_account_id(infra.region)}:model-customization-job/*" + ), + ), + ( + "bedrock:StopModelCustomizationJob", + lambda infra: ( + f"arn:aws:bedrock:{infra.region}:{get_caller_account_id(infra.region)}:model-customization-job/*" + ), + ), + ( + "bedrock:GetModelCustomizationJob", + lambda infra: ( + f"arn:aws:bedrock:{infra.region}:{get_caller_account_id(infra.region)}:model-customization-job/*" + ), + ), "iam:PassRole", ] ) @@ -2205,8 +2231,18 @@ def required_calling_role_permissions(cls, data_s3_path=None, output_s3_path=Non # Add SMTJ-specific permissions permissions.extend( [ - ("sagemaker:CreateTrainingJob", "*"), - ("sagemaker:DescribeTrainingJob", "*"), + ( + "sagemaker:CreateTrainingJob", + lambda infra: ( + f"arn:aws:sagemaker:{infra.region}:{get_caller_account_id(infra.region)}:training-job/*" + ), + ), + ( + "sagemaker:DescribeTrainingJob", + lambda infra: ( + f"arn:aws:sagemaker:{infra.region}:{get_caller_account_id(infra.region)}:training-job/*" + ), + ), "iam:GetRole", "iam:PassRole", "iam:GetPolicy", diff --git a/src/amzn_nova_forge/util/aws_utils.py b/src/amzn_nova_forge/util/aws_utils.py new file mode 100644 index 0000000..60b378d --- /dev/null +++ b/src/amzn_nova_forge/util/aws_utils.py @@ -0,0 +1,46 @@ +# 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. +"""AWS utility helpers shared across the SDK.""" + +import logging +from typing import Optional + +import boto3 + +logger = logging.getLogger(__name__) + +_account_id_cache: Optional[str] = None + + +def get_caller_account_id(region: str = "us-east-1") -> str: + """Return the AWS account ID of the caller, cached to avoid redundant STS calls. + + Only caches successful results — transient STS failures return "*" without poisoning + the cache, so subsequent calls will retry. + """ + global _account_id_cache + if _account_id_cache is None: + try: + _account_id_cache = boto3.client("sts", region_name=region).get_caller_identity()[ + "Account" + ] + except Exception: + logger.warning( + "Failed to retrieve caller account ID via STS in region %s; " + "falling back to wildcard '*'", + region, + exc_info=True, + ) + return "*" + return _account_id_cache diff --git a/src/amzn_nova_forge/validation/validator.py b/src/amzn_nova_forge/validation/validator.py index aa21f41..9f88deb 100644 --- a/src/amzn_nova_forge/validation/validator.py +++ b/src/amzn_nova_forge/validation/validator.py @@ -342,10 +342,11 @@ def _validate_calling_role_permissions( elif callable(resource_spec): # (api_string, resource_lambda) - call lambda with infra if infra is None: - errors.append( - f"Cannot evaluate resource lambda for {api_string}: infra is None" + raise ValueError( + f"Cannot evaluate resource ARN for {api_string}: " + "runtime manager is None. " + "Set validation_config={'iam': False} to skip IAM validation." ) - continue try: resource_arn = resource_spec(infra) diff --git a/tests/unit/util/test_aws_utils.py b/tests/unit/util/test_aws_utils.py new file mode 100644 index 0000000..a4534c6 --- /dev/null +++ b/tests/unit/util/test_aws_utils.py @@ -0,0 +1,86 @@ +# 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.aws_utils.""" + +import unittest +from unittest.mock import MagicMock, patch + +from botocore.exceptions import ClientError + +from amzn_nova_forge.util import aws_utils +from amzn_nova_forge.util.aws_utils import get_caller_account_id + + +class TestGetCallerAccountId(unittest.TestCase): + """Tests for get_caller_account_id caching and fallback behavior.""" + + def setUp(self): + # Reset module-level cache between tests + aws_utils._account_id_cache = None + + @patch("amzn_nova_forge.util.aws_utils.boto3.client") + def test_returns_account_id_on_success(self, mock_client): + mock_sts = MagicMock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + mock_client.return_value = mock_sts + + result = get_caller_account_id("us-east-1") + + self.assertEqual(result, "123456789012") + mock_client.assert_called_once_with("sts", region_name="us-east-1") + + @patch("amzn_nova_forge.util.aws_utils.boto3.client") + def test_caches_successful_result(self, mock_client): + mock_sts = MagicMock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + mock_client.return_value = mock_sts + + first = get_caller_account_id("us-east-1") + second = get_caller_account_id("us-west-2") + + self.assertEqual(first, "123456789012") + self.assertEqual(second, "123456789012") + # Only one STS call — second call uses cache + mock_sts.get_caller_identity.assert_called_once() + + @patch("amzn_nova_forge.util.aws_utils.boto3.client") + def test_returns_wildcard_on_exception(self, mock_client): + mock_sts = MagicMock() + mock_sts.get_caller_identity.side_effect = ClientError( + {"Error": {"Code": "ExpiredToken", "Message": "token expired"}}, + "GetCallerIdentity", + ) + mock_client.return_value = mock_sts + + result = get_caller_account_id("us-east-1") + + self.assertEqual(result, "*") + + @patch("amzn_nova_forge.util.aws_utils.boto3.client") + def test_does_not_cache_failed_result(self, mock_client): + mock_sts = MagicMock() + # First call fails + mock_sts.get_caller_identity.side_effect = [ + Exception("network error"), + {"Account": "123456789012"}, + ] + mock_client.return_value = mock_sts + + first = get_caller_account_id("us-east-1") + second = get_caller_account_id("us-east-1") + + self.assertEqual(first, "*") + self.assertEqual(second, "123456789012") + # Two STS calls — failure was not cached + self.assertEqual(mock_sts.get_caller_identity.call_count, 2) diff --git a/tests/unit/validation/test_validator.py b/tests/unit/validation/test_validator.py index eadaabf..cfa7c7f 100644 --- a/tests/unit/validation/test_validator.py +++ b/tests/unit/validation/test_validator.py @@ -1877,6 +1877,78 @@ def test_validate_calling_role_permissions_sts_error(self, mock_boto3_client): self.assertEqual(len(errors), 1) self.assertIn("Failed to validate calling role permissions: STS failed", errors[0]) + @patch("boto3.client") + def test_validate_calling_role_permissions_lambda_resource(self, mock_boto3_client): + """Test permission validation with lambda resource spec (the P449020299 fix).""" + mock_iam_client = MagicMock() + mock_sts_client = MagicMock() + + mock_boto3_client.side_effect = lambda service, **kwargs: { + "iam": mock_iam_client, + "sts": mock_sts_client, + }[service] + + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/TestRole/session", + "Account": "123456789012", + } + mock_iam_client.simulate_principal_policy.return_value = { + "EvaluationResults": [{"EvalDecision": "allowed"}] + } + + mock_infra = MagicMock() + mock_infra.region = "us-west-2" + + errors = [] + required_permissions = [ + ( + "sagemaker:CreateTrainingJob", + lambda infra: f"arn:aws:sagemaker:{infra.region}:123456789012:training-job/*", + ), + ] + + Validator._validate_calling_role_permissions( + errors, required_permissions, mock_infra, "us-west-2" + ) + + self.assertEqual(len(errors), 0) + mock_iam_client.simulate_principal_policy.assert_any_call( + PolicySourceArn="arn:aws:iam::123456789012:role/TestRole", + ActionNames=["sagemaker:CreateTrainingJob"], + ResourceArns=["arn:aws:sagemaker:us-west-2:123456789012:training-job/*"], + ) + + @patch("boto3.client") + def test_validate_calling_role_permissions_lambda_infra_none_raises(self, mock_boto3_client): + """Test that lambda resource spec with infra=None produces an error.""" + mock_iam_client = MagicMock() + mock_sts_client = MagicMock() + + mock_boto3_client.side_effect = lambda service, **kwargs: { + "iam": mock_iam_client, + "sts": mock_sts_client, + }[service] + + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/TestRole/session", + "Account": "123456789012", + } + + errors = [] + required_permissions = [ + ( + "sagemaker:CreateTrainingJob", + lambda infra: f"arn:aws:sagemaker:{infra.region}:*:training-job/*", + ), + ] + + Validator._validate_calling_role_permissions( + errors, required_permissions, None, "us-east-1" + ) + + self.assertEqual(len(errors), 1) + self.assertIn("runtime manager is None", errors[0]) + class TestPermissionValidationMethods(unittest.TestCase): """Test cases for permission validation helper methods and formats""" @@ -2411,7 +2483,7 @@ def test_runtime_managers_define_mixed_permission_validation_types( self.assertTrue(has_strings, "SMTJ should have string permissions") @patch("subprocess.run") - @patch("amzn_nova_forge.manager.runtime_manager._get_caller_account_id") + @patch("amzn_nova_forge.manager.runtime_manager.get_caller_account_id") def test_smhp_required_permissions_uses_caller_account_id(self, mock_get_account_id, mock_run): """Happy path: account ID from STS appears in generated ARNs.""" mock_run.return_value.stdout = "" @@ -2435,7 +2507,7 @@ def test_smhp_required_permissions_uses_caller_account_id(self, mock_get_account ) @patch("subprocess.run") - @patch("amzn_nova_forge.manager.runtime_manager._get_caller_account_id") + @patch("amzn_nova_forge.manager.runtime_manager.get_caller_account_id") def test_smhp_required_permissions_falls_back_to_wildcard_on_sts_error( self, mock_get_account_id, mock_run ): From a2b81fffca457b510383760787462273e70d9737 Mon Sep 17 00:00:00 2001 From: Lisa Ni Date: Wed, 19 Aug 2026 19:43:52 +0000 Subject: [PATCH 2/2] Address review: move What's Different summary to top, drop undefined evaluator ref - Move 'What's Different (Summary)' to the top of the migration guide (right after 'Why Migrate'), removing the duplicate at the end - Trim the log-streaming snippet to trainer.stream_logs() so no undefined 'evaluator' is referenced --- README.md | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 17797dc..8434791 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,15 @@ A comprehensive Python SDK for fine-tuning and customizing Amazon Nova models. T The `amzn-nova-forge` package is deprecated. Amazon Nova model customization functionality is available in the SageMaker Python SDK V3. +## What's Different (Summary) + +- Compute is a config object (`HyperPodCompute`, `TrainingJobCompute`), not a runtime manager +- Model is a string identifier (e.g. `"nova-textgeneration-lite-v2"`), not an enum; also accepts S3 checkpoint paths for iterative training +- Deployment uses `ModelBuilder`/`BedrockModelBuilder` pattern instead of `ForgeDeployer` +- Overrides use full recipe paths (e.g. `"recipes.training_config.trainer.lr"`); use `trainer.get_resolved_recipe()` to inspect the final merged recipe +- No `ForgeConfig` object — shared settings are passed directly to trainer constructors +- Job notifications currently support SMTJ only — pass a `notifications` dict with SNS topic and EventBridge event bus ARNs + ## Installation ```bash @@ -200,8 +209,6 @@ monitor.show_logs(limit=100) # Stream logs (works on both trainer and evaluator) trainer.stream_logs() trainer.stream_logs(tail_logs=50) # last 50 log entries - -evaluator.stream_logs() ``` #### Metrics Visualization @@ -486,15 +493,6 @@ trainer = SFTTrainer( trainer.train() ``` -## What's Different (Summary) - -- Compute is a config object (`HyperPodCompute`, `TrainingJobCompute`), not a runtime manager -- Model is a string identifier (e.g. `"nova-textgeneration-lite-v2"`), not an enum; also accepts S3 checkpoint paths for iterative training -- Deployment uses `ModelBuilder`/`BedrockModelBuilder` pattern instead of `ForgeDeployer` -- Overrides use full recipe paths (e.g. `"recipes.training_config.trainer.lr"`); use `trainer.get_resolved_recipe()` to inspect the final merged recipe -- No `ForgeConfig` object — shared settings are passed directly to trainer constructors -- Job notifications currently support SMTJ only — pass a `notifications` dict with SNS topic and EventBridge event bus ARNs - ## Support - SageMaker Python SDK docs: https://sagemaker.readthedocs.io/en/stable/