diff --git a/README.md b/README.md index f805b62..7b5fa33 100644 --- a/README.md +++ b/README.md @@ -493,6 +493,7 @@ result.enable_job_notifications( 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 diff --git a/docs/spec/dataset.md b/docs/spec/dataset.md index 3a359e6..bfaa996 100644 --- a/docs/spec/dataset.md +++ b/docs/spec/dataset.md @@ -414,6 +414,7 @@ def transform( - `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`. + - `platform` (Platform): Required for RFT Multiturn methods (`RFT_MULTITURN_LORA`, `RFT_MULTITURN_FULL`). Determines the output format: `Platform.SMTJServerless` produces flat `{"prompt": "..."}` format, `Platform.SMHP` produces nested `{"id": "...", "metadata": {"prompt": "..."}}` format. - `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. @@ -460,7 +461,7 @@ def validate( - `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()`. +- `platform` (Platform): Optional. The target platform (`Platform.SMTJServerless`, `Platform.SMHP`, `Platform.SMTJ`, or `Platform.BEDROCK`). Required for RFT Multiturn methods — routes to `RFTMultiturnServerlessValidator` when `Platform.SMTJServerless`, validating the flat `{"prompt": "..."}` format. **Returns:** - None diff --git a/docs/spec/model-customizer.md b/docs/spec/model-customizer.md index 9185f9f..e096c01 100644 --- a/docs/spec/model-customizer.md +++ b/docs/spec/model-customizer.md @@ -212,6 +212,7 @@ def train( - `loraplus_lr_ratio` (float): LoRA+ learning rate ratio - `global_batch_size` (int): Global batch size - `max_length` (int): Maximum sequence length + - `model_name_or_path` (str): S3 checkpoint path or a SageMaker Model Package ARN for iterative training - 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) @@ -225,8 +226,8 @@ def train( - `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. + - `output_s3_path` (str, Optional): Path to the metrics and output tar file. For MTRL jobs, fetched from the AgentRFT API. + - `output_model_arn` (str, Optional): Model package ARN for `SMTJServerless` jobs. Use as `model_arn` in `ForgeTrainer` for iterative training. - `model_type` (Model): Model type of the model being trained **Raises:** @@ -382,7 +383,7 @@ def deploy( * **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. +- `model_artifact_path` (Optional[str]): S3 path to the trained model checkpoint, or a SageMaker Model Package name/ARN. Model package ARNs are auto-detected for SageMaker deployments and passed via `ModelPackageName`. 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 @@ -432,6 +433,15 @@ sagemaker_deployment = customizer.deploy( print(f"Model deployed: {sagemaker_deployment.endpoint.uri}") print(f"Endpoint: {sagemaker_deployment.endpoint.endpoint_name}") print(f"Status: {sagemaker_deployment.status}") + +# Deploy to SageMaker using a model package ARN (auto-detected) +sagemaker_mp_deployment = customizer.deploy( + model_artifact_path="arn:aws:sagemaker:us-east-1:123456789012:model-package/my-group/1", + deploy_platform=DeployPlatform.SAGEMAKER, + unit_count=1, + endpoint_name="my-model-package-endpoint", +) +print(f"Model deployed: {sagemaker_mp_deployment.endpoint.uri}") ``` Optionally, you can provide a Bedrock execution role name to be used in deployment. @@ -454,12 +464,16 @@ create_bedrock_execution_role( ``` --- #### `create_custom_model()` -Creates a Bedrock custom model from S3 artifacts, decoupled from endpoint deployment. +Creates a Bedrock custom model from S3 artifacts or a model package ARN, 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. +Either `model_artifact_path` (maps to `modelSourceConfig`) or `custom_model_data_source` +(maps to `customModelDataSource`) must be provided, but not both. If neither is provided, +the method attempts to resolve from `job_result`. + **Signature:** ```python def create_custom_model( @@ -470,26 +484,28 @@ def create_custom_model( execution_role_name: Optional[str] = None, tags: Optional[List[Dict[str, str]]] = None, skip_model_reuse: bool = False, + custom_model_data_source: Optional[Dict[str, Any]] = None, ) -> ModelDeployResult ``` **Parameters:** -- `model_artifact_path` (Optional[str]): S3 path to trained model checkpoint. Takes precedence over `job_result` if both are provided. +- `model_artifact_path` (Optional[str]): S3 path to trained model checkpoint. Used to populate `modelSourceConfig.s3DataSource.s3Uri`. 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. +- `custom_model_data_source` (Optional[Dict[str, Any]]): Alternative data source configuration for the custom model. When provided, `modelSourceConfig` is omitted from the API call. **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 + - `escrow_uri` (str): S3 artifacts path or model package ARN 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`. +- `ValueError`: When neither `model_artifact_path`, `job_result`, nor `custom_model_data_source` is provided, or when both `model_artifact_path` and `custom_model_data_source` are provided. - `RuntimeError`: When IAM role creation or custom model creation fails. **Example:** @@ -505,6 +521,16 @@ publish_result.dump(file_path="./results/") # Or create from a training job result publish_result = customizer.create_custom_model(job_result=training_result) + +# Or create from a model package ARN +publish_result = customizer.create_custom_model( + custom_model_data_source={ + "modelPackageArnDataSource": { + "modelPackageArn": "arn:aws:sagemaker:us-east-1:123456789012:model-package/my-group/1" + } + } +) +print(f"Model ARN: {publish_result.model_arn}") ``` --- #### `deploy_to_bedrock()` diff --git a/docs/spec/rft-multiturn.md b/docs/spec/rft-multiturn.md index 7c0563e..10cc626 100644 --- a/docs/spec/rft-multiturn.md +++ b/docs/spec/rft-multiturn.md @@ -207,6 +207,16 @@ print(f"Uploaded to: {custom_env.s3_uri}") - `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). +**Note:** RFT multiturn supports SageMaker HyperPod (SMHP) and SMTJServerless platforms with Nova 2.0 models (NOVA_LITE_2). + +--- + +### Platform Requirements + +#### MLFlow Required for MTRL on SMTJServerless + +MTRL training and evaluation on `SMTJServerless` require `ForgeConfig.mlflow_monitor` with a valid `tracking_uri` — raises `ValueError` if missing. + +MLFlow is optional on SMHP. --- diff --git a/docs/spec/runtime-managers.md b/docs/spec/runtime-managers.md index 117a6c8..3ebafdf 100644 --- a/docs/spec/runtime-managers.md +++ b/docs/spec/runtime-managers.md @@ -426,7 +426,9 @@ def __init__( security_group_ids: Optional[list[str]] = None, max_job_runtime: Optional[int] = 86400, rft_lambda: Optional[str] = None, + agent_core_arn: Optional[str] = None, evaluator_name: Optional[str] = None, + intermediate_model_package_group_name: Optional[str] = None, ) ``` @@ -438,16 +440,25 @@ def __init__( - `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. +- `rft_lambda` (Optional[str]): Lambda ARN, SageMaker hub-content ARN, or local `.py` file path for the RFT reward function. Used for single-turn RFT or as an alternative agent environment for MTRL on serverless. - **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. +- `agent_core_arn` (Optional[str]): Bedrock AgentCore runtime ARN for MTRL training. If both `agent_core_arn` and `rft_lambda` are set, `agent_core_arn` takes priority as the agent environment. - `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. +- `intermediate_model_package_group_name` (Optional[str]): Model package group name for intermediate checkpoints during MTRL training. If not set, defaults to `{model_package_group_name}-checkpoints`. **Example:** ```python from amzn_nova_forge.manager import * +# With Bedrock AgentCore (for MTRL training) +infra = SMTJServerlessRuntimeManager( + model_package_group_name="nova-rft-serverless", + execution_role="arn:aws:iam::123456789012:role/my-role", + agent_core_arn="arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/my-agent", +) + # With a Lambda ARN (auto-registered as hub-content during train()) infra = SMTJServerlessRuntimeManager( model_package_group_name="nova-rft-serverless", diff --git a/docs/spec/service-classes.md b/docs/spec/service-classes.md index 8b0319f..a63e913 100644 --- a/docs/spec/service-classes.md +++ b/docs/spec/service-classes.md @@ -67,6 +67,7 @@ def __init__( infra: RuntimeManager, training_data_s3_path: Optional[str] = None, model_s3_path: Optional[str] = None, + model_arn: Optional[str] = None, data_mixing_enabled: bool = False, holdout_data_s3_path: Optional[str] = None, val_check_interval: Optional[int] = None, @@ -83,7 +84,8 @@ def __init__( - `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 +- `model_s3_path` (Optional[str]): S3 path for the base or previously trained model (SMHP/SMTJ). For SMTJServerless, use `model_arn` instead. +- `model_arn` (Optional[str]): Model package ARN for iterative training on SMTJServerless. Pass the `output_model_arn` from a previous job to train on top of it. Cannot be combined with `model_s3_path`. - `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. @@ -293,6 +295,59 @@ matched_file = trainer.trace_batch(result, step=42, output_path="/tmp/flagged.js ``` --- +##### `generate_training_metrics_csv()` + +Generates a `step_wise_training_metrics.csv` file from CloudWatch logs for a completed SMHP SFT training job and uploads it to S3. The CSV contains step-level metrics (step number, epoch number, training loss) in the same format produced by SMTJ/Bedrock SFT jobs. + +**Signature:** +```python +def generate_training_metrics_csv( + self, + job_result: Optional[SMHPTrainingResult] = None, + job_id: Optional[str] = None, + started_time=None, + end_time=None, + output_s3_path: Optional[str] = None, +) -> Optional[str] +``` + +**Parameters:** +- `job_result` (Optional[SMHPTrainingResult]): Result from a completed SMHP training job. If provided, `job_id`, `started_time`, and `output_s3_path` are extracted automatically. +- `job_id` (Optional[str]): The SMHP training job ID. Used if `job_result` is not provided. +- `started_time`: Start time for log filtering. Accepts a `datetime` object or an ISO date string (e.g., `"2025-05-26"`). Defaults to 7 days ago if not provided. +- `end_time`: Optional end time to bound the log search. Accepts a `datetime` object or an ISO date string. If not provided, searches up to the current time. Providing this significantly speeds up log retrieval for older jobs. +- `output_s3_path` (Optional[str]): S3 URI for the output. Defaults to the trainer's configured `output_s3_path`. + +**Returns:** +- `str | None`: S3 URI of the uploaded CSV (e.g., `s3://bucket/prefix/job-id/step_wise_training_metrics.csv`), or `None` if no metrics could be extracted. + +**Raises:** +- `ValueError`: If platform is not SMHP, method is not SFT_LORA/SFT_FULL, or required parameters are missing. + +**Example:** +```python +# From a job result (simplest) +result = trainer.train(job_name="my-sft-job") +csv_uri = trainer.generate_training_metrics_csv(job_result=result) + +# Standalone with job ID +csv_uri = trainer.generate_training_metrics_csv( + job_id="my-job-id", + started_time="2026-04-17", + end_time="2026-04-18", +) + +# Minimal — uses trainer's output_s3_path and defaults to 7-day lookback +csv_uri = trainer.generate_training_metrics_csv(job_id="my-job-id") +``` + +**Notes:** +- Only supported for SMHP platform with SFT_LORA or SFT_FULL training methods. +- For older jobs, providing `started_time` and `end_time` significantly reduces log retrieval time. +- If the job is still in progress or has failed, a warning is emitted but metrics are still extracted on a best-effort basis. + +--- + ### ForgeEvaluator Handles evaluation job configuration and execution for Nova models. @@ -361,7 +416,7 @@ def evaluate( - `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` +- `task_config` (Optional[EvalTaskConfig]): Task-specific configuration. Fields: `subtask`, `processor`, `rl_env`, `override_data_s3_path`, `evaluate_base_model` (MTRL only — when True, evaluates both base and fine-tuned model in one pipeline) - `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 @@ -609,11 +664,12 @@ def deploy( sagemaker_instance_type: str = "ml.p5.48xlarge", sagemaker_environment: Optional[SageMakerEndpointEnvironment] = None, skip_model_reuse: bool = False, + inference_component_configs: List[InferenceComponentConfig] = [], ) -> DeploymentResult ``` **Parameters:** -- `model_artifact_path` (str): S3 path to the trained model checkpoint +- `model_artifact_path` (str): S3 path to the trained model checkpoint, or a SageMaker Model Package name/ARN. Model package ARNs are auto-detected for SageMaker deployments and passed via `ModelPackageName` - `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 @@ -625,6 +681,7 @@ def deploy( - 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 +- `inference_component_configs` (List[InferenceComponentConfig]): List of inference component configs. When provided with `SAGEMAKER` platform, creates an IC-compatible endpoint and deploys the inference component(s) in one step. **Returns:** - `DeploymentResult`: Contains `endpoint` (EndpointInfo), `platform`, `endpoint_name`, `uri`, `model_artifact_path`, and `created_at` @@ -635,47 +692,75 @@ def deploy( **Example:** ```python +# Deploy from S3 artifacts to Bedrock 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}") + +# Deploy to SageMaker using a model package ARN (auto-detected) +deployment = deployer.deploy( + model_artifact_path="arn:aws:sagemaker:us-east-1:123456789012:model-package/my-group/1", + deploy_platform=DeployPlatform.SAGEMAKER, + unit_count=1, + sagemaker_instance_type="ml.p5.48xlarge", +) +print(f"Model deployed: {deployment.endpoint.uri}") ``` --- ##### `create_custom_model()` -Creates a Bedrock custom model from S3 artifacts without deploying to an endpoint. +Creates a Bedrock custom model from S3 artifacts or a model package ARN without deploying to an endpoint. + +Either `model_artifact_path` (maps to `modelSourceConfig`) or `custom_model_data_source` (maps to `customModelDataSource`) must be provided, but not both. **Signature:** ```python def create_custom_model( self, - model_artifact_path: str, + model_artifact_path: Optional[str] = None, endpoint_name: Optional[str] = None, execution_role_name: Optional[str] = None, tags: Optional[List[Dict[str, str]]] = None, skip_model_reuse: bool = False, + custom_model_data_source: Optional[Dict[str, Any]] = None, ) -> ModelDeployResult ``` **Parameters:** -- `model_artifact_path` (str): S3 path to trained model checkpoint +- `model_artifact_path` (Optional[str]): S3 path to trained model checkpoint. Used to populate `modelSourceConfig.s3DataSource.s3Uri` - `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 +- `custom_model_data_source` (Optional[Dict[str, Any]]): Alternative data source configuration for the custom model. When provided, `modelSourceConfig` is omitted from the API call **Returns:** - `ModelDeployResult`: Contains `model_arn`, `model_name`, `escrow_uri`, and `created_at` +**Raises:** +- `ValueError`: If neither or both of `model_artifact_path` and `custom_model_data_source` are provided + **Example:** ```python +# Create from S3 artifacts 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/") + +# Create from a model package ARN +publish_result = deployer.create_custom_model( + custom_model_data_source={ + "modelPackageArnDataSource": { + "modelPackageArn": "arn:aws:sagemaker:us-east-1:123456789012:model-package/my-group/1" + } + } +) +print(f"Model ARN: {publish_result.model_arn}") ``` --- @@ -818,6 +903,78 @@ def get_logs( --- +##### `create_inference_component()` +Creates an inference component on an existing SageMaker endpoint. Returns immediately without waiting for the component to become active. + +**Signature:** +```python +def create_inference_component( + self, + inference_component_name: str, + model_name: str, + num_cpus: int, + num_accelerators: int, + min_memory_in_mb: int, + endpoint_name: str, + variant_name: str = "primary", + copy_count: int = 1, +) -> DeploymentResult +``` + +**Parameters:** +- `inference_component_name` (str): Unique name for the inference component +- `model_name` (str): Name of the existing SageMaker model to use +- `num_cpus` (int): Number of vCPUs to allocate +- `num_accelerators` (int): Number of accelerators (GPUs) to allocate +- `min_memory_in_mb` (int): Minimum memory in MB to allocate +- `endpoint_name` (str): Name of the existing SageMaker endpoint (must be InService) +- `variant_name` (str): Production variant name on the endpoint. Default: `"primary"` +- `copy_count` (int): Number of model copies to deploy. Default: 1 + +**Returns:** +- `DeploymentResult`: Contains endpoint info with the inference component ARN as the URI and the deployer's `region` set on `EndpointInfo` + +**Raises:** +- `Exception`: If the endpoint does not exist, is not InService, or the API call fails + +**Example:** +```python +result = deployer.create_inference_component( + inference_component_name="my-model-ic", + model_name="my-sagemaker-model", + num_cpus=15, + num_accelerators=4, + min_memory_in_mb=25000, + endpoint_name="my-endpoint", +) +print(f"Inference component ARN: {result.endpoint.uri}") +``` +--- + +##### `monitor_inference_component()` +Polls an inference component until it reaches a terminal state (InService or Failed). + +**Signature:** +```python +def monitor_inference_component(self, inference_component_name: str) -> str +``` + +**Parameters:** +- `inference_component_name` (str): Name of the inference component to monitor + +**Returns:** +- `str`: Final status (`"InService"`) + +**Raises:** +- `Exception`: If the component reaches Failed status or the API call errors + +**Example:** +```python +status = deployer.monitor_inference_component(inference_component_name="my-model-ic") +print(f"Inference component is now: {status}") +``` +--- + ### ForgeInference Handles single and batch inference on trained Nova models. diff --git a/docs/spec/utilities.md b/docs/spec/utilities.md index ae4eabe..73af747 100644 --- a/docs/spec/utilities.md +++ b/docs/spec/utilities.md @@ -305,6 +305,87 @@ monitor.plot_metrics( --- +### MTRLLogMonitor + +Log monitor for MTRL (Multi-Turn RL) training and evaluation jobs. Auto-detects whether a job is a training job or an evaluation job by checking CloudWatch log groups. + +#### Constructor + +**Signature:** +```python +def __init__( + self, + job_id: str, + region: Optional[str] = None, + job_category: Optional[str] = None, +) +``` + +**Parameters:** +- `job_id` (str): The MTRL job name (training or evaluation) +- `region` (Optional[str]): AWS region +- `job_category` (Optional[str]): Job category override (`"AgentRFT"` for training, `"AgentRFTEvaluation"` for eval). Auto-detected if not provided. + +#### Class Methods + +##### `from_job_id()` + +**Signature:** +```python +@classmethod +def from_job_id( + cls, + job_id: str, + region: Optional[str] = None, + job_category: Optional[str] = None, +) -> "MTRLLogMonitor" +``` + +**Example:** +```python +from amzn_nova_forge.monitor import MTRLLogMonitor + +# Auto-detects training vs eval +monitor = MTRLLogMonitor.from_job_id( + job_id="my-mtrl-job", + region="us-east-1", +) +monitor.show_logs(limit=20) +``` + +#### Methods + +##### `show_logs()` + +Displays job logs. For training jobs, uses `AgentRFTJob.wait()` for live progress. For evaluation jobs, reads CloudWatch logs from `/aws/sagemaker/Job/AgentRFTEvaluation`. + +**Signature:** +```python +def show_logs( + self, + poll: int = 30, + timeout: int = 7200, + limit: Optional[int] = None, +) -> None +``` + +**Parameters:** +- `poll` (int): Polling interval in seconds (training jobs only) +- `timeout` (int): Maximum wait time in seconds (training jobs only) +- `limit` (Optional[int]): Maximum number of log events to display (eval jobs) + +**Example:** +```python +# Monitor an eval job +monitor = MTRLLogMonitor.from_job_id(job_id="my-eval-job", region="us-east-1") +monitor.show_logs(limit=50) + +# Also accessible via ForgeEvaluator.get_logs() +evaluator.get_logs(job_result=eval_result, limit=20) +``` + +--- + ### MLflowMonitor MLflow monitoring configuration for Nova model training. This class provides experiment tracking capabilities through MLflow integration. diff --git a/docs/user-guides/iam_setup.md b/docs/user-guides/iam_setup.md index b545e9a..5726b67 100644 --- a/docs/user-guides/iam_setup.md +++ b/docs/user-guides/iam_setup.md @@ -160,6 +160,14 @@ Please refer to the "Sid" of each statement to determine which policies you need "arn:aws:sagemaker:::model/*" ] }, + { + "Sid": "AccessModelPackageForDeployment", + "Effect": "Allow", + "Action": [ + "sagemaker:AccessModelPackage" + ], + "Resource": "arn:aws:sagemaker:::model-package/*" + }, { "Sid": "MLflowSageMaker", "Effect": "Allow", @@ -218,9 +226,95 @@ 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. +- [Model Package only] `AccessModelPackage` (`sagemaker:AccessModelPackage`) is required when deploying or evaluating a model using a SageMaker Model Package ARN. - [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. +### Multi-Turn RL (MTRL) on Serverless SMTJ + +If performing MTRL training or evaluation, your execution role needs the following additional permissions (beyond the basic SDK policies above): + +```json +{ + "Version": "2012-10-17", + "Statement": [ + { + "Sid": "MTRLJobManagement", + "Effect": "Allow", + "Action": [ + "sagemaker:CreateJob", + "sagemaker:DescribeJob", + "sagemaker:StopJob" + ], + "Resource": "arn:aws:sagemaker:::job/*" + }, + { + "Sid": "MTRLAgentCoreAccess", + "Effect": "Allow", + "Action": [ + "bedrock-agentcore:GetAgentRuntime", + "bedrock-agentcore:InvokeAgentRuntime" + ], + "Resource": "arn:aws:bedrock-agentcore:::runtime/*" + }, + { + "Sid": "MTRLModelPackageManagement", + "Effect": "Allow", + "Action": [ + "sagemaker:CreateModelPackageGroup", + "sagemaker:CreateModelPackage", + "sagemaker:DescribeModelPackage", + "sagemaker:DescribeModelPackageGroup", + "sagemaker:ListModelPackages", + "sagemaker:UpdateModelPackage" + ], + "Resource": [ + "arn:aws:sagemaker:::model-package/*", + "arn:aws:sagemaker:::model-package-group/*" + ] + }, + { + "Sid": "MTRLEvalPipeline", + "Effect": "Allow", + "Action": [ + "sagemaker:CreatePipeline", + "sagemaker:StartPipelineExecution", + "sagemaker:DescribePipelineExecution", + "sagemaker:ListPipelineExecutionSteps", + "sagemaker:StopPipelineExecution", + "sagemaker:UpdatePipeline" + ], + "Resource": "arn:aws:sagemaker:::pipeline/*" + } + ] +} +``` + +If using a **Lambda function as a 3P agent** (instead of Bedrock AgentCore directly), also add: +```json +{ + "Sid": "MTRLLambdaAgentInvoke", + "Effect": "Allow", + "Action": "lambda:InvokeFunction", + "Resource": "arn:aws:lambda:::function:" +} +``` + +**Trust Policy:** The execution role must include these service principals: +```json +{ + "Effect": "Allow", + "Principal": { + "Service": [ + "job.sagemaker.amazonaws.com", + "finetuning.sagemaker.amazonaws.com", + "bedrock-agentcore.amazonaws.com" + ] + }, + "Action": ["sts:AssumeRole", "sts:TagSession", "sts:SetSourceIdentity"] +} +``` + ### Job Monitoring via Email Notifications If you want to enable email notifications for SMTJ training jobs, your IAM role needs additional permissions to create and manage the notification infrastructure. The notification system uses CloudFormation to automatically provision resources including DynamoDB, SNS, Lambda, EventBridge, and IAM roles. diff --git a/docs/user-guides/instance_type_spec.md b/docs/user-guides/instance_type_spec.md index 9423dc9..66f448a 100644 --- a/docs/user-guides/instance_type_spec.md +++ b/docs/user-guides/instance_type_spec.md @@ -118,4 +118,4 @@ _All allow 1, 2, 4, 8, or 16 instances_ | Lite | ml.g6.12xlarge, ml.g6.24xlarge, ml.g6.48xlarge, ml.p5.48xlarge | | Lite 2.0 | ml.g6.48xlarge, ml.p5.48xlarge | -SageMaker Inference configuration for max_context_length and max_concurrency for each model/instance combination can be found in the [SageMaker Inference AWS documentation](https://docs.aws.amazon.com/nova/latest/userguide/nova-model-sagemaker-inference.html). \ No newline at end of file +SageMaker Inference configuration for max_context_length and max_concurrency for each model/instance combination, as well as inference component support, can be found in the [SageMaker Inference AWS documentation](https://docs.aws.amazon.com/nova/latest/userguide/nova-model-sagemaker-inference.html). \ No newline at end of file diff --git a/docs/user-guides/rft_multiturn.md b/docs/user-guides/rft_multiturn.md index 2894837..4dcee39 100644 --- a/docs/user-guides/rft_multiturn.md +++ b/docs/user-guides/rft_multiturn.md @@ -1,23 +1,369 @@ # RFT Multiturn -The Nova Forge SDK supports RFT (Reinforcement Fine-Tuning) multiturn training for multi-turn conversational tasks. This module provides infrastructure management and orchestration for running RFT training with custom reward environments. +The Nova Forge SDK supports RFT (Reinforcement Fine-Tuning) multiturn training for multi-turn conversational tasks. There are two deployment options: + +- **Serverless (SMTJServerless)**: Fully managed, no infrastructure setup required. Uses Bedrock AgentCore or a custom Lambda function as the agent environment. +- **HyperPod (SMHP)**: Self-managed infrastructure on SageMaker HyperPod with custom reward environments deployed on LOCAL, EC2, or ECS. ## Table of Contents -- [Overview](#overview) -- [Prerequisites](#prerequisites) -- [Quick Start](#quick-start) -- [Infrastructure Setup](#infrastructure-setup) +- [Serverless MTRL](#serverless-mtrl) + - [Prerequisites](#serverless-prerequisites) + - [Quick Start](#serverless-quick-start) + - [Training](#serverless-training) + - [Monitoring](#serverless-monitoring) + - [Iterative Training](#iterative-training) + - [Model Artifacts](#retrieving-model-artifacts) + - [Evaluation](#serverless-evaluation) + - [Save and Load Results](#save-and-load-results) - [Dataset Format](#dataset-format) -- [Training](#training) -- [Evaluation](#evaluation) -- [Monitoring](#monitoring) -- [Custom Environments](#custom-environments) -- [Helper Functions](#helper-functions) -- [Cleanup](#cleanup) +- [HyperPod MTRL](#hyperpod-mtrl) + - [Overview](#overview) + - [Prerequisites](#hyperpod-prerequisites) + - [Quick Start](#hyperpod-quick-start) + - [Infrastructure Setup](#infrastructure-setup) + - [Training](#hyperpod-training) + - [Evaluation](#hyperpod-evaluation) + - [Monitoring](#hyperpod-monitoring) + - [Custom Environments](#custom-environments) + - [Helper Functions](#helper-functions) + - [Cleanup](#cleanup) - [Platform Support](#platform-support) -## Overview +--- + +## Serverless MTRL + +Serverless MTRL uses `SMTJServerlessRuntimeManager` with a Bedrock AgentCore runtime or a custom Lambda function as the agent environment. No infrastructure setup, no HyperPod cluster — just configure and train. + +### Serverless Prerequisites + +- Python 3.12 +- AWS credentials configured +- A Bedrock AgentCore runtime deployed +- S3 bucket with training prompts (parquet format) +- IAM execution role with SageMaker permissions +- `sagemaker-train` wheel installed (`pip install sagemaker_train-*.whl sagemaker_core-*.whl`) + +### Serverless Quick Start + +This quickstart uses Bedrock AgentCore as the agent environment for MTRL training. + +```python +from amzn_nova_forge import * + +# 1. Configure runtime +runtime = SMTJServerlessRuntimeManager( + model_package_group_name="my-model-package-group", + execution_role="arn:aws:iam::123456789012:role/my-role", + agent_core_arn="arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/my-agent", +) + +# 2. Train +trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=runtime, + training_data_s3_path="s3://bucket/prompts/train.parquet", + config=ForgeConfig(output_s3_path="s3://bucket/output/"), +) + +result = trainer.train(job_name="my-mtrl-job") + +# 3. Wait and get output model +result.wait() +print(result.model_artifacts.output_model_arn) +``` + +### Serverless Training + +#### SMTJServerlessRuntimeManager + +```python +# Option 1: Using Bedrock AgentCore +runtime = SMTJServerlessRuntimeManager( + model_package_group_name="my-mpg", # Required: output model package group + execution_role="arn:aws:iam::123456789012:role/my-execution-role", # Required: IAM role + agent_core_arn="arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/my-agent", # Bedrock AgentCore runtime + intermediate_model_package_group_name="my-checkpoints", # Optional: for intermediate checkpoints + kms_key_id="arn:aws:kms:us-east-1:123456789012:key/my-key-id", # Optional: encryption + subnets=["subnet-123"], # Optional: VPC config + security_group_ids=["sg-123"], # Optional: VPC config +) + +# Option 2: Using a custom Lambda function +runtime = SMTJServerlessRuntimeManager( + model_package_group_name="my-mpg", + execution_role="arn:aws:iam::123456789012:role/my-execution-role", + rft_lambda="arn:aws:lambda:us-east-1:123456789012:function:my-reward-fn", # Lambda ARN +) +``` + +#### ForgeTrainer for MTRL + +```python +trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=runtime, + training_data_s3_path="s3://bucket/prompts/train.parquet", + model_arn=None, # Set for iterative training (model package ARN from previous MTRL job) + config=ForgeConfig( + output_s3_path="s3://bucket/output/", + mlflow_monitor=MLflowMonitor(tracking_uri="arn:aws:sagemaker:us-east-1:123456789012:mlflow-tracking-server/my-server"), + ), + region="us-east-1", +) +``` + +#### Training with Hyperparameter Overrides + +```python +result = trainer.train( + job_name="my-mtrl-job", + overrides={ + "global_batch_size": 16, + "max_steps": 50, + "rollout_timeout": 600, + # "learning_rate": 4e-5, + # "lora_rank": 32, + # "lora_alpha": 64, + # "advantage_method": "group_based", + # "group_size": 8, + # "rollout_max_concurrency": 96, + # "sampling_temperature": 0.7, + # "top_p": 0.95, + # "save_every": 10, + # "eval_every": 10, + }, +) +``` + +### Serverless Monitoring + +```python +from amzn_nova_forge.monitor import MTRLLogMonitor + +# Live progress panel (blocks while job is running) +monitor = MTRLLogMonitor.from_job_id(job_id=result.job_id, region="us-east-1") +monitor.show_logs() + +# Check status +status, raw_status = result.get_job_status() +print(f"Status: {status.value} ({raw_status})") + +# View per-step training metrics (from MLflow) +result.get_training_metrics() +``` + +### Iterative Training + +After an MTRL training job completes, continue training from the resulting model package using the `model_arn` parameter. This trains on top of the previous MTRL fine-tune rather than starting from the base model. The `model_arn` must be a model package ARN from a previous MTRL job (or an SFT checkpoint registered in a model package group). + +```python +# Get the output model ARN from a completed job +previous_model_arn = result.model_artifacts.output_model_arn + +# Create a new trainer with model_arn for iterative training +iterative_trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=runtime, + training_data_s3_path="s3://bucket/prompts/train.parquet", + model_arn=previous_model_arn, # Train on top of previous fine-tune + config=ForgeConfig(output_s3_path="s3://bucket/output/"), +) + +iter_result = iterative_trainer.train(job_name="my-iterative-job") +``` + +### Retrieving Model Artifacts + +Retrieve model artifacts from a completed MTRL job: + +```python +from amzn_nova_forge.util.sagemaker import get_model_artifacts + +artifacts = get_model_artifacts( + job_name="my-mtrl-job-id", + infra=runtime, + region="us-east-1", +) + +print(artifacts.output_model_arn) # Model package ARN +print(artifacts.output_s3_path) # S3 output path (from job config) +``` + +For MTRL jobs, `output_s3_path` is not required as a parameter — it is fetched from the AgentRFT job's `OutputDataConfig`. + +### Serverless Evaluation + +The evaluator supports three modes for MTRL on Serverless SMTJ: + +#### Base Model Only + +Evaluates the base Nova model without any fine-tuning: + +```python +evaluator = ForgeEvaluator( + model=Model.NOVA_LITE_2, + infra=runtime, + data_s3_path="s3://bucket/prompts/eval.parquet", + config=ForgeConfig(output_s3_path="s3://bucket/eval-output/"), +) + +eval_result = evaluator.evaluate( + job_name="eval-base-model", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, +) +``` + +#### Fine-Tuned Model Only + +Evaluates a fine-tuned model from a Restricted Model Package (RMP): + +```python +eval_result = evaluator.evaluate( + job_name="eval-finetuned", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + model_path="arn:aws:sagemaker:us-east-1:123456789012:model-package/my-rmp/1", +) +``` + +You can also pass a `job_result` from a completed training job to auto-resolve the model: + +```python +eval_result = evaluator.evaluate( + job_name="eval-finetuned", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + job_result=training_result, # Auto-resolves model checkpoint +) +``` + +#### Base + Fine-Tuned Comparison + +Evaluates both the base model and fine-tuned model in a single pipeline for side-by-side comparison: + +```python +eval_result = evaluator.evaluate( + job_name="eval-comparison", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + model_path="arn:aws:sagemaker:us-east-1:123456789012:model-package/my-rmp/1", + task_config=EvalTaskConfig(evaluate_base_model=True), +) +``` + +This creates a pipeline with both `EvaluateBaseModel` and `EvaluateFineTunedModel` steps. Results are tracked in MLflow for comparison. + +#### Evaluation Overrides + +All modes support hyperparameter overrides: + +```python +eval_result = evaluator.evaluate( + job_name="my-eval-job", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + model_path=model_package_arn, + overrides={ + "global_batch_size": 16, + "max_steps": 20, + }, +) + +print(f"Eval Job ID: {eval_result.job_id}") +``` + +### Save and Load Results + +```python +# Save result to file +result_path = result.dump() + +# Load in a new session +from amzn_nova_forge.core.result import TrainingResult +loaded_result = TrainingResult.load(result_path) + +# All operations work on loaded results +status, raw = loaded_result.get_job_status() +print(f"Output Model: {loaded_result.model_artifacts.output_model_arn}") +``` + +--- + +## Dataset Format + +RFT Multiturn training requires a dataset with specific fields. The SDK supports both flat and nested formats, as well as OpenAI message format for prompts. + +### Required Fields + +- `id` (str): Unique identifier for each sample +- `prompt` (str or list): The input prompt + - Can be a simple string: `"What is 2+2?"` + - Can be OpenAI message format: `[{"role": "user", "content": "What is 2+2?"}]` + +### Optional Fields + +- `answer` (str): Expected answer or completion +- `task` (str): Task category or type +- `info` (dict or str): Additional metadata + - Can be a dictionary: `{"difficulty": "easy"}` + - Can be a valid JSON string: `"{\"difficulty\": \"easy\"}"` + +**Important**: If any sample includes an optional field (answer, task, or info), ALL samples must include that field for consistency. + +**Loader initialization:** +```python +loader = CSVDatasetLoader(id="id", prompt="prompt", answer="answer", task="task", info="info") +``` + +### Dataset Validation + +```python +from amzn_nova_forge import JSONLDatasetLoader, TrainingMethod, Model, EvaluationTask + +# Load dataset +loader = JSONLDatasetLoader(id="id", prompt="prompt", answer="answer") +loader.load("data.jsonl") + +# Transform to RFT Multiturn format for training +loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + +# Transform to RFT Multiturn format for evaluation +# loader.transform(method=TrainingMethod.EVALUATION, eval_task=EvaluationTask.RFT_MULTITURN_EVAL, model=Model.NOVA_LITE_2) + + +# Validate dataset for training +loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + +# Validate dataset for evaluation +# loader.validate(method=TrainingMethod.EVALUATION, eval_task=EvaluationTask.RFT_MULTITURN_EVAL, model=Model.NOVA_LITE_2) + +# Upload to S3 +s3_path = loader.save_data("s3://my-bucket/data/training_data.jsonl") +print(f"Dataset uploaded to: {s3_path}") + +``` + +### Validation Rules + +- `id`: Must be unique across all samples +- `prompt`: Cannot be empty + - String prompts must be non-empty + - OpenAI format must have valid roles (system, user, assistant, tool, function) + - Tool messages must have `tool_call_id` and non-empty `content` + - Assistant messages with tool_calls must have valid structure +- `answer`: Optional, but if present in any sample, must be present in all samples +- `task`: Optional, but if present in any sample, must be present in all samples +- `info`: Optional, but if present in any sample, must be present in all samples + - Must be a dict or valid JSON string + +--- + +## HyperPod MTRL + +HyperPod MTRL uses `RFTMultiturnInfrastructure` to deploy custom reward environments on LOCAL, EC2, or ECS, and trains on a SageMaker HyperPod cluster. + +### Overview RFT Multiturn enables you to fine-tune Nova models using reinforcement learning with custom reward functions. The infrastructure can be deployed on three platforms: @@ -25,9 +371,9 @@ RFT Multiturn enables you to fine-tune Nova models using reinforcement learning - **EC2**: Runs on an AWS EC2 instance - **ECS**: Runs on AWS ECS Fargate -## Prerequisites +### HyperPod Prerequisites -### General Requirements +#### General Requirements - Python 3.12 - AWS credentials configured @@ -35,14 +381,14 @@ RFT Multiturn enables you to fine-tune Nova models using reinforcement learning dditional SSM and ECS permissions are required - see the "If performing RFT Multiturn training" section in the README - SageMaker HyperPod cluster (for training) -### Platform-Specific Requirements +#### Platform-Specific Requirements -#### For LOCAL Platform or SageMaker Notebook +##### For LOCAL Platform or SageMaker Notebook - Python 3.12 installed locally - Sufficient local compute resources -#### For EC2 Platform +##### For EC2 Platform Requirements: - EC2 instance launched with **Amazon Linux 2023** ([guide](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/LaunchingAndUsingInstances.html)) @@ -50,13 +396,13 @@ Requirements: - SSM access enabled on the instance - IAM permissions for SSM commands -#### For ECS Platform +##### For ECS Platform - ECS cluster with Fargate ([guide](https://docs.aws.amazon.com/AmazonECS/latest/developerguide/create-cluster-console-v2.html)) - VPC with subnets and security groups configured - IAM permissions for ECS task management -## Quick Start +### HyperPod Quick Start ```python from amzn_nova_forge import * @@ -97,9 +443,9 @@ training_result = trainer.train( rft_infra.cleanup(delete_stack=True) ``` -## Infrastructure Setup +### Infrastructure Setup -### RFTMultiturnInfrastructure Constructor +#### RFTMultiturnInfrastructure Constructor The `RFTMultiturnInfrastructure` class is the main entry point for managing RFT multiturn infrastructure. @@ -127,7 +473,7 @@ The platform is automatically detected based on `infrastructure_arn`: - **EC2**: `infrastructure_arn` which is ARN of EC2 instance or instance id starting with `i-` - **ECS**: `infrastructure_arn` which is ARN of ECS cluster -### LOCAL Platform +#### LOCAL Platform ```python from amzn_nova_forge import RFTMultiturnInfrastructure, VFEnvId @@ -142,7 +488,7 @@ rft_infra = RFTMultiturnInfrastructure( rft_infra.setup() ``` -### EC2 Platform +#### EC2 Platform ```python rft_infra = RFTMultiturnInfrastructure( @@ -156,7 +502,7 @@ rft_infra = RFTMultiturnInfrastructure( rft_infra.setup() ``` -### ECS Platform +#### ECS Platform ```python rft_infra = RFTMultiturnInfrastructure( @@ -175,7 +521,7 @@ rft_infra = RFTMultiturnInfrastructure( rft_infra.setup() ``` -### Custom Starter Kit Path +#### Custom Starter Kit Path You can provide a custom starter kit path instead of using the default AWS starter kit: @@ -209,7 +555,7 @@ rft_infra = RFTMultiturnInfrastructure( ) ``` -### Custom IAM Role +#### Custom IAM Role ```python rft_infra = RFTMultiturnInfrastructure( @@ -222,78 +568,9 @@ rft_infra = RFTMultiturnInfrastructure( ) ``` -## Dataset Format +### HyperPod Training -RFT Multiturn training requires a dataset with specific fields. The SDK supports both flat and nested formats, as well as OpenAI message format for prompts. - -### Required Fields - -- `id` (str): Unique identifier for each sample -- `prompt` (str or list): The input prompt - - Can be a simple string: `"What is 2+2?"` - - Can be OpenAI message format: `[{"role": "user", "content": "What is 2+2?"}]` - -### Optional Fields - -- `answer` (str): Expected answer or completion -- `task` (str): Task category or type -- `info` (dict or str): Additional metadata - - Can be a dictionary: `{"difficulty": "easy"}` - - Can be a valid JSON string: `"{\"difficulty\": \"easy\"}"` - -**Important**: If any sample includes an optional field (answer, task, or info), ALL samples must include that field for consistency. - -**Loader initialization:** -```python -loader = CSVDatasetLoader(id="id", prompt="prompt", answer="answer", task="task", info="info") -``` - -### Dataset Validation - -The SDK automatically validates your dataset: - -```python -from amzn_nova_forge import JSONLDatasetLoader, TrainingMethod, Model, EvaluationTask - -# Load dataset -loader = JSONLDatasetLoader(id="id", prompt="prompt", answer="answer") -loader.load("data.jsonl") - -# Transform to RFT Multiturn format for training -loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) - -# Transform to RFT Multiturn format for evaluation -# loader.transform(method=TrainingMethod.EVALUATION, eval_task=EvaluationTask.RFT_MULTITURN_EVAL, model=Model.NOVA_LITE_2) - - -# Validate dataset for training -loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) - -# Validate dataset for evaluation -# loader.validate(method=TrainingMethod.EVALUATION, eval_task=EvaluationTask.RFT_MULTITURN_EVAL, model=Model.NOVA_LITE_2) - -# Upload to S3 -s3_path = loader.save_data("s3://my-bucket/data/training_data.jsonl") -print(f"Dataset uploaded to: {s3_path}") - -``` - -### Validation Rules - -- `id`: Must be unique across all samples -- `prompt`: Cannot be empty - - String prompts must be non-empty - - OpenAI format must have valid roles (system, user, assistant, tool, function) - - Tool messages must have `tool_call_id` and non-empty `content` - - Assistant messages with tool_calls must have valid structure -- `answer`: Optional, but if present in any sample, must be present in all samples -- `task`: Optional, but if present in any sample, must be present in all samples -- `info`: Optional, but if present in any sample, must be present in all samples - - Must be a dict or valid JSON string - -## Training - -### setup() Method +#### setup() Method Deploys the SAM stack and validates platform requirements. @@ -320,7 +597,7 @@ This method takes no parameters. rft_infra.setup() ``` -### start_environment() Method +#### start_environment() Method Starts the training or evaluation environment on the configured platform using the unified environment client. @@ -369,7 +646,7 @@ rft_infra.start_environment( ``` -### start_training_environment() Method (DEPRECATED) +#### start_training_environment() Method (DEPRECATED) **DEPRECATED**: This method is deprecated and will be removed in a future version. Use `start_environment(env_type=EnvType.TRAIN, ...)` instead. @@ -395,7 +672,7 @@ rft_infra.start_environment( ) ``` -### start_evaluation_environment() Method (DEPRECATED) +#### start_evaluation_environment() Method (DEPRECATED) **DEPRECATED**: This method is deprecated and will be removed in a future version. Use `start_environment(env_type=EnvType.EVAL, ...)` instead. @@ -421,7 +698,7 @@ rft_infra.start_environment( ) ``` -### get_recipe_overrides() Method +#### get_recipe_overrides() Method Gets recipe parameter overrides for RFT multiturn training jobs. @@ -445,7 +722,7 @@ overrides = rft_infra.get_recipe_overrides() print(f"Lambda ARN: {overrides['rollout_request_arn']}") ``` -### Train with ForgeTrainer +#### Train with ForgeTrainer Use the `train()` method of `ForgeTrainer` with the `rft_multiturn_infra` parameter. @@ -489,9 +766,9 @@ training_result.wait() checkpoint_path = training_result.model_artifacts.checkpoint_s3_path ``` -## Evaluation +### HyperPod Evaluation -### Starting Evaluation Environment +#### Starting Evaluation Environment Use the `start_environment()` method with `env_type=EnvType.EVAL` to start the evaluation environment. @@ -524,11 +801,11 @@ rft_infra.start_environment( ) ``` -### start_evaluation_environment() Method (DEPRECATED) +#### start_evaluation_environment() Method (DEPRECATED) **DEPRECATED**: Use `start_environment(env_type=EnvType.EVAL, ...)` instead. See the [start_environment() documentation](#start_environment-method) for details. -### Evaluate with ForgeEvaluator +#### Evaluate with ForgeEvaluator Use the `evaluate()` method of `ForgeEvaluator` with the `rft_multiturn_infra` parameter. @@ -569,13 +846,13 @@ eval_result.wait() eval_result.show() ``` -## Monitoring +### HyperPod Monitoring -### Session Persistence +#### Session Persistence The SDK provides `dump()` and `load()` methods to save and restore infrastructure state across sessions (e.g., after notebook restarts). -#### dump() Method +##### dump() Method Saves infrastructure state to a JSON file for session recovery. @@ -608,7 +885,7 @@ state_file = rft_infra.dump( # Saves as: my_state.json ``` -#### load() Method (Class Method) +##### load() Method (Class Method) Loads infrastructure state from a file and reconnects to running processes. @@ -645,7 +922,7 @@ logs = rft_infra.get_logs(env_type=EnvType.TRAIN, limit=50) - Resume monitoring after disconnection - Debug issues by loading historical states -### get_logs() Method +#### get_logs() Method Retrieves logs from training, evaluation, or SAM deployment environments. @@ -697,7 +974,7 @@ logs = rft_infra.get_logs( ) ``` -### check_all_queues() Method +#### check_all_queues() Method Checks message counts in all SQS queues. @@ -721,7 +998,7 @@ for queue_name, counts in queue_status.items(): print(f" Last modified: {counts.last_receive_timestamp}") ``` -### flush_all_queues() Method +#### flush_all_queues() Method Purges all messages from all SQS queues. @@ -746,7 +1023,7 @@ This method takes no parameters. rft_infra.flush_all_queues() ``` -### get_configuration() Method +#### get_configuration() Method Gets complete infrastructure configuration. @@ -775,9 +1052,9 @@ print(f"Region: {config['region']}") print(f"Platform: {config['platform']}") ``` -## Custom Environments +### Custom Environments -### CustomEnvironment Class +#### CustomEnvironment Class The `CustomEnvironment` class allows you to create and package custom reward environments. @@ -865,7 +1142,7 @@ s3_uri = custom_env.package_and_upload( print(f"Uploaded to: {s3_uri}") ``` -### Create Custom Environment +#### Create Custom Environment ```python from amzn_nova_forge import CustomEnvironment @@ -886,7 +1163,7 @@ s3_uri = custom_env.package_and_upload() print(f"Environment uploaded to: {s3_uri}") ``` -### Use Custom Environment +#### Use Custom Environment ```python # For LOCAL platform @@ -908,11 +1185,11 @@ rft_infra = RFTMultiturnInfrastructure( ) ``` -### Built-in Environments +#### Built-in Environments The SDK provides two built-in environments via the `VFEnvId` enum: -#### VFEnvId.WORDLE +##### VFEnvId.WORDLE A Wordle game environment for word-guessing tasks. @@ -938,7 +1215,7 @@ rft_infra.start_environment( ) ``` -#### VFEnvId.TERMINAL_BENCH +##### VFEnvId.TERMINAL_BENCH A terminal benchmark environment for command-line tasks. @@ -961,9 +1238,9 @@ rft_infra.start_environment( ) ``` -## Helper Functions +### Helper Functions -### create_rft_execution_role() +#### create_rft_execution_role() Creates an IAM role with required permissions for RFT multiturn infrastructure. @@ -1025,7 +1302,7 @@ rft_infra = RFTMultiturnInfrastructure( ) ``` -### list_rft_stacks() +#### list_rft_stacks() Lists CloudFormation stacks related to RFT multiturn infrastructure. @@ -1054,9 +1331,9 @@ all_stacks = list_rft_stacks(region="us-east-1", all_stacks=True) print(f"Found {len(all_stacks)} total stacks") ``` -## Cleanup +### Cleanup -### kill_task() Method +#### kill_task() Method Stops a running training or evaluation task. @@ -1086,7 +1363,7 @@ rft_infra.kill_task(env_type=EnvType.TRAIN) rft_infra.kill_task(env_type=EnvType.EVAL) ``` -### cleanup() Method +#### cleanup() Method Cleans up infrastructure resources. @@ -1143,6 +1420,8 @@ rft_infra.cleanup(delete_stack=True) # cleanup_environment defaults to False - Environment cleanup is irreversible - you'll need to run `setup()` again - Use `delete_stack=False` if you plan to reuse the stack for another training/evaluation run +--- + ## Platform Support ### Supported Models @@ -1157,19 +1436,21 @@ rft_infra.cleanup(delete_stack=True) # cleanup_environment defaults to False ### Supported Platforms -- **Training**: SageMaker HyperPod (SMHP) only -- **Infrastructure**: LOCAL, EC2, or ECS +| Platform | Runtime Manager | Agent Environment | Infrastructure Required | +|------------|--------------------------------|-----------------------------|--------------------------------| +| Serverless | `SMTJServerlessRuntimeManager` | Bedrock AgentCore or Lambda | None | +| HyperPod | `SMHPRuntimeManager` | LOCAL / EC2 / ECS | `RFTMultiturnInfrastructure` | -### Platform Comparison +### HyperPod Infrastructure Comparison -| Feature | LOCAL | EC2 | ECS | -|-------------------|----------|----------|----------------------------------------------| -| Setup Complexity | Low | Medium | Medium with default network configs else High| -| Scalability | Limited | Medium | High | -| Cost | Low | Medium | Medium-High | -| python_venv_name | Required | Required | Optional | -| VPC Config | N/A | N/A | Required | -| CPU/Memory Config | N/A | N/A | Optional | +| Feature | LOCAL | EC2 | ECS | +|------------------|----------|----------|-----------------------------------------------| +| Setup Complexity | Low | Medium | Medium with default network configs else High | +| Scalability | Limited | Medium | High | +| Cost | Low | Medium | Medium-High | +| python_venv_name | Required | Required | Optional | +| VPC Config | N/A | N/A | Required | +| CPU/Memory Config| N/A | N/A | Optional | ## Examples @@ -1309,4 +1590,5 @@ get_logs(env_type=EnvType.TRAIN, tail=True) - [Main SDK Documentation](../README.md) - Complete SDK overview and getting started guide - [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 +- [RFT Multiturn Notebook](../samples/rft_multiturn_quickstart.ipynb) - RFT multiturn specific examples (HyperPod) +- [Serverless MTRL Notebook](../samples/rft_multiturn_serverless_quickstart.ipynb) - Serverless MTRL examples diff --git a/pyproject.toml b/pyproject.toml index 037792f..3c6b63a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,12 +17,13 @@ authors = [ license = "Apache-2.0" readme = "README.md" dependencies = [ - "boto3", + "boto3>=1.43.20", "filetype", "matplotlib", "pydantic>=2.0", "pyarrow>=14.0.1", - "sagemaker>=3.5.0", + "sagemaker>=3.13.0", + "tqdm", ] [project.optional-dependencies] video = ["pymediainfo"] diff --git a/samples/rft_multiturn_serverless_quickstart.ipynb b/samples/rft_multiturn_serverless_quickstart.ipynb new file mode 100644 index 0000000..a5a604f --- /dev/null +++ b/samples/rft_multiturn_serverless_quickstart.ipynb @@ -0,0 +1,628 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "7eeb1185", + "metadata": {}, + "source": [ + "# Amazon Nova Forge SDK - Serverless Multi-Turn RL Training (MTRL) Quick Start\n", + "\n", + "This notebook provides a walkthrough of the Amazon Nova Forge SDK for multi-turn reinforcement fine-tuning (MTRL) using SMTJ Serverless with a Bedrock AgentCore agent.\n", + "\n", + "## What You'll Learn\n", + "\n", + "1. Configuring a Bedrock AgentCore runtime as the agent environment\n", + "2. Launching an MTRL training job with hyperparameter overrides\n", + "3. Monitoring training progress, viewing metrics\n", + "4. Saving/loading job results across sessions\n", + "5. Running MTRL evaluation (base model, fine-tuned, or comparison)\n", + "6. Iterative training (SFT → MTRL)\n", + "7. Deploying to SageMaker and Bedrock\n", + "\n", + "## Table of Contents\n", + "- [Step 1: Import Required Modules](#step-1-import-required-modules)\n", + "- [Step 2: Configure Your AWS Resources](#step-2-configure-your-aws-resources)\n", + "- [Step 3: Configure Runtime Infrastructure](#step-3-configure-runtime-infrastructure)\n", + "- [Step 4: Initialize ForgeTrainer](#step-4-initialize-forgetrainer)\n", + "- [Step 5: Start Training](#step-5-start-training)\n", + "- [Step 6: Monitor Training Progress](#step-6-monitor-training-progress)\n", + "- [Step 7: Save and Load Job Results](#step-7-save-and-load-job-results)\n", + "- [Step 8: Evaluate Your Model](#step-8-evaluate-your-model-after-training-completes)\n", + "- [Step 9: Evaluation Modes](#step-9-evaluation-modes)\n", + "- [Step 10: Iterative Training (SFT → MTRL)](#step-10-iterative-training-sft--mtrl)\n", + "- [Step 11: Deploy Your Model](#step-11-deploy-your-model)\n", + "\n", + "## Prerequisites\n", + "\n", + "- AWS credentials configured\n", + "- A Bedrock AgentCore runtime deployed\n", + "- S3 bucket with training prompts (parquet format)\n", + "- IAM execution role with SageMaker permissions\n", + "- Nova Forge SDK installed per [its README](https://github.com/aws/nova-forge-sdk/blob/main/README.md#installation)" + ] + }, + { + "cell_type": "markdown", + "id": "2e9094fc", + "metadata": {}, + "source": [ + "## Step 1: Import Required Modules" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ede465fa", + "metadata": {}, + "outputs": [], + "source": [ + "!pip install amzn-nova-forge" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4acdd637", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "import boto3\n", + "from botocore.exceptions import ClientError, NoCredentialsError, ProfileNotFound\n", + "\n", + "\n", + "def load_credentials(profile=None):\n", + " \"\"\"Load AWS credentials with fallback behavior.\"\"\"\n", + " if profile:\n", + " try:\n", + " session = boto3.Session(profile_name=profile)\n", + " credentials = session.get_credentials()\n", + " if not credentials:\n", + " raise RuntimeError(f\"No credentials found for profile '{profile}'\")\n", + " except ProfileNotFound:\n", + " raise RuntimeError(f\"Profile '{profile}' not found in credentials file\")\n", + " else:\n", + " try:\n", + " session = boto3.Session()\n", + " credentials = session.get_credentials()\n", + " if not credentials:\n", + " raise RuntimeError(\"No credentials found in current AWS session\")\n", + " except NoCredentialsError:\n", + " raise RuntimeError(\"No AWS credentials configured\")\n", + "\n", + " # Validate credentials\n", + " try:\n", + " sts_client = session.client(\"sts\")\n", + " sts_client.get_caller_identity()\n", + " except ClientError as e:\n", + " raise RuntimeError(f\"Invalid AWS credentials: {e}\")\n", + "\n", + " return {\n", + " \"aws_access_key_id\": credentials.access_key,\n", + " \"aws_secret_access_key\": credentials.secret_key,\n", + " \"aws_session_token\": credentials.token,\n", + " \"region_name\": session.region_name or \"us-east-1\",\n", + " }" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5de14de6", + "metadata": {}, + "outputs": [], + "source": [ + "load_credentials()\n", + "print(\"AWS credentials validated successfully!\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "86f3d4ec", + "metadata": {}, + "outputs": [], + "source": [ + "from amzn_nova_forge import *\n", + "from amzn_nova_forge.core.result import BaseJobResult\n", + "\n", + "print(\"SDK imported successfully!\")" + ] + }, + { + "cell_type": "markdown", + "id": "0e8ae348", + "metadata": {}, + "source": [ + "## Step 2: Configure Your AWS Resources" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7a64c735", + "metadata": {}, + "outputs": [], + "source": [ + "# TODO: Update these values for your environment\n", + "ACCOUNT_ID = \"\"\n", + "REGION = \"us-east-1\"\n", + "\n", + "EXECUTION_ROLE = f\"arn:aws:iam::{ACCOUNT_ID}:role/\"\n", + "MODEL_PKG_GROUP = \"\"\n", + "MLFLOW_ARN = f\"arn:aws:sagemaker:{REGION}:{ACCOUNT_ID}:mlflow-tracking-server/\"\n", + "\n", + "# Agent environment: Bedrock AgentCore runtime ARN\n", + "AGENT_CORE_ARN = f\"arn:aws:bedrock-agentcore:{REGION}:{ACCOUNT_ID}:runtime/\"\n", + "\n", + "# Training and evaluation data\n", + "TRAINING_DATA = f\"s3:///prompts/train.parquet\"\n", + "EVAL_DATA = f\"s3:///prompts/eval.parquet\"\n", + "\n", + "# Output path\n", + "OUTPUT_S3 = f\"s3:///output/\"\n", + "\n", + "print(f\"Training Data: {TRAINING_DATA}\")\n", + "print(f\"Output Path: {OUTPUT_S3}\")" + ] + }, + { + "cell_type": "markdown", + "id": "6e7318e1", + "metadata": {}, + "source": [ + "## Step 3: Configure Runtime Infrastructure" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e13a074a", + "metadata": {}, + "outputs": [], + "source": [ + "# For MTRL, use agent_core_arn instead of rft_lambda\n", + "runtime = SMTJServerlessRuntimeManager(\n", + " model_package_group_name=MODEL_PKG_GROUP,\n", + " execution_role=EXECUTION_ROLE,\n", + " agent_core_arn=AGENT_CORE_ARN,\n", + ")\n", + "\n", + "print(\"Runtime configured for SMTJ Serverless (MTRL)\")" + ] + }, + { + "cell_type": "markdown", + "id": "23e06ab2", + "metadata": {}, + "source": [ + "## Step 4: Initialize ForgeTrainer" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4e9b3382", + "metadata": {}, + "outputs": [], + "source": [ + "# Create MLflow monitor (optional)\n", + "mlflow_monitor = MLflowMonitor(tracking_uri=MLFLOW_ARN)\n", + "\n", + "# Create trainer\n", + "trainer = ForgeTrainer(\n", + " model=Model.NOVA_LITE_2,\n", + " method=TrainingMethod.RFT_MULTITURN_LORA,\n", + " infra=runtime,\n", + " training_data_s3_path=TRAINING_DATA,\n", + " config=ForgeConfig(\n", + " output_s3_path=OUTPUT_S3,\n", + " mlflow_monitor=mlflow_monitor,\n", + " ),\n", + " region=REGION,\n", + ")\n", + "\n", + "print(\"ForgeTrainer initialized\")\n", + "print(f\" Model: Nova Lite 2.0\")\n", + "print(f\" Method: RFT Multi-Turn with LoRA\")" + ] + }, + { + "cell_type": "markdown", + "id": "a998a9be", + "metadata": {}, + "source": [ + "## Step 5: Start Training\n", + "\n", + "Overrides are applied directly to the `MultiTurnRLTrainer.hyperparameters` object. Use exact parameter names." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "00f85910", + "metadata": {}, + "outputs": [], + "source": [ + "# Define training hyperparameters\n", + "training_overrides = {\n", + " \"global_batch_size\": 16,\n", + " \"max_steps\": 50,\n", + " \"rollout_timeout\": 600,\n", + " # Other available overrides:\n", + " # \"learning_rate\": 4e-5,\n", + " # \"lora_rank\": 32,\n", + " # \"lora_alpha\": 64,\n", + " # \"advantage_method\": \"group_based\",\n", + " # \"group_size\": 8,\n", + " # \"rollout_max_concurrency\": 96,\n", + " # \"sampling_temperature\": 0.7,\n", + " # \"top_p\": 0.95,\n", + " # \"save_every\": 10,\n", + " # \"eval_every\": 10,\n", + "}\n", + "\n", + "# Start training\n", + "training_result = trainer.train(\n", + " job_name=\"\",\n", + " overrides=training_overrides,\n", + ")\n", + "\n", + "print(\"\\nTraining job started!\")\n", + "print(f\" Job ID: {training_result.job_id}\")\n", + "print(f\" Output Path: {training_result.model_artifacts.output_s3_path}\")" + ] + }, + { + "cell_type": "markdown", + "id": "b431de2c", + "metadata": {}, + "source": [ + "## Step 6: Monitor Training Progress" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "af254634", + "metadata": {}, + "outputs": [], + "source": [ + "# Block until the job reaches a terminal state (displays rich progress panel)\n", + "training_result.wait(poll=30, timeout=7200)" + ] + }, + { + "cell_type": "markdown", + "id": "7bafe857", + "metadata": {}, + "source": [ + "### Check Job Status and Metrics" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16437ed3", + "metadata": {}, + "outputs": [], + "source": [ + "status, raw_status = training_result.get_job_status()\n", + "print(f\"Status: {status.value} ({raw_status})\")\n", + "print(f\"Output Model Package: {training_result.model_artifacts.output_model_package_arn}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6bbece15", + "metadata": {}, + "outputs": [], + "source": [ + "# View per-step training metrics from MLflow\n", + "training_result.get_training_metrics()" + ] + }, + { + "cell_type": "markdown", + "id": "c6a11c59", + "metadata": {}, + "source": [ + "## Step 7: Save and Load Job Results" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "014c0fe7", + "metadata": {}, + "outputs": [], + "source": [ + "training_result_path = training_result.dump()\n", + "print(\"Training result saved\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "581729b2", + "metadata": {}, + "outputs": [], + "source": [ + "loaded_result = TrainingResult.load(training_result_path)\n", + "print(f\"Loaded - Job ID: {loaded_result.job_id}\")\n", + "\n", + "# All MTRL operations work on the loaded result\n", + "status, raw = loaded_result.get_job_status()\n", + "print(f\"Status: {status.value} ({raw})\")\n", + "print(f\"Output Model Package: {loaded_result.model_artifacts.output_model_package_arn}\")" + ] + }, + { + "cell_type": "markdown", + "id": "b023bb79", + "metadata": {}, + "source": [ + "## Step 8: Evaluate Your Model (After Training Completes)\n", + "\n", + "Run MTRL evaluation on the trained model using the same agent environment." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "20f4d5eb", + "metadata": {}, + "outputs": [], + "source": [ + "evaluator = ForgeEvaluator(\n", + " model=Model.NOVA_LITE_2,\n", + " infra=runtime,\n", + " data_s3_path=EVAL_DATA,\n", + " config=ForgeConfig(\n", + " output_s3_path=OUTPUT_S3,\n", + " mlflow_monitor=mlflow_monitor,\n", + " ),\n", + " region=REGION,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8c5ecd8d", + "metadata": {}, + "outputs": [], + "source": [ + "# Pass training_result to auto-resolve the model checkpoint\n", + "eval_result = evaluator.evaluate(\n", + " job_name=\"\",\n", + " eval_task=EvaluationTask.RFT_MULTITURN_EVAL,\n", + " job_result=training_result,\n", + " overrides={\n", + " \"global_batch_size\": 16,\n", + " \"max_steps\": 20,\n", + " },\n", + ")\n", + "\n", + "print(f\"Eval Job ID: {eval_result.job_id}\")\n", + "print(f\"Eval Output: {eval_result.eval_output_path}\")" + ] + }, + { + "cell_type": "markdown", + "id": "57335b1f", + "metadata": {}, + "source": [ + "## Step 9: Evaluation Modes\n", + "\n", + "You can evaluate the base model, fine-tuned model, or compare both in a single pipeline." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "867f3cff", + "metadata": {}, + "outputs": [], + "source": [ + "# Fine-tuned model only (pass model_path directly)\n", + "eval_finetuned = evaluator.evaluate(\n", + " job_name=\"eval-finetuned\",\n", + " eval_task=EvaluationTask.RFT_MULTITURN_EVAL,\n", + " model_path=training_result.model_artifacts.output_model_arn,\n", + ")\n", + "\n", + "# Base + fine-tuned comparison (both in one pipeline)\n", + "eval_comparison = evaluator.evaluate(\n", + " job_name=\"eval-comparison\",\n", + " eval_task=EvaluationTask.RFT_MULTITURN_EVAL,\n", + " model_path=training_result.model_artifacts.output_model_arn,\n", + " task_config=EvalTaskConfig(evaluate_base_model=True),\n", + ")\n", + "print(f\"Comparison eval: {eval_comparison.job_id}\")" + ] + }, + { + "cell_type": "markdown", + "id": "d43afcf5", + "metadata": {}, + "source": [ + "## Step 10: Iterative Training (SFT → MTRL)\n", + "\n", + "Use an SFT checkpoint as the starting point for MTRL. Register the checkpoint in a model package group, then pass it as `model_arn` to `ForgeTrainer`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ba1c2087", + "metadata": {}, + "outputs": [], + "source": [ + "import boto3\n", + "\n", + "sm = boto3.client(\"sagemaker\", region_name=REGION)\n", + "\n", + "# 1. Get SFT checkpoint path from manifest (after SFT job completes)\n", + "SFT_CHECKPOINT = \"s3://customer-escrow--smtj-//step_10\"\n", + "\n", + "# 2. Register in a model package group\n", + "resp = sm.create_model_package(\n", + " ModelPackageGroupName=\"\",\n", + " InferenceSpecification={\n", + " \"Containers\": [\n", + " {\n", + " \"ModelDataSource\": {\n", + " \"S3DataSource\": {\n", + " \"S3Uri\": SFT_CHECKPOINT,\n", + " \"S3DataType\": \"S3Prefix\",\n", + " \"CompressionType\": \"None\",\n", + " }\n", + " },\n", + " \"IsCheckpoint\": False,\n", + " \"BaseModel\": {\n", + " \"HubContentName\": \"nova-textgeneration-lite-v2\",\n", + " \"HubContentVersion\": \"3.48.0\",\n", + " },\n", + " }\n", + " ],\n", + " \"SupportedContentTypes\": [\"application/json\"],\n", + " \"SupportedResponseMIMETypes\": [\"application/json\"],\n", + " },\n", + " SkipModelValidation=\"All\",\n", + ")\n", + "SFT_MODEL_PACKAGE_ARN = resp[\"ModelPackageArn\"]\n", + "print(f\"SFT model package: {SFT_MODEL_PACKAGE_ARN}\")\n", + "\n", + "# 3. Launch MTRL with SFT checkpoint as base\n", + "trainer_iterative = ForgeTrainer(\n", + " model=Model.NOVA_LITE_2,\n", + " method=TrainingMethod.RFT_MULTITURN_LORA,\n", + " infra=runtime,\n", + " training_data_s3_path=TRAINING_DATA,\n", + " model_arn=SFT_MODEL_PACKAGE_ARN, # SFT checkpoint as starting point\n", + " config=ForgeConfig(output_s3_path=OUTPUT_S3, mlflow_monitor=mlflow_monitor),\n", + " region=REGION,\n", + ")\n", + "\n", + "iterative_result = trainer_iterative.train(\n", + " job_name=\"sft-to-mtrl\",\n", + " overrides={\"global_batch_size\": 8, \"max_steps\": 12, \"save_every\": 10},\n", + ")\n", + "print(f\"SFT→MTRL Job: {iterative_result.job_id}\")" + ] + }, + { + "cell_type": "markdown", + "id": "57a55ebb", + "metadata": {}, + "source": [ + "## Step 11: Deploy Your Model\n", + "\n", + "Deploy the trained model to SageMaker or Bedrock for inference." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "433fa57b", + "metadata": {}, + "outputs": [], + "source": [ + "MODEL_PACKAGE_ARN = training_result.model_artifacts.output_model_arn\n", + "\n", + "deployer = ForgeDeployer(region=REGION, model=Model.NOVA_LITE_2)\n", + "\n", + "# Option A: Deploy to SageMaker endpoint\n", + "smi_result = deployer.deploy(\n", + " model_artifact_path=MODEL_PACKAGE_ARN,\n", + " deploy_platform=DeployPlatform.SAGEMAKER,\n", + " sagemaker_instance_type=\"ml.g6.48xlarge\",\n", + " endpoint_name=\"my-mtrl-endpoint\",\n", + " execution_role_name=\"\",\n", + ")\n", + "print(f\"SageMaker endpoint: {smi_result.endpoint.uri}\")\n", + "\n", + "# Option B: Deploy to Bedrock On-Demand\n", + "bedrock_model = deployer.create_custom_model(\n", + " model_artifact_path=None,\n", + " custom_model_data_source={\"modelPackageArnDataSource\": {\"modelPackageArn\": MODEL_PACKAGE_ARN}},\n", + " endpoint_name=\"my-mtrl-bedrock\",\n", + " execution_role_name=\"\",\n", + ")\n", + "bedrock_result = deployer.deploy_to_bedrock(\n", + " model_deploy_result=bedrock_model,\n", + " deploy_platform=DeployPlatform.BEDROCK_OD,\n", + ")\n", + "print(f\"Bedrock deployment: {bedrock_result.endpoint.uri}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9aa28ac3", + "metadata": {}, + "outputs": [], + "source": [ + "# Invoke the deployed model\n", + "inference = ForgeInference(region=REGION)\n", + "\n", + "result = inference.invoke(\n", + " endpoint_arn=smi_result.endpoint.uri, # or bedrock_result.endpoint.uri\n", + " request_body={\n", + " \"messages\": [{\"role\": \"user\", \"content\": \"What is 25 * 4 + 10?\"}],\n", + " \"max_tokens\": 256,\n", + " },\n", + ")\n", + "result.show()" + ] + }, + { + "cell_type": "markdown", + "id": "50227b81", + "metadata": {}, + "source": [ + "---\n", + "## Summary\n", + "\n", + "| What | How |\n", + "|------|-----|\n", + "| Launch training | `trainer.train(job_name=..., overrides={...})` |\n", + "| Wait for completion | `training_result.wait()` |\n", + "| Check status | `training_result.get_job_status()` |\n", + "| View metrics | `training_result.get_training_metrics()` |\n", + "| Get model package | `training_result.model_artifacts.output_model_arn` |\n", + "| Save/load result | `training_result.dump()` / `TrainingResult.load(path)` |\n", + "| Eval (fine-tuned) | `evaluator.evaluate(model_path=arn, ...)` |\n", + "| Eval (base + FT) | `evaluator.evaluate(model_path=arn, task_config=EvalTaskConfig(evaluate_base_model=True))` |\n", + "| SFT → MTRL | Register SFT checkpoint in MPG, pass as `model_arn` to `ForgeTrainer` |\n", + "| Deploy to SMI | `deployer.deploy(model_artifact_path=arn, deploy_platform=DeployPlatform.SAGEMAKER)` |\n", + "| Deploy to Bedrock | `deployer.create_custom_model(custom_model_data_source=...) + deploy_to_bedrock()` |\n", + "| Invoke | `ForgeInference(region).invoke(endpoint_arn=..., request_body={...})` |" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/samples/serverless_quickstart.ipynb b/samples/serverless_quickstart.ipynb index 70bcce9..53b832a 100644 --- a/samples/serverless_quickstart.ipynb +++ b/samples/serverless_quickstart.ipynb @@ -516,13 +516,19 @@ "metadata": {}, "outputs": [], "source": [ + "# Use the model package ARN from the training result.\n", + "# For serverless, model_path must be a model package ARN (not an S3 path).\n", + "# The training result's output_model_arn contains this after a successful training job.\n", + "model_package_arn = loaded_training_result.model_artifacts.output_model_arn\n", + "print(f\"Model package ARN: {model_package_arn}\")\n", + "\n", "byod_eval_result = evaluator.evaluate(\n", " job_name=\"serverless-eval-byod\",\n", " eval_task=EvaluationTask.GEN_QA,\n", " task_config=EvalTaskConfig(\n", " override_data_s3_path=\"s3:///nova-customization/gen_qa.jsonl\", # TODO: Replace with your data path\n", " ),\n", - " # model_path='s3://customer-escrow-/your-model-path/' # TODO: Replace with your trained model path\n", + " model_path=model_package_arn,\n", " overrides={\"max_new_tokens\": 2048},\n", ")\n", "\n", diff --git a/src/amzn_nova_forge/__version__.py b/src/amzn_nova_forge/__version__.py index 699e2a5..d23c4ba 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.8" # pragma: no cover +VERSION = "1.4.9" # pragma: no cover diff --git a/src/amzn_nova_forge/core/__init__.py b/src/amzn_nova_forge/core/__init__.py index fec2e55..ded1d02 100644 --- a/src/amzn_nova_forge/core/__init__.py +++ b/src/amzn_nova_forge/core/__init__.py @@ -60,6 +60,7 @@ ModelArtifacts, ModelConfigDict, RecipeConfig, + ValidationConfig, validate_region, ) diff --git a/src/amzn_nova_forge/core/constants.py b/src/amzn_nova_forge/core/constants.py index 03a72ab..1895c2e 100644 --- a/src/amzn_nova_forge/core/constants.py +++ b/src/amzn_nova_forge/core/constants.py @@ -16,6 +16,7 @@ """ import os +import re from typing import Dict, List from amzn_nova_forge.core.enums import ( @@ -27,6 +28,15 @@ ) DEFAULT_REGION = "us-east-1" + +# MTRL CloudWatch log groups +MTRL_TRAIN_LOG_GROUP = "/aws/sagemaker/Job/AgentRFT" +MTRL_EVAL_LOG_GROUP = "/aws/sagemaker/Job/AgentRFTEvaluation" + +# MTRL pipeline execution ARN pattern +MTRL_PIPELINE_EXECUTION_RE = re.compile( + r"^arn:aws:sagemaker:[a-z0-9-]+:\d{12}:pipeline/.+/execution/.+" +) DEFAULT_JOB_CACHE_DIR = "~/.nova-forge/cache" DEFAULT_BATCH_TRACE_CACHE_DIR = "~/.nova-forge/batch_trace_cache" BATCH_TRACE_LOG_SUBDIR = "batch_tracing" diff --git a/src/amzn_nova_forge/core/result/__init__.py b/src/amzn_nova_forge/core/result/__init__.py index a08d200..c1f1465 100644 --- a/src/amzn_nova_forge/core/result/__init__.py +++ b/src/amzn_nova_forge/core/result/__init__.py @@ -46,6 +46,9 @@ from amzn_nova_forge.core.result.job_result import ( JobStatusManager as JobStatusManager, ) +from amzn_nova_forge.core.result.job_result import ( + MTRLStatusManager as MTRLStatusManager, +) from amzn_nova_forge.core.result.job_result import ( SMHPStatusManager as SMHPStatusManager, ) diff --git a/src/amzn_nova_forge/core/result/eval_result.py b/src/amzn_nova_forge/core/result/eval_result.py index 5bdd991..bbcc660 100644 --- a/src/amzn_nova_forge/core/result/eval_result.py +++ b/src/amzn_nova_forge/core/result/eval_result.py @@ -31,6 +31,7 @@ BedrockStatusManager, JobStatus, JobStatusManager, + MTRLStatusManager, SMHPStatusManager, SMTJStatusManager, ) @@ -199,6 +200,8 @@ def __del__(self): @dataclass class SMTJEvaluationResult(EvaluationResult): + _job_name: Optional[str] = None + def __init__( self, job_id: str, @@ -209,30 +212,43 @@ def __init__( s3_client=None, region: Optional[str] = None, ): + self._region = region + self._job_name: Optional[str] = None 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: + if self.eval_task == EvaluationTask.RFT_MULTITURN_EVAL: + return MTRLStatusManager(region=self._region) return SMTJStatusManager(self._sagemaker_client) def _to_dict(self): - return { + d = { "job_id": self.job_id, "started_time": self.started_time.isoformat(), "eval_task": self.eval_task.value, "eval_output_path": self.eval_output_path, } + if self._region: + d["region"] = self._region + if getattr(self, "_job_name", None): + d["job_name"] = self._job_name + return d @classmethod def _from_dict(cls, data) -> "SMTJEvaluationResult": - return cls( + result = cls( job_id=data["job_id"], started_time=datetime.fromisoformat(data["started_time"]), eval_task=EvaluationTask(data["eval_task"]), eval_output_path=data["eval_output_path"], + region=data.get("region"), ) + if data.get("job_name"): + result._job_name = data["job_name"] + return result @dataclass diff --git a/src/amzn_nova_forge/core/result/job_result.py b/src/amzn_nova_forge/core/result/job_result.py index b380fcc..6d7c0f5 100644 --- a/src/amzn_nova_forge/core/result/job_result.py +++ b/src/amzn_nova_forge/core/result/job_result.py @@ -24,6 +24,7 @@ import boto3 +from amzn_nova_forge.core.constants import MTRL_PIPELINE_EXECUTION_RE from amzn_nova_forge.core.enums import Platform from amzn_nova_forge.core.validation_patterns import ( validate_cluster_name, @@ -229,6 +230,88 @@ def resolve_start_time(self, job_id: str) -> datetime: raise ValueError(f"Cannot resolve start time for SMHP job {job_id}") +class MTRLStatusManager(JobStatusManager): + """Status manager for MTRL jobs (training via AgentRFT Job API, eval via Pipeline Execution).""" + + def __init__(self, region: Optional[str] = None): + super().__init__() + self._region = region + self._session = boto3.Session(region_name=region) if region else None + + def _is_pipeline_execution(self, job_id: str) -> bool: + return bool(MTRL_PIPELINE_EXECUTION_RE.match(job_id)) + + def _get_pipeline_client(self, job_id: str): + region = job_id.split(":")[3] if ":sagemaker:" in job_id else self._region + return 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: + return self._job_status, self._raw_status + + if self._is_pipeline_execution(job_id): + return self._get_pipeline_status(job_id) + return self._get_job_status(job_id) + + def _get_job_status(self, job_id: str) -> tuple[JobStatus, str]: + from sagemaker.train.agent_rft_job import AgentRFTJob + + rft_job = AgentRFTJob.get(job_id, session=self._session) + raw_status = rft_job.job_status + + status_mapping = { + "InProgress": JobStatus.IN_PROGRESS, + "Completed": JobStatus.COMPLETED, + "Failed": JobStatus.FAILED, + "Stopping": JobStatus.FAILED, + "Stopped": JobStatus.FAILED, + } + job_status = status_mapping.get(raw_status, JobStatus.IN_PROGRESS) + + self._job_status = job_status + self._raw_status = raw_status + return job_status, raw_status + + def _get_pipeline_status(self, job_id: str) -> tuple[JobStatus, str]: + client = self._get_pipeline_client(job_id) + resp = client.describe_pipeline_execution(PipelineExecutionArn=job_id) + raw_status = resp["PipelineExecutionStatus"] + + status_mapping = { + "Executing": JobStatus.IN_PROGRESS, + "Starting": JobStatus.IN_PROGRESS, + "Succeeded": JobStatus.COMPLETED, + "Failed": JobStatus.FAILED, + "Stopping": JobStatus.FAILED, + "Stopped": JobStatus.FAILED, + } + job_status = status_mapping.get(raw_status, JobStatus.IN_PROGRESS) + + self._job_status = job_status + self._raw_status = raw_status + return job_status, raw_status + + def resolve_start_time(self, job_id: str) -> datetime: + if self._is_pipeline_execution(job_id): + client = self._get_pipeline_client(job_id) + resp = client.describe_pipeline_execution(PipelineExecutionArn=job_id) + start_time = resp.get("CreationTime") + if start_time: + return ( + start_time + if isinstance(start_time, datetime) + else datetime.fromisoformat(str(start_time)) + ) + raise ValueError(f"Cannot resolve start time for MTRL eval pipeline {job_id}") + + from sagemaker.train.agent_rft_job import AgentRFTJob + + rft_job = AgentRFTJob.get(job_id, session=self._session) + if rft_job.creation_time: + return rft_job.creation_time + raise ValueError(f"Cannot resolve start time for MTRL job {job_id}") + + class BedrockStatusManager(JobStatusManager): # Injected by util/bedrock.py at import time so core/ has zero internal imports. _get_job_details: Optional[Callable] = None @@ -333,9 +416,13 @@ def __init__(self, job_id: str, started_time: Optional[datetime] = None): Platform.SMTJ if isinstance(self._status_manager, SMTJStatusManager) else ( - Platform.BEDROCK - if isinstance(self._status_manager, BedrockStatusManager) - else Platform.SMHP + Platform.SMTJServerless + if isinstance(self._status_manager, MTRLStatusManager) + else ( + Platform.BEDROCK + if isinstance(self._status_manager, BedrockStatusManager) + else Platform.SMHP + ) ) ) @@ -494,7 +581,15 @@ def dump(self, file_path: Optional[str] = None, file_name: Optional[str] = None) :param file_name: The file name of the result. Default to _.json if not provided :return: The full result file path """ - file_name = file_name or f"{self.job_id}_{self._platform.value}.json" + if not file_name: + # MTRL Pipeline execution ARNs contain "/" which is invalid in filenames + if "/" in self.job_id: + exec_id = self.job_id.rsplit("/", 1)[-1] + job_name = getattr(self, "_job_name", None) + name = f"{job_name}-{exec_id}" if job_name else exec_id + else: + name = self.job_id + file_name = f"{name}_{self._platform.value}.json" if file_path is None: full_path = Path(file_name) diff --git a/src/amzn_nova_forge/core/result/training_result.py b/src/amzn_nova_forge/core/result/training_result.py index 5295f31..bf7702b 100644 --- a/src/amzn_nova_forge/core/result/training_result.py +++ b/src/amzn_nova_forge/core/result/training_result.py @@ -23,6 +23,7 @@ BaseJobResult, BedrockStatusManager, JobStatusManager, + MTRLStatusManager, SMHPStatusManager, SMTJStatusManager, ) @@ -62,6 +63,8 @@ def show(self): @dataclass class SMTJTrainingResult(TrainingResult): + _MTRL_METHODS = (TrainingMethod.RFT_MULTITURN_LORA,) + def __init__( self, job_id: str, @@ -74,11 +77,66 @@ def __init__( ): self._region = region self._sagemaker_client = sagemaker_client or boto3.client("sagemaker", region_name=region) + self._is_mtrl = method in self._MTRL_METHODS + self._rft_job = None super().__init__(job_id, started_time, method, model_artifacts, model_type) def _create_status_manager(self) -> JobStatusManager: + if self._is_mtrl: + return MTRLStatusManager(region=self._region) return SMTJStatusManager(self._sagemaker_client, region=self._region) + def _get_rft_job(self): + """Lazily fetch the AgentRFTJob for MTRL jobs.""" + if self._rft_job is None: + from sagemaker.train.agent_rft_job import AgentRFTJob + + session = boto3.Session(region_name=self._region) if self._region else None + self._rft_job = AgentRFTJob.get(self.job_id, session=session) + # Populate output_model_arn if available and not already set + if not self.model_artifacts.output_model_arn: + arn = self._rft_job.output_model_package_arn + if arn: + self.model_artifacts.output_model_arn = arn + return self._rft_job + + def wait(self, poll: int = 30, timeout: int = 3600) -> None: + """Wait for an MTRL job to complete. + + Delegates to the AgentRFTJob.wait() which displays a rich + progress panel showing job status, metrics, and links. + + Args: + poll: Seconds between status polls (default 30). + timeout: Maximum seconds to wait (default 3600). + """ + if not self._is_mtrl: + raise NotImplementedError( + "wait() is only supported for MTRL jobs. " + "For standard SMTJ jobs, use get_job_status() to poll." + ) + rft_job = self._get_rft_job() + rft_job.wait(poll=poll, timeout=timeout) + # Populate model_artifacts.output_model_arn after completion + rft_job.refresh() + if rft_job.output_model_package_arn: + self.model_artifacts.output_model_arn = rft_job.output_model_package_arn + + def get_training_metrics(self) -> list: + """Fetch per-step training metrics from MLflow for MTRL jobs. + + Delegates to AgentRFTJob.get_training_metrics() which + retrieves reward/mean, turns/mean, total_tokens, and + num_trajectories for each training step. + + Returns: + List of dicts with per-step metrics. + """ + if not self._is_mtrl: + raise NotImplementedError("get_training_metrics() is only supported for MTRL jobs.") + rft_job = self._get_rft_job() + return rft_job.get_training_metrics() + def _to_dict(self): return { "job_id": self.job_id, diff --git a/src/amzn_nova_forge/core/training_overrides.py b/src/amzn_nova_forge/core/training_overrides.py index 2881228..1d68660 100644 --- a/src/amzn_nova_forge/core/training_overrides.py +++ b/src/amzn_nova_forge/core/training_overrides.py @@ -63,5 +63,18 @@ class TrainingOverrides(TypedDict, total=False): top_logprobs: int + # MTRL-specific overrides + advantage_method: str + lora_rank: int + loss_fn: str + max_tokens: int + group_size: int + rollout_max_concurrency: int + rollout_timeout: int + sampling_temperature: float + top_p: float + save_every: int + eval_every: int + NULLABLE_OVERRIDE_FIELDS: frozenset[str] = frozenset({"reasoning_effort"}) diff --git a/src/amzn_nova_forge/core/types.py b/src/amzn_nova_forge/core/types.py index bb1c300..2dfd598 100644 --- a/src/amzn_nova_forge/core/types.py +++ b/src/amzn_nova_forge/core/types.py @@ -72,8 +72,8 @@ class ModelConfigDict(TypedDict): @dataclass class ModelArtifacts: - checkpoint_s3_path: Optional[str] - output_s3_path: str + checkpoint_s3_path: Optional[str] = None + output_s3_path: Optional[str] = None output_model_arn: Optional[str] = None # Model package ARN for SMTJServerless jobs @@ -98,6 +98,7 @@ def escrow_uri(self) -> Optional[str]: return self.model_publish.escrow_uri if self.model_publish else None _status_checker: ClassVar[Optional[Callable]] = None + _sagemaker_status_checker: ClassVar[Optional[Callable]] = None @classmethod def _register_status_checker(cls, checker: Callable) -> None: @@ -108,8 +109,32 @@ def _register_status_checker(cls, checker: Callable) -> None: """ cls._status_checker = checker + @classmethod + def _register_sagemaker_status_checker(cls, checker: Callable) -> None: + """Register the function used to check SageMaker deployment status. + + Called by util/sagemaker.py at import time to wire up the status + property without core/ needing to import util/. + """ + cls._sagemaker_status_checker = checker + @property def status(self): + if self.endpoint.platform == DeployPlatform.SAGEMAKER: + if DeploymentResult._sagemaker_status_checker is None: + try: + import amzn_nova_forge.util.sagemaker # noqa: F401 + except ImportError: + pass + if DeploymentResult._sagemaker_status_checker is None: + raise RuntimeError( + "SageMaker status checker not available. " + "Ensure amzn_nova_forge.util.sagemaker is imported." + ) + return DeploymentResult._sagemaker_status_checker( + self.endpoint.uri, region=self.endpoint.region + ) + if DeploymentResult._status_checker is None: # Runtime fallback only — core/types.py does NOT import util.bedrock # at module load time. This triggers registration if the caller @@ -136,6 +161,76 @@ def validate_region(region: str) -> None: ) +@dataclass +class InferenceComponentConfig: + """Configuration for creating an inference component on a SageMaker endpoint. + + When passed to create_sagemaker_endpoint, the endpoint is created without a + ModelName in ProductionVariants (IC-compatible mode) and an inference component + is created after the endpoint reaches InService. The IC references the SageMaker + model (created during deploy) via ModelName. + + Args: + inference_component_name: Unique name for the inference component. + num_cpus: Number of vCPUs to allocate. + num_accelerators: Number of accelerators (GPUs) to allocate. + min_memory_in_mb: Minimum memory in MB to allocate. + copy_count: Number of model copies to deploy. Default: 1. + variant_name: Production variant name. Default: "primary". + """ + + inference_component_name: str + num_cpus: int + num_accelerators: int + min_memory_in_mb: int + copy_count: int = 1 + variant_name: str = "primary" + + +# Minimum compute resource requirements for inference components per model. +# Format: {Model: (min_cpus, min_memory_mb, min_gpus)} +_IC_MIN_COMPUTE_REQUIREMENTS: Dict[Model, tuple] = { + Model.NOVA_MICRO: (15, 25000, 4), + Model.NOVA_LITE: (20, 35000, 4), + Model.NOVA_LITE_2: (20, 100000, 4), +} + + +def validate_inference_component_resources(config: InferenceComponentConfig, model: Model) -> None: + """Validate that inference component compute resources meet minimum requirements. + + Args: + config: The inference component configuration to validate. + model: The Nova model being deployed. + + Raises: + ValueError: If any resource is below the minimum for the given model. + """ + requirements = _IC_MIN_COMPUTE_REQUIREMENTS.get(model) + if requirements is None: + return # No known requirements for this model, skip validation + + min_cpus, min_memory_mb, min_gpus = requirements + errors = [] + + if config.num_cpus < min_cpus: + errors.append(f"num_cpus={config.num_cpus} is below minimum {min_cpus} for {model.value}") + if config.min_memory_in_mb < min_memory_mb: + errors.append( + f"min_memory_in_mb={config.min_memory_in_mb} is below minimum {min_memory_mb} for {model.value}" + ) + if config.num_accelerators < min_gpus: + errors.append( + f"num_accelerators={config.num_accelerators} is below minimum {min_gpus} for {model.value}" + ) + + if errors: + raise ValueError( + f"Inference component resources do not meet minimum requirements for {model.value}: " + + "; ".join(errors) + ) + + @dataclass class JobConfig: job_name: str @@ -157,6 +252,7 @@ class JobConfig: method: Optional[TrainingMethod] = None # Training method (required for Bedrock) data_mixing_config: Optional[Dict[str, Any]] = None # Datamix percent fields (SMTJServerless) environment: Optional[Dict[str, str]] = None # Environment variables for the training container + model_name_or_path: Optional[str] = None # Model path or model package ARN # 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/core/validation_patterns.py b/src/amzn_nova_forge/core/validation_patterns.py index 21cd5bc..1567be6 100644 --- a/src/amzn_nova_forge/core/validation_patterns.py +++ b/src/amzn_nova_forge/core/validation_patterns.py @@ -28,6 +28,10 @@ # https://docs.aws.amazon.com/eks/latest/APIReference/API_CreateCluster.html#API_CreateCluster_RequestParameters CLUSTER_NAME_REGEX = re.compile(r"^[0-9A-Za-z][A-Za-z0-9\-_]{1,100}$") +MODEL_PACKAGE_ARN_REGEX = re.compile( + r"^arn:aws[\w-]*:sagemaker:[a-z0-9-]+:\d{12}:model-package/.+$" +) + def validate_job_name(job_name: str) -> None: if not JOB_NAME_REGEX.match(job_name): diff --git a/src/amzn_nova_forge/dataset/dataset_transformers.py b/src/amzn_nova_forge/dataset/dataset_transformers.py index 754256d..a2a72de 100644 --- a/src/amzn_nova_forge/dataset/dataset_transformers.py +++ b/src/amzn_nova_forge/dataset/dataset_transformers.py @@ -901,3 +901,59 @@ def convert_to_rft_multiturn(rec, column_mappings, transform_ctx=None, s3_client metadata["info"] = rec[info_col] return {"id": generated_id, "metadata": metadata} + + @staticmethod + def _extract_prompt(rec, column_mappings) -> str: + """Extract prompt value from a record, checking multiple input formats. + + Resolution order: + 1. Top-level "prompt" key + 2. Nested "metadata.prompt" + 3. Column mapping for "prompt" + + Returns the prompt as a string (serializes dicts/lists to JSON). + """ + if "prompt" in rec: + prompt_value = rec["prompt"] + elif ( + "metadata" in rec and isinstance(rec["metadata"], dict) and "prompt" in rec["metadata"] + ): + prompt_value = rec["metadata"]["prompt"] + else: + prompt_col = column_mappings.get("prompt") + if not prompt_col: + raise ValueError( + "'prompt' column mapping is required for RFT Multiturn Serverless.\n" + "Make sure to add prompt='your_column_name' when initializing DatasetLoader." + ) + if prompt_col not in rec: + raise ValueError( + f"'prompt' column '{prompt_col}' not found in record.\n" + f"Make sure the column exists in your data." + ) + prompt_value = rec[prompt_col] + + if isinstance(prompt_value, (dict, list)): + prompt_value = json.dumps(prompt_value) + + return prompt_value + + @staticmethod + def convert_to_rft_multiturn_serverless( + rec, column_mappings, transform_ctx=None, s3_client=None + ): + """Convert to flat RFT multiturn format for Serverless MTRL. + + The Serverless MTRL service reads only the `prompt` column and passes + it as-is to the agent/rollout server. This transformer outputs the flat + format: {"prompt": ""}. + + Args: + rec: A single dataset record (dict). Can have a "prompt" field directly, + or use column_mappings to locate the prompt column. + column_mappings: Dictionary mapping field names to column names. + + Returns: + dict: Record with flat structure: {"prompt": str} + """ + return {"prompt": DatasetTransformer._extract_prompt(rec, column_mappings)} diff --git a/src/amzn_nova_forge/dataset/dataset_validator/rft_multiturn_dataset_validator.py b/src/amzn_nova_forge/dataset/dataset_validator/rft_multiturn_dataset_validator.py index 959cd73..0e57321 100644 --- a/src/amzn_nova_forge/dataset/dataset_validator/rft_multiturn_dataset_validator.py +++ b/src/amzn_nova_forge/dataset/dataset_validator/rft_multiturn_dataset_validator.py @@ -220,9 +220,32 @@ def validate_id_not_empty(cls, v: str) -> str: return v.strip() +class RFTMultiturnServerlessSample(BaseModel): + """ + Represents an RFT Multiturn dataset sample for Serverless MTRL. + + Expected format: + { + "prompt": str # REQUIRED - the prompt string passed as-is to the agent + } + """ + + model_config = ConfigDict(extra="allow") + + prompt: str + + @field_validator("prompt") + @classmethod + def validate_prompt_not_empty(cls, v: str) -> str: + """Ensure prompt is not empty.""" + if not v or not v.strip(): + raise ValueError("'prompt' cannot be empty") + return v + + class RFTMultiturnDatasetValidator(BaseDatasetValidator): """ - Validator for RFT Multiturn datasets in SDK format. + Validator for RFT Multiturn datasets in SDK format (HyperPod). RFT Multiturn is only supported on Nova 2.0 and requires: - Unique sample IDs @@ -268,3 +291,29 @@ def get_optional_fields(self) -> List[str]: OPTIONAL_FIELDS: A list of all the optional fields for RFT Multiturn. """ return OPTIONAL_FIELDS + + +class RFTMultiturnServerlessValidator(BaseDatasetValidator): + """ + Validator for RFT Multiturn datasets in Serverless format. + + Validates flat format: {"prompt": str} + The service reads only the prompt column and passes it as-is to the agent. + """ + + def __init__(self, model: Model): + super().__init__() + if model != Model.NOVA_LITE_2: + raise ValueError( + f"RFT Multiturn is only supported on Nova 2.0 Lite (NOVA_LITE_2). " + f"Current model: {model}." + ) + + def get_sample_model(self): + return RFTMultiturnServerlessSample + + def get_success_message(self) -> str: + return f"Validation succeeded for {self.num_samples} samples on a Serverless RFT Multiturn dataset." + + def get_optional_fields(self) -> List[str]: + return [] diff --git a/src/amzn_nova_forge/dataset/operations/transform_operation.py b/src/amzn_nova_forge/dataset/operations/transform_operation.py index 6d76adc..ef94040 100644 --- a/src/amzn_nova_forge/dataset/operations/transform_operation.py +++ b/src/amzn_nova_forge/dataset/operations/transform_operation.py @@ -20,7 +20,7 @@ import boto3 import jsonschema -from ...core.enums import EvaluationTask, Model, TrainingMethod +from ...core.enums import EvaluationTask, Model, Platform, TrainingMethod from ...util.iterator_utils import peek from ...util.logging import logger from ..data_state import DataLocation, DataState @@ -41,6 +41,7 @@ class SchemaTransformOperation(NovaForgeTransformOperation): "convert_to_evaluation": DatasetTransformer.convert_to_evaluation, "convert_to_cpt": DatasetTransformer.convert_to_cpt, "convert_to_rft_multiturn": DatasetTransformer.convert_to_rft_multiturn, + "convert_to_rft_multiturn_serverless": DatasetTransformer.convert_to_rft_multiturn_serverless, } def execute(self, loader: Any, **kwargs) -> OperationResult: @@ -52,6 +53,7 @@ def execute(self, loader: Any, **kwargs) -> OperationResult: model: Optional[Model] = kwargs.get("model") eval_task: Optional[EvaluationTask] = kwargs.get("eval_task") region: Optional[str] = kwargs.get("region") + platform = kwargs.get("platform") if training_method is None or model is None: raise ValueError("training_method and model are required for schema transforms.") @@ -66,8 +68,34 @@ def execute(self, loader: Any, **kwargs) -> OperationResult: multimodal_data_s3_path = kwargs.get("multimodal_data_s3_path") multimodal_data_bucket_owner = kwargs.get("multimodal_data_bucket_owner") + # RFT Multiturn requires platform to determine output format + is_mtrl = training_method in ( + TrainingMethod.RFT_MULTITURN_LORA, + TrainingMethod.RFT_MULTITURN_FULL, + ) + if is_mtrl and platform is None: + raise ValueError( + "platform is required for RFT Multiturn transforms. " + "Pass platform=Platform.SMTJServerless or platform=Platform.SMHP. " + "Example: loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, " + "model=Model.NOVA_LITE_2, platform=Platform.SMTJServerless)" + ) + + if is_mtrl and platform == Platform.SMTJServerless: + transform_config = dict(transform_config) + transform_config["schema"] = None + transform_config["transformers"] = [ + { + "source_schema": None, + "method": "convert_to_rft_multiturn_serverless", + "msg": "Transforming to flat RFT Multiturn format for Serverless.", + }, + ] + # Already in target format — nothing to do - if self._validate_against_schema(loader.dataset, transform_config["schema"]): + if transform_config["schema"] is not None and self._validate_against_schema( + loader.dataset, transform_config["schema"] + ): logger.info("Transform: %s", transform_config["success_msg"]) return OperationResult(status="SUCCEEDED", output_state=state) diff --git a/src/amzn_nova_forge/dataset/operations/validate_operation.py b/src/amzn_nova_forge/dataset/operations/validate_operation.py index f57fc9d..31e2dba 100644 --- a/src/amzn_nova_forge/dataset/operations/validate_operation.py +++ b/src/amzn_nova_forge/dataset/operations/validate_operation.py @@ -51,7 +51,20 @@ def execute(self, loader: Any, **kwargs) -> None: "For RFT Multiturn evaluation, pass eval_task=EvaluationTask.RFT_MULTITURN_EVAL" ) - validator = resolve_schema_validator(training_method, model, eval_task) + # For Serverless MTRL, use the serverless validator + is_mtrl = training_method in ( + TrainingMethod.RFT_MULTITURN_LORA, + TrainingMethod.RFT_MULTITURN_FULL, + ) + if is_mtrl and platform == Platform.SMTJServerless: + from amzn_nova_forge.dataset.dataset_validator.rft_multiturn_dataset_validator import ( + RFTMultiturnServerlessValidator, + ) + + validator: Optional[Any] = RFTMultiturnServerlessValidator(model) + else: + validator = resolve_schema_validator(training_method, model, eval_task) + if validator is None: logger.info( "Validate: skipped — not available for model=%s / method=%s", diff --git a/src/amzn_nova_forge/deployer/forge_deployer.py b/src/amzn_nova_forge/deployer/forge_deployer.py index aacb7cd..92beec0 100644 --- a/src/amzn_nova_forge/deployer/forge_deployer.py +++ b/src/amzn_nova_forge/deployer/forge_deployer.py @@ -35,8 +35,11 @@ DeploymentResult, EndpointInfo, ForgeConfig, + InferenceComponentConfig, + validate_inference_component_resources, validate_region, ) +from amzn_nova_forge.core.validation_patterns import MODEL_PACKAGE_ARN_REGEX from amzn_nova_forge.iam.iam_role_creator import ( create_bedrock_execution_role, create_sagemaker_execution_role, @@ -57,10 +60,14 @@ from amzn_nova_forge.util.logging import logger from amzn_nova_forge.util.sagemaker import ( SAGEMAKER_EXECUTION_ROLE_NAME, + _get_sagemaker_inference_image, _validate_sagemaker_instance_type_for_model_deployment, + check_sagemaker_deployment_status, + create_inference_component, create_sagemaker_endpoint, create_sagemaker_model, find_sagemaker_model_by_tag, + monitor_inference_component, ) from amzn_nova_forge.validation.endpoint_validator import ( SAGEMAKER_ENDPOINT_ARN_REGEX, @@ -116,12 +123,15 @@ def deploy( sagemaker_instance_type: Optional[str] = "ml.p5.48xlarge", sagemaker_environment: Optional[SageMakerEndpointEnvironment] = None, skip_model_reuse: bool = False, + inference_component_configs: List[InferenceComponentConfig] = [], # noqa: B006 - never mutated ) -> DeploymentResult: """Deploy a model to Bedrock or SageMaker. Args: model_artifact_path: S3 path to the trained model checkpoint, - or an existing Bedrock custom model ARN. + an existing Bedrock custom model ARN, or a SageMaker Model + Package name/ARN. Model package ARNs are auto-detected for + SageMaker deployments and passed via ModelPackageName. deploy_platform: Target platform. endpoint_name: Custom endpoint name (auto-generated if omitted). unit_count: Instance/PT unit count. @@ -129,6 +139,9 @@ def deploy( sagemaker_instance_type: EC2 instance type for SageMaker. sagemaker_environment: Optional SageMaker endpoint environment config. skip_model_reuse: If True, always create a new model (skip tag-based discovery). + inference_component_configs: List of configs for creating inference components. + When provided with SAGEMAKER platform, creates an IC-compatible endpoint + and deploys the inference component(s) in one step. Returns: DeploymentResult with endpoint information. @@ -158,18 +171,24 @@ def deploy( sagemaker_instance_type, self.model, context_length, max_concurrency ) - artifact_path = ( - model_artifact_path - if model_artifact_path.endswith("/") - else model_artifact_path + "/" - ) - - if artifact_path.startswith("arn:aws:bedrock:"): + if model_artifact_path.startswith("arn:aws:bedrock:"): raise ValueError( "Cannot deploy Bedrock-customized models to SageMaker. " "Train on SageMaker first." ) + # Auto-detect model package ARN + model_package_name: Optional[str] = None + if MODEL_PACKAGE_ARN_REGEX.match(model_artifact_path): + model_package_name = model_artifact_path + artifact_path = model_artifact_path + else: + artifact_path = ( + model_artifact_path + if model_artifact_path.endswith("/") + else model_artifact_path + "/" + ) + return self._deploy_to_sagemaker( model_artifact_path=artifact_path, endpoint_name=endpoint_name, @@ -177,6 +196,9 @@ def deploy( unit_count=unit_count, sagemaker_environment=sagemaker_environment, execution_role_name=execution_role_name, + skip_model_reuse=skip_model_reuse, + inference_component_configs=inference_component_configs, + model_package_name=model_package_name, ) else: raise ValueError(f"Unsupported deployment platform: {deploy_platform}") @@ -202,7 +224,10 @@ 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, region=self.region) + if platform == DeployPlatform.SAGEMAKER: + status_str = check_sagemaker_deployment_status(endpoint_arn, region=self.region) + else: + status_str = check_deployment_status(endpoint_arn, platform, region=self.region) if status_str is None: return None try: @@ -244,9 +269,95 @@ def get_logs( else: platform = DeployPlatform.BEDROCK_OD - status = check_deployment_status(arn, platform, region=self.region) + if platform == DeployPlatform.SAGEMAKER: + status = check_sagemaker_deployment_status(arn, region=self.region) + else: + status = check_deployment_status(arn, platform, region=self.region) logger.info(f"Deployment status for {arn}: {status}") + @_telemetry_emitter( + Feature.DEPLOY, + "create_inference_component", + extra_info_fn=lambda self, *args, **kwargs: { + "model": self.model.value, + }, + ) + def create_inference_component( + self, + inference_component_name: str, + model_name: str, + num_cpus: int, + num_accelerators: int, + min_memory_in_mb: int, + endpoint_name: str, + variant_name: str = "primary", + copy_count: int = 1, + ) -> DeploymentResult: + """Create an inference component on an existing SageMaker endpoint. + + Validates the target endpoint is InService, then calls CreateInferenceComponent + using the specified SageMaker model and returns a DeploymentResult immediately + without waiting for the component to become active. + + Args: + inference_component_name: Unique name for the inference component. + model_name: Name of the existing SageMaker model to use. + num_cpus: Number of vCPUs to allocate. + num_accelerators: Number of accelerators (GPUs) to allocate. + min_memory_in_mb: Minimum memory in MB to allocate. + endpoint_name: Name of the existing SageMaker endpoint (must be InService). + variant_name: Production variant name on the endpoint. Default: "primary". + copy_count: Number of model copies to deploy. Default: 1. + + Returns: + DeploymentResult with endpoint info containing the inference component ARN. + + Raises: + Exception: If the endpoint does not exist, is not InService, + or the CreateInferenceComponent API call fails. + """ + sagemaker_client = boto3.client("sagemaker", region_name=self.region) + + return create_inference_component( + inference_component_name=inference_component_name, + endpoint_name=endpoint_name, + variant_name=variant_name, + model_name=model_name, + num_cpus=num_cpus, + num_accelerators=num_accelerators, + min_memory_in_mb=min_memory_in_mb, + copy_count=copy_count, + sagemaker_client=sagemaker_client, + region=self.region, + ) + + @_telemetry_emitter( + Feature.DEPLOY, + "monitor_inference_component", + extra_info_fn=lambda self, *args, **kwargs: { + "model": self.model.value, + }, + ) + def monitor_inference_component(self, inference_component_name: str) -> str: + """Monitor an inference component until it reaches a terminal state. + + Polls DescribeInferenceComponent every 30 seconds until InService or Failed. + + Args: + inference_component_name: Name of the inference component to monitor. + + Returns: + str: Final status ("InService"). + + Raises: + Exception: If the component reaches Failed status or the API call errors. + """ + sagemaker_client = boto3.client("sagemaker", region_name=self.region) + return monitor_inference_component( + inference_component_name=inference_component_name, + sagemaker_client=sagemaker_client, + ) + @_telemetry_emitter( Feature.DEPLOY, "find_published_model", @@ -298,34 +409,76 @@ def find_published_model( def create_custom_model( self, - model_artifact_path: str, + model_artifact_path: Optional[str] = None, endpoint_name: Optional[str] = None, execution_role_name: Optional[str] = None, tags: Optional[List[Dict[str, str]]] = None, skip_model_reuse: bool = False, + custom_model_data_source: Optional[Dict[str, Any]] = None, ) -> ModelDeployResult: - """Create a Bedrock custom model from S3 artifacts. + """Create a Bedrock custom model from S3 artifacts or a model package ARN. - Extracts the model-creation step from the deploy flow so users can - create a model independently of endpoint deployment. + Either ``model_artifact_path`` (maps to ``modelSourceConfig``) or + ``custom_model_data_source`` (maps to ``customModelDataSource``) must be + provided, but not both. Args: - model_artifact_path: S3 path to trained model checkpoint. + model_artifact_path: S3 path to trained model checkpoint. Used to + populate ``modelSourceConfig.s3DataSource.s3Uri``. endpoint_name: Optional name prefix for the model name. execution_role_name: Optional IAM role name. If None, the SDK creates and manages the default role. tags: Optional list of {"key": str, "value": str} dicts for source tracking. skip_model_reuse: If True, always create a new model (skip tag-based discovery). + custom_model_data_source: Alternative data source configuration for the + custom model. For example:: + + {"modelPackageArnDataSource": {"modelPackageArn": "MODEL_PACKAGE_ARN"}} + + When provided, ``modelSourceConfig`` is omitted from the API call. Returns: ModelDeployResult with model_arn, model_name, escrow_uri, etc. + + Raises: + ValueError: If neither or both of ``model_artifact_path`` and + ``custom_model_data_source`` are provided. """ + if model_artifact_path and custom_model_data_source: + raise ValueError( + "Only one of 'model_artifact_path' (modelSourceConfig) or " + "'custom_model_data_source' (customModelDataSource) may be provided, not both." + ) + if not model_artifact_path and not custom_model_data_source: + raise ValueError( + "Either 'model_artifact_path' (modelSourceConfig) or " + "'custom_model_data_source' (customModelDataSource) must be provided." + ) + # Determine the escrow path for model reuse tracking. + # For custom_model_data_source dicts, extract the ARN string to avoid + # dict repr characters ({, }, ') that violate Bedrock tag constraints. + if model_artifact_path: + escrow_path = model_artifact_path + else: + assert custom_model_data_source is not None + mp_ds = custom_model_data_source.get("modelPackageArnDataSource", {}) + escrow_path = mp_ds.get("modelPackageArn") + if not escrow_path: + logger.warning( + "Could not extract a tag-safe identifier from 'custom_model_data_source'. " + "Model reuse tracking via escrow tag will be skipped." + ) + # Check for existing model by escrow tag - existing = self.find_published_model("bedrock", model_artifact_path, skip_model_reuse) + existing = ( + self.find_published_model("bedrock", escrow_path, skip_model_reuse) + if escrow_path + else None + ) if existing: logger.info( f"Found existing Bedrock model {existing} for escrow URI " - f"'{model_artifact_path}'. Reusing instead of creating a duplicate." + f"'{escrow_path}'. Reusing instead of creating a duplicate." ) bedrock_client = boto3.client("bedrock", region_name=self.region) result = ModelDeployResult.from_arn(existing, bedrock_client) @@ -344,7 +497,10 @@ def create_custom_model( else: logger.warning("Model %s has unknown status. Creating new model.", existing) - bedrock_client = boto3.client("bedrock", region_name=self.region) + bedrock_client = boto3.client( + "bedrock", + region_name=self.region, + ) iam_client = boto3.client("iam", region_name=self.region) # Resolve execution role @@ -375,10 +531,14 @@ def create_custom_model( create_kwargs: Dict[str, Any] = { "modelName": model_name, - "modelSourceConfig": {"s3DataSource": {"s3Uri": model_artifact_path}}, "roleArn": bedrock_execution_role_arn, } + if model_artifact_path: + create_kwargs["modelSourceConfig"] = {"s3DataSource": {"s3Uri": model_artifact_path}} + else: + create_kwargs["customModelDataSource"] = custom_model_data_source + if self._config.kms_key_id: if self._config.kms_key_id.startswith("arn:aws:kms:"): kms_arn = self._config.kms_key_id @@ -392,12 +552,13 @@ def create_custom_model( create_kwargs["modelTags"] = tags # Inject escrow URI tag (Bedrock format: lowercase key/value) - escrow_tag = { - "key": ESCROW_URI_TAG_KEY, - "value": _escrow_tag_value(model_artifact_path), - } - all_tags = list(create_kwargs.get("modelTags", [])) + [escrow_tag] - create_kwargs["modelTags"] = all_tags + if escrow_path: + escrow_tag = { + "key": ESCROW_URI_TAG_KEY, + "value": _escrow_tag_value(escrow_path), + } + all_tags = list(create_kwargs.get("modelTags", [])) + [escrow_tag] + create_kwargs["modelTags"] = all_tags try: logger.info(f"Creating custom model '{model_name}'...") @@ -416,12 +577,13 @@ def create_custom_model( result = ModelDeployResult( model_arn=model["modelArn"], model_name=model_name, - escrow_uri=model_artifact_path, + escrow_uri=escrow_path or "", created_at=datetime.now(timezone.utc), ) self.last_model_publish = result - self._published_models.add(("bedrock", model["modelArn"], model_artifact_path)) + if escrow_path: + self._published_models.add(("bedrock", model["modelArn"], escrow_path)) logger.info(f"Custom model created: {model['modelArn']}") return result @@ -690,6 +852,8 @@ def _deploy_to_sagemaker( sagemaker_environment: Optional[SageMakerEndpointEnvironment] = None, execution_role_name: Optional[str] = None, skip_model_reuse: bool = False, + inference_component_configs: List[InferenceComponentConfig] = [], # noqa: B006 - never mutated + model_package_name: Optional[str] = None, ) -> DeploymentResult: env = sagemaker_environment or SageMakerEndpointEnvironment() env.validate_smi_config_bounds(model=self.model, instance_type=instance_type) @@ -743,12 +907,13 @@ def _deploy_to_sagemaker( model_arn = create_sagemaker_model( region=self.region, model_name=model_name, - model_s3_location=model_artifact_path, + model_s3_location=model_artifact_path if not model_package_name else None, sagemaker_execution_role_arn=sagemaker_role_arn, sagemaker_client=sagemaker_client, environment=env_vars, deployment_mode=self.deployment_mode, tags=[escrow_tag], + model_package_name=model_package_name, ) self._published_models.add(("sagemaker", model_arn, model_artifact_path)) @@ -760,16 +925,36 @@ def _deploy_to_sagemaker( ) self.last_model_publish = model_deploy + # Resolve inference component config defaults from deploy context + if inference_component_configs: + if len(inference_component_configs) > 1: + logger.warning( + f"Deploying {len(inference_component_configs)} inference components to endpoint " + f"'{endpoint_name}'. Total resource validation is not performed — ensure the " + f"combined resource requirements do not exceed the instance capacity." + ) + for ic_config in inference_component_configs: + validate_inference_component_resources(ic_config, self.model) + try: - endpoint_arn = create_sagemaker_endpoint( - model_name=model_name, - endpoint_config_name=endpoint_config_name, - endpoint_name=endpoint_name, - instance_type=instance_type, - sagemaker_client=sagemaker_client, - initial_instance_count=unit_count, - deployment_mode=self.deployment_mode, - ) + endpoint_kwargs: Dict[str, Any] = { + "model_name": model_name, + "endpoint_config_name": endpoint_config_name, + "endpoint_name": endpoint_name, + "instance_type": instance_type, + "sagemaker_client": sagemaker_client, + "initial_instance_count": unit_count, + "deployment_mode": self.deployment_mode, + **( + { + "inference_component_configs": inference_component_configs, + "execution_role_arn": sagemaker_role_arn, + } + if inference_component_configs + else {} + ), + } + endpoint_arn = create_sagemaker_endpoint(**endpoint_kwargs) except Exception as e: raise RuntimeError( f"SageMaker endpoint creation failed: {e}\n\n" diff --git a/src/amzn_nova_forge/evaluator/forge_evaluator.py b/src/amzn_nova_forge/evaluator/forge_evaluator.py index aa7a4e3..0fa9aaf 100644 --- a/src/amzn_nova_forge/evaluator/forge_evaluator.py +++ b/src/amzn_nova_forge/evaluator/forge_evaluator.py @@ -37,6 +37,7 @@ EvaluationResult, SMHPEvaluationResult, SMTJEvaluationResult, + SMTJTrainingResult, TrainingResult, ) from amzn_nova_forge.core.runtime import RuntimeManager @@ -46,7 +47,7 @@ validate_region, ) from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig -from amzn_nova_forge.manager.runtime_manager import SMHPRuntimeManager +from amzn_nova_forge.manager.runtime_manager import SMHPRuntimeManager, SMTJServerlessRuntimeManager from amzn_nova_forge.model.nova_model_customizer_util import ( requires_custom_eval_data, resolve_model_checkpoint_path, @@ -70,12 +71,24 @@ @dataclass class EvalTaskConfig: - """Per-task configuration for an evaluation job.""" + """Per-task configuration for an evaluation job. + + Attributes: + subtask: Subtask identifier for benchmark evaluations. + processor: Lambda processor config for BYOM evaluations. + rl_env: RL environment config for HyperPod evaluations. + override_data_s3_path: Override the default evaluation dataset path. + evaluate_base_model: When True and used with RFT_MULTITURN_EVAL, + evaluates both the base model and the fine-tuned model in a + single pipeline for side-by-side comparison. Only applies to + MTRL evaluation on SMTJServerless. Requires model_path to be set. + """ subtask: Optional[str] = None processor: Optional[Dict[str, Any]] = None rl_env: Optional[Dict[str, Any]] = None override_data_s3_path: Optional[str] = None + evaluate_base_model: bool = False class ForgeEvaluator: @@ -222,6 +235,24 @@ def evaluate( "Use SageMaker platforms (SMTJ, SMHP) instead." ) + # MTRL evaluation on SMTJServerless delegates to the MultiTurnRLEvaluator + if ( + eval_task == EvaluationTask.RFT_MULTITURN_EVAL + and self._platform == Platform.SMTJServerless + ): + evaluate_base_model = False + if task_config and hasattr(task_config, "evaluate_base_model"): + evaluate_base_model = task_config.evaluate_base_model + return self._execute_mtrl_eval( + job_name=job_name, + model_path=model_path, + job_result=job_result, + task_config=task_config, + overrides=overrides, + dry_run=dry_run, + evaluate_base_model=evaluate_base_model, + ) + # Check job cache cached = load_existing_result( self._cache_context, @@ -354,6 +385,7 @@ def evaluate( recipe_path=resolved_recipe_path, input_s3_data_type="S3Prefix", method=TrainingMethod.EVALUATION, + model_name_or_path=resolved_model_path, ) ) @@ -774,6 +806,88 @@ def _upload_config_to_s3(self, config_yaml: str, s3_uri: str) -> None: ) logger.info(f"Uploaded InspectLens config to s3://{bucket}/{key}") + def _execute_mtrl_eval( + self, + job_name: str, + model_path: Optional[str] = None, + job_result: Optional[TrainingResult] = None, + task_config: Optional[EvalTaskConfig] = None, + overrides: Optional[Dict[str, Any]] = None, + dry_run: bool = False, + evaluate_base_model: bool = False, + ) -> Optional[EvaluationResult]: + """Delegate MTRL evaluation to the MultiTurnRLEvaluator. + + This is called when eval_task is RFT_MULTITURN_EVAL on the + SMTJServerless platform. It uses the sagemaker.train.evaluate + MultiTurnRLEvaluator which creates a SageMaker Pipeline with + the AgentRFTEvaluation job type. + """ + infra = cast(SMTJServerlessRuntimeManager, self.infra) + + # Resolve model path from job_result if not explicitly provided + resolved_model_path = model_path + if resolved_model_path is None and job_result is not None: + # For MTRL, use the model package ARN directly + + if isinstance(job_result, SMTJTrainingResult) and job_result._is_mtrl: + resolved_model_path = job_result.model_artifacts.output_model_arn + else: + resolved_model_path = resolve_model_checkpoint_path( + model_path=None, + job_result=job_result, + customizer_job_id=None, + customizer_output_s3_path=self.output_s3_path, + customizer_model_path=None, + ) + + # Resolve MLflow tracking URI (required for AgentRFT evaluation jobs) + if not self._config.mlflow_monitor or not self._config.mlflow_monitor.tracking_uri: + raise ValueError( + "MLflow configuration is required for AgentRFT evaluation jobs. " + "Please provide an mlflow_monitor with a valid tracking_uri when " + "using RFT_MULTITURN_EVAL on the SMTJServerless platform." + ) + mlflow_uri = self._config.mlflow_monitor.tracking_uri + + if dry_run: + logger.info( + f"[dry_run] Would launch MTRL evaluation '{job_name}' " + f"with model_path={resolved_model_path}" + ) + return None + + # Pass training job name so the evaluator can attach and resolve model artifacts + training_job_name = job_result.job_id if job_result is not None else None + + execution = infra.execute_mtrl_eval( + model=self.model, + data_s3_path=self.data_s3_path, + output_s3_path=self.output_s3_path, + mlflow_tracking_uri=mlflow_uri, + model_path=resolved_model_path, + overrides=overrides, + training_job_name=training_job_name, + evaluate_base_model=evaluate_base_model, + ) + + start_time = datetime.now(timezone.utc) + eval_output_s3_path = f"{self.output_s3_path.rstrip('/')}/{job_name}/" + + evaluation_result = SMTJEvaluationResult( + job_id=execution.arn, + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + started_time=start_time, + eval_output_path=eval_output_s3_path, + region=self.region, + ) + evaluation_result._job_name = job_name # type: ignore[attr-defined] + # Attach the execution object for downstream access (wait, show_results, etc.) + evaluation_result._mtrl_execution = execution # type: ignore[attr-defined] + + logger.info(f"Started MTRL evaluation pipeline. Execution ARN: {execution.arn}") + return evaluation_result + @_telemetry_emitter(Feature.EVAL, "get_logs") def get_logs( self, @@ -798,6 +912,23 @@ def get_logs( "No job reference provided. Pass either a job_result or explicit job_id and started_time." ) + is_mtrl_eval = ( + job_result + and hasattr(job_result, "eval_task") + and job_result.eval_task == EvaluationTask.RFT_MULTITURN_EVAL + ) + + if is_mtrl_eval: + from amzn_nova_forge.monitor import MTRLLogMonitor + + monitor = MTRLLogMonitor.from_job_id( + job_id=resolved_job_id, + region=self.region, + job_category="AgentRFTEvaluation", + ) + monitor.show_logs(limit=limit) + return + kwargs: Dict[str, Any] = {} if self._platform == Platform.SMHP: kwargs["cluster_name"] = cast(SMHPRuntimeManager, self.infra).cluster_name diff --git a/src/amzn_nova_forge/iam/bedrock_policies.json b/src/amzn_nova_forge/iam/bedrock_policies.json index 440a814..708aa1d 100644 --- a/src/amzn_nova_forge/iam/bedrock_policies.json +++ b/src/amzn_nova_forge/iam/bedrock_policies.json @@ -49,5 +49,17 @@ "Resource": "LAMBDA_ARN_PLACEHOLDER" } ] + }, + "sagemaker_model_package_policy": { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "sagemaker:AccessModelPackage" + ], + "Resource": "*" + } + ] } } \ No newline at end of file diff --git a/src/amzn_nova_forge/iam/iam_role_creator.py b/src/amzn_nova_forge/iam/iam_role_creator.py index bb791d3..3aef8d2 100644 --- a/src/amzn_nova_forge/iam/iam_role_creator.py +++ b/src/amzn_nova_forge/iam/iam_role_creator.py @@ -163,7 +163,7 @@ def create_bedrock_execution_role( policies["s3_read_policy"]["Statement"][0]["Resource"] = "*" # Create and attach policies - for policy_name in ["bedrock_policy", "s3_read_policy"]: + for policy_name in ["bedrock_policy", "s3_read_policy", "sagemaker_model_package_policy"]: try: policy_arn = iam_client.create_policy( PolicyName=f"{role_name}{policy_name.title()}", @@ -293,6 +293,7 @@ def create_sagemaker_execution_role( "s3_read_policy", "kms_policy", "ec2_policy", + "sagemaker_model_package_policy", ]: try: logger.info(f"{json.dumps(policies[policy_name])}.") diff --git a/src/amzn_nova_forge/iam/sagemaker_policies.json b/src/amzn_nova_forge/iam/sagemaker_policies.json index b6d4bb3..213a940 100644 --- a/src/amzn_nova_forge/iam/sagemaker_policies.json +++ b/src/amzn_nova_forge/iam/sagemaker_policies.json @@ -120,5 +120,17 @@ "Resource": "*" } ] + }, + "sagemaker_model_package_policy": { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "sagemaker:AccessModelPackage" + ], + "Resource": "*" + } + ] } } \ No newline at end of file diff --git a/src/amzn_nova_forge/inference/forge_inference.py b/src/amzn_nova_forge/inference/forge_inference.py index 7838bda..bbcf402 100644 --- a/src/amzn_nova_forge/inference/forge_inference.py +++ b/src/amzn_nova_forge/inference/forge_inference.py @@ -101,12 +101,15 @@ def invoke( self, endpoint_arn: str, request_body: Dict[str, Any], + inference_component_name: Optional[str] = None, ) -> Any: """Invoke real-time inference against an endpoint. Args: endpoint_arn: ARN of the deployed endpoint. request_body: Inference request body. + inference_component_name: Optional inference component to target on a + SageMaker endpoint. Required for IC-enabled endpoints. Returns: Inference result. @@ -117,7 +120,12 @@ def invoke( runtime_client = boto3.client("sagemaker-runtime", region_name=self.region) endpoint_name = endpoint_arn.split("/")[-1] logger.info(f"Invoking SageMaker endpoint: {endpoint_name}") - return invoke_sagemaker_inference(request_body, endpoint_name, runtime_client) + return invoke_sagemaker_inference( + request_body, + endpoint_name, + runtime_client, + inference_component_name=inference_component_name, + ) else: runtime_client = boto3.client("bedrock-runtime", region_name=self.region) return invoke_model( diff --git a/src/amzn_nova_forge/manager/mtrl_manager.py b/src/amzn_nova_forge/manager/mtrl_manager.py new file mode 100644 index 0000000..d0dad24 --- /dev/null +++ b/src/amzn_nova_forge/manager/mtrl_manager.py @@ -0,0 +1,260 @@ +# 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. +from typing import Any, Dict, List, Optional + +import boto3 + +from amzn_nova_forge.core.enums import Model +from amzn_nova_forge.telemetry.constants import Feature +from amzn_nova_forge.telemetry.telemetry_logging import _telemetry_emitter +from amzn_nova_forge.util.logging import logger + + +class MTRLOperations: + """Multi-turn RL training and evaluation operations for serverless runtime. + + Mixin class — expects to be composed with SMTJServerlessRuntimeManager + which provides the attributes declared below. + """ + + model_package_group_arn: str + execution_role: Optional[str] + subnets: Optional[List[str]] + security_group_ids: Optional[List[str]] + kms_key_id: Optional[str] + agent_core_arn: Optional[str] + rft_lambda: Optional[str] + region: str + + def _get_or_create_checkpoint_model_package_group_arn(self) -> str: + raise NotImplementedError + + def _mtrl_common_kwargs( + self, + output_s3_path: Optional[str] = None, + mlflow_tracking_uri: Optional[str] = None, + ) -> Dict[str, Any]: + """Build common kwargs shared between MTRL training and evaluation.""" + kwargs: Dict[str, Any] = {} + if output_s3_path: + kwargs["s3_output_path"] = output_s3_path + if self.model_package_group_arn: + kwargs["output_model_package_group"] = self.model_package_group_arn + kwargs["intermediate_checkpoint_model_package_group"] = ( + self._get_or_create_checkpoint_model_package_group_arn() + ) + if mlflow_tracking_uri: + kwargs["mlflow_app_arn"] = mlflow_tracking_uri + if self.execution_role: + kwargs["role"] = self.execution_role + if self.subnets or self.security_group_ids: + from sagemaker.core.shapes import VpcConfig + + kwargs["networking"] = VpcConfig( + security_group_ids=self.security_group_ids or [], + subnets=self.subnets or [], + ) + if self.kms_key_id: + kwargs["kms_key_arn"] = self.kms_key_id + return kwargs + + @staticmethod + def _apply_mtrl_overrides(obj, overrides: Optional[Dict[str, Any]], label: str = "") -> None: + """Apply overrides to trainer/evaluator hyperparameters object.""" + if not overrides: + return + # obj.hyperparameters is a lazy property that fetches valid parameter specs + # from the SageMaker Public Hub. We need it before applying overrides so that + # setattr can validate each key/value against the recipe's allowed parameters. + try: + hp = obj.hyperparameters + except (ValueError, RuntimeError) as e: + skipped = ", ".join(f"{k}={v}" for k, v in overrides.items()) + logger.warning( + f"Overrides not supported currently — " + f"skipping {label}hyperparameters: [{skipped}]. " + f"The job will proceed with default values." + ) + return + for key, value in overrides.items(): + try: + setattr(hp, key, value) + except AttributeError: + logger.warning( + f"Skipping unsupported {label}hyperparameter '{key}' — not in recipe." + ) + except (TypeError, ValueError) as e: + raise ValueError(f"Invalid {label}hyperparameter '{key}={value}': {e}") from e + + @_telemetry_emitter(Feature.TRAINING, "execute_mtrl") + def execute_mtrl( + self, + model: Model, + job_name: str, + data_s3_path: Optional[str] = None, + output_s3_path: Optional[str] = None, + mlflow_tracking_uri: Optional[str] = None, + overrides: Optional[Dict[str, Any]] = None, + model_path: Optional[str] = None, + ) -> str: + """Execute a multi-turn RFT job via sagemaker.train MultiTurnRLTrainer. + + Args: + model: The Nova model enum. + job_name: Base job name. + data_s3_path: S3 path to training data. + output_s3_path: S3 output path. + mlflow_tracking_uri: MLflow tracking server ARN. + overrides: Optional hyperparameter overrides. + model_path: Model package ARN from a previous job for iterative training. + + Returns: + The job name string from the SDK. + """ + # TODO: Move to top-level import when MTRL is available in the PySdk + from sagemaker.core.resources import ModelPackage + from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer as SmMTRL + + agent_arn = self.agent_core_arn or self.rft_lambda + + resolved_model: Any = model.hub_content_name + if model_path and ":model-package/" in model_path: + resolved_model = ModelPackage.get(model_path) + + kwargs = { + "model": resolved_model, + "agent_env": agent_arn, + "accept_eula": True, + **self._mtrl_common_kwargs(output_s3_path, mlflow_tracking_uri), + } + if data_s3_path: + kwargs["training_dataset"] = data_s3_path + + sm_trainer = SmMTRL(**kwargs) + self._apply_mtrl_overrides(sm_trainer, overrides) + sm_trainer.base_job_name = job_name + + sm_job = sm_trainer.train(wait=False) + logger.info(f"MTRL job created: {sm_job.job_name} ({sm_job.job_arn})") + return sm_job.job_name + + @_telemetry_emitter(Feature.EVAL, "execute_mtrl_eval") + def execute_mtrl_eval( + self, + model: Model, + data_s3_path: Optional[str] = None, + output_s3_path: Optional[str] = None, + mlflow_tracking_uri: Optional[str] = None, + model_path: Optional[str] = None, + overrides: Optional[Dict[str, Any]] = None, + training_job_name: Optional[str] = None, + evaluate_base_model: bool = False, + ): + """Execute MTRL evaluation via sagemaker.train MultiTurnRLEvaluator. + + Constructs a MultiTurnRLTrainer, attaches the completed training job, + and passes the trainer to MultiTurnRLEvaluator so it can auto-resolve + the model package and agent config. + + Args: + model: The Nova model enum. + data_s3_path: S3 path to evaluation dataset. + output_s3_path: S3 output path for results. + mlflow_tracking_uri: MLflow tracking server ARN. + model_path: Model package ARN of a fine-tuned model to evaluate. + overrides: Optional hyperparameter overrides for evaluation. + training_job_name: Name of the completed MTRL training job to attach. + + Returns: + The MTRLEvaluationExecution object from the SDK. + """ + # TODO: Move to top-level import when MTRL is available in the PySdk + from sagemaker.train.evaluate import MultiTurnRLEvaluator + from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer as SmMTRL + + agent_arn = self.agent_core_arn or self.rft_lambda + common_kwargs = self._mtrl_common_kwargs(output_s3_path, mlflow_tracking_uri) + + # The evaluator uses 'mlflow_resource_arn' not 'mlflow_app_arn', and accepts + # 'kms_key_id' and 'region' which the trainer does not (BaseEvaluator fields). + eval_common_kwargs = dict(common_kwargs) + if "mlflow_app_arn" in eval_common_kwargs: + eval_common_kwargs["mlflow_resource_arn"] = eval_common_kwargs.pop("mlflow_app_arn") + if self.kms_key_id: + eval_common_kwargs["kms_key_id"] = self.kms_key_id + if self.region: + eval_common_kwargs["region"] = self.region + + # Build a trainer instance and attach the completed job so the evaluator + # can auto-resolve model artifacts and agent config. + if training_job_name and model_path: + trainer_kwargs = { + "model": model.hub_content_name, + "agent_env": agent_arn, + "accept_eula": True, + **common_kwargs, + } + if data_s3_path: + trainer_kwargs["training_dataset"] = data_s3_path + + sm_trainer = SmMTRL(**trainer_kwargs) + job = SmMTRL.attach(training_job_name) + sm_trainer._latest_job = job + sm_trainer.base_model_arn = sm_trainer._model_arn + sm_trainer.base_model_name = model.hub_content_name + sm_trainer.agent_config = agent_arn + + eval_kwargs = { + "model": sm_trainer, + "dataset": data_s3_path, + **eval_common_kwargs, + } + else: + # Fallback: pass model directly with agent_config. + # If model_path is a model-package ARN (Restricted MPG), pass the + # base hub-content name as `model` and inject the ARN post-init to + # avoid pydantic validation failure on missing s3_uri. + is_model_package_arn = model_path and ":model-package/" in model_path + eval_kwargs = { + "model": model.hub_content_name + if is_model_package_arn + else (model_path or model.hub_content_name), + "dataset": data_s3_path, + "agent_config": agent_arn, + **eval_common_kwargs, + } + + if evaluate_base_model: + eval_kwargs["evaluate_base_model"] = True + + evaluator = MultiTurnRLEvaluator(**eval_kwargs) + + if not training_job_name and model_path and ":model-package/" in model_path: + evaluator._source_model_package_arn_cache = model_path + + self._apply_mtrl_overrides(evaluator, overrides, label="eval ") + + execution = evaluator.evaluate() + logger.info(f"MTRL evaluation pipeline started: {execution.arn}") + return execution + + def _cleanup_mtrl(self, job_name: str) -> None: + """Stop an MTRL job.""" + # TODO: Move to top-level import when MTRL is available in the PySdk + from sagemaker.train.agent_rft_job import AgentRFTJob + + session = boto3.Session(region_name=self.region) + rft_job = AgentRFTJob.get(job_name, session=session) + rft_job.stop() + logger.info(f"Stopped MTRL job '{job_name}'") diff --git a/src/amzn_nova_forge/manager/runtime_manager.py b/src/amzn_nova_forge/manager/runtime_manager.py index aedf8f9..17286fb 100644 --- a/src/amzn_nova_forge/manager/runtime_manager.py +++ b/src/amzn_nova_forge/manager/runtime_manager.py @@ -21,15 +21,16 @@ import textwrap import time import zipfile -from urllib.parse import urlparse from dataclasses import dataclass, field from typing import Any, Dict, List, Optional +from urllib.parse import urlparse import boto3 import yaml from botocore.exceptions import ClientError, NoRegionError from sagemaker.core.helper.session_helper import Session, get_execution_role from sagemaker.core.shapes import ( + ModelPackageConfig, OutputDataConfig, S3DataSource, TensorBoardOutputConfig, @@ -50,6 +51,8 @@ from amzn_nova_forge.core.enums import Model, Platform, TrainingMethod from amzn_nova_forge.core.runtime import RuntimeManager as RuntimeManagerBase from amzn_nova_forge.core.types import JobConfig +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.bedrock import ( get_customization_type, @@ -73,6 +76,7 @@ TrainingMethod.DPO_LORA: ("DPO", "LORA"), TrainingMethod.DPO_FULL: ("DPO", None), TrainingMethod.RFT_LORA: ("RLVR", "LORA"), + TrainingMethod.RFT_MULTITURN_LORA: ("RLVR", "LORA"), # TrainingMethod.RFT_FULL: ("RLVR", None), # TODO: Add RLVR full support } DEFAULT_SMTJ_JOB_MAX_RUNTIME = 86400 # 1 day @@ -742,6 +746,45 @@ def execute(self, job_config: JobConfig) -> str: if job_config.trainer_config_hyperparameters: trainer_config["hyperparameters"] = job_config.trainer_config_hyperparameters + # If the model_name_or_path is a model package ARN, + # ModelTrainer.from_recipe() requires a ModelPackageConfig with model_package_group_arn. + model_package_config = None + if ( + job_config.method == TrainingMethod.EVALUATION + and job_config.model_name_or_path + and MODEL_PACKAGE_ARN_REGEX.match(job_config.model_name_or_path) + ): + # Extract model package group name from the ARN + _parts = job_config.model_name_or_path.split("/") + if len(_parts) < 2: + raise ValueError( + f"Could not parse model package group name from ARN: '{job_config.model_name_or_path}'. " + f"Expected format: arn:aws:sagemaker:region:account:model-package/group-name/version" + ) + _group_name = _parts[1] + # Look up the existing model package group ARN + try: + _resp = self.sagemaker_client.describe_model_package_group( + ModelPackageGroupName=_group_name + ) + _group_arn = _resp["ModelPackageGroupArn"] + except ClientError as e: + if e.response["Error"]["Code"] in ( + "ValidationException", + "ResourceNotFoundException", + ): + raise ValueError( + f"Model package group '{_group_name}' does not exist. " + f"Ensure the model package ARN '{job_config.model_name_or_path}' refers to an existing group." + ) from e + raise # Re-raise unexpected errors + model_package_config = ModelPackageConfig( + model_package_group_arn=_group_arn, + ) + + if model_package_config: + trainer_config["model_package_config"] = model_package_config + model_trainer = ModelTrainer.from_recipe( **trainer_config ).with_tensorboard_output_config(tensorboard_output_config) @@ -1398,6 +1441,8 @@ def _is_standard_ecr_image(image_uri: str) -> bool: """Return True if the image is a standard private ECR image (account.dkr.ecr.region.amazonaws.com/...).""" parsed = urlparse(image_uri) host = parsed.hostname + # Handle schemeless image URIs such as: + # 123456789012.dkr.ecr.us-west-2.amazonaws.com/repo:tag if not host: host = parsed.path.split("/", 1)[0] if not host: @@ -2120,7 +2165,7 @@ def required_calling_role_permissions(cls, data_s3_path=None, output_s3_path=Non return permissions -class SMTJServerlessRuntimeManager(RuntimeManager): +class SMTJServerlessRuntimeManager(MTRLOperations, RuntimeManager): def __init__( self, model_package_group_name: str, @@ -2131,13 +2176,18 @@ def __init__( security_group_ids: Optional[list[str]] = None, max_job_runtime: Optional[int] = DEFAULT_SMTJ_JOB_MAX_RUNTIME, # 1 day rft_lambda: Optional[str] = None, + agent_core_arn: Optional[str] = None, evaluator_name: Optional[str] = None, + hub_content_version: Optional[str] = None, + intermediate_model_package_group_name: Optional[str] = None, ): # NOTE: Not setting execution_role directly due to issues with mypy type inference self._execution_role = execution_role self.model_package_group_name = model_package_group_name + self.intermediate_model_package_group_name = intermediate_model_package_group_name + self.agent_core_arn = agent_core_arn self.evaluator_name = evaluator_name - self.hub_content_version: Optional[str] = None + self.hub_content_version: Optional[str] = hub_content_version self.subnets = subnets self.security_group_ids = security_group_ids self.encrypt_inter_container_traffic = encrypt_inter_container_traffic @@ -2204,6 +2254,23 @@ def setup(self) -> None: ) self.model_package_group_arn = resp["ModelPackageGroupArn"] + def _get_or_create_checkpoint_model_package_group_arn(self) -> str: + """Get or create a separate model package group for intermediate checkpoints.""" + checkpoint_group_name = ( + self.intermediate_model_package_group_name + or f"{self.model_package_group_name}-checkpoints"[:63] + ) + try: + resp = self.sagemaker_client.describe_model_package_group( + ModelPackageGroupName=checkpoint_group_name + ) + return resp["ModelPackageGroupArn"] + except self.sagemaker_client.exceptions.ClientError: + resp = self.sagemaker_client.create_model_package_group( + ModelPackageGroupName=checkpoint_group_name + ) + return resp["ModelPackageGroupArn"] + def _resolve_base_model_arn(self, model: Model) -> str: """Resolve the BaseModelArn from SageMaker Hub for the given model.""" hub_content = get_hub_content( @@ -2357,9 +2424,13 @@ def _build_serverless_job_config( config["CustomizationTechnique"] = technique if peft: config["Peft"] = peft - # EvaluatorArn: reward function hub-content ARN for RLVR training. + # EvaluatorArn: reward function hub-content ARN for RLVR training, + # or agent runtime ARN for multiturn RLVR. # Lambda ARNs are passed via HyperParameters (reward_lambda_arn) instead. - if _is_hub_content_arn(evaluator_arn): + if evaluator_arn and ( + _is_hub_content_arn(evaluator_arn) + or method in (TrainingMethod.RFT_MULTITURN_LORA, TrainingMethod.RFT_MULTITURN_FULL) + ): config["EvaluatorArn"] = evaluator_arn return config @@ -2452,16 +2523,32 @@ def execute(self, job_config: JobConfig) -> str: eval_task=recipe.get("evaluation", {}).get("task"), evaluator_arn=evaluator_arn, ), - "ModelPackageConfig": { + } + + # ModelPackageConfig: + # - Training jobs: always include ModelPackageGroupArn. + # - Eval jobs: only include ModelPackageConfig when model_name_or_path is a + # SageMaker ARN (fine-tuned model eval). Base model evals omit it entirely. + model_name_or_path = recipe.get("run", {}).get("model_name_or_path", "") + if job_config.method == TrainingMethod.EVALUATION: + if is_sagemaker_arn(model_name_or_path): + model_package_config = { + "ModelPackageGroupArn": self.model_package_group_arn, + "SourceModelPackageArn": model_name_or_path, + } + else: + model_package_config = {} + else: + model_package_config = { "ModelPackageGroupArn": self.model_package_group_arn, - # model_name_or_path is a model package ARN for iterative training **( - {"SourceModelPackageArn": recipe["run"]["model_name_or_path"]} - if is_sagemaker_arn(recipe.get("run", {}).get("model_name_or_path", "")) + {"SourceModelPackageArn": model_name_or_path} + if is_sagemaker_arn(model_name_or_path) else {} ), - }, - } + } + if model_package_config: + create_params["ModelPackageConfig"] = model_package_config if self.kms_key_id: create_params["OutputDataConfig"]["KmsKeyId"] = self.kms_key_id @@ -2531,10 +2618,13 @@ def execute(self, job_config: JobConfig) -> str: logger.error(f"Failed to start training job: {str(e)}") raise - def cleanup(self, job_name: str) -> None: + def cleanup(self, job_name: str, is_mtrl: bool = False) -> None: try: - self.sagemaker_client.stop_training_job(TrainingJobName=job_name) - self.sagemaker_client.close() + if is_mtrl: + self._cleanup_mtrl(job_name) + else: + self.sagemaker_client.stop_training_job(TrainingJobName=job_name) + self.sagemaker_client.close() except Exception as e: logger.error(f"Failed to cleanup job {job_name}: {str(e)}") raise diff --git a/src/amzn_nova_forge/model/nova_model_customizer.py b/src/amzn_nova_forge/model/nova_model_customizer.py index 1b2f1dd..28b83d3 100644 --- a/src/amzn_nova_forge/model/nova_model_customizer.py +++ b/src/amzn_nova_forge/model/nova_model_customizer.py @@ -1027,8 +1027,9 @@ def create_custom_model( execution_role_name: Optional[str] = None, tags: Optional[List[Dict[str, str]]] = None, skip_model_reuse: bool = False, + custom_model_data_source: Optional[Dict[str, Any]] = None, ) -> ModelDeployResult: - """Create a Bedrock custom model from S3 artifacts. + """Create a Bedrock custom model from S3 artifacts or a model package ARN. Delegates to ForgeDeployer.create_custom_model(). Handles job_result resolution at the facade level before delegating. @@ -1039,8 +1040,14 @@ def create_custom_model( DeprecationWarning, stacklevel=2, ) - # Resolve model_artifact_path from job_result (facade-level concern) - if model_artifact_path is None and job_result is not None: + # Resolve model_artifact_path from job_result (facade-level concern). + # Skip extraction when custom_model_data_source is already provided, + # since it serves as a valid alternative to model_artifact_path. + if ( + model_artifact_path is None + and job_result is not None + and custom_model_data_source is None + ): model_artifact_path = extract_checkpoint_path_from_job_output( output_s3_path=job_result.model_artifacts.output_s3_path, job_result=job_result, @@ -1048,11 +1055,13 @@ def create_custom_model( if model_artifact_path is None: raise ValueError( f"Could not resolve checkpoint path from job result '{job_result.job_id}'. " - f"Provide model_artifact_path explicitly." + f"Provide model_artifact_path or custom_model_data_source explicitly." ) - if model_artifact_path is None: - raise ValueError("Either model_artifact_path or job_result must be provided.") + if model_artifact_path is None and custom_model_data_source is None: + raise ValueError( + "Either model_artifact_path, job_result, or custom_model_data_source must be provided." + ) deployer = self._build_deployer() result = deployer.create_custom_model( @@ -1061,6 +1070,7 @@ def create_custom_model( execution_role_name=execution_role_name, tags=tags, skip_model_reuse=skip_model_reuse, + custom_model_data_source=custom_model_data_source, ) # Sync state back diff --git a/src/amzn_nova_forge/monitor/__init__.py b/src/amzn_nova_forge/monitor/__init__.py index 4a7d77c..b9f6d05 100644 --- a/src/amzn_nova_forge/monitor/__init__.py +++ b/src/amzn_nova_forge/monitor/__init__.py @@ -11,10 +11,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from .log_monitor import CloudWatchLogMonitor +from .log_monitor import CloudWatchLogMonitor, MTRLLogMonitor from .mlflow_monitor import MLflowMonitor __all__ = [ "MLflowMonitor", "CloudWatchLogMonitor", + "MTRLLogMonitor", ] diff --git a/src/amzn_nova_forge/monitor/log_monitor.py b/src/amzn_nova_forge/monitor/log_monitor.py index 32144d7..2bce91f 100644 --- a/src/amzn_nova_forge/monitor/log_monitor.py +++ b/src/amzn_nova_forge/monitor/log_monitor.py @@ -19,6 +19,11 @@ import pandas from matplotlib import pyplot +from amzn_nova_forge.core.constants import ( + MTRL_EVAL_LOG_GROUP, + MTRL_PIPELINE_EXECUTION_RE, + MTRL_TRAIN_LOG_GROUP, +) from amzn_nova_forge.core.enums import Platform, TrainingMethod from amzn_nova_forge.core.result.job_result import ( BaseJobResult, @@ -545,3 +550,164 @@ def plot_metrics( pyplot.grid(True) pyplot.style.use("seaborn-v0_8-white") pyplot.show() + + +class MTRLLogMonitor: + """Log monitor for MTRL jobs. + + Provides the same ``from_job_id`` / ``show_logs`` interface as + ``CloudWatchLogMonitor`` so the user experience is consistent. + Internally delegates to the ``AgentRFTJob.wait()`` for live + progress and ``get_training_metrics()`` for completed jobs. + Auto-detects whether the job is a training or evaluation job. + """ + + def __init__( + self, job_id: str, region: Optional[str] = None, job_category: Optional[str] = None + ): + self.job_id = job_id + self._region = region + self._job_category = job_category + self._rft_job = None + + @classmethod + @_telemetry_emitter(Feature.MONITOR, "mtrl_from_job_id") + def from_job_id( + cls, job_id: str, region: Optional[str] = None, job_category: Optional[str] = None, **kwargs + ) -> "MTRLLogMonitor": + return cls(job_id=job_id, region=region, job_category=job_category) + + def _detect_job_category(self) -> str: + """Auto-detect whether this is a training or evaluation job.""" + if self._job_category: + return self._job_category + + # Pipeline execution ARNs are always evaluation jobs + if MTRL_PIPELINE_EXECUTION_RE.match(self.job_id): + self._job_category = "AgentRFTEvaluation" + return self._job_category + + try: + logs_client = boto3.client("logs", region_name=self._region) + + for category, log_group in [ + ("AgentRFT", MTRL_TRAIN_LOG_GROUP), + ("AgentRFTEvaluation", MTRL_EVAL_LOG_GROUP), + ]: + try: + resp = logs_client.describe_log_streams( + logGroupName=log_group, + logStreamNamePrefix=self.job_id, + limit=1, + ) + if resp.get("logStreams"): + self._job_category = category + return category + except logs_client.exceptions.ResourceNotFoundException: + continue + except Exception: + pass + + self._job_category = "AgentRFT" + return self._job_category + + def _get_rft_job(self): + if self._rft_job is None: + from sagemaker.train.agent_rft_job import AgentRFTJob + + session = boto3.Session(region_name=self._region) if self._region else None + self._rft_job = AgentRFTJob.get(self.job_id, session=session) + return self._rft_job + + @_telemetry_emitter(Feature.MONITOR, "mtrl_show_logs") + def show_logs( + self, poll: int = 30, timeout: int = 7200, limit: Optional[int] = None, **kwargs + ) -> None: + """Show MTRL job progress. + + If the job is still running, blocks and displays a live progress panel. + If the job is completed, prints training metrics. + For eval jobs, reads CloudWatch logs directly. + """ + category = self._detect_job_category() + + if category == "AgentRFTEvaluation": + self._show_eval_logs(limit=limit) + return + + rft_job = self._get_rft_job() + rft_job.refresh() + if rft_job.job_status in ("Completed", "Failed", "Stopped"): + print(f"Job '{self.job_id}' status: {rft_job.job_status}") + if rft_job.job_status == "Completed": + rft_job.get_training_metrics() + else: + rft_job.wait(poll=poll, timeout=timeout) + + def _resolve_eval_job_names(self) -> List[tuple]: + """Resolve the underlying eval job name(s) from a pipeline execution ARN. + + Returns list of (step_name, job_name) tuples. + """ + if not MTRL_PIPELINE_EXECUTION_RE.match(self.job_id): + return [("", self.job_id)] + + region = self.job_id.split(":")[3] + sm_client = boto3.client("sagemaker", region_name=region) + steps = sm_client.list_pipeline_execution_steps(PipelineExecutionArn=self.job_id) + jobs = [] + for step in steps.get("PipelineExecutionSteps", []): + job_meta = step.get("Metadata", {}).get("Job", {}) + if job_meta: + job_name = job_meta["Arn"].rsplit("/", 1)[-1] + step_name = step.get("StepName", "") + jobs.append((step_name, job_name)) + return jobs or [("", self.job_id)] + + def _show_eval_logs(self, limit: Optional[int] = None) -> None: + """Read and display CloudWatch logs for an MTRL evaluation job.""" + jobs = self._resolve_eval_job_names() + region = ( + self.job_id.split(":")[3] + if MTRL_PIPELINE_EXECUTION_RE.match(self.job_id) + else self._region + ) + logs_client = boto3.client("logs", region_name=region) + + for step_name, job_name in jobs: + if step_name: + print(f"\n--- {step_name} ---") + + try: + resp = logs_client.describe_log_streams( + logGroupName=MTRL_EVAL_LOG_GROUP, + logStreamNamePrefix=job_name, + limit=5, + ) + except logs_client.exceptions.ResourceNotFoundException: + print(f"No log group found: {MTRL_EVAL_LOG_GROUP}") + return + + streams = resp.get("logStreams", []) + if not streams: + print(f"No log stream found for job '{job_name}'") + continue + + stream_name = streams[0]["logStreamName"] + params: Dict[str, Any] = { + "logGroupName": MTRL_EVAL_LOG_GROUP, + "logStreamName": stream_name, + "startFromHead": False, + } + if limit: + params["limit"] = limit + + events_resp = logs_client.get_log_events(**params) + events = events_resp.get("events", []) + + if not events: + print(f"No logs available yet for job '{job_name}'") + continue + + for event in events: + print(event["message"].strip()) diff --git a/src/amzn_nova_forge/recipe/recipe_builder.py b/src/amzn_nova_forge/recipe/recipe_builder.py index 664c9e7..49cc2a9 100644 --- a/src/amzn_nova_forge/recipe/recipe_builder.py +++ b/src/amzn_nova_forge/recipe/recipe_builder.py @@ -42,6 +42,9 @@ from amzn_nova_forge.core.runtime import RuntimeManager from amzn_nova_forge.core.training_overrides import NULLABLE_OVERRIDE_FIELDS from amzn_nova_forge.core.types import ConfigParameter, RecipeConfig, ValidationConfig +from amzn_nova_forge.core.validation_patterns import ( + MODEL_PACKAGE_ARN_REGEX, +) from amzn_nova_forge.monitor import MLflowMonitor from amzn_nova_forge.rft_multiturn import RFTMultiturnInfrastructure from amzn_nova_forge.util.checkpoint_util import validate_checkpoint_uri @@ -133,7 +136,10 @@ def __init__( self.mlflow_tracking_uri = None # RFT - if method == TrainingMethod.RFT_LORA or method == TrainingMethod.RFT_FULL: + if method in ( + TrainingMethod.RFT_LORA, + TrainingMethod.RFT_FULL, + ): self.rft_lambda_arn = rft_lambda_arn elif rft_lambda_arn is not None: logger.info("'rft_lambda_arn' is only required for RFT. Will ignore.") @@ -149,10 +155,10 @@ def __init__( f"{method} method is only supported on Nova 2.0. " f"You provided {model}. Please use Model.NOVA_LITE_2." ) - if not rft_multiturn_infra: + if not rft_multiturn_infra and platform != Platform.SMTJServerless: raise ValueError("'rft_multiturn_infra' is required for RFT multiturn training") elif method == TrainingMethod.EVALUATION and eval_task == EvaluationTask.RFT_MULTITURN_EVAL: - if not rft_multiturn_infra: + if not rft_multiturn_infra and platform != Platform.SMTJServerless: raise ValueError("'rft_multiturn_infra' is required for RFT multiturn evaluation") # Validation data handling @@ -365,8 +371,11 @@ def apply_user_provided_inputs_into_overrides_template(): 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: + # RFT (single-turn only — MTRL uses agent_core_arn directly) + if self.method in ( + TrainingMethod.RFT_LORA, + TrainingMethod.RFT_FULL, + ): overrides_template.setdefault("reward_lambda_arn", {})["default"] = ( self.rft_lambda_arn ) @@ -474,28 +483,28 @@ def handle_edge_case_keys(key: str) -> bool: return False if key == "model_name_or_path": if key in overrides: - if not str(overrides[key]).startswith("s3://"): + if str(overrides[key]).startswith("s3://"): + validate_checkpoint_uri( + checkpoint_uri=str(overrides[key]), region=self.region + ) + elif not MODEL_PACKAGE_ARN_REGEX.match(str(overrides[key])): logger.warning( f"Override for '{key}' will be ignored. If you wish to use a different model than {self.model.name}, please update your NovaModelCustomizer object." ) return False - else: + elif key in input_recipe_key_values: + if str(input_recipe_key_values[key]).startswith("s3://"): validate_checkpoint_uri( - checkpoint_uri=str(overrides[key]), region=self.region + checkpoint_uri=str(input_recipe_key_values[key]), + region=self.region, ) - elif key in input_recipe_key_values: - if not str(input_recipe_key_values[key]).startswith( - "s3://" + elif not MODEL_PACKAGE_ARN_REGEX.match( + str(input_recipe_key_values[key]) ) and self.model.model_path != str(input_recipe_key_values[key]): logger.warning( f"{key} '{str(input_recipe_key_values[key])}' will be ignored from your input recipe. If you wish to use a different model than {self.model.name}, please update your NovaModelCustomizer object." ) return False - elif str(input_recipe_key_values[key]).startswith("s3://"): - validate_checkpoint_uri( - checkpoint_uri=str(input_recipe_key_values[key]), - region=self.region, - ) # ignoring data_s3_path as we use it to determine modality of data for getting Datamixing recipes if key == "data_s3_path" and self.is_multimodal: if key in overrides or key in input_recipe_key_values: @@ -988,6 +997,15 @@ def build_and_validate( if "rl_env" in final_recipe_dict and "reward_lambda_arn" in final_recipe_dict["rl_env"]: del final_recipe_dict["rl_env"]["reward_lambda_arn"] + if self.mlflow_tracking_uri: + final_recipe_dict.setdefault("run", {})["mlflow_tracking_uri"] = ( + self.mlflow_tracking_uri + ) + if self.mlflow_experiment_name: + final_recipe_dict["run"]["mlflow_experiment_name"] = self.mlflow_experiment_name + if self.mlflow_run_name: + final_recipe_dict["run"]["mlflow_run_name"] = self.mlflow_run_name + if self.enable_batch_sample_tracing: final_recipe_dict.setdefault("training_config", {})["enable_batch_sample_tracing"] = ( True diff --git a/src/amzn_nova_forge/trainer/forge_trainer.py b/src/amzn_nova_forge/trainer/forge_trainer.py index d31cf3a..9cbab18 100644 --- a/src/amzn_nova_forge/trainer/forge_trainer.py +++ b/src/amzn_nova_forge/trainer/forge_trainer.py @@ -40,6 +40,7 @@ SMTJTrainingResult, TrainingResult, ) +from amzn_nova_forge.core.result.job_result import JobStatus from amzn_nova_forge.core.runtime import RuntimeManager from amzn_nova_forge.core.training_overrides import TrainingOverrides from amzn_nova_forge.core.types import ( @@ -49,7 +50,7 @@ RecipeConfig, validate_region, ) -from amzn_nova_forge.manager.runtime_manager import SMHPRuntimeManager +from amzn_nova_forge.manager.runtime_manager import SMHPRuntimeManager, SMTJServerlessRuntimeManager from amzn_nova_forge.model.nova_model_customizer_util import set_output_s3_path from amzn_nova_forge.monitor.log_monitor import CloudWatchLogMonitor from amzn_nova_forge.recipe.recipe_builder import RecipeBuilder @@ -58,6 +59,10 @@ from amzn_nova_forge.util.data_mixing import DataMixing from amzn_nova_forge.util.data_utils import is_multimodal_data from amzn_nova_forge.util.logging import logger +from amzn_nova_forge.util.metric_util import ( + _build_and_upload_training_metrics_csv, + _parse_user_time, +) from amzn_nova_forge.util.recipe import load_recipe_templates from amzn_nova_forge.util.sagemaker import get_model_artifacts from amzn_nova_forge.validation.endpoint_validator import is_sagemaker_arn @@ -82,6 +87,7 @@ def __init__( infra: RuntimeManager, training_data_s3_path: Optional[str] = None, model_s3_path: Optional[str] = None, + model_arn: Optional[str] = None, data_mixing_enabled: bool = False, holdout_data_s3_path: Optional[str] = None, val_check_interval: Optional[int] = None, @@ -95,6 +101,7 @@ def __init__( self.method = method self.infra = infra self.training_data_s3_path = training_data_s3_path + self.model_arn = model_arn self.model_s3_path = model_s3_path self.holdout_data_s3_path = holdout_data_s3_path self.hub_content_version = hub_content_version @@ -126,11 +133,23 @@ def __init__( self._platform = infra.platform - if self._platform == Platform.SMTJServerless and self.model_s3_path is not None: - if not is_sagemaker_arn(self.model_s3_path): + if self._platform == Platform.SMTJServerless: + if self.model_s3_path is not None and self.model_arn is not None: + raise ValueError( + "Cannot specify both model_s3_path and model_arn. " + "For SMTJServerless, use model_arn with a model package ARN." + ) + if self.model_s3_path is not None: + if not is_sagemaker_arn(self.model_s3_path): + raise ValueError( + f"For SMTJServerless, model_s3_path must be a SageMaker model package ARN, " + f"got: '{self.model_s3_path}'. Use model_arn instead." + ) + self.model_arn = self.model_s3_path + self.model_s3_path = None + if self.model_arn is not None and not is_sagemaker_arn(self.model_arn): raise ValueError( - f"For SMTJServerless, model_path must be a SageMaker model package ARN, " - f"got: '{self.model_s3_path}'." + f"model_arn must be a SageMaker model package ARN, got: '{self.model_arn}'." ) if self._platform == Platform.BEDROCK and self.model_s3_path is not None: logger.warning( @@ -207,7 +226,7 @@ def __init__( model=model, method=method, data_s3_path=training_data_s3_path, - model_path=model_s3_path, + model_path=self.model_arn or self.model_s3_path, output_s3_path=self.output_s3_path, instance_type=infra.instance_type, instance_count=infra.instance_count, @@ -265,6 +284,65 @@ def train( validate_rft_lambda_name(rft_lambda_arn.split(":")[-1], self._platform) logger.info(f"Using reward lambda: {rft_lambda_arn}") + # MTRL serverless: skip recipe generation, call directly + is_mtrl_serverless = self._platform == Platform.SMTJServerless and self.method in ( + TrainingMethod.RFT_MULTITURN_LORA, + TrainingMethod.RFT_MULTITURN_FULL, + ) + if is_mtrl_serverless: + if not self._config.mlflow_monitor or not self._config.mlflow_monitor.tracking_uri: + raise ValueError( + "MLflow configuration is required for AgentRFT jobs. " + "Please provide an mlflow_monitor with a valid tracking_uri when " + "using RFT_MULTITURN methods on the SMTJServerless platform." + ) + + if dry_run: + return None + + mtrl_infra = cast(SMTJServerlessRuntimeManager, self.infra) + mlflow_uri = self._config.mlflow_monitor.tracking_uri + + unique_job_name = f"{job_name}-{uuid.uuid4().hex[:8]}"[:48].rstrip("-") + start_time = datetime.now(timezone.utc) + + job_id = mtrl_infra.execute_mtrl( + model=self.model, + job_name=unique_job_name, + data_s3_path=self.training_data_s3_path, + output_s3_path=self.output_s3_path, + mlflow_tracking_uri=mlflow_uri, + overrides=dict(overrides) if overrides else None, + model_path=self.model_arn, + ) + + training_result: TrainingResult = SMTJTrainingResult( + job_id=job_id, + started_time=start_time, + method=self.method, + model_type=self.model, + model_artifacts=get_model_artifacts( + job_name=job_id, + infra=self.infra, + region=self.region, + ), + region=self.region, + ) + + logger.info(f"Started MTRL job '{training_result.job_id}'.") + if training_result.model_artifacts.output_s3_path: + logger.info(f"Output S3 path is: {training_result.model_artifacts.output_s3_path}.") + + persist_result( + self._cache_context, + training_result, + job_name=job_name, + job_type="train", + recipe_path=recipe_path, + overrides=overrides or {}, + ) + return training_result + recipe_builder = RecipeBuilder( region=self.region, job_name=job_name, @@ -276,7 +354,7 @@ def train( infra=self.infra, data_s3_path=self.training_data_s3_path, output_s3_path=self.output_s3_path, - model_path=self.model_s3_path, + model_path=self.model_arn or 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, @@ -334,7 +412,6 @@ def train( job_id = self.infra.execute(job_config=JobConfig(**job_config_params)) - training_result: TrainingResult if self._platform in (Platform.SMTJ, Platform.SMTJServerless): training_result = SMTJTrainingResult( job_id=job_id, @@ -434,7 +511,7 @@ def get_config( infra=self.infra, data_s3_path=self.training_data_s3_path, output_s3_path=self.output_s3_path, - model_path=self.model_s3_path, + model_path=self.model_arn or self.model_s3_path, validation_data_s3_path=self.holdout_data_s3_path, val_check_interval=self.val_check_interval, data_mixing_instance=self.data_mixing, @@ -454,8 +531,15 @@ def get_logs( limit: Optional[int] = None, start_from_head: bool = False, end_time: Optional[int] = None, + poll: int = 30, + timeout: int = 7200, ) -> None: - """Stream CloudWatch logs for a training job. + """Stream logs or wait for a training job. + + For MTRL jobs, delegates to ``job_result.wait()`` which displays a + rich progress panel with status, metrics, and links. + + For all other jobs, streams CloudWatch logs. Provide either a ``job_result`` or explicit ``job_id`` + ``started_time``. @@ -463,6 +547,15 @@ def get_logs( ValueError: If neither ``job_result`` nor both ``job_id`` and ``started_time`` are provided. """ + # MTRL: delegate to wait() which shows progress + if ( + job_result is not None + and isinstance(job_result, SMTJTrainingResult) + and job_result._is_mtrl + ): + job_result.wait(poll=poll, timeout=timeout) + return + resolved_job_id = job_result.job_id if job_result else job_id resolved_started = job_result.started_time if job_result else started_time @@ -543,3 +636,131 @@ def trace_batch( s3_client=self._s3_client, cache_dir=cache_dir, ) + + @_telemetry_emitter(Feature.TRAINING, "generate_training_metrics_csv") + def generate_training_metrics_csv( + self, + job_result: Optional[SMHPTrainingResult] = None, + job_id: Optional[str] = None, + started_time=None, + end_time=None, + output_s3_path: Optional[str] = None, + ) -> Optional[str]: + """Generate step-wise training metrics CSV for an SMHP SFT job. + + Fetches CloudWatch logs for the specified job, extracts step-level + training metrics (step number, epoch number, training loss), and + uploads a ``step_wise_training_metrics.csv`` to the job's output S3 path. + + Args: + job_result: An SMHPTrainingResult from a completed training job. + job_id: The SMHP training job ID (used if job_result isn't provided). + started_time: Job start time for log filtering. Accepts a datetime + object or an ISO date str (e.g. "2025-05-26"). Defaults to 7 days. + end_time: Optional end time to bound the log search. If not provided, + searches up to the current time. + output_s3_path: S3 URI for output. Defaults to ForgeTrainer's output path. + + Returns: + S3 URI of the uploaded CSV, or None (no logs/metrics found). + + Raises: + ValueError: If platform is not SMHP, method is not SFT_LORA/SFT_FULL, + or required parameters are missing. + """ + # Validate platform + if self._platform != Platform.SMHP: + raise ValueError( + "generate_training_metrics_csv is only supported for SMHP platform. " + f"Current platform: {self._platform.value}" + ) + + # Validate method + if self.method not in (TrainingMethod.SFT_LORA, TrainingMethod.SFT_FULL): + raise ValueError( + "generate_training_metrics_csv is only supported for SFT training methods " + f"(SFT_LORA, SFT_FULL). Current method: {self.method.value}" + ) + + # Extract parameters from job_result or use standalone params + resolved_output_s3_path: Optional[str] = None + resolved_job_id: Optional[str] = None + if job_result is not None: + resolved_job_id = job_result.job_id + resolved_started_time = job_result.started_time + resolved_cluster_name = job_result.cluster_name + resolved_namespace = job_result.namespace + resolved_output_s3_path = output_s3_path or job_result.model_artifacts.output_s3_path + else: + resolved_job_id = job_id + resolved_started_time = started_time + resolved_cluster_name = cast(SMHPRuntimeManager, self.infra).cluster_name + resolved_namespace = cast(SMHPRuntimeManager, self.infra).namespace + resolved_output_s3_path = output_s3_path or self.output_s3_path + + if not resolved_job_id: + raise ValueError("No job_id provided. Pass either a job_result object or a job_id.") + + if not resolved_output_s3_path: + raise ValueError( + "output_s3_path is required but was not provided and could not be " + "extracted from job_result." + ) + + # Check job status and emit warnings + if job_result is not None: + try: + job_status, raw_status = job_result.get_job_status() + if job_status == JobStatus.IN_PROGRESS: + logger.warning( + "Job '%s' is still in progress (status: %s). " + "Training metrics may be incomplete.", + resolved_job_id, + raw_status, + ) + elif job_status == JobStatus.FAILED: + logger.warning( + "Job '%s' has failed (status: %s). Training metrics may be partial.", + resolved_job_id, + raw_status, + ) + except Exception: + logger.warning( + "Could not determine job status for '%s'. " + "Proceeding with best-effort log parsing.", + resolved_job_id, + ) + + # Fetch logs from CloudWatch + resolved_started_time_dt = _parse_user_time(resolved_started_time) + started_time_ms = int(resolved_started_time_dt.timestamp() * 1000) + + end_time_ms = None + if end_time is not None: + resolved_end_time_dt = _parse_user_time(end_time) + end_time_ms = int(resolved_end_time_dt.timestamp() * 1000) + + # Create a monitor object and retrieve relevant job logs + logger.info( + "Fetching CloudWatch logs for job '%s' - this can take a few minutes.", + resolved_job_id, + ) + monitor = CloudWatchLogMonitor( + job_id=resolved_job_id, + platform=Platform.SMHP, + started_time=started_time_ms, + region=self.region, + cluster_name=resolved_cluster_name, + namespace=resolved_namespace, + ) + log_events = monitor.get_logs(end_time=end_time_ms) + logger.info("Retrieved %d log events.", len(log_events) if log_events else 0) + + logger.info("Parsing training metrics and generating CSV. ") + return _build_and_upload_training_metrics_csv( + job_id=resolved_job_id, + log_events=log_events, + output_s3_path=resolved_output_s3_path, + training_method=self.method, + region=self.region, + ) diff --git a/src/amzn_nova_forge/util/metric_util.py b/src/amzn_nova_forge/util/metric_util.py index 13f3adb..55a324b 100644 --- a/src/amzn_nova_forge/util/metric_util.py +++ b/src/amzn_nova_forge/util/metric_util.py @@ -11,13 +11,23 @@ # 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 os import re +import tempfile +from datetime import datetime, timedelta, timezone from typing import Dict, List, Optional +import boto3 import pandas from amzn_nova_forge.core.enums import Platform, TrainingMethod +from amzn_nova_forge.util.s3_utils import parse_s3_uri +logger = logging.getLogger(__name__) + +DEFAULT_LOOKBACK_DAYS = 7 +EPOCH_IDX_PATTERN = re.compile(r'"EpochIdx"\s*:\s*(\d+)') GLOBAL_STEP_REGEX = r"global_step[=:]\s*([\d.]+)" TRAINING_LOSS_REGEX = r"reduced_train_loss[=:]\s*(-?[\d.]+(?:[eE][+-]?\d+)?)" SMHP_RFT_REWARD_SCORE_REGEX = r"train_rm_score:\s*(-?[\d.]+(?:[eE][+-]?\d+)?)" @@ -75,3 +85,191 @@ def get_metrics( pass return pandas.DataFrame(all_metrics, columns=["global_step"] + metrics) + + +def _extract_epoch_boundaries(log_events: list[dict]) -> list[tuple[int, int]]: + """ + Extract epoch start timestamps from CloudWatch log events. + Returns a list of (epoch_index, timestamp_ms) tuples sorted by timestamp. + """ + boundaries: list[tuple[int, int]] = [] + for event in log_events: + message = event.get("message", "") + match = EPOCH_IDX_PATTERN.search(message) + if match: + epoch_index = int(match.group(1)) + timestamp_ms = event.get("timestamp", 0) + boundaries.append((epoch_index, timestamp_ms)) + boundaries.sort(key=lambda x: x[1]) + return boundaries + + +def _assign_epochs_to_steps( + metrics_df: pandas.DataFrame, + epoch_boundaries: list[tuple[int, int]], + log_events: list[dict], +) -> pandas.DataFrame: + """Adds epoch_number column to the metrics DataFrame. + + Compares the timestamp of the step to the most recent epoch boundary + that precedes the step. Defaults to epoch 0 if an epoch doesn't precede + a step value. + """ + if metrics_df.empty: + metrics_df = metrics_df.copy() + metrics_df["epoch_number"] = pandas.Series(dtype=int) + return metrics_df + + # Build a mapping from global_step to timestamp by scanning log events + step_timestamps: dict[int, int] = {} + for event in log_events: + message = event.get("message", "") + for line in message.splitlines(): + step_match = re.search(GLOBAL_STEP_REGEX, line) + if step_match: + step_value = int(float(step_match.group(1))) + if step_value not in step_timestamps: + step_timestamps[step_value] = event.get("timestamp", 0) + + # Assign epoch numbers based on most recent preceding epoch boundary + epoch_numbers: list[int] = [] + for _, row in metrics_df.iterrows(): + step = int(row["global_step"]) + step_ts = step_timestamps.get(step, 0) + + # Find the most recent epoch boundary at or before this step's timestamp + assigned_epoch = 0 + for epoch_index, boundary_ts in epoch_boundaries: + if boundary_ts <= step_ts: + assigned_epoch = epoch_index + else: + break + + epoch_numbers.append(assigned_epoch) + + result_df = metrics_df.copy() + result_df["epoch_number"] = epoch_numbers + return result_df + + +def _parse_user_time(started_time) -> datetime: + """Parse start or end time into a datetime. + + Accepts: datetime object (returned as-is), ISO date string (e.g. "2025-05-26") + If "None" is provided, defaults to 7 days. + """ + if started_time is None: + return datetime.now(tz=timezone.utc) - timedelta(days=DEFAULT_LOOKBACK_DAYS) + + if isinstance(started_time, datetime): + return started_time + + if isinstance(started_time, str): + try: + parsed = datetime.fromisoformat(started_time.replace("Z", "+00:00")) + return parsed + except ValueError: + raise ValueError( + f"Cannot parse started_time '{started_time}'. " + "Use ISO date format, e.g., '2025-05-26'." + ) + + raise TypeError( + f"started_time must be a datetime, ISO date string, or None. " + f"Got {type(started_time).__name__}." + ) + + +def _build_and_upload_training_metrics_csv( + job_id: str, + log_events: List[Dict], + output_s3_path: str, + training_method: TrainingMethod, + region: Optional[str] = None, + s3_client=None, +) -> Optional[str]: + """Generate step_wise_training_metrics.csv from SMHP SFT log events. + + Parses the provided log events to extract step-level training metrics, + writes them to a CSV, and uploads to S3. + + Args: + job_id: The SMHP training job ID. + log_events: CloudWatch log events (list of dicts with 'message' and 'timestamp'). + output_s3_path: S3 URI for output (e.g., s3://bucket/prefix). + training_method: The training method (SFT_LORA or SFT_FULL). + region: AWS region (used if s3_client is not provided). + s3_client: Optional boto3 S3 client. + + Returns: + S3 URI of the uploaded CSV, or None if no metrics could be extracted. + + Raises: + ValueError: If output_s3_path is not provided. + """ + if not output_s3_path: + raise ValueError("output_s3_path is required but was not provided.") + + if not log_events: + logger.warning( + "No CloudWatch logs found for job %s. Cannot generate training metrics CSV.", + job_id, + ) + return None + + metrics_df = get_metrics( + platform=Platform.SMHP, + training_method=training_method, + logs=log_events, + ) + if metrics_df.empty: + logger.warning( + "No training metrics could be extracted from logs for job %s. " + "Cannot generate training metrics CSV.", + job_id, + ) + return None + + # Extract epoch boundaries and assign epochs to steps + epoch_boundaries = _extract_epoch_boundaries(log_events) + metrics_df = _assign_epochs_to_steps(metrics_df, epoch_boundaries, log_events) + + # Construct final DataFrame with expected column names + final_df = pandas.DataFrame( + { + "step_number": metrics_df["global_step"].astype(int), + "epoch_number": metrics_df["epoch_number"].astype(int), + "training_loss": metrics_df["training_loss"], + } + ) + + # Write CSV to a temp file + csv_filename = "step_wise_training_metrics.csv" + tmp_file = None + try: + tmp_file = tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) + final_df.to_csv(tmp_file.name, index=False) + tmp_file.close() + + # Parse S3 path for upload + bucket, base_key = parse_s3_uri(output_s3_path) + base_key = base_key.strip("/") + + # Construct the full S3 key + s3_key = f"{base_key}/{job_id}/{csv_filename}" if base_key else f"{job_id}/{csv_filename}" + + # Upload to S3 + if s3_client is None: + s3_client = boto3.client("s3", region_name=region) + + s3_client.upload_file(tmp_file.name, bucket, s3_key) + + # Construct and return the S3 URI + s3_uri = f"s3://{bucket}/{s3_key}" + logger.info("Training metrics CSV uploaded to %s", s3_uri) + return s3_uri + + finally: + # Clean up temp file + if tmp_file and os.path.exists(tmp_file.name): + os.unlink(tmp_file.name) diff --git a/src/amzn_nova_forge/util/recipe.py b/src/amzn_nova_forge/util/recipe.py index ac55ade..fa7fae6 100644 --- a/src/amzn_nova_forge/util/recipe.py +++ b/src/amzn_nova_forge/util/recipe.py @@ -756,14 +756,14 @@ def load_recipe_templates( and rft_multiturn_infra ) - if is_rft_multiturn_training or is_rft_multiturn_eval: + if (is_rft_multiturn_training or is_rft_multiturn_eval) and platform != Platform.SMTJServerless: if not rft_multiturn_infra: raise ValueError( f"rft_multiturn_infra is required for RFT multiturn {'training' if is_rft_multiturn_training else 'evaluation'}" ) if platform != Platform.SMHP: raise ValueError( - f"RFT multiturn {'training' if is_rft_multiturn_training else 'eval'} is only supported on HyperPod (SMHP)" + f"RFT multiturn {'training' if is_rft_multiturn_training else 'eval'} is only supported on HyperPod (SMHP) or SMTJServerless" ) # RFT multiturn requires instance_type (only supported on SMHP) diff --git a/src/amzn_nova_forge/util/sagemaker.py b/src/amzn_nova_forge/util/sagemaker.py index a5a4135..0d9e91c 100644 --- a/src/amzn_nova_forge/util/sagemaker.py +++ b/src/amzn_nova_forge/util/sagemaker.py @@ -18,7 +18,8 @@ import json import re import time -from datetime import datetime, timezone +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone from typing import Any, Dict, List, Optional import boto3 @@ -31,12 +32,21 @@ SUPPORTED_SMI_CONFIGS, _escrow_tag_value, ) -from amzn_nova_forge.core.enums import DeploymentMode, Model, Platform +from amzn_nova_forge.core.enums import DeploymentMode, DeployPlatform, Model, Platform from amzn_nova_forge.core.result.inference_result import ( SingleInferenceResult, ) from amzn_nova_forge.core.runtime import RuntimeManager -from amzn_nova_forge.core.types import ModelArtifacts +from amzn_nova_forge.core.types import ( + _IC_MIN_COMPUTE_REQUIREMENTS, + DeploymentResult, + EndpointInfo, + InferenceComponentConfig, + ModelArtifacts, + validate_inference_component_resources, +) +from amzn_nova_forge.telemetry.constants import Feature +from amzn_nova_forge.telemetry.telemetry_logging import _telemetry_emitter from amzn_nova_forge.validation.endpoint_validator import ( validate_s3_uri_prefix, ) @@ -204,8 +214,12 @@ def _get_sagemaker_inference_image(region: str) -> str: return f"{REGION_TO_ESCROW_ACCOUNT_MAPPING[region]}.dkr.ecr.{region}.amazonaws.com/nova-inference-repo:SM-Inference-latest" +@_telemetry_emitter(Feature.TRAINING, "get_model_artifacts") def get_model_artifacts( - job_name: str, infra: RuntimeManager, output_s3_path: str, region: Optional[str] = None + job_name: str, + infra: RuntimeManager, + output_s3_path: Optional[str] = None, + region: Optional[str] = None, ) -> ModelArtifacts: """ Retrieve model artifacts for a job @@ -213,7 +227,7 @@ def get_model_artifacts( Args: job_name: Name of the job infra: Infrastructure of the job - output_s3_path: Output S3 path of the job (only necessary for HyperPod) + output_s3_path: Output S3 path of the job (required for HyperPod) Returns: ModelArtifacts: Model artifact S3 paths @@ -221,10 +235,31 @@ def get_model_artifacts( Raises: Exception: If unable to obtain job artifact information """ - sagemaker_client = boto3.client("sagemaker", region_name=region) + # Use the infra's sagemaker client if available — it may be configured with a custom + # endpoint (e.g. gamma) that the job was submitted to. + sagemaker_client = getattr(infra, "sagemaker_client", None) or boto3.client( + "sagemaker", region_name=region + ) if infra.platform in (Platform.SMTJ, Platform.SMTJServerless): - response = sagemaker_client.describe_training_job(TrainingJobName=job_name) + try: + response = sagemaker_client.describe_training_job(TrainingJobName=job_name) + except ClientError as e: + error_code = e.response["Error"]["Code"] + if error_code not in ("ValidationException", "ResourceNotFound"): + raise + if infra.platform != Platform.SMTJServerless: + raise + # MTRL jobs are AgentRFT jobs, not standard SageMaker training jobs + from sagemaker.train.agent_rft_job import AgentRFTJob + + session = boto3.Session(region_name=region) + rft_job = AgentRFTJob.get(job_name, session=session) + return ModelArtifacts( + checkpoint_s3_path=None, + output_s3_path=rft_job.s3_output_path, + output_model_arn=rft_job.output_model_package_arn, + ) # Serverless jobs populate OutputModelPackageArn; use it to get the checkpoint S3 URI # SMTJ jobs use CheckpointConfig.S3Uri checkpoint_s3_path = None @@ -258,6 +293,8 @@ def get_model_artifacts( output_model_arn=model_package_arn, ) elif infra.platform == Platform.SMHP: + if not output_s3_path: + raise ValueError("output_s3_path is required for HyperPod jobs") try: cluster_name = infra.cluster_name # type: ignore[attr-defined] except AttributeError: @@ -427,33 +464,50 @@ def _validate_sagemaker_instance_type_for_model_deployment( def create_sagemaker_model( region: str, model_name: str, - model_s3_location: str, sagemaker_execution_role_arn: str, sagemaker_client: BaseClient, - environment: Dict[str, Any] = {}, + model_s3_location: Optional[str] = None, + environment: Dict[str, Any] = {}, # noqa: B006 - never mutated network_isolation: bool = True, deployment_mode: DeploymentMode = DeploymentMode.FAIL_IF_EXISTS, tags: Optional[List[Dict[str, str]]] = None, + model_package_name: Optional[str] = None, ) -> str: """Create a SageMaker model resource. + Supports two mutually exclusive model source modes: + 1. S3 data source: provide ``model_s3_location`` + 2. Model package: provide ``model_package_name`` (name or ARN of a SageMaker Model Package). + Args: region: AWS region model_name: Name of the SageMaker model - model_s3_location: S3 URI where model artifacts are stored + model_s3_location: S3 URI where model artifacts are stored. + Required when model_package_name is not provided. sagemaker_execution_role_arn: IAM role ARN for SageMaker execution sagemaker_client: SageMaker client environment: Environment variables for the model network_isolation: Enable network isolation deployment_mode: How to handle existing model + tags: Optional resource tags + model_package_name: Name or ARN of a SageMaker Model Package to use as the model source. Returns: str: Model ARN Raises: + ValueError: If neither model_s3_location nor model_package_name is provided, or if both are provided. Exception: If model already exists (FAIL_IF_EXISTS) or creation fails """ - validate_s3_uri_prefix(s3_uri=model_s3_location) + if model_package_name and model_s3_location: + raise ValueError( + "Only one of model_s3_location or model_package_name may be provided, not both." + ) + if not model_package_name and not model_s3_location: + raise ValueError("Either model_s3_location or model_package_name must be provided.") + + if model_s3_location: + validate_s3_uri_prefix(s3_uri=model_s3_location) if deployment_mode in [ DeploymentMode.FAIL_IF_EXISTS, @@ -467,19 +521,27 @@ def create_sagemaker_model( raise logger.info(f"Creating model: {model_name}...") + + # Build PrimaryContainer based on model source + primary_container: Dict[str, Any] = { + "Image": _get_sagemaker_inference_image(region), + "Environment": environment, + } + + if model_package_name: + primary_container["ModelPackageName"] = model_package_name + else: + primary_container["ModelDataSource"] = { + "S3DataSource": { + "S3Uri": model_s3_location, + "S3DataType": "S3Prefix", + "CompressionType": "None", + } + } + create_kwargs = { "ModelName": model_name, - "PrimaryContainer": { - "Image": _get_sagemaker_inference_image(region), - "ModelDataSource": { - "S3DataSource": { - "S3Uri": model_s3_location, - "S3DataType": "S3Prefix", - "CompressionType": "None", - } - }, - "Environment": environment, - }, + "PrimaryContainer": primary_container, "ExecutionRoleArn": sagemaker_execution_role_arn, "EnableNetworkIsolation": network_isolation, } @@ -537,17 +599,27 @@ def create_sagemaker_endpoint( sagemaker_client: BaseClient, initial_instance_count: int = 1, deployment_mode: DeploymentMode = DeploymentMode.FAIL_IF_EXISTS, + inference_component_configs: List[InferenceComponentConfig] = [], # noqa: B006 - never mutated + execution_role_arn: Optional[str] = None, ) -> str: """Create a SageMaker endpoint config and endpoint. + When inference_component_configs is provided, creates an IC-compatible endpoint + (no ModelName in ProductionVariants, uses RoutingConfig) and then creates the + inference component(s) after the endpoint reaches InService. + Args: - model_name: Name of the existing SageMaker model + model_name: Name of the existing SageMaker model (used in ProductionVariants for standard + endpoints, or in InferenceComponent Specification when inference_component_configs is provided) endpoint_config_name: Name for the endpoint configuration endpoint_name: Name for the endpoint instance_type: EC2 instance type sagemaker_client: SageMaker client initial_instance_count: Number of instances deployment_mode: How to handle existing resources + inference_component_configs: List of configs for creating inference components. + When provided, the endpoint is created in IC-compatible mode. + execution_role_arn: IAM execution role ARN. Required when inference_component_configs is provided. Returns: str: Endpoint ARN @@ -573,18 +645,47 @@ def create_sagemaker_endpoint( if e.response["Error"]["Code"] != "ValidationException": raise + # Build endpoint config based on whether we're using inference components + if inference_component_configs: + if not execution_role_arn: + raise ValueError( + "execution_role_arn is required when creating an inference component endpoint." + ) + # IC-enabled endpoints support a single production variant; all ICs must + # target the same variant. + variant_names = {ic.variant_name for ic in inference_component_configs} + if len(variant_names) > 1: + raise ValueError( + f"All inference component configs must use the same variant_name when " + f"deploying to a single endpoint, but found: {variant_names}" + ) + variant_name = inference_component_configs[0].variant_name + production_variant = { + "VariantName": variant_name, + "InstanceType": instance_type, + "InitialInstanceCount": initial_instance_count, + "RoutingConfig": {"RoutingStrategy": "LEAST_OUTSTANDING_REQUESTS"}, + } + config_kwargs: Dict[str, Any] = { + "EndpointConfigName": endpoint_config_name, + "ExecutionRoleArn": execution_role_arn, + "ProductionVariants": [production_variant], + "Tags": [{"Key": "sagemaker:nova-inference-component", "Value": "true"}], + } + else: + production_variant = { + "VariantName": "primary", + "ModelName": model_name, + "InitialInstanceCount": initial_instance_count, + "InstanceType": instance_type, + } + config_kwargs = { + "EndpointConfigName": endpoint_config_name, + "ProductionVariants": [production_variant], + } + logger.info(f"Creating endpoint configuration: {endpoint_config_name}...") - config_response = sagemaker_client.create_endpoint_config( - EndpointConfigName=endpoint_config_name, - ProductionVariants=[ - { - "VariantName": "primary", - "ModelName": model_name, - "InitialInstanceCount": initial_instance_count, - "InstanceType": instance_type, - } - ], - ) + config_response = sagemaker_client.create_endpoint_config(**config_kwargs) logger.info( f"Endpoint configuration created successfully: {config_response['EndpointConfigArn']}" ) @@ -600,6 +701,29 @@ def create_sagemaker_endpoint( except Exception as e: raise Exception(f"Failed to create deployment {endpoint_name}: {e}") + # If inference component configs are provided, create ICs after endpoint is InService + if inference_component_configs: + for ic_config in inference_component_configs: + logger.info(f"Creating inference component: {ic_config.inference_component_name}...") + ic_response = sagemaker_client.create_inference_component( + InferenceComponentName=ic_config.inference_component_name, + EndpointName=endpoint_name, + VariantName=ic_config.variant_name, + Specification={ + "ModelName": model_name, + "ComputeResourceRequirements": { + "NumberOfCpuCoresRequired": ic_config.num_cpus, + "NumberOfAcceleratorDevicesRequired": ic_config.num_accelerators, + "MinMemoryRequiredInMb": ic_config.min_memory_in_mb, + }, + }, + RuntimeConfig={ + "CopyCount": ic_config.copy_count, + }, + ) + ic_arn = ic_response["InferenceComponentArn"] + logger.info(f"Triggered inference component creation: {ic_arn}") + return endpoint_response["EndpointArn"] @@ -607,6 +731,7 @@ def invoke_sagemaker_inference( request_body: Dict[str, Any], endpoint_name: str, sagemaker_client: BaseClient, + inference_component_name: Optional[str] = None, ) -> SingleInferenceResult: """ Invoke Sagemaker inference and return result @@ -615,6 +740,8 @@ def invoke_sagemaker_inference( request_body (Dict[str, Any]): The payload to send to the inference endpoint. endpoint_name (str): Name of the SageMaker inference endpoint. sagemaker_client (BaseClient): Sagemaker client + inference_component_name (Optional[str]): Optional inference component to target. + When provided, adds InferenceComponentName to the API call. Returns: - Generator[str, None, None] for streaming responses @@ -629,9 +756,15 @@ def invoke_sagemaker_inference( logger.info(f"Invoking endpoint ({'streaming' if is_streaming else 'non-streaming'})...") if is_streaming: - response = sagemaker_client.invoke_endpoint_with_response_stream( - EndpointName=endpoint_name, ContentType="application/json", Body=body - ) + stream_kwargs: Dict[str, Any] = { + "EndpointName": endpoint_name, + "ContentType": "application/json", + "Body": body, + } + if inference_component_name is not None: + stream_kwargs["InferenceComponentName"] = inference_component_name + + response = sagemaker_client.invoke_endpoint_with_response_stream(**stream_kwargs) event_stream = response["Body"] @@ -651,12 +784,16 @@ def stream_generator(): nonstreaming_response=None, ) else: - response = sagemaker_client.invoke_endpoint( - EndpointName=endpoint_name, - ContentType="application/json", - Accept="application/json", - Body=body, - ) + invoke_kwargs: Dict[str, Any] = { + "EndpointName": endpoint_name, + "ContentType": "application/json", + "Accept": "application/json", + "Body": body, + } + if inference_component_name is not None: + invoke_kwargs["InferenceComponentName"] = inference_component_name + + response = sagemaker_client.invoke_endpoint(**invoke_kwargs) body_content = json.loads(response["Body"].read().decode("utf-8")) @@ -670,3 +807,241 @@ def stream_generator(): except Exception as e: raise Exception(f"Error invoking endpoint {endpoint_name}: {str(e)}") + + +def create_inference_component( + inference_component_name: str, + endpoint_name: str, + variant_name: str, + model_name: str, + num_cpus: int, + num_accelerators: int, + min_memory_in_mb: int, + copy_count: int, + sagemaker_client: BaseClient, + region: Optional[str] = None, +) -> DeploymentResult: + """Create a SageMaker inference component on an existing endpoint. + + Validates the target endpoint is InService and IC-compatible, then calls + CreateInferenceComponent using the specified SageMaker model name and returns + a DeploymentResult immediately without waiting for the component to become active. + + Args: + inference_component_name: Unique name for the inference component. + endpoint_name: Name of the existing SageMaker endpoint (must be InService). + variant_name: Production variant name on the endpoint. + model_name: Name of the existing SageMaker model to use. + num_cpus: Number of vCPUs to allocate. + num_accelerators: Number of accelerators (GPUs) to allocate. + min_memory_in_mb: Minimum memory in MB to allocate. + copy_count: Number of model copies to deploy. + sagemaker_client: Boto3 SageMaker client. + region: Optional AWS region name. When provided, stored in the returned EndpointInfo. + + Returns: + DeploymentResult: Contains endpoint info with the inference component ARN + as the URI. Use .status to check current deployment state. + + Raises: + ValueError: If required parameters are missing or invalid. + Exception: If the endpoint does not exist, is not InService, + or the CreateInferenceComponent API call fails. + """ + # Validate required string parameters + required_str_params = { + "inference_component_name": inference_component_name, + "endpoint_name": endpoint_name, + "variant_name": variant_name, + "model_name": model_name, + } + for param_name, param_value in required_str_params.items(): + if not param_value: + raise ValueError( + f"Parameter '{param_name}' is required for inference component creation" + ) + + # Validate endpoint exists and is InService + try: + endpoint_response = sagemaker_client.describe_endpoint(EndpointName=endpoint_name) + except ClientError as e: + error_code = e.response["Error"]["Code"] + if error_code in ("ValidationException", "ResourceNotFound"): + raise Exception(f"Endpoint '{endpoint_name}' not found") from e + raise + + endpoint_status = endpoint_response["EndpointStatus"] + if endpoint_status != "InService": + raise Exception( + f"Endpoint '{endpoint_name}' is not InService (current status: {endpoint_status})" + ) + + # Validate endpoint config is IC-compatible (no ModelName, has RoutingConfig) + endpoint_config_name = endpoint_response.get("EndpointConfigName") + if endpoint_config_name: + try: + config_response = sagemaker_client.describe_endpoint_config( + EndpointConfigName=endpoint_config_name + ) + variants = config_response.get("ProductionVariants", []) + for variant in variants: + if "ModelName" in variant: + raise Exception( + f"Endpoint '{endpoint_name}' is not configured for inference components. " + f"The endpoint config '{endpoint_config_name}' has 'ModelName' set on " + f"variant '{variant.get('VariantName')}'. Inference components require an " + f"endpoint config without ModelName and with RoutingConfig set to " + f"LEAST_OUTSTANDING_REQUESTS." + ) + routing_config = variant.get("RoutingConfig", {}) + routing_strategy = routing_config.get("RoutingStrategy") + if routing_strategy != "LEAST_OUTSTANDING_REQUESTS": + raise Exception( + f"Endpoint '{endpoint_name}' is not configured for inference components. " + f"The endpoint config '{endpoint_config_name}' variant " + f"'{variant.get('VariantName')}' has RoutingStrategy " + f"'{routing_strategy}' but inference components require " + f"'LEAST_OUTSTANDING_REQUESTS'." + ) + except ClientError as e: + logger.warning( + f"Could not validate endpoint config '{endpoint_config_name}': {e}. " + f"Proceeding with inference component creation." + ) + + create_request = { + "InferenceComponentName": inference_component_name, + "EndpointName": endpoint_name, + "VariantName": variant_name, + "Specification": { + "ModelName": model_name, + "ComputeResourceRequirements": { + "NumberOfCpuCoresRequired": num_cpus, + "NumberOfAcceleratorDevicesRequired": num_accelerators, + "MinMemoryRequiredInMb": min_memory_in_mb, + }, + }, + "RuntimeConfig": { + "CopyCount": copy_count, + }, + } + + try: + response = sagemaker_client.create_inference_component(**create_request) + except ClientError as e: + error_code = e.response["Error"]["Code"] + if error_code == "ResourceInUse": + raise Exception(f"Inference component '{inference_component_name}' already exists") + raise Exception(f"Failed to create inference component '{inference_component_name}': {e}") + + inference_component_arn = response["InferenceComponentArn"] + + return DeploymentResult( + endpoint=EndpointInfo( + platform=DeployPlatform.SAGEMAKER, + endpoint_name=endpoint_name, + uri=inference_component_arn, + model_artifact_path=model_name, + region=region, + ), + created_at=datetime.now(timezone.utc), + ) + + +def monitor_inference_component( + inference_component_name: str, + sagemaker_client: BaseClient, +) -> str: + """Monitor an inference component until it reaches a terminal state. + + Polls DescribeInferenceComponent every 60 seconds until the status is + InService (success) or Failed (raises exception). Times out after 1 hour + if the component remains in a non-terminal state. + + Args: + inference_component_name: Name of the inference component to monitor. + sagemaker_client: Boto3 SageMaker client. + + Returns: + str: Final status ("InService"). + + Raises: + Exception: If the component reaches Failed status, times out, or the API call errors. + """ + MAX_WAIT_SECONDS = 3600 # 1 hour + POLL_INTERVAL_SECONDS = 60 + start_time = datetime.now(timezone.utc) + + while True: + try: + response = sagemaker_client.describe_inference_component( + InferenceComponentName=inference_component_name + ) + except ClientError as e: + raise Exception( + f"Error describing inference component '{inference_component_name}': " + f"{e.response['Error']['Code']} - {e.response['Error']['Message']}" + ) + + status = response["InferenceComponentStatus"] + elapsed = (datetime.now(timezone.utc) - start_time).total_seconds() + + logger.info( + f"Inference component '{inference_component_name}' status: {status} | " + f"Elapsed: {timedelta(seconds=int(elapsed))}" + ) + + if status == "InService": + logger.info( + f"Inference component '{inference_component_name}' is now InService. " + f"Total time elapsed: {timedelta(seconds=int(elapsed))}" + ) + return "InService" + elif status == "Failed": + raise Exception(f"Inference component '{inference_component_name}' failed to deploy") + + if elapsed > MAX_WAIT_SECONDS: + raise Exception( + f"Inference component '{inference_component_name}' did not reach a terminal state " + f"within {MAX_WAIT_SECONDS} seconds (last status: {status})" + ) + + time.sleep(POLL_INTERVAL_SECONDS) + + +def check_sagemaker_deployment_status( + deployment_arn: str, region: Optional[str] = None +) -> Optional[str]: + """Check the current status of a SageMaker deployment (endpoint or inference component). + + For inference component ARNs (containing "inference-component/"), calls + DescribeInferenceComponent and returns InferenceComponentStatus. + For endpoint ARNs, calls DescribeEndpoint and returns EndpointStatus. + + Args: + deployment_arn: The ARN of the SageMaker deployment to check. + region: Optional AWS region for the client. + + Returns: + str: Current status of the deployment. + + Raises: + Exception: If unable to check deployment status. + """ + sagemaker_client = boto3.client("sagemaker", region_name=region) + try: + if "inference-component/" in deployment_arn: + component_name = deployment_arn.split("/")[-1] + response = sagemaker_client.describe_inference_component( + InferenceComponentName=component_name + ) + return response["InferenceComponentStatus"] + else: + endpoint_name = deployment_arn.split("/")[-1] + response = sagemaker_client.describe_endpoint(EndpointName=endpoint_name) + return response["EndpointStatus"] + except Exception as e: + raise Exception(f"Failed to check deployment status: {e}") + + +DeploymentResult._register_sagemaker_status_checker(check_sagemaker_deployment_status) diff --git a/src/amzn_nova_forge/util/subprocess_utils.py b/src/amzn_nova_forge/util/subprocess_utils.py index 8b5ac7b..4313797 100644 --- a/src/amzn_nova_forge/util/subprocess_utils.py +++ b/src/amzn_nova_forge/util/subprocess_utils.py @@ -28,6 +28,7 @@ r"DeprecationWarning:", r"FutureWarning:", r"InsecureRequestWarning:", + r"RequestsDependencyWarning:", r"ResourceWarning:", r"UserWarning:", r"urllib3\.[a-zA-Z_]+Warning:", diff --git a/src/amzn_nova_forge/validation/validator.py b/src/amzn_nova_forge/validation/validator.py index b47d6ec..aa21f41 100644 --- a/src/amzn_nova_forge/validation/validator.py +++ b/src/amzn_nova_forge/validation/validator.py @@ -859,6 +859,12 @@ def validate_rft( elif rft_lambda_arn is not None and not is_lambda_arn(rft_lambda_arn): if is_hub_content_arn(rft_lambda_arn) and platform == Platform.SMTJServerless: pass # valid — hub-content ARN used as EvaluatorArn on SMTJServerless + elif ( + platform == Platform.SMTJServerless + and method == TrainingMethod.RFT_MULTITURN_LORA + and rft_lambda_arn.startswith("arn:") + ): + pass # valid — AgentCore runtime ARN for serverless multiturn else: errors.append( "'rft_lambda_arn' must be a valid Lambda function ARN" @@ -962,6 +968,8 @@ def get_recipe_value(data: Dict[str, Any], key_to_find: str) -> Any: if method in [TrainingMethod.RFT_LORA, TrainingMethod.RFT_FULL]: validate_rft(rft_lambda_arn=rft_lambda_arn, rft_lambda_source=rft_lambda_source) + elif method == TrainingMethod.RFT_MULTITURN_LORA and platform == Platform.SMTJServerless: + pass # MTRL serverless uses agent_core_arn — no Lambda required elif method in [TrainingMethod.EVALUATION]: assert eval_task is not None validate_eval( @@ -972,10 +980,20 @@ def get_recipe_value(data: Dict[str, Any], key_to_find: str) -> Any: rl_env_config=rl_env_config, ) + # MTRL serverless jobs use the SDK's own validation (FineTuningOptions). + # The JumpStart overrides template has constraints for single-turn RFT that don't + # apply to MTRL (e.g., min global_batch_size=16, required reward_lambda_arn). + _is_mtrl_serverless = ( + method == TrainingMethod.RFT_MULTITURN_LORA and platform == Platform.SMTJServerless + ) + for key, override_metadata in overrides_template.items(): # Skip HyperPod specific key since it's not actually present within recipes if key == "namespace": continue + # Skip min/max/enum/type validation for MTRL serverless — validates these + if _is_mtrl_serverless and key not in ("output_s3_path", "data_s3_path", "model_type"): + continue # TODO: Need to figure out what this actually refers to within the recipe. Until then, we won't validate it. elif key == "max_context_length": continue diff --git a/tests/unit/dataset/test_rft_multiturn_serverless.py b/tests/unit/dataset/test_rft_multiturn_serverless.py new file mode 100644 index 0000000..bb75b67 --- /dev/null +++ b/tests/unit/dataset/test_rft_multiturn_serverless.py @@ -0,0 +1,252 @@ +# 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. +""" +Test cases for RFT Multiturn Serverless dataset transformation and validation. + +Tests the platform-specific behavior: +- Platform.SMTJServerless: flat {"prompt": "..."} format +- Platform.SMHP: nested {"id": "...", "metadata": {"prompt": "..."}} format +- No platform: raises ValueError for RFT Multiturn +""" + +import json +import tempfile +from pathlib import Path + +import pytest + +from amzn_nova_forge.core.enums import Model, Platform, TrainingMethod +from amzn_nova_forge.dataset import JSONLDatasetLoader +from amzn_nova_forge.dataset.dataset_validator.rft_multiturn_dataset_validator import ( + RFTMultiturnServerlessSample, + RFTMultiturnServerlessValidator, +) + + +def _write_jsonl(data): + """Write data to a temp jsonl file and return the path.""" + f = tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) + for record in data: + f.write(json.dumps(record) + "\n") + f.close() + return f.name + + +def _read_jsonl(path): + """Read jsonl file and return list of dicts.""" + with open(path) as f: + return [json.loads(line) for line in f if line.strip()] + + +class TestServerlessTransform: + """Test transform with Platform.SMTJServerless produces flat format.""" + + def test_flat_input_produces_flat_output(self): + data = [ + {"id": "q1", "prompt": "What is 2+2?", "answer": "4"}, + {"id": "q2", "prompt": "What is 3+3?", "answer": "6"}, + ] + path = _write_jsonl(data) + try: + loader = JSONLDatasetLoader(id="id", prompt="prompt", answer="answer") + loader.load(path) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMTJServerless, + ) + output = tempfile.mktemp(suffix=".jsonl") + loader.save(output) + results = _read_jsonl(output) + + assert len(results) == 2 + assert results[0] == {"prompt": "What is 2+2?"} + assert results[1] == {"prompt": "What is 3+3?"} + finally: + Path(path).unlink() + + def test_json_string_prompt_preserved(self): + data = [ + { + "id": "q1", + "prompt": '{"instance": "Solve 3+3", "reward_spec": {"ground_truth": "6"}}', + }, + ] + path = _write_jsonl(data) + try: + loader = JSONLDatasetLoader(id="id", prompt="prompt") + loader.load(path) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMTJServerless, + ) + output = tempfile.mktemp(suffix=".jsonl") + loader.save(output) + results = _read_jsonl(output) + + assert len(results) == 1 + assert ( + results[0]["prompt"] + == '{"instance": "Solve 3+3", "reward_spec": {"ground_truth": "6"}}' + ) + finally: + Path(path).unlink() + + def test_nested_input_extracts_prompt(self): + """When input has metadata.prompt (nested format), serverless transform extracts it.""" + data = [ + {"metadata": {"prompt": "Hello world"}}, + ] + path = _write_jsonl(data) + try: + loader = JSONLDatasetLoader(prompt="prompt") + loader.load(path) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMTJServerless, + ) + output = tempfile.mktemp(suffix=".jsonl") + loader.save(output) + results = _read_jsonl(output) + + assert results[0] == {"prompt": "Hello world"} + finally: + Path(path).unlink() + + +class TestHyperPodTransform: + """Test transform with Platform.SMHP produces nested format.""" + + def test_produces_nested_format(self): + data = [ + {"id": "q1", "prompt": "What is 2+2?", "answer": "4"}, + ] + path = _write_jsonl(data) + try: + loader = JSONLDatasetLoader(id="id", prompt="prompt", answer="answer") + loader.load(path) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) + output = tempfile.mktemp(suffix=".jsonl") + loader.save(output) + results = _read_jsonl(output) + + assert results[0] == {"id": "q1", "metadata": {"prompt": "What is 2+2?", "answer": "4"}} + finally: + Path(path).unlink() + + +class TestPlatformRequired: + """Test that platform is required for RFT Multiturn transforms.""" + + def test_no_platform_raises_error(self): + data = [{"id": "q1", "prompt": "test"}] + path = _write_jsonl(data) + try: + loader = JSONLDatasetLoader(id="id", prompt="prompt") + loader.load(path) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + ) + with pytest.raises(ValueError, match="platform is required"): + loader.save(tempfile.mktemp(suffix=".jsonl")) + finally: + Path(path).unlink() + + def test_sft_does_not_require_platform(self): + """Non-MTRL methods should not require platform.""" + data = [ + { + "messages": [ + {"role": "user", "content": [{"text": "hi"}]}, + {"role": "assistant", "content": [{"text": "hello"}]}, + ] + }, + ] + path = _write_jsonl(data) + try: + loader = JSONLDatasetLoader() + loader.load(path) + # SFT should not raise even without platform + loader.transform( + method=TrainingMethod.SFT_LORA, + model=Model.NOVA_LITE_2, + ) + # Just verify no ValueError about platform — don't save (avoids STS calls) + loader.execute() + finally: + Path(path).unlink() + + +class TestServerlessValidator: + """Test RFTMultiturnServerlessValidator.""" + + def test_valid_sample(self): + sample = {"prompt": "What is 2+2?"} + RFTMultiturnServerlessSample.model_validate(sample) + + def test_empty_prompt_fails(self): + with pytest.raises(Exception, match="cannot be empty"): + RFTMultiturnServerlessSample.model_validate({"prompt": ""}) + + def test_missing_prompt_fails(self): + with pytest.raises(Exception): + RFTMultiturnServerlessSample.model_validate({"id": "q1"}) + + def test_extra_fields_allowed(self): + sample = {"prompt": "test", "extra_field": "value"} + validated = RFTMultiturnServerlessSample.model_validate(sample) + assert validated.prompt == "test" + + def test_validator_init_rejects_non_lite2(self): + with pytest.raises(ValueError, match="NOVA_LITE_2"): + RFTMultiturnServerlessValidator(Model.NOVA_LITE) + + def test_validator_init_accepts_lite2(self): + validator = RFTMultiturnServerlessValidator(Model.NOVA_LITE_2) + assert validator is not None + + +class TestServerlessValidateOperation: + """Test validate() with platform=Platform.SMTJServerless.""" + + def test_validate_serverless_format(self): + data = [ + {"prompt": "What is 2+2?"}, + {"prompt": "Explain gravity"}, + ] + path = _write_jsonl(data) + try: + loader = JSONLDatasetLoader(prompt="prompt") + loader.load(path) + loader.validate( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMTJServerless, + ) + loader.execute() + finally: + Path(path).unlink() + + def test_validate_rejects_empty_prompt(self): + """Validator rejects samples with empty/whitespace-only prompt.""" + sample = {"prompt": " "} + with pytest.raises(Exception, match="cannot be empty"): + RFTMultiturnServerlessSample.model_validate(sample) diff --git a/tests/unit/dataset/test_rft_multiturn_tool_calling.py b/tests/unit/dataset/test_rft_multiturn_tool_calling.py index bfe67d6..2909bef 100644 --- a/tests/unit/dataset/test_rft_multiturn_tool_calling.py +++ b/tests/unit/dataset/test_rft_multiturn_tool_calling.py @@ -26,7 +26,7 @@ import pytest -from amzn_nova_forge.core.enums import EvaluationTask, Model, TrainingMethod +from amzn_nova_forge.core.enums import EvaluationTask, Model, Platform, TrainingMethod from amzn_nova_forge.dataset import JSONLDatasetLoader @@ -75,7 +75,11 @@ def test_single_tool_call_valid(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) finally: Path(temp_path).unlink() @@ -138,7 +142,11 @@ def test_multiple_tool_calls_valid(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) finally: Path(temp_path).unlink() @@ -166,7 +174,11 @@ def test_prompt_as_simple_string_valid(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt", answer="answer") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) finally: Path(temp_path).unlink() @@ -213,7 +225,11 @@ def test_assistant_with_both_content_and_tool_calls_valid(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) finally: Path(temp_path).unlink() @@ -243,7 +259,11 @@ def test_tool_message_missing_tool_call_id_fails(self): with pytest.raises(ValueError): loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) finally: Path(temp_path).unlink() @@ -272,7 +292,11 @@ def test_assistant_without_content_or_tool_calls_fails(self): with pytest.raises(ValueError): loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) finally: Path(temp_path).unlink() @@ -295,7 +319,11 @@ def test_empty_string_prompt_fails(self): with pytest.raises(ValueError): loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) finally: Path(temp_path).unlink() @@ -333,7 +361,11 @@ def test_tool_call_invalid_type_fails(self): with pytest.raises(ValueError): loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) finally: Path(temp_path).unlink() @@ -359,7 +391,11 @@ def test_extra_fields_forbidden_in_metadata(self): with pytest.raises(ValueError, match="Extra inputs are not permitted"): loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) finally: Path(temp_path).unlink() @@ -385,7 +421,11 @@ def test_extra_fields_forbidden_at_top_level(self): with pytest.raises(ValueError, match="Extra inputs are not permitted"): loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) finally: Path(temp_path).unlink() @@ -416,6 +456,7 @@ def test_evaluation_method_with_rft_multiturn_eval_task(self): method=TrainingMethod.EVALUATION, eval_task=EvaluationTask.RFT_MULTITURN_EVAL, model=Model.NOVA_LITE_2, + platform=Platform.SMHP, ) # Validate with EVALUATION method and RFT_MULTITURN_EVAL task @@ -477,6 +518,7 @@ def test_evaluation_method_with_rft_multiturn_eval_openai_format(self): method=TrainingMethod.EVALUATION, eval_task=EvaluationTask.RFT_MULTITURN_EVAL, model=Model.NOVA_LITE_2, + platform=Platform.SMHP, ) # Validate with EVALUATION method and RFT_MULTITURN_EVAL task diff --git a/tests/unit/dataset/test_rft_multiturn_validation.py b/tests/unit/dataset/test_rft_multiturn_validation.py index 9eb19d5..8390ed4 100644 --- a/tests/unit/dataset/test_rft_multiturn_validation.py +++ b/tests/unit/dataset/test_rft_multiturn_validation.py @@ -24,7 +24,7 @@ import pytest -from amzn_nova_forge.core.enums import EvaluationTask, Model, TrainingMethod +from amzn_nova_forge.core.enums import EvaluationTask, Model, Platform, TrainingMethod from amzn_nova_forge.dataset import ( CSVDatasetLoader, JSONDatasetLoader, @@ -52,7 +52,11 @@ def test_flat_format_minimal_valid(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) @@ -88,7 +92,11 @@ def test_flat_format_all_fields_valid(self): id="id", prompt="prompt", answer="answer", task="task", info="info" ) loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) @@ -126,7 +134,11 @@ def test_nested_format_valid(self): try: loader = JSONLDatasetLoader(id="id", metadata="metadata") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) @@ -169,7 +181,11 @@ def test_openai_messages_format_valid(self): id="id", prompt="prompt", answer="answer", task="task", info="info" ) loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) @@ -190,7 +206,11 @@ def test_empty_info_dict_valid(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt", info="info") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) @@ -211,7 +231,11 @@ def test_json_format_valid(self): try: loader = JSONDatasetLoader(id="id", prompt="prompt", info="info") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) @@ -234,7 +258,11 @@ def test_csv_format_valid(self): id="id", prompt="prompt", answer="answer", task="task", info="info" ) loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) @@ -262,7 +290,11 @@ def test_info_as_invalid_json_string_fails(self): with pytest.raises(Exception): loader = JSONLDatasetLoader(id="id", prompt="prompt", info="info") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) @@ -284,7 +316,11 @@ def test_empty_prompt_fails(self): with pytest.raises(Exception): loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) @@ -306,7 +342,11 @@ def test_missing_prompt_fails(self): with pytest.raises(Exception): loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) @@ -333,7 +373,11 @@ def test_invalid_openai_role_fails(self): with pytest.raises(Exception): loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) @@ -360,7 +404,11 @@ def test_openai_missing_content_fails(self): with pytest.raises(Exception): loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) @@ -387,7 +435,11 @@ def test_empty_openai_content_fails(self): with pytest.raises(Exception): loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) @@ -412,7 +464,11 @@ def test_id_auto_generation_when_missing(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() @@ -440,7 +496,11 @@ def test_id_auto_generation_when_empty_string(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() @@ -467,7 +527,11 @@ def test_id_auto_generation_nested_format(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() @@ -496,7 +560,11 @@ def test_id_counter_persistence_across_records(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() @@ -525,7 +593,11 @@ def test_answer_type_conversion_integer(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt", answer="answer") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() @@ -550,7 +622,11 @@ def test_answer_type_conversion_float(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt", answer="answer") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() @@ -575,7 +651,11 @@ def test_task_type_conversion_integer(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt", task="task") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() @@ -605,7 +685,11 @@ def test_answer_and_task_type_conversion_combined(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt", answer="answer", task="task") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() @@ -632,7 +716,11 @@ def test_answer_empty_string_preserved(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt", answer="answer") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() @@ -660,7 +748,11 @@ def test_answer_with_special_characters(self): try: loader = JSONLDatasetLoader(id="id", prompt="prompt", answer="answer") loader.load(temp_path) - loader.transform(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) + loader.transform( + method=TrainingMethod.RFT_MULTITURN_LORA, + model=Model.NOVA_LITE_2, + platform=Platform.SMHP, + ) loader.execute() loader.execute() loader.validate(method=TrainingMethod.RFT_MULTITURN_LORA, model=Model.NOVA_LITE_2) diff --git a/tests/unit/deployer/test_forge_deployer.py b/tests/unit/deployer/test_forge_deployer.py index 398c1fc..8b94878 100644 --- a/tests/unit/deployer/test_forge_deployer.py +++ b/tests/unit/deployer/test_forge_deployer.py @@ -15,6 +15,7 @@ from datetime import datetime, timezone from unittest.mock import MagicMock, patch +from amzn_nova_forge.core.constants import ESCROW_URI_TAG_KEY from amzn_nova_forge.core.enums import ( DeploymentMode, DeployPlatform, @@ -803,13 +804,12 @@ def test_get_logs_with_job_result(self, mock_check_status, mock_validate_region) region="us-east-1", ) - @patch(f"{PATCH_PREFIX}.check_deployment_status", return_value="InProgress") + @patch(f"{PATCH_PREFIX}.check_sagemaker_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:123456789012:endpoint/ep") mock_check_status.assert_called_once_with( "arn:aws:sagemaker:us-east-1:123456789012:endpoint/ep", - DeployPlatform.SAGEMAKER, region="us-east-1", ) @@ -821,6 +821,402 @@ def test_get_logs_no_arn_raises_value_error(self, mock_validate_region): self.assertIn("endpoint_arn", str(ctx.exception)) +@patch(f"{PATCH_PREFIX}.validate_region") +class TestCreateInferenceComponent(unittest.TestCase): + """Tests for ForgeDeployer.create_inference_component().""" + + def _make_deployer(self, **kwargs): + defaults = dict(region="us-east-1", model=Model.NOVA_MICRO) + defaults.update(kwargs) + return ForgeDeployer(**defaults) + + @patch(f"{PATCH_PREFIX}.create_inference_component") + @patch("boto3.client") + def test_successful_create_inference_component( + self, mock_boto_client, mock_create_ic, mock_validate_region + ): + mock_sagemaker = MagicMock() + mock_boto_client.return_value = mock_sagemaker + + expected_result = DeploymentResult( + endpoint=EndpointInfo( + platform=DeployPlatform.SAGEMAKER, + endpoint_name="my-endpoint", + uri="arn:aws:sagemaker:us-east-1:123456789012:inference-component/my-ic", + model_artifact_path="my-model", + ), + created_at=datetime.now(timezone.utc), + ) + mock_create_ic.return_value = expected_result + + deployer = self._make_deployer() + result = deployer.create_inference_component( + inference_component_name="my-ic", + model_name="my-model", + num_cpus=4, + num_accelerators=1, + min_memory_in_mb=8192, + endpoint_name="my-endpoint", + variant_name="primary", + copy_count=1, + ) + + self.assertEqual(result, expected_result) + mock_boto_client.assert_called_once_with("sagemaker", region_name="us-east-1") + mock_create_ic.assert_called_once_with( + inference_component_name="my-ic", + endpoint_name="my-endpoint", + variant_name="primary", + model_name="my-model", + num_cpus=4, + num_accelerators=1, + min_memory_in_mb=8192, + copy_count=1, + sagemaker_client=mock_sagemaker, + region="us-east-1", + ) + + @patch(f"{PATCH_PREFIX}.create_inference_component") + @patch("boto3.client") + def test_create_inference_component_uses_default_variant_and_copy_count( + self, mock_boto_client, mock_create_ic, mock_validate_region + ): + mock_sagemaker = MagicMock() + mock_boto_client.return_value = mock_sagemaker + mock_create_ic.return_value = MagicMock(spec=DeploymentResult) + + deployer = self._make_deployer() + deployer.create_inference_component( + inference_component_name="my-ic", + model_name="my-model", + num_cpus=2, + num_accelerators=4, + min_memory_in_mb=16384, + endpoint_name="my-endpoint", + ) + + call_kwargs = mock_create_ic.call_args[1] + self.assertEqual(call_kwargs["variant_name"], "primary") + self.assertEqual(call_kwargs["copy_count"], 1) + + +@patch(f"{PATCH_PREFIX}.validate_region") +class TestMonitorInferenceComponent(unittest.TestCase): + def _make_deployer(self, **kwargs): + defaults = dict(region="us-east-1", model=Model.NOVA_MICRO) + defaults.update(kwargs) + return ForgeDeployer(**defaults) + + @patch(f"{PATCH_PREFIX}.monitor_inference_component") + @patch("boto3.client") + def test_successful_monitor_returns_status( + self, mock_boto_client, mock_monitor_ic, mock_validate_region + ): + mock_sagemaker = MagicMock() + mock_boto_client.return_value = mock_sagemaker + mock_monitor_ic.return_value = "InService" + + deployer = self._make_deployer() + result = deployer.monitor_inference_component( + inference_component_name="my-ic", + ) + + self.assertEqual(result, "InService") + mock_boto_client.assert_called_once_with("sagemaker", region_name="us-east-1") + mock_monitor_ic.assert_called_once_with( + inference_component_name="my-ic", + sagemaker_client=mock_sagemaker, + ) + + @patch(f"{PATCH_PREFIX}.monitor_inference_component") + @patch("boto3.client") + def test_monitor_propagates_exception_on_failure( + self, mock_boto_client, mock_monitor_ic, mock_validate_region + ): + mock_sagemaker = MagicMock() + mock_boto_client.return_value = mock_sagemaker + mock_monitor_ic.side_effect = Exception("Inference component 'my-ic' reached Failed status") + + deployer = self._make_deployer() + with self.assertRaises(Exception) as ctx: + deployer.monitor_inference_component(inference_component_name="my-ic") + self.assertIn("Failed status", str(ctx.exception)) + + +@patch(f"{PATCH_PREFIX}.find_sagemaker_model_by_tag", return_value=None) +@patch(f"{PATCH_PREFIX}.validate_region") +class TestDeploySageMakerWithInferenceComponentConfigs(unittest.TestCase): + """Tests for deploy() with inference_component_configs on SageMaker.""" + + def _make_deployer(self, **kwargs): + defaults = dict(region="us-east-1", model=Model.NOVA_MICRO) + defaults.update(kwargs) + return ForgeDeployer(**defaults) + + def _make_ic_config(self, **overrides): + from amzn_nova_forge.util.sagemaker import InferenceComponentConfig + + defaults = dict( + inference_component_name="my-ic", + num_cpus=15, + num_accelerators=4, + min_memory_in_mb=25000, + ) + defaults.update(overrides) + return InferenceComponentConfig(**defaults) + + @patch(f"{PATCH_PREFIX}.create_sagemaker_endpoint") + @patch(f"{PATCH_PREFIX}.create_sagemaker_model") + @patch(f"{PATCH_PREFIX}._validate_sagemaker_instance_type_for_model_deployment") + @patch(f"{PATCH_PREFIX}.validate_inference_component_resources") + @patch(f"{PATCH_PREFIX}.create_sagemaker_execution_role") + @patch("boto3.client") + def test_validate_inference_component_resources_called_for_each_config( + self, + mock_boto_client, + mock_create_role, + mock_validate_ic_resources, + mock_validate_instance, + mock_create_model, + mock_create_endpoint, + mock_validate_region, + mock_find_by_tag, + ): + mock_iam = MagicMock() + + def client_side_effect(service, **kwargs): + if service == "iam": + return mock_iam + return MagicMock() + + mock_boto_client.side_effect = client_side_effect + mock_create_role.return_value = { + "Role": {"Arn": "arn:aws:iam::123456789012:role/SageMakerRole"} + } + mock_create_model.return_value = ( + "arn:aws:sagemaker:us-east-1:123456789012:model/my-ep-model" + ) + mock_create_endpoint.return_value = ( + "arn:aws:sagemaker:us-east-1:123456789012:endpoint/my-ep" + ) + + ic_config_1 = self._make_ic_config(inference_component_name="ic-1") + ic_config_2 = self._make_ic_config(inference_component_name="ic-2") + + deployer = self._make_deployer() + deployer.deploy( + model_artifact_path="s3://bucket/model", + deploy_platform=DeployPlatform.SAGEMAKER, + inference_component_configs=[ic_config_1, ic_config_2], + ) + + self.assertEqual(mock_validate_ic_resources.call_count, 2) + mock_validate_ic_resources.assert_any_call(ic_config_1, Model.NOVA_MICRO) + mock_validate_ic_resources.assert_any_call(ic_config_2, Model.NOVA_MICRO) + + @patch(f"{PATCH_PREFIX}.create_sagemaker_endpoint") + @patch(f"{PATCH_PREFIX}.create_sagemaker_model") + @patch(f"{PATCH_PREFIX}._validate_sagemaker_instance_type_for_model_deployment") + @patch(f"{PATCH_PREFIX}.validate_inference_component_resources") + @patch(f"{PATCH_PREFIX}.create_sagemaker_execution_role") + @patch(f"{PATCH_PREFIX}.logger") + @patch("boto3.client") + def test_multi_ic_warning_logged_when_more_than_one_config( + self, + mock_boto_client, + mock_logger, + mock_create_role, + mock_validate_ic_resources, + mock_validate_instance, + mock_create_model, + mock_create_endpoint, + mock_validate_region, + mock_find_by_tag, + ): + mock_iam = MagicMock() + + def client_side_effect(service, **kwargs): + if service == "iam": + return mock_iam + return MagicMock() + + mock_boto_client.side_effect = client_side_effect + mock_create_role.return_value = { + "Role": {"Arn": "arn:aws:iam::123456789012:role/SageMakerRole"} + } + mock_create_model.return_value = ( + "arn:aws:sagemaker:us-east-1:123456789012:model/my-ep-model" + ) + mock_create_endpoint.return_value = ( + "arn:aws:sagemaker:us-east-1:123456789012:endpoint/my-ep" + ) + + ic_configs = [ + self._make_ic_config(inference_component_name="ic-1"), + self._make_ic_config(inference_component_name="ic-2"), + self._make_ic_config(inference_component_name="ic-3"), + ] + + deployer = self._make_deployer() + deployer.deploy( + model_artifact_path="s3://bucket/model", + deploy_platform=DeployPlatform.SAGEMAKER, + inference_component_configs=ic_configs, + ) + + mock_logger.warning.assert_called() + warning_msg = mock_logger.warning.call_args[0][0] + self.assertIn("3", warning_msg) + self.assertIn("inference components", warning_msg) + + @patch(f"{PATCH_PREFIX}.create_sagemaker_endpoint") + @patch(f"{PATCH_PREFIX}.create_sagemaker_model") + @patch(f"{PATCH_PREFIX}._validate_sagemaker_instance_type_for_model_deployment") + @patch(f"{PATCH_PREFIX}.validate_inference_component_resources") + @patch(f"{PATCH_PREFIX}.create_sagemaker_execution_role") + @patch(f"{PATCH_PREFIX}.logger") + @patch("boto3.client") + def test_no_multi_ic_warning_for_single_config( + self, + mock_boto_client, + mock_logger, + mock_create_role, + mock_validate_ic_resources, + mock_validate_instance, + mock_create_model, + mock_create_endpoint, + mock_validate_region, + mock_find_by_tag, + ): + mock_iam = MagicMock() + + def client_side_effect(service, **kwargs): + if service == "iam": + return mock_iam + return MagicMock() + + mock_boto_client.side_effect = client_side_effect + mock_create_role.return_value = { + "Role": {"Arn": "arn:aws:iam::123456789012:role/SageMakerRole"} + } + mock_create_model.return_value = ( + "arn:aws:sagemaker:us-east-1:123456789012:model/my-ep-model" + ) + mock_create_endpoint.return_value = ( + "arn:aws:sagemaker:us-east-1:123456789012:endpoint/my-ep" + ) + + ic_configs = [self._make_ic_config(inference_component_name="ic-1")] + + deployer = self._make_deployer() + deployer.deploy( + model_artifact_path="s3://bucket/model", + deploy_platform=DeployPlatform.SAGEMAKER, + inference_component_configs=ic_configs, + ) + + # No warning should be logged about multiple inference components + for call in mock_logger.warning.call_args_list: + self.assertNotIn("inference components", call[0][0]) + + @patch(f"{PATCH_PREFIX}.create_sagemaker_endpoint") + @patch(f"{PATCH_PREFIX}.create_sagemaker_model") + @patch(f"{PATCH_PREFIX}._validate_sagemaker_instance_type_for_model_deployment") + @patch(f"{PATCH_PREFIX}.validate_inference_component_resources") + @patch(f"{PATCH_PREFIX}.create_sagemaker_execution_role") + @patch("boto3.client") + def test_create_sagemaker_endpoint_receives_ic_configs_and_role_arn( + self, + mock_boto_client, + mock_create_role, + mock_validate_ic_resources, + mock_validate_instance, + mock_create_model, + mock_create_endpoint, + mock_validate_region, + mock_find_by_tag, + ): + mock_iam = MagicMock() + + def client_side_effect(service, **kwargs): + if service == "iam": + return mock_iam + return MagicMock() + + mock_boto_client.side_effect = client_side_effect + mock_create_role.return_value = { + "Role": {"Arn": "arn:aws:iam::123456789012:role/SageMakerRole"} + } + mock_create_model.return_value = ( + "arn:aws:sagemaker:us-east-1:123456789012:model/my-ep-model" + ) + mock_create_endpoint.return_value = ( + "arn:aws:sagemaker:us-east-1:123456789012:endpoint/my-ep" + ) + + ic_configs = [self._make_ic_config(inference_component_name="ic-1")] + + deployer = self._make_deployer() + deployer.deploy( + model_artifact_path="s3://bucket/model", + deploy_platform=DeployPlatform.SAGEMAKER, + inference_component_configs=ic_configs, + ) + + mock_create_endpoint.assert_called_once() + call_kwargs = mock_create_endpoint.call_args[1] + self.assertEqual(call_kwargs["inference_component_configs"], ic_configs) + self.assertEqual( + call_kwargs["execution_role_arn"], + "arn:aws:iam::123456789012:role/SageMakerRole", + ) + + @patch(f"{PATCH_PREFIX}.create_sagemaker_endpoint") + @patch(f"{PATCH_PREFIX}.create_sagemaker_model") + @patch(f"{PATCH_PREFIX}._validate_sagemaker_instance_type_for_model_deployment") + @patch(f"{PATCH_PREFIX}.create_sagemaker_execution_role") + @patch("boto3.client") + def test_create_sagemaker_endpoint_no_ic_kwargs_when_configs_none( + self, + mock_boto_client, + mock_create_role, + mock_validate_instance, + mock_create_model, + mock_create_endpoint, + mock_validate_region, + mock_find_by_tag, + ): + mock_iam = MagicMock() + + def client_side_effect(service, **kwargs): + if service == "iam": + return mock_iam + return MagicMock() + + mock_boto_client.side_effect = client_side_effect + mock_create_role.return_value = { + "Role": {"Arn": "arn:aws:iam::123456789012:role/SageMakerRole"} + } + mock_create_model.return_value = ( + "arn:aws:sagemaker:us-east-1:123456789012:model/my-ep-model" + ) + mock_create_endpoint.return_value = ( + "arn:aws:sagemaker:us-east-1:123456789012:endpoint/my-ep" + ) + + deployer = self._make_deployer() + deployer.deploy( + model_artifact_path="s3://bucket/model", + deploy_platform=DeployPlatform.SAGEMAKER, + inference_component_configs=[], + ) + + mock_create_endpoint.assert_called_once() + call_kwargs = mock_create_endpoint.call_args[1] + self.assertNotIn("inference_component_configs", call_kwargs) + self.assertNotIn("execution_role_arn", call_kwargs) + + class TestDeploymentResultRegion(unittest.TestCase): """Tests for region propagation through DeploymentResult.status.""" @@ -839,5 +1235,326 @@ def test_deployment_result_status_passes_region(self): mock_checker.assert_called_once_with("arn:test", DeployPlatform.BEDROCK_OD, "eu-west-1") +@patch(f"{PATCH_PREFIX}.find_sagemaker_model_by_tag", return_value=None) +@patch(f"{PATCH_PREFIX}.validate_region") +class TestDeploySageMakerModelPackageArn(unittest.TestCase): + """Tests for deploy() auto-detecting model package ARNs for SageMaker.""" + + def _make_deployer(self, **kwargs): + defaults = dict(region="us-east-1", model=Model.NOVA_MICRO) + defaults.update(kwargs) + return ForgeDeployer(**defaults) + + @patch(f"{PATCH_PREFIX}.create_sagemaker_endpoint") + @patch(f"{PATCH_PREFIX}.create_sagemaker_model") + @patch(f"{PATCH_PREFIX}._validate_sagemaker_instance_type_for_model_deployment") + @patch(f"{PATCH_PREFIX}.create_sagemaker_execution_role") + @patch("boto3.client") + def test_model_package_arn_sets_model_package_name( + self, + mock_boto_client, + mock_create_role, + mock_validate_instance, + mock_create_model, + mock_create_endpoint, + mock_validate_region, + mock_find_by_tag, + ): + """A model package ARN should be auto-detected and passed as model_package_name.""" + mock_iam = MagicMock() + + def client_side_effect(service, **kwargs): + if service == "iam": + return mock_iam + return MagicMock() + + mock_boto_client.side_effect = client_side_effect + + mock_create_role.return_value = { + "Role": {"Arn": "arn:aws:iam::123456789012:role/SageMakerRole"} + } + mock_create_model.return_value = "arn:aws:sagemaker:us-east-1:123456789012:model/my-model" + mock_create_endpoint.return_value = ( + "arn:aws:sagemaker:us-east-1:123456789012:endpoint/my-ep" + ) + + model_package_arn = "arn:aws:sagemaker:us-east-1:123456789012:model-package/my-package/1" + deployer = self._make_deployer() + result = deployer.deploy( + model_artifact_path=model_package_arn, + deploy_platform=DeployPlatform.SAGEMAKER, + ) + + self.assertIsInstance(result, DeploymentResult) + self.assertEqual(result.endpoint.platform, DeployPlatform.SAGEMAKER) + + # model_package_name should be the ARN itself + call_kwargs = mock_create_model.call_args.kwargs + self.assertEqual(call_kwargs["model_package_name"], model_package_arn) + # model_s3_location should be None when model_package_name is set + self.assertIsNone(call_kwargs["model_s3_location"]) + + @patch(f"{PATCH_PREFIX}.create_sagemaker_endpoint") + @patch(f"{PATCH_PREFIX}.create_sagemaker_model") + @patch(f"{PATCH_PREFIX}._validate_sagemaker_instance_type_for_model_deployment") + @patch(f"{PATCH_PREFIX}.create_sagemaker_execution_role") + @patch("boto3.client") + def test_model_package_arn_skips_trailing_slash_normalization( + self, + mock_boto_client, + mock_create_role, + mock_validate_instance, + mock_create_model, + mock_create_endpoint, + mock_validate_region, + mock_find_by_tag, + ): + """Model package ARN should NOT get a trailing slash appended.""" + mock_iam = MagicMock() + + def client_side_effect(service, **kwargs): + if service == "iam": + return mock_iam + return MagicMock() + + mock_boto_client.side_effect = client_side_effect + + mock_create_role.return_value = { + "Role": {"Arn": "arn:aws:iam::123456789012:role/SageMakerRole"} + } + mock_create_model.return_value = "arn:aws:sagemaker:us-east-1:123456789012:model/my-model" + mock_create_endpoint.return_value = ( + "arn:aws:sagemaker:us-east-1:123456789012:endpoint/my-ep" + ) + + model_package_arn = "arn:aws:sagemaker:us-west-2:123456789012:model-package/nova-pkg" + deployer = self._make_deployer() + deployer.deploy( + model_artifact_path=model_package_arn, + deploy_platform=DeployPlatform.SAGEMAKER, + ) + + # The artifact_path passed to _deploy_to_sagemaker should be the ARN unchanged + # (no trailing slash appended) + call_kwargs = mock_create_model.call_args.kwargs + self.assertFalse( + call_kwargs.get("model_package_name", "").endswith("/"), + "Model package ARN should not have a trailing slash", + ) + self.assertEqual(call_kwargs["model_package_name"], model_package_arn) + + @patch(f"{PATCH_PREFIX}.create_sagemaker_endpoint") + @patch(f"{PATCH_PREFIX}.create_sagemaker_model") + @patch(f"{PATCH_PREFIX}._validate_sagemaker_instance_type_for_model_deployment") + @patch(f"{PATCH_PREFIX}.create_sagemaker_execution_role") + @patch("boto3.client") + def test_s3_path_does_not_set_model_package_name( + self, + mock_boto_client, + mock_create_role, + mock_validate_instance, + mock_create_model, + mock_create_endpoint, + mock_validate_region, + mock_find_by_tag, + ): + """A regular S3 path should NOT set model_package_name (remains None).""" + mock_iam = MagicMock() + + def client_side_effect(service, **kwargs): + if service == "iam": + return mock_iam + return MagicMock() + + mock_boto_client.side_effect = client_side_effect + + mock_create_role.return_value = { + "Role": {"Arn": "arn:aws:iam::123456789012:role/SageMakerRole"} + } + mock_create_model.return_value = "arn:aws:sagemaker:us-east-1:123456789012:model/my-model" + mock_create_endpoint.return_value = ( + "arn:aws:sagemaker:us-east-1:123456789012:endpoint/my-ep" + ) + + deployer = self._make_deployer() + deployer.deploy( + model_artifact_path="s3://bucket/model/path", + deploy_platform=DeployPlatform.SAGEMAKER, + ) + + call_kwargs = mock_create_model.call_args.kwargs + self.assertIsNone(call_kwargs.get("model_package_name")) + # S3 path should get trailing slash normalization + self.assertEqual(call_kwargs["model_s3_location"], "s3://bucket/model/path/") + + +@patch(f"{PATCH_PREFIX}.find_bedrock_model_by_tag", return_value=None) +@patch(f"{PATCH_PREFIX}.validate_region") +class TestCreateCustomModel(unittest.TestCase): + """Tests for create_custom_model() validation and data source handling.""" + + def _make_deployer(self, **kwargs): + defaults = dict(region="us-east-1", model=Model.NOVA_MICRO) + defaults.update(kwargs) + return ForgeDeployer(**defaults) + + @patch(f"{PATCH_PREFIX}.monitor_model_create") + @patch(f"{PATCH_PREFIX}.create_bedrock_execution_role") + @patch("boto3.client") + def test_custom_model_data_source_happy_path( + self, + mock_boto_client, + mock_create_role, + mock_monitor, + mock_validate_region, + mock_find_by_tag, + ): + """Passing custom_model_data_source alone should use customModelDataSource in API call.""" + mock_bedrock = MagicMock() + mock_iam = MagicMock() + + def client_side_effect(service, **kwargs): + if service == "bedrock": + return mock_bedrock + if service == "iam": + return mock_iam + return MagicMock() + + mock_boto_client.side_effect = client_side_effect + + mock_create_role.return_value = { + "Role": {"Arn": "arn:aws:iam::123456789012:role/BedrockRole"} + } + model_arn = "arn:aws:bedrock:us-east-1:123456789012:custom-model/my-model" + mock_bedrock.create_custom_model.return_value = {"modelArn": model_arn} + + data_source = { + "modelPackageArnDataSource": { + "modelPackageArn": "arn:aws:sagemaker:us-east-1:123456789012:model-package/pkg/1" + } + } + + deployer = self._make_deployer() + result = deployer.create_custom_model(custom_model_data_source=data_source) + + self.assertEqual(result.model_arn, model_arn) + create_kwargs = mock_bedrock.create_custom_model.call_args[1] + # Should use customModelDataSource, not modelSourceConfig + self.assertEqual(create_kwargs["customModelDataSource"], data_source) + self.assertNotIn("modelSourceConfig", create_kwargs) + + def test_both_model_artifact_and_data_source_raises( + self, mock_validate_region, mock_find_by_tag + ): + """Passing both model_artifact_path and custom_model_data_source should raise ValueError.""" + deployer = self._make_deployer() + with self.assertRaises(ValueError) as ctx: + deployer.create_custom_model( + model_artifact_path="s3://bucket/model", + custom_model_data_source={"modelPackageArnDataSource": {"modelPackageArn": "arn"}}, + ) + self.assertIn("not both", str(ctx.exception)) + + def test_neither_model_artifact_nor_data_source_raises( + self, mock_validate_region, mock_find_by_tag + ): + """Passing neither model_artifact_path nor custom_model_data_source should raise ValueError.""" + deployer = self._make_deployer() + with self.assertRaises(ValueError) as ctx: + deployer.create_custom_model() + self.assertIn("must be provided", str(ctx.exception)) + + @patch(f"{PATCH_PREFIX}.monitor_model_create") + @patch(f"{PATCH_PREFIX}.create_bedrock_execution_role") + @patch("boto3.client") + def test_escrow_path_extracted_from_model_package_arn_data_source( + self, + mock_boto_client, + mock_create_role, + mock_monitor, + mock_validate_region, + mock_find_by_tag, + ): + """Escrow path should be extracted from modelPackageArnDataSource.modelPackageArn.""" + mock_bedrock = MagicMock() + mock_iam = MagicMock() + + def client_side_effect(service, **kwargs): + if service == "bedrock": + return mock_bedrock + if service == "iam": + return mock_iam + return MagicMock() + + mock_boto_client.side_effect = client_side_effect + + mock_create_role.return_value = { + "Role": {"Arn": "arn:aws:iam::123456789012:role/BedrockRole"} + } + model_arn = "arn:aws:bedrock:us-east-1:123456789012:custom-model/my-model" + mock_bedrock.create_custom_model.return_value = {"modelArn": model_arn} + + package_arn = "arn:aws:sagemaker:us-east-1:123456789012:model-package/my-pkg/1" + data_source = {"modelPackageArnDataSource": {"modelPackageArn": package_arn}} + + deployer = self._make_deployer() + result = deployer.create_custom_model(custom_model_data_source=data_source) + + # The escrow_uri on the result should be the extracted model package ARN + self.assertEqual(result.escrow_uri, package_arn) + + @patch(f"{PATCH_PREFIX}.monitor_model_create") + @patch(f"{PATCH_PREFIX}.create_bedrock_execution_role") + @patch("boto3.client") + def test_escrow_path_fallback_for_unknown_data_source_structure( + self, + mock_boto_client, + mock_create_role, + mock_monitor, + mock_validate_region, + mock_find_by_tag, + ): + """When modelPackageArnDataSource is absent, a warning is emitted and no escrow tag is added.""" + mock_bedrock = MagicMock() + mock_iam = MagicMock() + + def client_side_effect(service, **kwargs): + if service == "bedrock": + return mock_bedrock + if service == "iam": + return mock_iam + return MagicMock() + + mock_boto_client.side_effect = client_side_effect + + mock_create_role.return_value = { + "Role": {"Arn": "arn:aws:iam::123456789012:role/BedrockRole"} + } + model_arn = "arn:aws:bedrock:us-east-1:123456789012:custom-model/my-model" + mock_bedrock.create_custom_model.return_value = {"modelArn": model_arn} + + data_source = {"someOtherSource": {"key": "value"}} + + deployer = self._make_deployer() + with self.assertLogs("nova_forge_sdk", level="WARNING") as log: + result = deployer.create_custom_model(custom_model_data_source=data_source) + + # Warning should mention inability to extract tag-safe identifier + warning_found = any("tag-safe identifier" in msg for msg in log.output) + self.assertTrue(warning_found, f"Expected warning not found in: {log.output}") + + # escrow_uri should be empty since no tag-safe value was available + self.assertEqual(result.escrow_uri, "") + + # find_published_model should not have been called (no escrow_path) + mock_find_by_tag.assert_not_called() + + # No escrow tag should be in the create_custom_model call + call_kwargs = mock_bedrock.create_custom_model.call_args[1] + model_tags = call_kwargs.get("modelTags", []) + escrow_tags = [t for t in model_tags if t.get("key") == ESCROW_URI_TAG_KEY] + self.assertEqual(len(escrow_tags), 0) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/evaluator/test_forge_evaluator.py b/tests/unit/evaluator/test_forge_evaluator.py index e7cef14..f83f65c 100644 --- a/tests/unit/evaluator/test_forge_evaluator.py +++ b/tests/unit/evaluator/test_forge_evaluator.py @@ -11,6 +11,8 @@ # 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 os +import tempfile import unittest from datetime import datetime, timezone from unittest.mock import MagicMock, PropertyMock, create_autospec, patch @@ -19,14 +21,18 @@ from amzn_nova_forge.core.result import ( SMHPEvaluationResult, SMTJEvaluationResult, + SMTJTrainingResult, ) from amzn_nova_forge.core.result.training_result import TrainingResult from amzn_nova_forge.core.types import ForgeConfig, ModelArtifacts from amzn_nova_forge.evaluator.forge_evaluator import EvalTaskConfig, ForgeEvaluator +from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig from amzn_nova_forge.manager.runtime_manager import ( SMHPRuntimeManager, SMTJRuntimeManager, + SMTJServerlessRuntimeManager, ) +from amzn_nova_forge.monitor import MLflowMonitor class TestForgeEvaluatorInit(unittest.TestCase): @@ -767,8 +773,6 @@ class TestForgeEvaluatorInspectLens(unittest.TestCase): """Tests for ForgeEvaluator InspectLens path.""" def _make_evaluator(self, image_uri=None): - from unittest.mock import PropertyMock - mock_infra = create_autospec(SMTJRuntimeManager) mock_infra.kms_key_id = None mock_infra.instance_type = "ml.m5.large" @@ -800,8 +804,6 @@ def _make_evaluator(self, image_uri=None): return evaluator def test_dry_run_returns_none(self): - from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig - evaluator = self._make_evaluator() config = InspectLensConfig( benchmarks_path="s3://bucket/benchmarks/boolq/", @@ -832,8 +834,6 @@ def test_missing_inspect_lens_config_raises(self): self.assertIn("inspect_lens_config is required", str(ctx.exception)) def test_overrides_warning_for_non_decoding_keys(self): - from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig - evaluator = self._make_evaluator() config = InspectLensConfig( benchmarks_path="s3://bucket/benchmarks/boolq/", @@ -857,8 +857,6 @@ def test_overrides_warning_for_non_decoding_keys(self): self.assertTrue(any("benchmarks_path" in w for w in warning_calls)) def test_valid_decoding_overrides_no_warning(self): - from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig - evaluator = self._make_evaluator() config = InspectLensConfig( benchmarks_path="s3://bucket/benchmarks/boolq/", @@ -882,8 +880,6 @@ def test_valid_decoding_overrides_no_warning(self): self.assertFalse(any("overrides" in w for w in warning_calls)) def test_inference_provider_bedrock_default(self): - from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig - evaluator = self._make_evaluator() config = InspectLensConfig( benchmarks_path="s3://bucket/benchmarks/", @@ -894,8 +890,6 @@ def test_inference_provider_bedrock_default(self): self.assertIn("us.amazon.nova-micro-v1:0", provider["bedrock"]["model_id"]) def test_inference_provider_existing_endpoint(self): - from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig - evaluator = self._make_evaluator() config = InspectLensConfig( benchmarks_path="s3://bucket/benchmarks/", @@ -906,8 +900,6 @@ def test_inference_provider_existing_endpoint(self): self.assertEqual(provider["sagemaker_endpoint"]["endpoint_name"], "my-endpoint") def test_inference_provider_model_path_overrides_bedrock(self): - from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig - evaluator = self._make_evaluator() config = InspectLensConfig( benchmarks_path="s3://bucket/benchmarks/", @@ -921,8 +913,6 @@ def test_inference_provider_model_path_overrides_bedrock(self): def test_cache_hit_returns_cached_result(self): """When job caching is enabled and a matching result exists, return it without submitting.""" - from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig - evaluator = self._make_evaluator() config = InspectLensConfig( benchmarks_path="s3://bucket/benchmarks/boolq/", @@ -957,11 +947,6 @@ def test_cache_hit_returns_cached_result(self): def test_mlflow_tracking_injected_into_config_dict(self): """When ForgeConfig.mlflow_monitor is set, tracking section must appear in the YAML dict.""" - from unittest.mock import PropertyMock - - from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig - from amzn_nova_forge.monitor import MLflowMonitor - mock_infra = create_autospec(SMTJRuntimeManager) mock_infra.kms_key_id = None mock_infra.instance_type = "ml.m5.large" @@ -1041,8 +1026,6 @@ class TestInspectLensS3Paths(unittest.TestCase): FIXED_RUN_ID = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee" def _make_evaluator(self, output_s3_path="s3://bucket/output/"): - from unittest.mock import PropertyMock - mock_infra = create_autospec(SMTJRuntimeManager) mock_infra.kms_key_id = None mock_infra.instance_type = "ml.m5.large" @@ -1092,11 +1075,6 @@ def _run_evaluate(self, evaluator, inspect_lens_config, mock_sm_client): def test_local_benchmarks_path_raises_valueerror(self): """Local benchmarks_path should raise ValueError — user must upload first.""" - import os - import tempfile - - from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig - evaluator = self._make_evaluator(output_s3_path="s3://bucket/output/") with tempfile.TemporaryDirectory() as tmpdir: @@ -1127,9 +1105,6 @@ def test_local_benchmarks_path_raises_valueerror(self): def test_upload_benchmarks(self): """upload_benchmarks() uploads .py files to S3 and returns the S3 URI.""" - import os - import tempfile - evaluator = self._make_evaluator(output_s3_path="s3://bucket/output/") with tempfile.TemporaryDirectory() as tmpdir: @@ -1160,8 +1135,6 @@ def test_upload_benchmarks_invalid_local_dir_raises(self): def test_upload_benchmarks_invalid_s3_path_raises(self): """upload_benchmarks() raises ValueError for non-S3 path.""" - import tempfile - evaluator = self._make_evaluator() with tempfile.TemporaryDirectory() as tmpdir: with self.assertRaises(ValueError) as ctx: @@ -1170,8 +1143,6 @@ def test_upload_benchmarks_invalid_s3_path_raises(self): def test_s3_benchmarks_config_colocated_with_benchmarks(self): """S3 benchmarks_path in a different bucket → config and output still under output_s3_path.""" - from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig - evaluator = self._make_evaluator(output_s3_path="s3://bucket/output/") cfg = InspectLensConfig( @@ -1210,8 +1181,6 @@ def test_s3_benchmarks_config_colocated_with_benchmarks(self): def test_run_id_in_job_name(self): """unique_job_name should contain the run_id UUID.""" - from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig - evaluator = self._make_evaluator() cfg = InspectLensConfig( @@ -1241,8 +1210,6 @@ def test_run_id_in_job_name(self): def test_s3_benchmarks_bucket_only_path(self): """S3 benchmarks_path in a separate bucket → output still goes under output_s3_path.""" - from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig - evaluator = self._make_evaluator(output_s3_path="s3://bucket/output/") cfg = InspectLensConfig( @@ -1274,5 +1241,231 @@ def test_s3_benchmarks_bucket_only_path(self): self.assertEqual(output_path, f"s3://bucket/output/{self.FIXED_RUN_ID}/output/") +class TestForgeEvaluatorMTRL(unittest.TestCase): + """Tests for ForgeEvaluator._execute_mtrl_eval via evaluate().""" + + def setUp(self): + self.model = Model.NOVA_LITE_2 + self.mock_infra = create_autospec(SMTJServerlessRuntimeManager) + self.mock_infra.kms_key_id = None + self.mock_infra.instance_type = None + self.mock_infra.instance_count = None + self.mock_infra.platform = Platform.SMTJServerless + self.mock_infra.rft_lambda_arn = None + self.mock_infra.hub_content_version = None + + # Mock execute_mtrl_eval return value + mock_execution = MagicMock() + mock_execution.arn = ( + "arn:aws:sagemaker:us-east-1:123456789012:pipeline-execution/mtrl-eval-exec" + ) + self.mock_infra.execute_mtrl_eval.return_value = mock_execution + + self._patcher_set_output = patch( + "amzn_nova_forge.evaluator.forge_evaluator.set_output_s3_path", + return_value="s3://bucket/output", + ) + self._patcher_session = patch("boto3.session.Session") + self._patcher_client = patch("boto3.client") + + self._patcher_set_output.start() + mock_session = self._patcher_session.start() + self._patcher_client.start() + + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + + mock_mlflow = MagicMock(spec=MLflowMonitor) + mock_mlflow.tracking_uri = ( + "arn:aws:sagemaker:us-east-1:123456789012:mlflow-tracking-server/my-server" + ) + + self.evaluator = ForgeEvaluator( + model=self.model, + infra=self.mock_infra, + data_s3_path="s3://bucket/data", + config=ForgeConfig(mlflow_monitor=mock_mlflow), + ) + + def tearDown(self): + self._patcher_client.stop() + self._patcher_session.stop() + self._patcher_set_output.stop() + + def test_mtrl_eval_resolves_model_path_from_job_result(self): + """_execute_mtrl_eval resolves model_path from MTRL job_result's output_model_arn.""" + mock_job_result = MagicMock(spec=SMTJTrainingResult) + mock_job_result._is_mtrl = True + mock_job_result.job_id = "mtrl-train-123" + mock_job_result.model_artifacts = ModelArtifacts( + output_s3_path="s3://bucket/output/", + output_model_arn="arn:aws:sagemaker:us-east-1:123456789012:model-package/my-group/1", + ) + + result = self.evaluator.evaluate( + job_name="test-mtrl-eval", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + job_result=mock_job_result, + ) + + self.mock_infra.execute_mtrl_eval.assert_called_once() + call_kwargs = self.mock_infra.execute_mtrl_eval.call_args[1] + self.assertEqual( + call_kwargs["model_path"], + "arn:aws:sagemaker:us-east-1:123456789012:model-package/my-group/1", + ) + self.assertEqual(call_kwargs["training_job_name"], "mtrl-train-123") + self.assertIsInstance(result, SMTJEvaluationResult) + + def test_mtrl_eval_uses_explicit_model_path(self): + """_execute_mtrl_eval uses explicit model_path when provided directly.""" + explicit_arn = "arn:aws:sagemaker:us-east-1:123456789012:model-package/explicit-group/2" + + result = self.evaluator.evaluate( + job_name="test-mtrl-eval-explicit", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + model_path=explicit_arn, + ) + + self.mock_infra.execute_mtrl_eval.assert_called_once() + call_kwargs = self.mock_infra.execute_mtrl_eval.call_args[1] + self.assertEqual(call_kwargs["model_path"], explicit_arn) + self.assertIsNone(call_kwargs["training_job_name"]) + self.assertIsInstance(result, SMTJEvaluationResult) + + def test_mtrl_eval_dry_run_returns_none(self): + """_execute_mtrl_eval with dry_run=True returns None without calling infra.""" + result = self.evaluator.evaluate( + job_name="test-mtrl-eval-dryrun", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + dry_run=True, + ) + + self.assertIsNone(result) + self.mock_infra.execute_mtrl_eval.assert_not_called() + + def test_mtrl_eval_evaluate_base_model_passed_to_infra(self): + """evaluate_base_model from EvalTaskConfig is passed to execute_mtrl_eval.""" + model_path = "arn:aws:sagemaker:us-east-1:123456789012:model-package/rmp/1" + + result = self.evaluator.evaluate( + job_name="test-mtrl-eval-both", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + model_path=model_path, + task_config=EvalTaskConfig(evaluate_base_model=True), + ) + + self.mock_infra.execute_mtrl_eval.assert_called_once() + call_kwargs = self.mock_infra.execute_mtrl_eval.call_args[1] + self.assertEqual(call_kwargs["model_path"], model_path) + self.assertTrue(call_kwargs["evaluate_base_model"]) + self.assertIsInstance(result, SMTJEvaluationResult) + + def test_mtrl_eval_evaluate_base_model_defaults_to_false(self): + """evaluate_base_model defaults to False when not in task_config.""" + model_path = "arn:aws:sagemaker:us-east-1:123456789012:model-package/rmp/1" + + result = self.evaluator.evaluate( + job_name="test-mtrl-eval-ft-only", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + model_path=model_path, + ) + + self.mock_infra.execute_mtrl_eval.assert_called_once() + call_kwargs = self.mock_infra.execute_mtrl_eval.call_args[1] + self.assertFalse(call_kwargs["evaluate_base_model"]) + + def test_evaluate_base_model_only_used_for_mtrl(self): + """evaluate_base_model in EvalTaskConfig is only extracted for RFT_MULTITURN_EVAL.""" + # When using MTRL without evaluate_base_model, it should default to False + self.evaluator.evaluate( + job_name="test-mtrl-no-base", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + ) + + call_kwargs = self.mock_infra.execute_mtrl_eval.call_args[1] + self.assertFalse(call_kwargs["evaluate_base_model"]) + + # When using MTRL WITH evaluate_base_model=True, it should be True + self.mock_infra.execute_mtrl_eval.reset_mock() + self.evaluator.evaluate( + job_name="test-mtrl-with-base", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + model_path="arn:aws:sagemaker:us-east-1:123:model-package/rmp/1", + task_config=EvalTaskConfig(evaluate_base_model=True), + ) + + call_kwargs = self.mock_infra.execute_mtrl_eval.call_args[1] + self.assertTrue(call_kwargs["evaluate_base_model"]) + + @patch( + "amzn_nova_forge.evaluator.forge_evaluator.set_output_s3_path", + return_value="s3://bucket/output", + ) + @patch("boto3.client") + @patch("boto3.session.Session") + def test_mtrl_eval_raises_without_mlflow_config(self, mock_session, _mock_client, _mock_output): + """AgentRFT eval jobs must have an MLflow config; raise ValueError if missing.""" + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + mock_infra = create_autospec(SMTJServerlessRuntimeManager) + mock_infra.kms_key_id = None + mock_infra.instance_type = None + mock_infra.instance_count = None + mock_infra.platform = Platform.SMTJServerless + mock_infra.rft_lambda_arn = None + mock_infra.hub_content_version = None + + evaluator = ForgeEvaluator( + model=Model.NOVA_LITE_2, + infra=mock_infra, + data_s3_path="s3://bucket/data", + config=ForgeConfig(), # No mlflow_monitor + ) + + with self.assertRaises(ValueError) as ctx: + evaluator.evaluate( + job_name="test-mtrl-eval-no-mlflow", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + model_path="arn:aws:sagemaker:us-east-1:123:model-package/rmp/1", + ) + self.assertIn("MLflow configuration is required", str(ctx.exception)) + + @patch( + "amzn_nova_forge.evaluator.forge_evaluator.set_output_s3_path", + return_value="s3://bucket/output", + ) + @patch("boto3.client") + @patch("boto3.session.Session") + def test_mtrl_eval_raises_without_mlflow_tracking_uri( + self, mock_session, _mock_client, _mock_output + ): + """AgentRFT eval jobs must have a tracking_uri; raise ValueError if None.""" + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + mock_infra = create_autospec(SMTJServerlessRuntimeManager) + mock_infra.kms_key_id = None + mock_infra.instance_type = None + mock_infra.instance_count = None + mock_infra.platform = Platform.SMTJServerless + mock_infra.rft_lambda_arn = None + mock_infra.hub_content_version = None + + mock_monitor = MagicMock(spec=MLflowMonitor) + mock_monitor.tracking_uri = None + + evaluator = ForgeEvaluator( + model=Model.NOVA_LITE_2, + infra=mock_infra, + data_s3_path="s3://bucket/data", + config=ForgeConfig(mlflow_monitor=mock_monitor), + ) + + with self.assertRaises(ValueError) as ctx: + evaluator.evaluate( + job_name="test-mtrl-eval-no-tracking-uri", + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + model_path="arn:aws:sagemaker:us-east-1:123:model-package/rmp/1", + ) + self.assertIn("MLflow configuration is required", str(ctx.exception)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/iam/test_iam_role_creator.py b/tests/unit/iam/test_iam_role_creator.py index 48b0a29..bfe8416 100644 --- a/tests/unit/iam/test_iam_role_creator.py +++ b/tests/unit/iam/test_iam_role_creator.py @@ -917,6 +917,165 @@ def test_creates_exactly_two_policies(self, mock_boto_client): ) +class TestSagemakerModelPackagePolicy(unittest.TestCase): + """Tests that sagemaker_model_package_policy is created and attached by both role creators.""" + + @classmethod + def setUpClass(cls): + from importlib import resources as importlib_resources + + with ( + importlib_resources.files("amzn_nova_forge.iam") + .joinpath("sagemaker_policies.json") + .open() as f + ): + sagemaker_policies = json.load(f) + with ( + importlib_resources.files("amzn_nova_forge.iam") + .joinpath("bedrock_policies.json") + .open() as f + ): + bedrock_policies = json.load(f) + + cls.sagemaker_model_package_policy = sagemaker_policies["sagemaker_model_package_policy"] + cls.bedrock_model_package_policy = bedrock_policies["sagemaker_model_package_policy"] + + @patch("boto3.client") + def test_bedrock_role_creates_sagemaker_model_package_policy(self, mock_boto_client): + """create_bedrock_execution_role creates and attaches sagemaker_model_package_policy.""" + mock_sts = MagicMock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + mock_boto_client.return_value = mock_sts + + role_name = "test-bedrock-role" + 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=role_name) + + mock_iam.create_policy.assert_any_call( + PolicyName=f"{role_name}Sagemaker_Model_Package_Policy", + PolicyDocument=json.dumps(self.bedrock_model_package_policy), + ) + mock_iam.attach_role_policy.assert_any_call( + RoleName=role_name, + PolicyArn="arn:aws:iam::123456789012:policy/foo", + ) + + @patch("boto3.client") + def test_sagemaker_role_creates_sagemaker_model_package_policy(self, mock_boto_client): + """create_sagemaker_execution_role creates and attaches sagemaker_model_package_policy.""" + mock_sts = MagicMock() + mock_sts.get_caller_identity.return_value = {"Account": "123456789012"} + mock_boto_client.return_value = mock_sts + + role_name = "test-sagemaker-role" + 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_sagemaker_execution_role(iam_client=mock_iam, role_name=role_name) + + mock_iam.create_policy.assert_any_call( + PolicyName=f"{role_name}Sagemaker_Model_Package_Policy", + PolicyDocument=json.dumps(self.sagemaker_model_package_policy), + ) + mock_iam.attach_role_policy.assert_any_call( + RoleName=role_name, + PolicyArn="arn:aws:iam::123456789012:policy/foo", + ) + + @patch("boto3.client") + def test_bedrock_role_handles_existing_model_package_policy(self, mock_boto_client): + """create_bedrock_execution_role handles EntityAlreadyExistsException for sagemaker_model_package_policy.""" + account_id = "123456789012" + mock_sts = MagicMock() + mock_sts.get_caller_identity.return_value = {"Account": account_id} + mock_boto_client.return_value = mock_sts + + role_name = "test-bedrock-role" + 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") + + # First two policies succeed, sagemaker_model_package_policy already exists + existing_policy_arn = ( + f"arn:aws:iam::{account_id}:policy/{role_name}Sagemaker_Model_Package_Policy" + ) + mock_iam.create_policy.side_effect = [ + {"Policy": {"Arn": f"arn:aws:iam::{account_id}:policy/{role_name}Bedrock_Policy"}}, + {"Policy": {"Arn": f"arn:aws:iam::{account_id}:policy/{role_name}S3_Read_Policy"}}, + mock_iam.exceptions.EntityAlreadyExistsException("already exists"), + ] + mock_iam.get_policy.return_value = {"Policy": {"Arn": existing_policy_arn}} + + create_bedrock_execution_role(iam_client=mock_iam, role_name=role_name) + + mock_iam.get_policy.assert_called_with(PolicyArn=existing_policy_arn) + mock_iam.attach_role_policy.assert_any_call( + RoleName=role_name, + PolicyArn=existing_policy_arn, + ) + + @patch("boto3.client") + def test_sagemaker_role_handles_existing_model_package_policy(self, mock_boto_client): + """create_sagemaker_execution_role handles EntityAlreadyExistsException for sagemaker_model_package_policy.""" + account_id = "123456789012" + mock_sts = MagicMock() + mock_sts.get_caller_identity.return_value = {"Account": account_id} + mock_boto_client.return_value = mock_sts + + role_name = "test-sagemaker-role" + 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") + + # First 8 policies succeed, sagemaker_model_package_policy (9th) already exists + existing_policy_arn = ( + f"arn:aws:iam::{account_id}:policy/{role_name}Sagemaker_Model_Package_Policy" + ) + successful_result = {"Policy": {"Arn": f"arn:aws:iam::{account_id}:policy/foo"}} + mock_iam.create_policy.side_effect = [ + successful_result, # cloudwatch_ec2_ec2_policy + successful_result, # cloudwatch_metric_policy + successful_result, # cloudwatch_logstream_policy + successful_result, # cloudwatch_loggroup_policy + successful_result, # ecr_read_policy + successful_result, # s3_read_policy + successful_result, # kms_policy + successful_result, # ec2_policy + mock_iam.exceptions.EntityAlreadyExistsException("already exists"), + ] + mock_iam.get_policy.return_value = {"Policy": {"Arn": existing_policy_arn}} + + create_sagemaker_execution_role(iam_client=mock_iam, role_name=role_name) + + mock_iam.get_policy.assert_called_with(PolicyArn=existing_policy_arn) + mock_iam.attach_role_policy.assert_any_call( + RoleName=role_name, + PolicyArn=existing_policy_arn, + ) + + class TestRegionPropagation(unittest.TestCase): """Tests that the region parameter is propagated to the STS client.""" diff --git a/tests/unit/inference/test_forge_inference.py b/tests/unit/inference/test_forge_inference.py index da65cc5..d309651 100644 --- a/tests/unit/inference/test_forge_inference.py +++ b/tests/unit/inference/test_forge_inference.py @@ -102,7 +102,7 @@ def test_invoke_sagemaker_endpoint( mock_regex.match.assert_called_once_with(arn) mock_boto_client.assert_called_once_with("sagemaker-runtime", region_name="us-east-1") mock_invoke_sm.assert_called_once_with( - {"prompt": "hello"}, "my-endpoint", mock_runtime_client + {"prompt": "hello"}, "my-endpoint", mock_runtime_client, inference_component_name=None ) self.assertEqual(result, {"result": "ok"}) diff --git a/tests/unit/manager/test_runtime_manager.py b/tests/unit/manager/test_runtime_manager.py index 278df75..c26401a 100644 --- a/tests/unit/manager/test_runtime_manager.py +++ b/tests/unit/manager/test_runtime_manager.py @@ -20,6 +20,7 @@ import zipfile from unittest.mock import MagicMock, patch +from amzn_nova_forge.core.enums import TrainingMethod from amzn_nova_forge.manager.runtime_manager import ( DataPrepJobConfig, JobConfig, @@ -163,6 +164,112 @@ def test_execute_without_optional_params( mock_model_trainer.train.assert_called_once_with(wait=False, logs=False) self.assertEqual(job_id, "test-job-suffix") + @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_execute_eval_with_model_package_arn_passes_model_package_config( + self, mock_setup, mock_model_trainer_cls, mock_boto_client + ): + """Eval job with a model package ARN should pass ModelPackageConfig to from_recipe.""" + 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": "eval-job-suffix"}] + } + manager.sagemaker_client.describe_model_package_group.return_value = { + "ModelPackageGroupArn": "arn:aws:sagemaker:us-east-1:123456789012:model-package-group/my-group" + } + + model_package_arn = "arn:aws:sagemaker:us-east-1:123456789012:model-package/my-group/3" + job_config = JobConfig( + job_name="eval-job", + image_uri="123456789012.dkr.ecr.us-east-1.amazonaws.com/eval-image:latest", + recipe_path="/path/to/recipe", + output_s3_path="s3://output-bucket/eval-output", + method=TrainingMethod.EVALUATION, + model_name_or_path=model_package_arn, + ) + + job_id = manager.execute(job_config) + + call_kwargs = mock_model_trainer_cls.from_recipe.call_args.kwargs + self.assertIn("model_package_config", call_kwargs) + manager.sagemaker_client.describe_model_package_group.assert_called_once_with( + ModelPackageGroupName="my-group" + ) + self.assertEqual(job_id, "eval-job-suffix") + + @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_execute_eval_with_s3_path_omits_model_package_config( + self, mock_setup, mock_model_trainer_cls, mock_boto_client + ): + """Eval job with an S3 path should not pass ModelPackageConfig.""" + 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": "eval-job-suffix"}] + } + + job_config = JobConfig( + job_name="eval-job", + image_uri="123456789012.dkr.ecr.us-east-1.amazonaws.com/eval-image:latest", + recipe_path="/path/to/recipe", + output_s3_path="s3://output-bucket/eval-output", + method=TrainingMethod.EVALUATION, + model_name_or_path="s3://my-bucket/checkpoint/step_10", + ) + + job_id = manager.execute(job_config) + + call_kwargs = mock_model_trainer_cls.from_recipe.call_args.kwargs + self.assertNotIn("model_package_config", call_kwargs) + self.assertEqual(job_id, "eval-job-suffix") + + @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_execute_eval_with_nonexistent_model_package_group_raises( + self, mock_setup, mock_model_trainer_cls, mock_boto_client + ): + """Eval job with a model package ARN pointing to a nonexistent group should raise.""" + from botocore.exceptions import ClientError + + 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.describe_model_package_group.side_effect = ClientError( + {"Error": {"Code": "ValidationException", "Message": "not found"}}, + "DescribeModelPackageGroup", + ) + + model_package_arn = "arn:aws:sagemaker:us-east-1:123456789012:model-package/missing-group/1" + job_config = JobConfig( + job_name="eval-job", + image_uri="123456789012.dkr.ecr.us-east-1.amazonaws.com/eval-image:latest", + recipe_path="/path/to/recipe", + output_s3_path="s3://output-bucket/eval-output", + method=TrainingMethod.EVALUATION, + model_name_or_path=model_package_arn, + ) + + with self.assertRaises(ValueError) as ctx: + manager.execute(job_config) + + self.assertIn("does not exist", str(ctx.exception)) + @patch.object(SMTJRuntimeManager, "setup", return_value=None) def test_cleanup_success(self, mock_setup): manager = self._create_manager() diff --git a/tests/unit/manager/test_serverless_runtime_manager.py b/tests/unit/manager/test_serverless_runtime_manager.py index b1619a2..0e4a558 100644 --- a/tests/unit/manager/test_serverless_runtime_manager.py +++ b/tests/unit/manager/test_serverless_runtime_manager.py @@ -444,8 +444,6 @@ def test_build_serverless_job_config_evaluator_arn_not_set_for_benchmark(self): self.assertEqual(config["EvaluationType"], "BenchmarkEvaluation") self.assertNotIn("EvaluatorArn", config) - # --- _resolve_base_model_arn tests --- - @patch( "amzn_nova_forge.manager.runtime_manager.get_hub_content", return_value={"HubContentArn": "arn:hub:content"}, @@ -462,8 +460,6 @@ def test_resolve_base_model_arn(self, mock_hub): hub_content_version=None, ) - # --- execute tests --- - @patch("amzn_nova_forge.manager.runtime_manager.get_hub_content") @patch("sagemaker.ai_registry.dataset.DataSet") def test_execute_success(self, mock_dataset_cls, mock_hub_content): @@ -666,6 +662,85 @@ def test_execute_benchmark_eval_skips_input_data(self, mock_dataset_cls, mock_hu self.assertNotIn("InputDataConfig", call_kwargs) mock_dataset_cls.create.assert_not_called() + @patch("amzn_nova_forge.manager.runtime_manager.get_hub_content") + @patch("sagemaker.ai_registry.dataset.DataSet") + def test_execute_eval_with_sagemaker_arn_includes_source_model_package( + self, mock_dataset_cls, mock_hub_content + ): + """Eval job with a SageMaker ARN sets both ModelPackageGroupArn and SourceModelPackageArn.""" + manager = self._create_manager() + mock_hub_content.return_value = {"HubContentArn": self.mock_hub_content_arn} + manager.sagemaker_client.create_training_job.return_value = { + "TrainingJobArn": "arn:aws:sagemaker:us-east-1:123456789012:training-job/test-job" + } + model_package_arn = "arn:aws:sagemaker:us-east-1:123456789012:model-package/group/1" + recipe = { + "run": { + "model_type": "amazon.nova-2-lite-v1:0:256k", + "model_name_or_path": model_package_arn, + }, + "evaluation": {"task": "mmlu"}, + } + 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-eval-job", + image_uri="", + recipe_path=recipe_path, + output_s3_path="s3://output-bucket/output", + method=TrainingMethod.EVALUATION, + ) + manager.execute(job_config) + finally: + os.unlink(recipe_path) + + call_kwargs = manager.sagemaker_client.create_training_job.call_args.kwargs + self.assertIn("ModelPackageConfig", call_kwargs) + self.assertEqual( + call_kwargs["ModelPackageConfig"]["SourceModelPackageArn"], + model_package_arn, + ) + self.assertIn("ModelPackageGroupArn", call_kwargs["ModelPackageConfig"]) + self.assertEqual( + call_kwargs["ModelPackageConfig"]["ModelPackageGroupArn"], + manager.model_package_group_arn, + ) + + @patch("amzn_nova_forge.manager.runtime_manager.get_hub_content") + @patch("sagemaker.ai_registry.dataset.DataSet") + def test_execute_eval_without_sagemaker_arn_omits_model_package_config( + self, mock_dataset_cls, mock_hub_content + ): + """Eval job without a SageMaker ARN omits ModelPackageConfig entirely.""" + manager = self._create_manager() + mock_hub_content.return_value = {"HubContentArn": self.mock_hub_content_arn} + 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"}, + "evaluation": {"task": "mmlu"}, + } + 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-eval-job", + image_uri="", + recipe_path=recipe_path, + output_s3_path="s3://output-bucket/output", + method=TrainingMethod.EVALUATION, + ) + manager.execute(job_config) + finally: + os.unlink(recipe_path) + + call_kwargs = manager.sagemaker_client.create_training_job.call_args.kwargs + self.assertNotIn("ModelPackageConfig", call_kwargs) + @patch("amzn_nova_forge.manager.runtime_manager.get_hub_content") @patch("sagemaker.ai_registry.dataset.DataSet") def test_execute_hub_content_arn_skips_registration(self, mock_dataset_cls, mock_hub_content): @@ -976,8 +1051,6 @@ def test_execute_raises_on_api_error(self, mock_hub_content): finally: os.unlink(recipe_path) - # --- cleanup tests --- - def test_cleanup_success(self): manager = self._create_manager() manager.cleanup("test-job") @@ -986,6 +1059,29 @@ def test_cleanup_success(self): ) manager.sagemaker_client.close.assert_called_once() + def test_cleanup_is_mtrl_false_calls_stop_training_job(self): + manager = self._create_manager() + manager.cleanup("test-job", is_mtrl=False) + manager.sagemaker_client.stop_training_job.assert_called_once_with( + TrainingJobName="test-job" + ) + manager.sagemaker_client.close.assert_called_once() + + @patch.object(SMTJServerlessRuntimeManager, "_cleanup_mtrl") + def test_cleanup_is_mtrl_true_calls_cleanup_mtrl(self, mock_cleanup_mtrl): + manager = self._create_manager() + manager.cleanup("mtrl-job-123", is_mtrl=True) + mock_cleanup_mtrl.assert_called_once_with("mtrl-job-123") + manager.sagemaker_client.stop_training_job.assert_not_called() + + @patch.object(SMTJServerlessRuntimeManager, "_cleanup_mtrl") + def test_cleanup_is_mtrl_true_raises_on_error(self, mock_cleanup_mtrl): + manager = self._create_manager() + mock_cleanup_mtrl.side_effect = Exception("MTRL cleanup failed") + with self.assertRaises(Exception) as ctx: + manager.cleanup("mtrl-job-123", is_mtrl=True) + self.assertEqual(str(ctx.exception), "MTRL cleanup failed") + def test_cleanup_raises_on_error(self): manager = self._create_manager() manager.sagemaker_client.stop_training_job.side_effect = Exception("Cleanup failed") @@ -993,8 +1089,6 @@ def test_cleanup_raises_on_error(self): manager.cleanup("test-job") self.assertEqual(str(ctx.exception), "Cleanup failed") - # --- required_calling_role_permissions tests --- - def test_required_calling_role_permissions(self): perms = SMTJServerlessRuntimeManager.required_calling_role_permissions( data_s3_path="s3://input/data.jsonl", @@ -1211,6 +1305,7 @@ def test_all_expected_methods_present(self): TrainingMethod.DPO_LORA, TrainingMethod.DPO_FULL, TrainingMethod.RFT_LORA, + TrainingMethod.RFT_MULTITURN_LORA, # TrainingMethod.RFT_FULL, } self.assertEqual(set(_METHOD_TO_SERVERLESS_CONFIG.keys()), expected) @@ -1220,6 +1315,7 @@ def test_lora_methods_have_peft(self): TrainingMethod.SFT_LORA, TrainingMethod.DPO_LORA, TrainingMethod.RFT_LORA, + TrainingMethod.RFT_MULTITURN_LORA, ): _, peft = _METHOD_TO_SERVERLESS_CONFIG[method] self.assertEqual(peft, "LORA") @@ -1612,5 +1708,349 @@ def test_empty_datamix_config_does_not_error(self): self.assertNotIn("customer_data_percent", result) +class TestSMTJServerlessMTRL(unittest.TestCase): + """Tests for MTRL operations on SMTJServerlessRuntimeManager.""" + + MOCK_AGENT_CORE_ARN = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/my-agent" + MOCK_LAMBDA_ARN = "arn:aws:lambda:us-east-1:123456789012:function:my-reward-fn" + MOCK_MODEL_PACKAGE_ARN = "arn:aws:sagemaker:us-east-1:123456789012:model-package/grp/1" + + @patch.object(SMTJServerlessRuntimeManager, "setup", return_value=None) + def test_init_with_agent_core_arn(self, mock_setup): + manager = SMTJServerlessRuntimeManager( + model_package_group_name="test-group", + agent_core_arn=self.MOCK_AGENT_CORE_ARN, + ) + self.assertEqual(manager.agent_core_arn, self.MOCK_AGENT_CORE_ARN) + + @patch.object(SMTJServerlessRuntimeManager, "setup", return_value=None) + def test_init_with_intermediate_model_package_group(self, mock_setup): + manager = SMTJServerlessRuntimeManager( + model_package_group_name="test-group", + intermediate_model_package_group_name="test-checkpoints", + ) + self.assertEqual(manager.intermediate_model_package_group_name, "test-checkpoints") + + @patch.object(SMTJServerlessRuntimeManager, "setup", return_value=None) + def test_init_with_rft_lambda(self, mock_setup): + manager = SMTJServerlessRuntimeManager( + model_package_group_name="test-group", + rft_lambda=self.MOCK_LAMBDA_ARN, + ) + self.assertEqual(manager.rft_lambda, self.MOCK_LAMBDA_ARN) + + def _create_mtrl_ops(self, agent_core_arn=None, rft_lambda=None): + from amzn_nova_forge.manager.mtrl_manager import MTRLOperations + + ops = MTRLOperations.__new__(MTRLOperations) + ops.agent_core_arn = agent_core_arn + ops.rft_lambda = rft_lambda + ops.model_package_group_arn = ( + "arn:aws:sagemaker:us-east-1:123456789012:model-package-group/grp" + ) + ops.execution_role = "arn:aws:iam::123456789012:role/role" + ops.subnets = None + ops.security_group_ids = None + ops.kms_key_id = None + ops.region = "us-east-1" + return ops + + def _mock_sagemaker_modules(self): + import types + + mock_sm_train = types.ModuleType("sagemaker.train") + mock_sm_train_mtrl = types.ModuleType("sagemaker.train.multi_turn_rl_trainer") + mock_sm_core = types.ModuleType("sagemaker.core") + mock_sm_core_resources = types.ModuleType("sagemaker.core.resources") + + self._mock_trainer_cls = MagicMock() + self._mock_trainer_instance = MagicMock() + self._mock_trainer_cls.return_value = self._mock_trainer_instance + self._mock_model_package_cls = MagicMock() + + mock_sm_train_mtrl.MultiTurnRLTrainer = self._mock_trainer_cls + mock_sm_core_resources.ModelPackage = self._mock_model_package_cls + + return { + "sagemaker.train": mock_sm_train, + "sagemaker.train.multi_turn_rl_trainer": mock_sm_train_mtrl, + "sagemaker.core": mock_sm_core, + "sagemaker.core.resources": mock_sm_core_resources, + } + + @patch("amzn_nova_forge.manager.mtrl_manager.logger") + def test_execute_mtrl_with_agent_core(self, _mock_logger): + ops = self._create_mtrl_ops(agent_core_arn=self.MOCK_AGENT_CORE_ARN) + modules = self._mock_sagemaker_modules() + + mock_job = MagicMock() + mock_job.job_name = "test-job-123" + mock_job.job_arn = "arn:aws:sagemaker:us-east-1:123456789012:job/test-job-123" + self._mock_trainer_instance.train.return_value = mock_job + + with patch.dict("sys.modules", modules): + with patch.object( + ops, + "_get_or_create_checkpoint_model_package_group_arn", + return_value="arn:checkpoint", + ): + result = ops.execute_mtrl( + model=Model.NOVA_LITE_2, + job_name="test-job", + data_s3_path="s3://bucket/data", + output_s3_path="s3://bucket/output", + ) + + self.assertEqual(result, "test-job-123") + call_kwargs = self._mock_trainer_cls.call_args[1] + self.assertEqual(call_kwargs["agent_env"], self.MOCK_AGENT_CORE_ARN) + self.assertEqual(call_kwargs["model"], Model.NOVA_LITE_2.hub_content_name) + + @patch("amzn_nova_forge.manager.mtrl_manager.logger") + def test_execute_mtrl_with_lambda(self, _mock_logger): + ops = self._create_mtrl_ops(rft_lambda=self.MOCK_LAMBDA_ARN) + modules = self._mock_sagemaker_modules() + + mock_job = MagicMock() + mock_job.job_name = "lambda-job" + mock_job.job_arn = "arn:aws:sagemaker:us-east-1:123456789012:job/lambda-job" + self._mock_trainer_instance.train.return_value = mock_job + + with patch.dict("sys.modules", modules): + with patch.object( + ops, + "_get_or_create_checkpoint_model_package_group_arn", + return_value="arn:checkpoint", + ): + result = ops.execute_mtrl( + model=Model.NOVA_LITE_2, + job_name="lambda-job", + data_s3_path="s3://bucket/data", + ) + + call_kwargs = self._mock_trainer_cls.call_args[1] + self.assertEqual(call_kwargs["agent_env"], self.MOCK_LAMBDA_ARN) + + @patch("amzn_nova_forge.manager.mtrl_manager.logger") + def test_execute_mtrl_with_model_path_resolves_model_package(self, _mock_logger): + ops = self._create_mtrl_ops(agent_core_arn=self.MOCK_AGENT_CORE_ARN) + modules = self._mock_sagemaker_modules() + + mock_model_package = MagicMock() + self._mock_model_package_cls.get.return_value = mock_model_package + + mock_job = MagicMock() + mock_job.job_name = "iter-job" + mock_job.job_arn = "arn:aws:sagemaker:us-east-1:123456789012:job/iter-job" + self._mock_trainer_instance.train.return_value = mock_job + + with patch.dict("sys.modules", modules): + with patch.object( + ops, + "_get_or_create_checkpoint_model_package_group_arn", + return_value="arn:checkpoint", + ): + result = ops.execute_mtrl( + model=Model.NOVA_LITE_2, + job_name="iter-job", + data_s3_path="s3://bucket/data", + model_path=self.MOCK_MODEL_PACKAGE_ARN, + ) + + self.assertEqual(result, "iter-job") + self._mock_model_package_cls.get.assert_called_once_with(self.MOCK_MODEL_PACKAGE_ARN) + call_kwargs = self._mock_trainer_cls.call_args[1] + self.assertEqual(call_kwargs["model"], mock_model_package) + + @patch("amzn_nova_forge.manager.mtrl_manager.logger") + def test_execute_mtrl_without_model_path_uses_base_model(self, _mock_logger): + ops = self._create_mtrl_ops(agent_core_arn=self.MOCK_AGENT_CORE_ARN) + modules = self._mock_sagemaker_modules() + + mock_job = MagicMock() + mock_job.job_name = "base-job" + mock_job.job_arn = "arn:aws:sagemaker:us-east-1:123456789012:job/base-job" + self._mock_trainer_instance.train.return_value = mock_job + + with patch.dict("sys.modules", modules): + with patch.object( + ops, + "_get_or_create_checkpoint_model_package_group_arn", + return_value="arn:checkpoint", + ): + ops.execute_mtrl( + model=Model.NOVA_LITE_2, + job_name="base-job", + data_s3_path="s3://bucket/data", + ) + + self._mock_model_package_cls.get.assert_not_called() + call_kwargs = self._mock_trainer_cls.call_args[1] + self.assertEqual(call_kwargs["model"], Model.NOVA_LITE_2.hub_content_name) + + @patch("amzn_nova_forge.manager.mtrl_manager.logger") + def test_execute_mtrl_agent_core_takes_priority_over_lambda(self, _mock_logger): + ops = self._create_mtrl_ops( + agent_core_arn=self.MOCK_AGENT_CORE_ARN, + rft_lambda=self.MOCK_LAMBDA_ARN, + ) + modules = self._mock_sagemaker_modules() + + mock_job = MagicMock() + mock_job.job_name = "priority-job" + mock_job.job_arn = "arn:aws:sagemaker:us-east-1:123456789012:job/priority-job" + self._mock_trainer_instance.train.return_value = mock_job + + with patch.dict("sys.modules", modules): + with patch.object( + ops, + "_get_or_create_checkpoint_model_package_group_arn", + return_value="arn:checkpoint", + ): + ops.execute_mtrl( + model=Model.NOVA_LITE_2, + job_name="priority-job", + data_s3_path="s3://bucket/data", + ) + + call_kwargs = self._mock_trainer_cls.call_args[1] + self.assertEqual(call_kwargs["agent_env"], self.MOCK_AGENT_CORE_ARN) + + +class TestSMTJServerlessMTRLEval(unittest.TestCase): + """Tests for MTRLOperations.execute_mtrl_eval.""" + + MOCK_AGENT_CORE_ARN = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/my-agent" + MOCK_MODEL_PACKAGE_ARN = "arn:aws:sagemaker:us-east-1:123456789012:model-package/grp/1" + + def _create_mtrl_ops(self): + from amzn_nova_forge.manager.mtrl_manager import MTRLOperations + + ops = MTRLOperations.__new__(MTRLOperations) + ops.agent_core_arn = self.MOCK_AGENT_CORE_ARN + ops.rft_lambda = None + ops.model_package_group_arn = ( + "arn:aws:sagemaker:us-east-1:123456789012:model-package-group/grp" + ) + ops.execution_role = "arn:aws:iam::123456789012:role/role" + ops.subnets = None + ops.security_group_ids = None + ops.kms_key_id = None + ops.region = "us-east-1" + return ops + + def _mock_sagemaker_modules(self): + import types + + mock_sm_train = types.ModuleType("sagemaker.train") + mock_sm_train_mtrl = types.ModuleType("sagemaker.train.multi_turn_rl_trainer") + mock_sm_train_eval = types.ModuleType("sagemaker.train.evaluate") + mock_sm_core = types.ModuleType("sagemaker.core") + mock_sm_core_resources = types.ModuleType("sagemaker.core.resources") + + self._mock_trainer_cls = MagicMock() + self._mock_trainer_instance = MagicMock() + self._mock_trainer_cls.return_value = self._mock_trainer_instance + + self._mock_evaluator_cls = MagicMock() + self._mock_evaluator_instance = MagicMock() + self._mock_evaluator_cls.return_value = self._mock_evaluator_instance + + self._mock_execution = MagicMock() + self._mock_execution.arn = "arn:aws:sagemaker:us-east-1:123456789012:pipeline/eval-exec" + self._mock_evaluator_instance.evaluate.return_value = self._mock_execution + + mock_sm_train_mtrl.MultiTurnRLTrainer = self._mock_trainer_cls + mock_sm_train_eval.MultiTurnRLEvaluator = self._mock_evaluator_cls + mock_sm_core_resources.ModelPackage = MagicMock() + + return { + "sagemaker.train": mock_sm_train, + "sagemaker.train.multi_turn_rl_trainer": mock_sm_train_mtrl, + "sagemaker.train.evaluate": mock_sm_train_eval, + "sagemaker.core": mock_sm_core, + "sagemaker.core.resources": mock_sm_core_resources, + } + + @patch("amzn_nova_forge.manager.mtrl_manager.logger") + def test_execute_mtrl_eval_with_training_job_name(self, _mock_logger): + ops = self._create_mtrl_ops() + modules = self._mock_sagemaker_modules() + + mock_attached_job = MagicMock() + self._mock_trainer_cls.attach.return_value = mock_attached_job + + with patch.dict("sys.modules", modules): + with patch.object( + ops, + "_get_or_create_checkpoint_model_package_group_arn", + return_value="arn:checkpoint", + ): + result = ops.execute_mtrl_eval( + model=Model.NOVA_LITE_2, + data_s3_path="s3://bucket/eval-data", + output_s3_path="s3://bucket/output", + model_path=self.MOCK_MODEL_PACKAGE_ARN, + training_job_name="completed-training-job", + ) + + self.assertEqual(result, self._mock_execution) + # Should have built a trainer and attached the job + self._mock_trainer_cls.assert_called_once() + self._mock_trainer_cls.attach.assert_called_once_with("completed-training-job") + # Evaluator should receive the trainer as model + eval_call_kwargs = self._mock_evaluator_cls.call_args[1] + self.assertEqual(eval_call_kwargs["model"], self._mock_trainer_instance) + + @patch("amzn_nova_forge.manager.mtrl_manager.logger") + def test_execute_mtrl_eval_fallback_without_training_job(self, _mock_logger): + ops = self._create_mtrl_ops() + modules = self._mock_sagemaker_modules() + + with patch.dict("sys.modules", modules): + with patch.object( + ops, + "_get_or_create_checkpoint_model_package_group_arn", + return_value="arn:checkpoint", + ): + result = ops.execute_mtrl_eval( + model=Model.NOVA_LITE_2, + data_s3_path="s3://bucket/eval-data", + model_path=self.MOCK_MODEL_PACKAGE_ARN, + ) + + self.assertEqual(result, self._mock_execution) + # Should NOT attach a job — uses fallback path + self._mock_trainer_cls.attach.assert_not_called() + # Evaluator should receive model name (base hub content) since model_path is ARN + eval_call_kwargs = self._mock_evaluator_cls.call_args[1] + self.assertEqual(eval_call_kwargs["model"], Model.NOVA_LITE_2.hub_content_name) + # Should inject model_package_arn post-init + self.assertEqual( + self._mock_evaluator_instance._source_model_package_arn_cache, + self.MOCK_MODEL_PACKAGE_ARN, + ) + + @patch("amzn_nova_forge.manager.mtrl_manager.logger") + def test_execute_mtrl_eval_applies_overrides(self, _mock_logger): + ops = self._create_mtrl_ops() + modules = self._mock_sagemaker_modules() + + with patch.dict("sys.modules", modules): + with patch.object( + ops, + "_get_or_create_checkpoint_model_package_group_arn", + return_value="arn:checkpoint", + ): + ops.execute_mtrl_eval( + model=Model.NOVA_LITE_2, + data_s3_path="s3://bucket/eval-data", + overrides={"global_batch_size": 8}, + ) + + # Verify evaluate was called (overrides applied without error) + self._mock_evaluator_instance.evaluate.assert_called_once() + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/model/result/test_eval_result.py b/tests/unit/model/result/test_eval_result.py index e7358fc..a1196d8 100644 --- a/tests/unit/model/result/test_eval_result.py +++ b/tests/unit/model/result/test_eval_result.py @@ -462,6 +462,7 @@ def test_load_with_class_name(self, mock_open, mock_json_load): started_time=self.started_time, eval_task=self.eval_task, eval_output_path=self.eval_output_path, + region=None, ) @patch("amzn_nova_forge.core.result.job_result.json.load") @@ -530,6 +531,7 @@ def test_load_with_unknown_fields_succeed(self, mock_open, mock_json_load): started_time=self.started_time, eval_task=self.eval_task, eval_output_path=self.eval_output_path, + region=None, ) def test_get_method_returns_dict(self): @@ -637,5 +639,185 @@ def test_smtj_evaluation_result_passes_region_to_sagemaker(self, mock_boto3): mock_boto3.client.assert_called_with("sagemaker", region_name="eu-west-1") +class TestSMTJEvaluationResultMTRL(unittest.TestCase): + """Tests for MTRL eval pipeline status polling.""" + + @patch("boto3.client") + def test_mtrl_eval_uses_pipeline_status_manager(self, mock_boto3): + mock_boto3.return_value = Mock() + result = SMTJEvaluationResult( + job_id="arn:aws:sagemaker:us-east-1:123456789012:pipeline/SagemakerEvaluation-MTRLEvaluation/execution/abc123", + started_time=datetime.now(), + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + eval_output_path="s3://bucket/output/", + s3_client=Mock(), + ) + from amzn_nova_forge.core.result.job_result import MTRLStatusManager + + self.assertIsInstance(result._status_manager, MTRLStatusManager) + + @patch("boto3.client") + def test_mtrl_eval_polls_pipeline_execution(self, mock_boto3): + mock_sm_client = Mock() + mock_sm_client.describe_pipeline_execution.return_value = { + "PipelineExecutionStatus": "Executing" + } + mock_boto3.return_value = mock_sm_client + + result = SMTJEvaluationResult( + job_id="arn:aws:sagemaker:us-east-1:123456789012:pipeline/Test/execution/xyz", + started_time=datetime.now(), + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + eval_output_path="s3://bucket/output/", + s3_client=Mock(), + ) + + status, raw = result.get_job_status() + self.assertEqual(status, JobStatus.IN_PROGRESS) + self.assertEqual(raw, "Executing") + + @patch("boto3.client") + def test_mtrl_eval_pipeline_succeeded(self, mock_boto3): + mock_sm_client = Mock() + mock_sm_client.describe_pipeline_execution.return_value = { + "PipelineExecutionStatus": "Succeeded" + } + mock_boto3.return_value = mock_sm_client + + result = SMTJEvaluationResult( + job_id="arn:aws:sagemaker:us-east-1:123456789012:pipeline/Test/execution/xyz", + started_time=datetime.now(), + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + eval_output_path="s3://bucket/output/", + s3_client=Mock(), + ) + + status, raw = result.get_job_status() + self.assertEqual(status, JobStatus.COMPLETED) + self.assertEqual(raw, "Succeeded") + + @patch("boto3.client") + def test_mtrl_eval_pipeline_failed(self, mock_boto3): + mock_sm_client = Mock() + mock_sm_client.describe_pipeline_execution.return_value = { + "PipelineExecutionStatus": "Failed" + } + mock_boto3.return_value = mock_sm_client + + result = SMTJEvaluationResult( + job_id="arn:aws:sagemaker:us-east-1:123456789012:pipeline/Test/execution/xyz", + started_time=datetime.now(), + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + eval_output_path="s3://bucket/output/", + s3_client=Mock(), + ) + + status, raw = result.get_job_status() + self.assertEqual(status, JobStatus.FAILED) + + @patch("boto3.client") + def test_non_mtrl_eval_uses_smtj_status_manager(self, mock_boto3): + mock_boto3.return_value = Mock() + result = SMTJEvaluationResult( + job_id="some-training-job", + started_time=datetime.now(), + eval_task=EvaluationTask.MMLU, + eval_output_path="s3://bucket/output/", + sagemaker_client=Mock(), + s3_client=Mock(), + ) + from amzn_nova_forge.core.result.job_result import SMTJStatusManager + + self.assertIsInstance(result._status_manager, SMTJStatusManager) + + @patch("boto3.client") + def test_mtrl_eval_extracts_region_from_arn(self, mock_boto3): + mock_sm_client = Mock() + mock_sm_client.describe_pipeline_execution.return_value = { + "PipelineExecutionStatus": "Executing" + } + mock_boto3.return_value = mock_sm_client + + result = SMTJEvaluationResult( + job_id="arn:aws:sagemaker:eu-west-1:123456789012:pipeline/Test/execution/xyz", + started_time=datetime.now(), + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + eval_output_path="s3://bucket/output/", + s3_client=Mock(), + region="us-east-1", + ) + + result.get_job_status() + mock_boto3.assert_any_call("sagemaker", region_name="eu-west-1") + + @patch("boto3.client") + def test_mtrl_eval_resolve_start_time_pipeline(self, mock_boto3): + from datetime import timezone + + mock_sm_client = Mock() + mock_sm_client.describe_pipeline_execution.return_value = { + "CreationTime": datetime(2026, 6, 3, 18, 0, 0, tzinfo=timezone.utc) + } + mock_boto3.return_value = mock_sm_client + + from amzn_nova_forge.core.result.job_result import MTRLStatusManager + + mgr = MTRLStatusManager(region="us-east-1") + result = mgr.resolve_start_time( + "arn:aws:sagemaker:us-east-1:123456789012:pipeline/Test/execution/xyz" + ) + self.assertEqual(result, datetime(2026, 6, 3, 18, 0, 0, tzinfo=timezone.utc)) + + @patch("boto3.client") + def test_dump_pipeline_arn_with_job_name(self, mock_boto3): + import os + + mock_boto3.return_value = Mock() + result = SMTJEvaluationResult( + job_id="arn:aws:sagemaker:us-east-1:123456789012:pipeline/Eval/execution/abc123", + started_time=datetime.now(), + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + eval_output_path="s3://bucket/output/", + s3_client=Mock(), + ) + result._job_name = "my-eval" + path = result.dump() + self.assertEqual(str(path), "my-eval-abc123_SMTJServerless.json") + os.remove(str(path)) + + @patch("boto3.client") + def test_dump_pipeline_arn_without_job_name(self, mock_boto3): + import os + + mock_boto3.return_value = Mock() + result = SMTJEvaluationResult( + job_id="arn:aws:sagemaker:us-east-1:123456789012:pipeline/Eval/execution/xyz789", + started_time=datetime.now(), + eval_task=EvaluationTask.RFT_MULTITURN_EVAL, + eval_output_path="s3://bucket/output/", + s3_client=Mock(), + ) + path = result.dump() + self.assertEqual(str(path), "xyz789_SMTJServerless.json") + os.remove(str(path)) + + @patch("boto3.client") + def test_dump_regular_job_id(self, mock_boto3): + import os + + mock_boto3.return_value = Mock() + result = SMTJEvaluationResult( + job_id="my-training-job-123", + started_time=datetime.now(), + eval_task=EvaluationTask.MMLU, + eval_output_path="s3://bucket/output/", + sagemaker_client=Mock(), + s3_client=Mock(), + ) + path = result.dump() + self.assertEqual(str(path), "my-training-job-123_SMTJ.json") + os.remove(str(path)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/model/result/test_training_result.py b/tests/unit/model/result/test_training_result.py index 8f433bb..8d1ecf2 100644 --- a/tests/unit/model/result/test_training_result.py +++ b/tests/unit/model/result/test_training_result.py @@ -438,5 +438,241 @@ def test_bedrock_training_result_passes_region(self): mock_boto3.assert_called_once_with("bedrock", region_name="eu-west-1") +class TestSMTJTrainingResultMTRL(unittest.TestCase): + """Tests for MTRL-specific methods on SMTJTrainingResult.""" + + def setUp(self): + self.model_artifacts = ModelArtifacts( + output_s3_path="s3://bucket/output/", + output_model_arn="arn:aws:sagemaker:us-east-1:123456789012:model-package/grp/1", + ) + + @patch("amzn_nova_forge.core.result.training_result.boto3.client") + def test_wait_raises_for_non_mtrl(self, _mock_client): + result = SMTJTrainingResult( + job_id="sft-job", + started_time=datetime(2026, 1, 1), + method=TrainingMethod.SFT_LORA, + model_artifacts=self.model_artifacts, + model_type=Model.NOVA_LITE_2, + ) + with self.assertRaises(NotImplementedError): + result.wait() + + @patch("amzn_nova_forge.core.result.training_result.boto3.client") + def test_get_training_metrics_raises_for_non_mtrl(self, _mock_client): + result = SMTJTrainingResult( + job_id="sft-job", + started_time=datetime(2026, 1, 1), + method=TrainingMethod.SFT_LORA, + model_artifacts=self.model_artifacts, + model_type=Model.NOVA_LITE_2, + ) + with self.assertRaises(NotImplementedError): + result.get_training_metrics() + + @patch("amzn_nova_forge.core.result.training_result.boto3.client") + def test_wait_delegates_to_rft_job(self, _mock_client): + import sys + import types + + mock_agent_module = types.ModuleType("sagemaker.train.agent_rft_job") + mock_rft_job = Mock() + mock_rft_job.output_model_package_arn = ( + "arn:aws:sagemaker:us-east-1:123456789012:model-package/grp/2" + ) + mock_agent_cls = Mock() + mock_agent_cls.get.return_value = mock_rft_job + mock_agent_module.AgentRFTJob = mock_agent_cls + + result = SMTJTrainingResult( + job_id="mtrl-job", + started_time=datetime(2026, 1, 1), + method=TrainingMethod.RFT_MULTITURN_LORA, + model_artifacts=ModelArtifacts(output_s3_path="s3://bucket/output/"), + model_type=Model.NOVA_LITE_2, + region="us-east-1", + ) + + with patch.dict("sys.modules", {"sagemaker.train.agent_rft_job": mock_agent_module}): + result.wait(poll=5, timeout=60) + + mock_rft_job.wait.assert_called_once_with(poll=5, timeout=60) + mock_rft_job.refresh.assert_called_once() + self.assertEqual( + result.model_artifacts.output_model_arn, + "arn:aws:sagemaker:us-east-1:123456789012:model-package/grp/2", + ) + + @patch("amzn_nova_forge.core.result.training_result.boto3.client") + def test_get_training_metrics_delegates_to_rft_job(self, _mock_client): + import sys + import types + + mock_agent_module = types.ModuleType("sagemaker.train.agent_rft_job") + mock_rft_job = Mock() + mock_rft_job.output_model_package_arn = None + mock_rft_job.get_training_metrics.return_value = [{"step": 1, "reward": 0.5}] + mock_agent_cls = Mock() + mock_agent_cls.get.return_value = mock_rft_job + mock_agent_module.AgentRFTJob = mock_agent_cls + + result = SMTJTrainingResult( + job_id="mtrl-job", + started_time=datetime(2026, 1, 1), + method=TrainingMethod.RFT_MULTITURN_LORA, + model_artifacts=ModelArtifacts(output_s3_path="s3://bucket/output/"), + model_type=Model.NOVA_LITE_2, + region="us-east-1", + ) + + with patch.dict("sys.modules", {"sagemaker.train.agent_rft_job": mock_agent_module}): + metrics = result.get_training_metrics() + + self.assertEqual(metrics, [{"step": 1, "reward": 0.5}]) + + +class TestMTRLStatusManager(unittest.TestCase): + """Tests for MTRLStatusManager.""" + + def test_get_job_status_in_progress(self): + import sys + import types + + from amzn_nova_forge.core.result.job_result import MTRLStatusManager + + mock_agent_module = types.ModuleType("sagemaker.train.agent_rft_job") + mock_rft_job = Mock() + mock_rft_job.job_status = "InProgress" + mock_agent_cls = Mock() + mock_agent_cls.get.return_value = mock_rft_job + mock_agent_module.AgentRFTJob = mock_agent_cls + + manager = MTRLStatusManager(region="us-east-1") + + with patch.dict("sys.modules", {"sagemaker.train.agent_rft_job": mock_agent_module}): + status, raw = manager.get_job_status("mtrl-job-123") + + self.assertEqual(status, JobStatus.IN_PROGRESS) + self.assertEqual(raw, "InProgress") + + def test_get_job_status_completed(self): + import sys + import types + + from amzn_nova_forge.core.result.job_result import MTRLStatusManager + + mock_agent_module = types.ModuleType("sagemaker.train.agent_rft_job") + mock_rft_job = Mock() + mock_rft_job.job_status = "Completed" + mock_agent_cls = Mock() + mock_agent_cls.get.return_value = mock_rft_job + mock_agent_module.AgentRFTJob = mock_agent_cls + + manager = MTRLStatusManager(region="us-east-1") + + with patch.dict("sys.modules", {"sagemaker.train.agent_rft_job": mock_agent_module}): + status, raw = manager.get_job_status("mtrl-job-123") + + self.assertEqual(status, JobStatus.COMPLETED) + self.assertEqual(raw, "Completed") + + def test_get_job_status_caches_terminal_state(self): + import sys + import types + + from amzn_nova_forge.core.result.job_result import MTRLStatusManager + + mock_agent_module = types.ModuleType("sagemaker.train.agent_rft_job") + mock_rft_job = Mock() + mock_rft_job.job_status = "Completed" + mock_agent_cls = Mock() + mock_agent_cls.get.return_value = mock_rft_job + mock_agent_module.AgentRFTJob = mock_agent_cls + + manager = MTRLStatusManager(region="us-east-1") + + with patch.dict("sys.modules", {"sagemaker.train.agent_rft_job": mock_agent_module}): + manager.get_job_status("mtrl-job-123") + status, raw = manager.get_job_status("mtrl-job-123") + + # Should only call API once due to caching + mock_agent_cls.get.assert_called_once() + self.assertEqual(status, JobStatus.COMPLETED) + + def test_resolve_start_time(self): + import sys + import types + + from amzn_nova_forge.core.result.job_result import MTRLStatusManager + + mock_agent_module = types.ModuleType("sagemaker.train.agent_rft_job") + mock_rft_job = Mock() + mock_rft_job.creation_time = datetime(2026, 5, 20, 10, 0, 0) + mock_agent_cls = Mock() + mock_agent_cls.get.return_value = mock_rft_job + mock_agent_module.AgentRFTJob = mock_agent_cls + + manager = MTRLStatusManager(region="us-east-1") + + with patch.dict("sys.modules", {"sagemaker.train.agent_rft_job": mock_agent_module}): + start_time = manager.resolve_start_time("mtrl-job-123") + + self.assertEqual(start_time, datetime(2026, 5, 20, 10, 0, 0)) + + +class TestMTRLLogMonitor(unittest.TestCase): + """Tests for MTRLLogMonitor.""" + + def test_from_job_id(self): + from amzn_nova_forge.monitor.log_monitor import MTRLLogMonitor + + monitor = MTRLLogMonitor.from_job_id(job_id="mtrl-job-123", region="us-east-1") + self.assertEqual(monitor.job_id, "mtrl-job-123") + self.assertEqual(monitor._region, "us-east-1") + + def test_show_logs_completed_job(self): + import sys + import types + + from amzn_nova_forge.monitor.log_monitor import MTRLLogMonitor + + mock_agent_module = types.ModuleType("sagemaker.train.agent_rft_job") + mock_rft_job = Mock() + mock_rft_job.job_status = "Completed" + mock_agent_cls = Mock() + mock_agent_cls.get.return_value = mock_rft_job + mock_agent_module.AgentRFTJob = mock_agent_cls + + monitor = MTRLLogMonitor(job_id="mtrl-job-123", region="us-east-1") + + with patch.dict("sys.modules", {"sagemaker.train.agent_rft_job": mock_agent_module}): + monitor.show_logs() + + mock_rft_job.refresh.assert_called_once() + mock_rft_job.get_training_metrics.assert_called_once() + mock_rft_job.wait.assert_not_called() + + def test_show_logs_running_job_calls_wait(self): + import sys + import types + + from amzn_nova_forge.monitor.log_monitor import MTRLLogMonitor + + mock_agent_module = types.ModuleType("sagemaker.train.agent_rft_job") + mock_rft_job = Mock() + mock_rft_job.job_status = "InProgress" + mock_agent_cls = Mock() + mock_agent_cls.get.return_value = mock_rft_job + mock_agent_module.AgentRFTJob = mock_agent_cls + + monitor = MTRLLogMonitor(job_id="mtrl-job-123", region="us-east-1") + + with patch.dict("sys.modules", {"sagemaker.train.agent_rft_job": mock_agent_module}): + monitor.show_logs(poll=10, timeout=120) + + mock_rft_job.wait.assert_called_once_with(poll=10, timeout=120) + + 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 bfa55c0..0dd312e 100644 --- a/tests/unit/model/test_nova_model_customizer.py +++ b/tests/unit/model/test_nova_model_customizer.py @@ -215,7 +215,7 @@ def test_serverless_model_path_must_be_arn(self, mock_setup): model=Model.NOVA_MICRO, method=TrainingMethod.SFT_LORA, infra=runtime, - model_path="arn:aws:sagemaker:us-east-1:123:model-package/group/1", + model_path="arn:aws:sagemaker:us-east-1:123456789012:model-package/group/1", ) # None should not raise (no iterative training) @@ -1816,7 +1816,7 @@ def test_deploy_sagemaker_success( mock_boto_client, ): mock_sagemaker_role_creation.return_value = {"Role": {"Arn": "sagemaker:role:arn"}} - mock_create_model.return_value = "arn:aws:sagemaker:us-east-1:123:model/test-model" + mock_create_model.return_value = "arn:aws:sagemaker:us-east-1:123456789012:model/test-model" mock_create_endpoint.return_value = "sagemaker:endpoint:arn" result = self.customizer.deploy( @@ -1833,7 +1833,7 @@ def test_deploy_sagemaker_success( self.assertIsNotNone(result.model_publish) self.assertEqual( result.model_publish.model_arn, - "arn:aws:sagemaker:us-east-1:123:model/test-model", + "arn:aws:sagemaker:us-east-1:123456789012:model/test-model", ) @patch("boto3.client") @@ -1889,6 +1889,7 @@ def test_deploy_sagemaker_with_job_success( "Value": "s3://xn---checkpointbucket/ckpt/", } ], + model_package_name=None, ) mock_create_endpoint.assert_called_once_with( model_name="nova-micro-sft-lora-sagemaker-model", @@ -1947,7 +1948,7 @@ def test_deploy_sagemaker_failure( mock_boto_client, ): mock_sagemaker_role_creation.return_value = {"Role": {"Arn": "sagemaker:role:arn"}} - mock_create_model.return_value = "arn:aws:sagemaker:us-east-1:123:model/test-model" + mock_create_model.return_value = "arn:aws:sagemaker:us-east-1:123456789012:model/test-model" mock_create_endpoint.side_effect = Exception("Failed to create deployment") with self.assertRaises(Exception) as context: @@ -1972,7 +1973,7 @@ def test_deploy_to_sagemaker_with_model_deploy_result( ): """deploy_to_sagemaker with model_deploy_result skips model creation.""" mock_sagemaker_role_creation.return_value = {"Role": {"Arn": "sagemaker:role:arn"}} - mock_create_model.return_value = "arn:aws:sagemaker:us-east-1:123:model/test-model" + mock_create_model.return_value = "arn:aws:sagemaker:us-east-1:123456789012:model/test-model" mock_create_endpoint.return_value = "sagemaker:endpoint:arn" mock_boto_client.side_effect = self._make_boto_client_dispatcher() @@ -2007,7 +2008,7 @@ def test_deploy_to_sagemaker_endpoint_failure_shows_retry_hint( ): """Endpoint failure error message includes model ARN and retry hint.""" mock_sagemaker_role_creation.return_value = {"Role": {"Arn": "sagemaker:role:arn"}} - mock_create_model.return_value = "arn:aws:sagemaker:us-east-1:123:model/test-model" + mock_create_model.return_value = "arn:aws:sagemaker:us-east-1:123456789012:model/test-model" mock_create_endpoint.side_effect = Exception("Endpoint creation failed") mock_boto_client.side_effect = self._make_boto_client_dispatcher() @@ -2070,7 +2071,7 @@ def test_deploy_to_sagemaker_skip_model_reuse( ): """deploy_to_sagemaker with skip_model_reuse=True skips tag-based model discovery.""" mock_sagemaker_role_creation.return_value = {"Role": {"Arn": "sagemaker:role:arn"}} - mock_create_model.return_value = "arn:aws:sagemaker:us-east-1:123:model/test-model" + mock_create_model.return_value = "arn:aws:sagemaker:us-east-1:123456789012:model/test-model" mock_create_endpoint.return_value = "sagemaker:endpoint:arn" mock_boto_client.side_effect = self._make_boto_client_dispatcher() @@ -2898,7 +2899,7 @@ def test_invoke_inference_sagemaker_endpoint(self, mock_invoke_sagemaker, mock_b # Assert mock_boto3_client.assert_called_once_with("sagemaker-runtime", region_name="us-east-1") mock_invoke_sagemaker.assert_called_once_with( - request_body, "test-endpoint", mock_runtime_client + request_body, "test-endpoint", mock_runtime_client, inference_component_name=None ) assert result == "Inference Result" @@ -4847,5 +4848,46 @@ def test_get_config_passes_overrides(self, mock_init, mock_get_config): mock_get_config.assert_called_once_with(overrides=overrides) +class TestCreateCustomModelDataSource(TestNovaModelCustomizer): + """Tests for facade-level create_custom_model() with custom_model_data_source.""" + + def test_raises_when_no_source_provided(self): + """Validation at facade level requires at least one source.""" + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + with self.assertRaises(ValueError) as ctx: + self.customizer.create_custom_model() + + self.assertIn("custom_model_data_source", str(ctx.exception)) + + def test_delegates_custom_model_data_source_to_deployer(self): + """custom_model_data_source is passed through to ForgeDeployer.""" + import warnings + + data_source = {"modelPackageArnDataSource": {"modelPackageArn": "arn:pkg"}} + + with patch.object(self.customizer, "_build_deployer") as mock_build: + mock_deployer = MagicMock() + mock_deployer.create_custom_model.return_value = MagicMock( + model_arn="arn:test", model_name="test" + ) + mock_build.return_value = mock_deployer + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + self.customizer.create_custom_model(custom_model_data_source=data_source) + + mock_deployer.create_custom_model.assert_called_once_with( + model_artifact_path=None, + endpoint_name=None, + execution_role_name=None, + tags=None, + skip_model_reuse=False, + custom_model_data_source=data_source, + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/monitor/test_log_monitor.py b/tests/unit/monitor/test_log_monitor.py index 19e6d34..269ce60 100644 --- a/tests/unit/monitor/test_log_monitor.py +++ b/tests/unit/monitor/test_log_monitor.py @@ -1036,6 +1036,137 @@ def test_plot_metrics_bedrock_get_metrics_raises_not_implemented(self): monitor.plot_metrics(TrainingMethod.SFT_FULL) +class TestMTRLLogMonitor(unittest.TestCase): + def test_from_job_id_no_category(self): + from amzn_nova_forge.monitor.log_monitor import MTRLLogMonitor + + monitor = MTRLLogMonitor.from_job_id(job_id="test-job", region="us-east-1") + self.assertEqual(monitor.job_id, "test-job") + self.assertIsNone(monitor._job_category) + + def test_from_job_id_with_category(self): + from amzn_nova_forge.monitor.log_monitor import MTRLLogMonitor + + monitor = MTRLLogMonitor.from_job_id( + job_id="test-job", region="us-east-1", job_category="AgentRFTEvaluation" + ) + self.assertEqual(monitor._job_category, "AgentRFTEvaluation") + + @patch("boto3.client") + def test_detect_job_category_finds_eval(self, mock_boto3): + from amzn_nova_forge.monitor.log_monitor import MTRLLogMonitor + + mock_logs_client = Mock() + mock_logs_client.describe_log_streams.side_effect = [ + {"logStreams": []}, + {"logStreams": [{"logStreamName": "test-eval-job/"}]}, + ] + mock_boto3.return_value = mock_logs_client + + monitor = MTRLLogMonitor(job_id="test-eval-job", region="us-east-1") + category = monitor._detect_job_category() + + self.assertEqual(category, "AgentRFTEvaluation") + + @patch("boto3.client") + def test_detect_job_category_finds_training(self, mock_boto3): + from amzn_nova_forge.monitor.log_monitor import MTRLLogMonitor + + mock_logs_client = Mock() + mock_logs_client.describe_log_streams.side_effect = [ + {"logStreams": [{"logStreamName": "test-train-job/"}]}, + ] + mock_boto3.return_value = mock_logs_client + + monitor = MTRLLogMonitor(job_id="test-train-job", region="us-east-1") + category = monitor._detect_job_category() + + self.assertEqual(category, "AgentRFT") + + @patch("boto3.client") + def test_detect_job_category_skips_if_already_set(self, mock_boto3): + from amzn_nova_forge.monitor.log_monitor import MTRLLogMonitor + + monitor = MTRLLogMonitor( + job_id="test-job", region="us-east-1", job_category="AgentRFTEvaluation" + ) + category = monitor._detect_job_category() + + self.assertEqual(category, "AgentRFTEvaluation") + mock_boto3.assert_not_called() + + @patch("builtins.print") + @patch("boto3.client") + def test_show_eval_logs(self, mock_boto3, mock_print): + from amzn_nova_forge.monitor.log_monitor import MTRLLogMonitor + + mock_logs_client = Mock() + mock_logs_client.describe_log_streams.return_value = { + "logStreams": [{"logStreamName": "eval-job-123/"}] + } + mock_logs_client.get_log_events.return_value = { + "events": [ + {"message": "Progress: 5/50\n"}, + {"message": "Progress: 10/50\n"}, + ] + } + mock_boto3.return_value = mock_logs_client + + monitor = MTRLLogMonitor( + job_id="eval-job-123", region="us-east-1", job_category="AgentRFTEvaluation" + ) + monitor.show_logs(limit=5) + + mock_print.assert_any_call("Progress: 5/50") + mock_print.assert_any_call("Progress: 10/50") + + @patch("builtins.print") + @patch("boto3.client") + def test_show_eval_logs_no_streams(self, mock_boto3, mock_print): + from amzn_nova_forge.monitor.log_monitor import MTRLLogMonitor + + mock_logs_client = Mock() + mock_logs_client.describe_log_streams.return_value = {"logStreams": []} + mock_boto3.return_value = mock_logs_client + + monitor = MTRLLogMonitor( + job_id="no-such-job", region="us-east-1", job_category="AgentRFTEvaluation" + ) + monitor.show_logs() + + mock_print.assert_called_once_with("No log stream found for job 'no-such-job'") + + @patch("amzn_nova_forge.monitor.log_monitor.MTRLLogMonitor.from_job_id") + def test_forge_evaluator_get_logs_routes_to_mtrl_monitor(self, mock_from_job_id): + """ForgeEvaluator.get_logs() routes MTRL eval results to MTRLLogMonitor.""" + from datetime import datetime, timezone + from unittest.mock import patch as mock_patch + + from amzn_nova_forge.core.enums import EvaluationTask + + mock_monitor = Mock() + mock_from_job_id.return_value = mock_monitor + + mock_result = Mock() + mock_result.job_id = "arn:aws:sagemaker:us-east-1:123456789012:pipeline/Eval/execution/abc" + mock_result.eval_task = EvaluationTask.RFT_MULTITURN_EVAL + mock_result.started_time = datetime(2026, 6, 3, tzinfo=timezone.utc) + + from amzn_nova_forge.evaluator.forge_evaluator import ForgeEvaluator + + with mock_patch.object(ForgeEvaluator, "__init__", return_value=None): + evaluator = ForgeEvaluator.__new__(ForgeEvaluator) + evaluator.region = "us-east-1" + evaluator.get_logs(job_result=mock_result, limit=10) + + mock_from_job_id.assert_called_once_with( + job_id=mock_result.job_id, + region="us-east-1", + job_category="AgentRFTEvaluation", + ) + mock_monitor.show_logs.assert_called_once_with(limit=10) + + class TestRegionPropagation(unittest.TestCase): @patch("boto3.client") def test_cloudwatch_log_monitor_passes_region_to_logs_client(self, mock_boto_client): diff --git a/tests/unit/recipe/test_recipe_builder.py b/tests/unit/recipe/test_recipe_builder.py index 82fcb0f..ad1f4ef 100644 --- a/tests/unit/recipe/test_recipe_builder.py +++ b/tests/unit/recipe/test_recipe_builder.py @@ -2202,6 +2202,115 @@ def test_model_name_or_path_s3_override_validated( self.assertEqual(config["run"]["model_name_or_path"], s3_checkpoint) + @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_model_name_or_path_model_package_arn_override_accepted( + self, mock_validator, mock_download, mock_metadata + ): + mock_metadata.return_value = {"recipe_uri": "s3://bucket/recipe"} + + recipe_template = { + "run": { + "name": "{{name}}", + "model_name_or_path": "{{model_name_or_path}}", + } + } + + overrides_template = { + "name": {"default": "", "type": "string"}, + "model_name_or_path": {"default": "models/test", "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") + + mp_arn = "arn:aws:sagemaker:us-east-1:123456789012:model-package/my-mpg/1" + overrides = {"model_name_or_path": mp_arn} + recipe_path, *_ = builder.build_and_validate( + overrides=overrides, + output_recipe_path=output_path, + ) + + with open(recipe_path, "r") as f: + config = yaml.safe_load(f) + + self.assertEqual(config["run"]["model_name_or_path"], mp_arn) + + @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") + @patch("amzn_nova_forge.recipe.recipe_builder.load_file_as_string") + def test_model_name_or_path_model_package_arn_from_input_recipe_accepted( + self, mock_load_file, mock_validator, mock_download, mock_metadata + ): + mock_metadata.return_value = {"recipe_uri": "s3://bucket/recipe"} + + recipe_template = { + "run": { + "name": "{{name}}", + "model_name_or_path": "{{model_name_or_path}}", + } + } + + overrides_template = { + "name": {"default": "", "type": "string"}, + "model_name_or_path": {"default": "models/test", "type": "string"}, + } + + mock_download.return_value = (recipe_template, overrides_template, "image_uri") + + mp_arn = "arn:aws:sagemaker:us-east-1:123456789012:model-package/my-mpg/1" + input_recipe = { + "run": { + "name": "test", + "model_name_or_path": mp_arn, + } + } + mock_load_file.return_value = yaml.dump(input_recipe) + + 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") + input_recipe_path = os.path.join(tmpdir, "input_recipe.yaml") + + recipe_path, *_ = builder.build_and_validate( + input_recipe_path=input_recipe_path, + output_recipe_path=output_path, + ) + + with open(recipe_path, "r") as f: + config = yaml.safe_load(f) + + self.assertEqual(config["run"]["model_name_or_path"], mp_arn) + @patch("amzn_nova_forge.recipe.recipe_builder.logger") @patch("amzn_nova_forge.util.checkpoint_util.validate_checkpoint_uri") @patch("amzn_nova_forge.util.recipe.get_hub_recipe_metadata") diff --git a/tests/unit/trainer/test_forge_trainer.py b/tests/unit/trainer/test_forge_trainer.py index 256c20f..75290ae 100644 --- a/tests/unit/trainer/test_forge_trainer.py +++ b/tests/unit/trainer/test_forge_trainer.py @@ -30,6 +30,7 @@ SMHPRuntimeManager, SMTJRuntimeManager, ) +from amzn_nova_forge.monitor.mlflow_monitor import MLflowMonitor from amzn_nova_forge.trainer.forge_trainer import ForgeTrainer FIXED_OUTPUT_PATH = "s3://sagemaker-nova-123456789012-us-east-1/output" @@ -1249,5 +1250,420 @@ def test_get_config_rft_multiturn_raises(self): self.assertIn("rft_multiturn_infra", str(ctx.exception)) +class TestForgeTrainerModelArn(unittest.TestCase): + """Tests for model_arn parameter on ForgeTrainer for SMTJServerless.""" + + MOCK_MODEL_PACKAGE_ARN = "arn:aws:sagemaker:us-east-1:123456789012:model-package/my-group/1" + + 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 + infra.rft_lambda_arn = None + infra.hub_content_version = None + return infra + + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_model_arn_accepted_for_serverless(self, mock_session, _mock_output): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + infra = self._make_serverless_infra() + + trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + model_arn=self.MOCK_MODEL_PACKAGE_ARN, + ) + + self.assertEqual(trainer.model_arn, self.MOCK_MODEL_PACKAGE_ARN) + self.assertIsNone(trainer.model_s3_path) + + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_model_s3_path_migrated_to_model_arn_for_serverless(self, mock_session, _mock_output): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + infra = self._make_serverless_infra() + + trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + model_s3_path=self.MOCK_MODEL_PACKAGE_ARN, + ) + + self.assertEqual(trainer.model_arn, self.MOCK_MODEL_PACKAGE_ARN) + self.assertIsNone(trainer.model_s3_path) + + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_both_model_arn_and_model_s3_path_raises(self, mock_session, _mock_output): + 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_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + model_s3_path=self.MOCK_MODEL_PACKAGE_ARN, + model_arn=self.MOCK_MODEL_PACKAGE_ARN, + ) + self.assertIn("Cannot specify both", 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_model_arn_invalid_raises(self, mock_session, _mock_output): + 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_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + model_arn="s3://bucket/not-an-arn/", + ) + self.assertIn("model_arn must be a SageMaker model package ARN", 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_no_model_arn_trains_from_base(self, mock_session, _mock_output): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + infra = self._make_serverless_infra() + + trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + ) + + self.assertIsNone(trainer.model_arn) + self.assertIsNone(trainer.model_s3_path) + + +class TestForgeTrainerMTRLTrain(unittest.TestCase): + """Tests for the MTRL serverless training path in ForgeTrainer.train().""" + + MOCK_MLFLOW_ARN = "arn:aws:sagemaker:us-east-1:123456789012:mlflow-tracking-server/my-server" + + 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 + infra.rft_lambda_arn = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/agent" + infra.hub_content_version = None + infra.execute_mtrl.return_value = "mtrl-job-id-123" + return infra + + def _make_mlflow_config(self): + mock_monitor = MagicMock(spec=MLflowMonitor) + mock_monitor.tracking_uri = self.MOCK_MLFLOW_ARN + return ForgeConfig(mlflow_monitor=mock_monitor) + + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_mtrl_train_raises_without_mlflow_config(self, mock_session, _mock_output): + """AgentRFT jobs must have an MLflow config; raise ValueError if missing.""" + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + infra = self._make_serverless_infra() + + trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + ) + + with self.assertRaises(ValueError) as ctx: + trainer.train(job_name="my-mtrl-job") + self.assertIn("MLflow configuration is required", 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_mtrl_train_raises_without_mlflow_tracking_uri(self, mock_session, _mock_output): + """AgentRFT jobs must have a tracking_uri; raise ValueError if None.""" + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + infra = self._make_serverless_infra() + + mock_monitor = MagicMock(spec=MLflowMonitor) + mock_monitor.tracking_uri = None + + trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + config=ForgeConfig(mlflow_monitor=mock_monitor), + ) + + with self.assertRaises(ValueError) as ctx: + trainer.train(job_name="my-mtrl-job") + self.assertIn("MLflow configuration is required", 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_mtrl_train_raises_for_rft_multiturn_full_without_mlflow( + self, mock_session, _mock_output + ): + """RFT_MULTITURN_FULL also requires MLflow config.""" + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + infra = self._make_serverless_infra() + + trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_FULL, + infra=infra, + training_data_s3_path="s3://bucket/data", + ) + + with self.assertRaises(ValueError) as ctx: + trainer.train(job_name="my-mtrl-full-job") + self.assertIn("MLflow configuration is required", str(ctx.exception)) + + @patch("amzn_nova_forge.trainer.forge_trainer.get_model_artifacts") + @patch("amzn_nova_forge.trainer.forge_trainer.persist_result") + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_mtrl_serverless_train_returns_smtj_training_result( + self, mock_session, _mock_output, mock_persist, mock_get_artifacts + ): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + infra = self._make_serverless_infra() + + mock_get_artifacts.return_value = ModelArtifacts( + output_s3_path="s3://bucket/output/", + output_model_arn=None, + ) + + trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + config=self._make_mlflow_config(), + ) + + with patch("amzn_nova_forge.trainer.forge_trainer.SMTJTrainingResult") as MockResult: + mock_result_instance = MagicMock() + mock_result_instance.job_id = "mtrl-job-id-123" + mock_result_instance.model_artifacts = ModelArtifacts( + output_s3_path="s3://bucket/output/" + ) + MockResult.return_value = mock_result_instance + + result = trainer.train(job_name="my-mtrl-job") + + infra.execute_mtrl.assert_called_once() + call_kwargs = infra.execute_mtrl.call_args[1] + self.assertEqual(call_kwargs["model"], Model.NOVA_LITE_2) + self.assertEqual(call_kwargs["data_s3_path"], "s3://bucket/data") + self.assertEqual(call_kwargs["output_s3_path"], FIXED_OUTPUT_PATH) + self.assertIsNone(call_kwargs["model_path"]) + self.assertEqual(result, mock_result_instance) + + @patch("amzn_nova_forge.trainer.forge_trainer.get_model_artifacts") + @patch("amzn_nova_forge.trainer.forge_trainer.persist_result") + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_mtrl_serverless_train_passes_model_arn( + self, mock_session, _mock_output, mock_persist, mock_get_artifacts + ): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + infra = self._make_serverless_infra() + model_arn = "arn:aws:sagemaker:us-east-1:123456789012:model-package/grp/1" + + mock_get_artifacts.return_value = ModelArtifacts( + output_s3_path="s3://bucket/output/", + output_model_arn=None, + ) + + trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + model_arn=model_arn, + config=self._make_mlflow_config(), + ) + + with patch("amzn_nova_forge.trainer.forge_trainer.SMTJTrainingResult") as MockResult: + mock_result_instance = MagicMock() + mock_result_instance.job_id = "mtrl-job-id-456" + mock_result_instance.model_artifacts = ModelArtifacts( + output_s3_path="s3://bucket/output/" + ) + MockResult.return_value = mock_result_instance + + trainer.train(job_name="my-iterative-job") + + call_kwargs = infra.execute_mtrl.call_args[1] + self.assertEqual(call_kwargs["model_path"], model_arn) + + @patch("amzn_nova_forge.trainer.forge_trainer.get_model_artifacts") + @patch("amzn_nova_forge.trainer.forge_trainer.persist_result") + @patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value=FIXED_OUTPUT_PATH, + ) + @patch("boto3.session.Session") + def test_mtrl_serverless_train_passes_overrides( + self, mock_session, _mock_output, mock_persist, mock_get_artifacts + ): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + infra = self._make_serverless_infra() + + mock_get_artifacts.return_value = ModelArtifacts( + output_s3_path="s3://bucket/output/", + ) + + trainer = ForgeTrainer( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + config=self._make_mlflow_config(), + ) + + with patch("amzn_nova_forge.trainer.forge_trainer.SMTJTrainingResult") as MockResult: + mock_result_instance = MagicMock() + mock_result_instance.job_id = "mtrl-job-id-789" + mock_result_instance.model_artifacts = ModelArtifacts( + output_s3_path="s3://bucket/output/" + ) + MockResult.return_value = mock_result_instance + + trainer.train( + job_name="my-overrides-job", + overrides={"global_batch_size": 16, "max_steps": 50}, + ) + + call_kwargs = infra.execute_mtrl.call_args[1] + self.assertEqual(call_kwargs["overrides"], {"global_batch_size": 16, "max_steps": 50}) + + +class TestForgeTrainerGetLogsMTRL(unittest.TestCase): + """Tests for ForgeTrainer.get_logs() with MTRL SMTJTrainingResult.""" + + def _make_trainer(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 + infra.rft_lambda_arn = None + infra.hub_content_version = None + + 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( + model=Model.NOVA_LITE_2, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=infra, + training_data_s3_path="s3://bucket/data", + ) + + @patch("amzn_nova_forge.trainer.forge_trainer.CloudWatchLogMonitor") + def test_get_logs_mtrl_delegates_to_wait(self, MockMonitor): + """When job_result is an MTRL SMTJTrainingResult, get_logs() should call wait() instead of streaming CloudWatch logs.""" + trainer = self._make_trainer() + + mock_result = MagicMock(spec=SMTJTrainingResult) + mock_result._is_mtrl = True + mock_result.job_id = "mtrl-job-abc" + mock_result.started_time = datetime(2025, 5, 1, tzinfo=timezone.utc) + + trainer.get_logs(job_result=mock_result) + + # Should call wait() on the job_result + mock_result.wait.assert_called_once_with(poll=30, timeout=7200) + # Should NOT create a CloudWatch monitor + MockMonitor.assert_not_called() + + @patch("amzn_nova_forge.trainer.forge_trainer.CloudWatchLogMonitor") + def test_get_logs_mtrl_passes_custom_poll_and_timeout(self, MockMonitor): + """MTRL get_logs() should pass through custom poll and timeout values.""" + trainer = self._make_trainer() + + mock_result = MagicMock(spec=SMTJTrainingResult) + mock_result._is_mtrl = True + mock_result.job_id = "mtrl-job-def" + mock_result.started_time = datetime(2025, 5, 1, tzinfo=timezone.utc) + + trainer.get_logs(job_result=mock_result, poll=60, timeout=7200) + + mock_result.wait.assert_called_once_with(poll=60, timeout=7200) + MockMonitor.assert_not_called() + + @patch("amzn_nova_forge.trainer.forge_trainer.CloudWatchLogMonitor") + def test_get_logs_non_mtrl_smtj_result_still_streams_cloudwatch(self, MockMonitor): + """Non-MTRL SMTJTrainingResult should still stream CloudWatch logs normally.""" + trainer = self._make_trainer() + + mock_result = MagicMock(spec=SMTJTrainingResult) + mock_result._is_mtrl = False + mock_result.job_id = "smtj-job-xyz" + mock_result.started_time = datetime(2025, 5, 1, tzinfo=timezone.utc) + + trainer.get_logs(job_result=mock_result) + + # Should NOT call wait() + mock_result.wait.assert_not_called() + # Should create a CloudWatch monitor + MockMonitor.assert_called_once() + MockMonitor.return_value.show_logs.assert_called_once() + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/util/test_metric_util.py b/tests/unit/util/test_metric_util.py new file mode 100644 index 0000000..c77effc --- /dev/null +++ b/tests/unit/util/test_metric_util.py @@ -0,0 +1,216 @@ +# 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. +"""Tests for SFT training metrics CSV generation. + +Covers epoch extraction, CSV orchestration, and ForgeTrainer integration. +""" + +import logging +from datetime import datetime, timezone +from unittest.mock import MagicMock, PropertyMock, create_autospec, patch + +import pandas +import pytest + +from amzn_nova_forge.core.enums import Model, Platform, TrainingMethod +from amzn_nova_forge.core.result.job_result import JobStatus +from amzn_nova_forge.manager.runtime_manager import SMHPRuntimeManager +from amzn_nova_forge.trainer.forge_trainer import ForgeTrainer +from amzn_nova_forge.util.metric_util import ( + _assign_epochs_to_steps, + _extract_epoch_boundaries, + get_metrics, +) + + +def _epoch_event(epoch_idx: int, ts: int) -> dict: + return { + "timestamp": ts, + "message": f'[job] [INFO] {{"EpochIdx": {epoch_idx}}}', + } + + +def _step_event(step: int, loss: float, ts: int) -> dict: + return { + "timestamp": ts, + "message": f"[job] [INFO] global_step: {step} | reduced_train_loss: {loss}", + } + + +SAMPLE_LOG_EVENTS = [ + _epoch_event(0, 1000), + _step_event(1, 3.711, 2000), + _step_event(2, 3.589, 3000), + _epoch_event(1, 4000), + _step_event(3, 3.456, 5000), +] + + +def _make_trainer(method=TrainingMethod.SFT_LORA, platform=Platform.SMHP): + infra = create_autospec(SMHPRuntimeManager) + infra.instance_type = "ml.p5.48xlarge" + infra.instance_count = 2 + infra.kms_key_id = None + infra.platform = platform + infra.cluster_name = "my-cluster" + infra.namespace = "kubeflow" + infra.rft_lambda_arn = None + with ( + patch( + "amzn_nova_forge.trainer.forge_trainer.set_output_s3_path", + return_value="s3://bucket/output", + ), + patch("boto3.session.Session") as mock_session, + ): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + return ForgeTrainer( + model=Model.NOVA_MICRO, + method=method, + infra=infra, + training_data_s3_path="s3://bucket/data", + ) + + +class TestEpochExtraction: + """Epoch boundary extraction and step assignment.""" + + def test_extract_boundaries_and_assign_epochs(self): + """Epoch boundaries are extracted and steps are assigned correctly.""" + boundaries = _extract_epoch_boundaries(SAMPLE_LOG_EVENTS) + assert boundaries == [(0, 1000), (1, 4000)] + + metrics_df = get_metrics( + platform=Platform.SMHP, + training_method=TrainingMethod.SFT_LORA, + logs=SAMPLE_LOG_EVENTS, + ) + result = _assign_epochs_to_steps(metrics_df, boundaries, SAMPLE_LOG_EVENTS) + assert list(result["epoch_number"]) == [0, 0, 1] + + def test_steps_before_any_epoch_default_to_zero(self): + """Steps with no preceding epoch boundary get epoch 0.""" + events = [ + _step_event(1, 3.9, 1000), + _step_event(2, 3.8, 2000), + _epoch_event(1, 3000), + _step_event(3, 3.7, 4000), + ] + metrics_df = get_metrics( + platform=Platform.SMHP, + training_method=TrainingMethod.SFT_LORA, + logs=events, + ) + boundaries = _extract_epoch_boundaries(events) + result = _assign_epochs_to_steps(metrics_df, boundaries, events) + assert list(result["epoch_number"]) == [0, 0, 1] + + def test_empty_logs_return_empty(self): + """Empty log list produces empty boundaries and handles empty DataFrame.""" + assert _extract_epoch_boundaries([]) == [] + empty_df = pandas.DataFrame(columns=["global_step", "training_loss"]) + result = _assign_epochs_to_steps(empty_df, [], []) + assert "epoch_number" in result.columns + assert len(result) == 0 + + +class TestBuildAndUploadTrainingMetricsCsv: + """The _build_and_upload_training_metrics_csv() utility function.""" + + def test_successful_generation(self): + """Returns S3 URI and uploads CSV with correct columns.""" + from amzn_nova_forge.util.metric_util import _build_and_upload_training_metrics_csv + + mock_s3 = MagicMock() + + result = _build_and_upload_training_metrics_csv( + job_id="job-123", + log_events=SAMPLE_LOG_EVENTS, + output_s3_path="s3://bucket/prefix", + training_method=TrainingMethod.SFT_LORA, + s3_client=mock_s3, + ) + + assert result == "s3://bucket/prefix/job-123/step_wise_training_metrics.csv" + mock_s3.upload_file.assert_called_once() + _, bucket, key = mock_s3.upload_file.call_args[0] + assert bucket == "bucket" + assert key == "prefix/job-123/step_wise_training_metrics.csv" + + def test_no_logs_returns_none(self): + """Returns None when no logs are provided.""" + from amzn_nova_forge.util.metric_util import _build_and_upload_training_metrics_csv + + result = _build_and_upload_training_metrics_csv( + job_id="job-123", + log_events=[], + output_s3_path="s3://bucket/prefix", + training_method=TrainingMethod.SFT_LORA, + ) + assert result is None + + def test_missing_output_path_raises(self): + """Raises ValueError when output_s3_path is empty.""" + from amzn_nova_forge.util.metric_util import _build_and_upload_training_metrics_csv + + with pytest.raises(ValueError, match="output_s3_path"): + _build_and_upload_training_metrics_csv( + job_id="job-123", + log_events=SAMPLE_LOG_EVENTS, + output_s3_path="", + training_method=TrainingMethod.SFT_LORA, + ) + + +class TestForgeTrainerIntegration: + """ForgeTrainer.generate_training_metrics_csv() validation.""" + + def test_non_smhp_platform_raises_error(self): + """Non-SMHP platform raises ValueError.""" + trainer = _make_trainer() + trainer._platform = Platform.SMTJ + with pytest.raises(ValueError, match="SMHP"): + trainer.generate_training_metrics_csv(job_id="j", output_s3_path="s3://b/o") + + def test_non_sft_method_raises_error(self): + """Non-SFT method raises ValueError.""" + trainer = _make_trainer(method=TrainingMethod.CPT) + with pytest.raises(ValueError, match="SFT"): + trainer.generate_training_metrics_csv(job_id="j", output_s3_path="s3://b/o") + + @patch("amzn_nova_forge.trainer.forge_trainer.CloudWatchLogMonitor") + @patch("amzn_nova_forge.trainer.forge_trainer._build_and_upload_training_metrics_csv") + def test_delegates_to_metric_util_with_correct_params(self, mock_gen, mock_monitor_cls): + """Fetches logs and delegates to metric_util function.""" + mock_monitor_cls.return_value.get_logs.return_value = SAMPLE_LOG_EVENTS + mock_gen.return_value = "s3://bucket/output/job-abc/step_wise_training_metrics.csv" + trainer = _make_trainer() + + job_result = MagicMock() + job_result.job_id = "job-abc" + job_result.started_time = datetime(2025, 6, 1, tzinfo=timezone.utc) + job_result.cluster_name = "hp-cluster" + job_result.namespace = "ns" + job_result.model_artifacts.output_s3_path = "s3://bucket/output" + job_result.get_job_status.return_value = (JobStatus.COMPLETED, "Succeeded") + + result = trainer.generate_training_metrics_csv(job_result=job_result) + + assert result == "s3://bucket/output/job-abc/step_wise_training_metrics.csv" + mock_gen.assert_called_once_with( + job_id="job-abc", + log_events=SAMPLE_LOG_EVENTS, + output_s3_path="s3://bucket/output", + training_method=TrainingMethod.SFT_LORA, + region="us-east-1", + ) diff --git a/tests/unit/util/test_sagemaker.py b/tests/unit/util/test_sagemaker.py index abccf41..453b8de 100644 --- a/tests/unit/util/test_sagemaker.py +++ b/tests/unit/util/test_sagemaker.py @@ -12,17 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. import json +import time import unittest from datetime import datetime from unittest.mock import MagicMock, patch from botocore.exceptions import ClientError -from amzn_nova_forge.core.enums import DeploymentMode, Model, Platform +from amzn_nova_forge.core.enums import DeploymentMode, DeployPlatform, Model, Platform from amzn_nova_forge.core.result.inference_result import ( SingleInferenceResult, ) -from amzn_nova_forge.core.types import ModelArtifacts +from amzn_nova_forge.core.types import DeploymentResult, EndpointInfo, ModelArtifacts from amzn_nova_forge.manager.runtime_manager import ( RuntimeManager, SMHPRuntimeManager, @@ -30,14 +31,41 @@ SMTJServerlessRuntimeManager, ) from amzn_nova_forge.util.sagemaker import ( + _IC_MIN_COMPUTE_REQUIREMENTS, + InferenceComponentConfig, _get_sagemaker_inference_image, _validate_sagemaker_instance_type_for_model_deployment, + create_inference_component, create_sagemaker_endpoint, create_sagemaker_model, get_model_artifacts, invoke_sagemaker_inference, + monitor_inference_component, + validate_inference_component_resources, ) +IC_NAME = "my-ic-component" +IC_ENDPOINT_NAME = "my-endpoint" +IC_VARIANT_NAME = "AllTraffic" +IC_IMAGE_URI = "123456789012.dkr.ecr.us-east-1.amazonaws.com/my-repo:latest" +IC_MODEL_ARTIFACT_PATH = "s3://my-bucket/model/artifacts" +IC_ENVIRONMENT = {"MODEL_NAME": "nova-micro"} +IC_NUM_CPUS = 4 +IC_NUM_ACCELERATORS = 1 +IC_MIN_MEMORY_IN_MB = 8192 +IC_COPY_COUNT = 1 + +DEFAULT_IC_PARAMS = { + "inference_component_name": IC_NAME, + "endpoint_name": IC_ENDPOINT_NAME, + "variant_name": IC_VARIANT_NAME, + "model_name": "my-sagemaker-model", + "num_cpus": IC_NUM_CPUS, + "num_accelerators": IC_NUM_ACCELERATORS, + "min_memory_in_mb": IC_MIN_MEMORY_IN_MB, + "copy_count": IC_COPY_COUNT, +} + class TestSagemaker(unittest.TestCase): def setUp(self): @@ -57,7 +85,7 @@ def test_create_sagemaker_model_success(self): {"Error": {"Code": "ValidationException"}}, "" ) mock_client.create_model.return_value = { - "ModelArn": "arn:aws:sagemaker:us-east-1:123:model/test-model" + "ModelArn": "arn:aws:sagemaker:us-east-1:123456789012:model/test-model" } result = create_sagemaker_model( @@ -68,7 +96,7 @@ def test_create_sagemaker_model_success(self): sagemaker_client=mock_client, ) - self.assertEqual(result, "arn:aws:sagemaker:us-east-1:123:model/test-model") + self.assertEqual(result, "arn:aws:sagemaker:us-east-1:123456789012:model/test-model") mock_client.create_model.assert_called_once() def test_create_sagemaker_model_already_exists(self): @@ -95,6 +123,57 @@ def test_create_sagemaker_model_invalid_s3(self): sagemaker_client=MagicMock(), ) + def test_create_sagemaker_model_with_model_package_name(self): + """Model package name is passed as ModelPackageName in PrimaryContainer.""" + mock_client = MagicMock() + mock_client.describe_model.side_effect = ClientError( + {"Error": {"Code": "ValidationException"}}, "" + ) + mock_client.create_model.return_value = { + "ModelArn": "arn:aws:sagemaker:us-east-1:123456789012:model/test-model" + } + + model_package_arn = "arn:aws:sagemaker:us-east-1:123456789012:model-package/my-group/1" + result = create_sagemaker_model( + region=self.region, + model_name=self.model_name, + sagemaker_execution_role_arn=self.sagemaker_execution_role_arn, + sagemaker_client=mock_client, + model_package_name=model_package_arn, + ) + + self.assertEqual(result, "arn:aws:sagemaker:us-east-1:123456789012:model/test-model") + call_kwargs = mock_client.create_model.call_args[1] + primary_container = call_kwargs["PrimaryContainer"] + self.assertEqual(primary_container["ModelPackageName"], model_package_arn) + self.assertNotIn("ModelDataSource", primary_container) + # Image should still be present + self.assertIn("Image", primary_container) + + def test_create_sagemaker_model_both_s3_and_package_raises(self): + """Providing both model_s3_location and model_package_name raises ValueError.""" + with self.assertRaises(ValueError) as ctx: + create_sagemaker_model( + region=self.region, + model_name=self.model_name, + model_s3_location=self.model_s3_location, + sagemaker_execution_role_arn=self.sagemaker_execution_role_arn, + sagemaker_client=MagicMock(), + model_package_name="arn:aws:sagemaker:us-east-1:123456789012:model-package/group/1", + ) + self.assertIn("Only one of", str(ctx.exception)) + + def test_create_sagemaker_model_neither_s3_nor_package_raises(self): + """Providing neither model_s3_location nor model_package_name raises ValueError.""" + with self.assertRaises(ValueError) as ctx: + create_sagemaker_model( + region=self.region, + model_name=self.model_name, + sagemaker_execution_role_arn=self.sagemaker_execution_role_arn, + sagemaker_client=MagicMock(), + ) + self.assertIn("must be provided", str(ctx.exception)) + @patch("amzn_nova_forge.util.sagemaker._monitor_endpoint_creation") def test_create_sagemaker_endpoint_success(self, mock_monitor): mock_client = MagicMock() @@ -137,6 +216,76 @@ def test_create_sagemaker_endpoint_already_exists(self): ) self.assertIn("already exists", str(ctx.exception)) + @patch("amzn_nova_forge.util.sagemaker._monitor_endpoint_creation") + def test_create_sagemaker_endpoint_with_ic_returns_endpoint_arn(self, mock_monitor): + mock_client = MagicMock() + mock_client.describe_endpoint_config.side_effect = ClientError( + {"Error": {"Code": "ValidationException"}}, "" + ) + mock_client.describe_endpoint.side_effect = ClientError( + {"Error": {"Code": "ValidationException"}}, "" + ) + mock_client.create_endpoint_config.return_value = {"EndpointConfigArn": "test-config-arn"} + mock_client.create_endpoint.return_value = { + "EndpointArn": "arn:aws:sagemaker:us-east-1:123456789012:endpoint/test-endpoint" + } + mock_client.create_inference_component.return_value = { + "InferenceComponentArn": "arn:aws:sagemaker:us-east-1:123456789012:inference-component/test-ic" + } + mock_monitor.return_value = "InService" + + ic_config = InferenceComponentConfig( + inference_component_name="test-ic", + num_cpus=15, + num_accelerators=4, + min_memory_in_mb=25000, + ) + + result = create_sagemaker_endpoint( + model_name=self.model_name, + endpoint_config_name=self.endpoint_config_name, + endpoint_name=self.endpoint_name, + instance_type="ml.g5.4xlarge", + sagemaker_client=mock_client, + inference_component_configs=[ic_config], + execution_role_arn="arn:aws:iam::123456789012:role/test-role", + ) + + self.assertEqual(result, "arn:aws:sagemaker:us-east-1:123456789012:endpoint/test-endpoint") + mock_client.create_inference_component.assert_called_once() + + def test_create_sagemaker_endpoint_ic_without_execution_role_raises(self): + """Providing inference_component_configs without execution_role_arn must raise ValueError.""" + mock_client = MagicMock() + mock_client.describe_endpoint_config.side_effect = ClientError( + {"Error": {"Code": "ValidationException"}}, "" + ) + mock_client.describe_endpoint.side_effect = ClientError( + {"Error": {"Code": "ValidationException"}}, "" + ) + + ic_config = InferenceComponentConfig( + inference_component_name="test-ic", + num_cpus=15, + num_accelerators=4, + min_memory_in_mb=25000, + ) + + with self.assertRaises(ValueError) as ctx: + create_sagemaker_endpoint( + model_name=self.model_name, + endpoint_config_name=self.endpoint_config_name, + endpoint_name=self.endpoint_name, + instance_type="ml.g5.4xlarge", + sagemaker_client=mock_client, + inference_component_configs=[ic_config], + # execution_role_arn intentionally omitted + ) + + self.assertIn("execution_role_arn", str(ctx.exception)) + mock_client.create_endpoint_config.assert_not_called() + mock_client.create_endpoint.assert_not_called() + def test_get_sagemaker_inference_image_unsupported_region(self): with self.assertRaises(ValueError): _get_sagemaker_inference_image(region="unsupported_region") @@ -330,7 +479,7 @@ def test_get_model_artifacts_smhp_client_error(self, mock_boto_client): def test_get_model_artifacts_serverless_with_model_package(self, mock_boto_client): """Serverless jobs return output_model_arn and checkpoint_s3_path from model package.""" job_name = "test-serverless-job" - model_package_arn = "arn:aws:sagemaker:us-east-1:123:model-package/group/1" + model_package_arn = "arn:aws:sagemaker:us-east-1:123456789012:model-package/group/1" output_s3_path = "s3://my-bucket/output/" checkpoint_s3_path = "s3://customer-escrow/job/384/" @@ -526,6 +675,492 @@ def test_invalid_instance_types(self): with self.assertRaises(ValueError): _validate_sagemaker_instance_type_for_model_deployment(instance_type, model) + @patch("amzn_nova_forge.util.sagemaker.boto3.client") + def test_get_model_artifacts_mtrl_fallback_on_validation_exception(self, mock_boto_client): + """SMTJServerless jobs fall through to AgentRFTJob on ValidationException.""" + import sys + import types + + mock_sagemaker = MagicMock() + mock_sagemaker.describe_training_job.side_effect = ClientError( + error_response={"Error": {"Code": "ValidationException", "Message": "not found"}}, + operation_name="DescribeTrainingJob", + ) + mock_boto_client.return_value = mock_sagemaker + + infra = MagicMock() + infra.platform = Platform.SMTJServerless + infra.sagemaker_client = mock_sagemaker + + mock_rft_job = MagicMock() + mock_rft_job.s3_output_path = "s3://bucket/output/" + mock_rft_job.output_model_package_arn = ( + "arn:aws:sagemaker:us-east-1:123456789012:model-package/grp/1" + ) + + mock_agent_module = types.ModuleType("sagemaker.train.agent_rft_job") + mock_agent_cls = MagicMock() + mock_agent_cls.get.return_value = mock_rft_job + mock_agent_module.AgentRFTJob = mock_agent_cls + + with patch.dict("sys.modules", {"sagemaker.train.agent_rft_job": mock_agent_module}): + result = get_model_artifacts( + job_name="mtrl-job-123", + infra=infra, + region="us-east-1", + ) + + self.assertEqual( + result.output_model_arn, "arn:aws:sagemaker:us-east-1:123456789012:model-package/grp/1" + ) + self.assertEqual(result.output_s3_path, "s3://bucket/output/") + self.assertIsNone(result.checkpoint_s3_path) + + @patch("amzn_nova_forge.util.sagemaker.boto3.client") + def test_get_model_artifacts_mtrl_no_output_s3_path_param_needed(self, mock_boto_client): + """output_s3_path parameter is optional for MTRL — fetched from AgentRFTJob.""" + import sys + import types + + mock_sagemaker = MagicMock() + mock_sagemaker.describe_training_job.side_effect = ClientError( + error_response={"Error": {"Code": "ValidationException", "Message": "not found"}}, + operation_name="DescribeTrainingJob", + ) + mock_boto_client.return_value = mock_sagemaker + + infra = MagicMock() + infra.platform = Platform.SMTJServerless + infra.sagemaker_client = mock_sagemaker + + mock_rft_job = MagicMock() + mock_rft_job.s3_output_path = "s3://from-api/output/" + mock_rft_job.output_model_package_arn = None + + mock_agent_module = types.ModuleType("sagemaker.train.agent_rft_job") + mock_agent_cls = MagicMock() + mock_agent_cls.get.return_value = mock_rft_job + mock_agent_module.AgentRFTJob = mock_agent_cls + + with patch.dict("sys.modules", {"sagemaker.train.agent_rft_job": mock_agent_module}): + result = get_model_artifacts( + job_name="mtrl-job-456", + infra=infra, + ) + + self.assertEqual(result.output_s3_path, "s3://from-api/output/") + self.assertIsNone(result.output_model_arn) + + @patch("amzn_nova_forge.util.sagemaker.boto3.client") + def test_get_model_artifacts_smtj_does_not_fallback_to_agent_rft(self, mock_boto_client): + """Platform.SMTJ should re-raise ResourceNotFound, not fallback to AgentRFT.""" + mock_sagemaker = MagicMock() + mock_sagemaker.describe_training_job.side_effect = ClientError( + error_response={"Error": {"Code": "ResourceNotFound", "Message": "not found"}}, + operation_name="DescribeTrainingJob", + ) + mock_boto_client.return_value = mock_sagemaker + + infra = MagicMock() + infra.platform = Platform.SMTJ + infra.sagemaker_client = mock_sagemaker + + with self.assertRaises(ClientError): + get_model_artifacts( + job_name="smtj-job-123", + infra=infra, + output_s3_path="s3://bucket/output/", + ) + + def test_model_artifacts_all_fields_optional(self): + """ModelArtifacts can be constructed with no arguments.""" + artifacts = ModelArtifacts() + self.assertIsNone(artifacts.checkpoint_s3_path) + self.assertIsNone(artifacts.output_s3_path) + self.assertIsNone(artifacts.output_model_arn) + + +def _mock_client_for_successful_creation(ic_name="my-ic-component", endpoint_name="my-endpoint"): + mock_client = MagicMock() + mock_client.describe_endpoint.return_value = { + "EndpointName": endpoint_name, + "EndpointStatus": "InService", + "EndpointConfigName": "my-endpoint-config", + } + mock_client.describe_endpoint_config.return_value = { + "ProductionVariants": [ + { + "VariantName": "AllTraffic", + "RoutingConfig": {"RoutingStrategy": "LEAST_OUTSTANDING_REQUESTS"}, + } + ] + } + mock_client.create_inference_component.return_value = { + "InferenceComponentArn": ( + f"arn:aws:sagemaker:us-east-1:123456789012:inference-component/{ic_name}" + ), + } + return mock_client + + +class TestCreateInferenceComponentNonBlocking(unittest.TestCase): + def test_returns_deployment_result(self): + params = {**DEFAULT_IC_PARAMS} + mock_client = _mock_client_for_successful_creation() + + with patch("amzn_nova_forge.util.sagemaker.time.sleep") as mock_sleep: + result = create_inference_component(**params, sagemaker_client=mock_client) + + self.assertIsInstance(result, DeploymentResult) + mock_client.describe_inference_component.assert_not_called() + mock_sleep.assert_not_called() + + def test_returns_deployment_result_various_names(self): + names = ["component-alpha", "ic_beta_123", "X", "a-very-long-component-name-here"] + for name in names: + params = {**DEFAULT_IC_PARAMS} + params["inference_component_name"] = name + mock_client = _mock_client_for_successful_creation(ic_name=name) + + with patch("amzn_nova_forge.util.sagemaker.time.sleep") as mock_sleep: + result = create_inference_component(**params, sagemaker_client=mock_client) + + self.assertIsInstance(result, DeploymentResult) + mock_client.describe_inference_component.assert_not_called() + mock_sleep.assert_not_called() + + +class TestCreateInferenceComponentValidation(unittest.TestCase): + def _assert_invalid_param_raises(self, param_name, value): + params = {**DEFAULT_IC_PARAMS, param_name: value} + mock_client = MagicMock() + + with self.assertRaises(ValueError) as ctx: + create_inference_component(**params, sagemaker_client=mock_client) + + self.assertIn(param_name, str(ctx.exception)) + mock_client.create_inference_component.assert_not_called() + + def test_empty_string_params_raise(self): + for param in ( + "inference_component_name", + "endpoint_name", + "variant_name", + "model_name", + ): + with self.subTest(param=param): + self._assert_invalid_param_raises(param, "") + + def test_none_params_raise(self): + for param in ("inference_component_name", "endpoint_name", "model_name"): + with self.subTest(param=param): + self._assert_invalid_param_raises(param, None) + + +class TestCreateInferenceComponentEndpointValidation(unittest.TestCase): + def test_non_existent_endpoint_raises(self): + params = {**DEFAULT_IC_PARAMS} + params["endpoint_name"] = "missing-endpoint" + mock_client = MagicMock() + mock_client.describe_endpoint.side_effect = ClientError( + {"Error": {"Code": "ValidationException", "Message": "Could not find endpoint"}}, + "DescribeEndpoint", + ) + + with self.assertRaises(Exception) as ctx: + create_inference_component(**params, sagemaker_client=mock_client) + + self.assertIn("not found", str(ctx.exception).lower()) + self.assertIn("missing-endpoint", str(ctx.exception)) + mock_client.create_inference_component.assert_not_called() + + def test_non_inservice_endpoint_raises(self): + params = {**DEFAULT_IC_PARAMS} + status = "Creating" + mock_client = MagicMock() + mock_client.describe_endpoint.return_value = { + "EndpointName": params["endpoint_name"], + "EndpointStatus": status, + } + + with self.assertRaises(Exception) as ctx: + create_inference_component(**params, sagemaker_client=mock_client) + + self.assertIn(status, str(ctx.exception)) + mock_client.create_inference_component.assert_not_called() + + +class TestCreateInferenceComponentEndpointInfo(unittest.TestCase): + def test_endpoint_info_fields(self): + params = {**DEFAULT_IC_PARAMS} + mock_client = _mock_client_for_successful_creation( + ic_name=params["inference_component_name"], + endpoint_name=params["endpoint_name"], + ) + + result = create_inference_component(**params, sagemaker_client=mock_client) + + self.assertEqual(result.endpoint.platform, DeployPlatform.SAGEMAKER) + self.assertEqual(result.endpoint.endpoint_name, params["endpoint_name"]) + self.assertIn(params["inference_component_name"], result.endpoint.uri) + self.assertEqual(result.endpoint.model_artifact_path, params["model_name"]) + + def test_endpoint_info_with_different_params(self): + params = {**DEFAULT_IC_PARAMS} + params["inference_component_name"] = "custom-ic-name" + params["endpoint_name"] = "custom-endpoint" + params["model_name"] = "custom-sagemaker-model" + mock_client = _mock_client_for_successful_creation( + ic_name="custom-ic-name", + endpoint_name="custom-endpoint", + ) + + result = create_inference_component(**params, sagemaker_client=mock_client) + + self.assertEqual(result.endpoint.platform, DeployPlatform.SAGEMAKER) + self.assertEqual(result.endpoint.endpoint_name, "custom-endpoint") + self.assertIn("custom-ic-name", result.endpoint.uri) + self.assertEqual(result.endpoint.model_artifact_path, "custom-sagemaker-model") + + +class TestMonitorInferenceComponentPolling(unittest.TestCase): + @patch("amzn_nova_forge.util.sagemaker.time.sleep") + def test_polling_terminates_on_terminal_state(self, mock_sleep): + cases = [ + # (side_effect, expected_calls, expected_sleeps, terminal_status, should_raise) + ( + [{"InferenceComponentName": "my-ic", "InferenceComponentStatus": "InService"}], + 1, + 0, + "InService", + False, + ), + ( + [ + {"InferenceComponentName": "my-ic", "InferenceComponentStatus": "Creating"}, + {"InferenceComponentName": "my-ic", "InferenceComponentStatus": "Creating"}, + {"InferenceComponentName": "my-ic", "InferenceComponentStatus": "InService"}, + ], + 3, + 2, + "InService", + False, + ), + ( + [ + {"InferenceComponentName": "my-ic", "InferenceComponentStatus": "Creating"}, + {"InferenceComponentName": "my-ic", "InferenceComponentStatus": "Failed"}, + ], + 2, + 1, + None, + True, + ), + ] + + mock_client = MagicMock() + + for side_effect, expected_calls, expected_sleeps, terminal_status, should_raise in cases: + with self.subTest(side_effect=side_effect): + mock_client.reset_mock() + mock_sleep.reset_mock() + mock_client.describe_inference_component.side_effect = side_effect + + if should_raise: + with self.assertRaises(Exception) as ctx: + monitor_inference_component("my-ic", mock_client) + self.assertIn("my-ic", str(ctx.exception)) + else: + result = monitor_inference_component("my-ic", mock_client) + self.assertEqual(result, terminal_status) + + self.assertEqual( + mock_client.describe_inference_component.call_count, expected_calls + ) + self.assertEqual(mock_sleep.call_count, expected_sleeps) + + +class TestInvokeInferenceComponentName(unittest.TestCase): + def test_non_streaming_includes_inference_component_name(self): + mock_client = MagicMock() + mock_client.invoke_endpoint.return_value = { + "Body": MagicMock( + read=MagicMock( + return_value=json.dumps( + {"id": "test-id", "created": 1234567890, "choices": []} + ).encode() + ) + ), + "ResponseMetadata": {"RequestId": "test-request-id"}, + } + + invoke_sagemaker_inference( + request_body={"inputs": "test"}, + endpoint_name="my-endpoint", + sagemaker_client=mock_client, + inference_component_name="my-ic", + ) + + call_kwargs = mock_client.invoke_endpoint.call_args[1] + self.assertIn("InferenceComponentName", call_kwargs) + self.assertEqual(call_kwargs["InferenceComponentName"], "my-ic") + + def test_streaming_includes_inference_component_name(self): + mock_client = MagicMock() + mock_client.invoke_endpoint_with_response_stream.return_value = { + "Body": [], + "ResponseMetadata": {"RequestId": "test-request-id"}, + } + + invoke_sagemaker_inference( + request_body={"inputs": "test", "stream": True}, + endpoint_name="my-endpoint", + sagemaker_client=mock_client, + inference_component_name="my-ic", + ) + + call_kwargs = mock_client.invoke_endpoint_with_response_stream.call_args[1] + self.assertIn("InferenceComponentName", call_kwargs) + self.assertEqual(call_kwargs["InferenceComponentName"], "my-ic") + + +class TestCheckDeploymentStatusSagemaker(unittest.TestCase): + @patch("amzn_nova_forge.util.sagemaker.boto3.client") + def test_inference_component_status_returned(self, mock_boto_client): + from amzn_nova_forge.util.sagemaker import check_sagemaker_deployment_status + + mock_sm_client = MagicMock() + mock_sm_client.describe_inference_component.return_value = { + "InferenceComponentName": "my-ic", + "InferenceComponentStatus": "InService", + } + mock_boto_client.return_value = mock_sm_client + + arn = "arn:aws:sagemaker:us-east-1:123456789012:inference-component/my-ic" + status = check_sagemaker_deployment_status(arn) + + self.assertEqual(status, "InService") + mock_boto_client.assert_any_call("sagemaker", region_name=None) + mock_sm_client.describe_inference_component.assert_called_once_with( + InferenceComponentName="my-ic" + ) + + @patch("amzn_nova_forge.util.sagemaker.boto3.client") + def test_inference_component_creating_status(self, mock_boto_client): + from amzn_nova_forge.util.sagemaker import check_sagemaker_deployment_status + + mock_sm_client = MagicMock() + mock_sm_client.describe_inference_component.return_value = { + "InferenceComponentName": "ic-creating", + "InferenceComponentStatus": "Creating", + } + mock_boto_client.return_value = mock_sm_client + + arn = "arn:aws:sagemaker:us-east-1:123456789012:inference-component/ic-creating" + status = check_sagemaker_deployment_status(arn) + + self.assertEqual(status, "Creating") + + @patch("amzn_nova_forge.util.sagemaker.boto3.client") + def test_inference_component_failed_status(self, mock_boto_client): + from amzn_nova_forge.util.sagemaker import check_sagemaker_deployment_status + + mock_sm_client = MagicMock() + mock_sm_client.describe_inference_component.return_value = { + "InferenceComponentName": "ic-failed", + "InferenceComponentStatus": "Failed", + } + mock_boto_client.return_value = mock_sm_client + + arn = "arn:aws:sagemaker:us-east-1:123456789012:inference-component/ic-failed" + status = check_sagemaker_deployment_status(arn) + + self.assertEqual(status, "Failed") + + @patch("amzn_nova_forge.util.sagemaker.boto3.client") + def test_deployment_result_status_property(self, mock_boto_client): + mock_sm_client = MagicMock() + mock_sm_client.describe_inference_component.return_value = { + "InferenceComponentName": "my-ic", + "InferenceComponentStatus": "InService", + } + mock_boto_client.return_value = mock_sm_client + + result = DeploymentResult( + endpoint=EndpointInfo( + platform=DeployPlatform.SAGEMAKER, + endpoint_name="my-endpoint", + uri="arn:aws:sagemaker:us-east-1:123456789012:inference-component/my-ic", + model_artifact_path="s3://bucket/model", + ), + created_at=datetime.now(), + ) + + status = result.status + self.assertEqual(status, "InService") + + +class TestCheckDeploymentStatusEndpoint(unittest.TestCase): + @patch("amzn_nova_forge.util.sagemaker.boto3.client") + def test_returns_endpoint_status_in_service(self, mock_boto_client): + from amzn_nova_forge.util.sagemaker import check_sagemaker_deployment_status + + mock_sm_client = MagicMock() + mock_sm_client.describe_endpoint.return_value = {"EndpointStatus": "InService"} + mock_boto_client.return_value = mock_sm_client + + arn = "arn:aws:sagemaker:us-east-1:123456789012:endpoint/my-endpoint" + status = check_sagemaker_deployment_status(arn) + + self.assertEqual(status, "InService") + mock_boto_client.assert_any_call("sagemaker", region_name=None) + mock_sm_client.describe_endpoint.assert_called_once_with(EndpointName="my-endpoint") + + @patch("amzn_nova_forge.util.sagemaker.boto3.client") + def test_returns_endpoint_status_creating(self, mock_boto_client): + from amzn_nova_forge.util.sagemaker import check_sagemaker_deployment_status + + mock_sm_client = MagicMock() + mock_sm_client.describe_endpoint.return_value = {"EndpointStatus": "Creating"} + mock_boto_client.return_value = mock_sm_client + + arn = "arn:aws:sagemaker:us-east-1:123456789012:endpoint/my-endpoint" + status = check_sagemaker_deployment_status(arn) + + self.assertEqual(status, "Creating") + + +class TestValidateInferenceComponentResources(unittest.TestCase): + def test_sufficient_resources_passes(self): + config = InferenceComponentConfig( + inference_component_name="test-ic", + num_cpus=15, + num_accelerators=4, + min_memory_in_mb=25000, + ) + validate_inference_component_resources(config, Model.NOVA_MICRO) + + def test_resources_below_minimums_raises_valueerror(self): + config = InferenceComponentConfig( + inference_component_name="test-ic", + num_cpus=4, + num_accelerators=1, + min_memory_in_mb=8192, + ) + with self.assertRaises(ValueError): + validate_inference_component_resources(config, Model.NOVA_MICRO) + + def test_model_not_in_requirements_passes(self): + self.assertNotIn(Model.NOVA_PRO, _IC_MIN_COMPUTE_REQUIREMENTS) + + config = InferenceComponentConfig( + inference_component_name="test-ic", + num_cpus=1, + num_accelerators=0, + min_memory_in_mb=512, + ) + validate_inference_component_resources(config, Model.NOVA_PRO) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/validation/test_validator.py b/tests/unit/validation/test_validator.py index 894502a..eadaabf 100644 --- a/tests/unit/validation/test_validator.py +++ b/tests/unit/validation/test_validator.py @@ -2849,6 +2849,125 @@ def test_none_value_required_str_field_without_null_enum_still_raises(self): self.assertIn("NoneType", errors[0]) +class TestMTRLServerlessValidation(unittest.TestCase): + """Tests for MTRL serverless validation logic in Validator._validate_recipe.""" + + def setUp(self): + self.mock_infra = Mock(spec=SMTJRuntimeManager) + self.mock_infra.instance_type = "ml.p5.48xlarge" + self.mock_infra.region = "us-east-1" + + def test_mtrl_serverless_skips_override_constraint_validation(self): + """MTRL serverless jobs skip min/max/enum/type validation for most keys.""" + recipe = { + "training_config": { + "global_batch_size": 2, # Would fail min=16 constraint for single-turn RFT + "max_steps": 5, + }, + "run": { + "output_s3_path": "s3://bucket/output/", + "data_s3_path": "s3://bucket/data/", + }, + } + overrides_template = { + "global_batch_size": {"type": "integer", "min": 16, "max": 256}, + "max_steps": {"type": "integer", "min": 10, "max": 100000}, + "output_s3_path": {"type": "string", "required": True}, + "data_s3_path": {"type": "string", "required": True}, + } + + errors = [] + Validator._validate_recipe( + recipe=recipe, + overrides_template=overrides_template, + instance_type=None, + errors=errors, + method=TrainingMethod.RFT_MULTITURN_LORA, + platform=Platform.SMTJServerless, + ) + + # Should have NO errors because MTRL serverless skips constraint validation + self.assertEqual(len(errors), 0) + + def test_non_mtrl_serverless_still_validates_constraints(self): + """Non-MTRL methods on SMTJServerless should still validate constraints.""" + recipe = {"training_config": {"max_steps": 2}} + overrides_template = {"max_steps": {"type": "integer", "min": 4}} + + errors = [] + Validator._validate_recipe( + recipe=recipe, + overrides_template=overrides_template, + instance_type=None, + errors=errors, + method=TrainingMethod.SFT_LORA, + platform=Platform.SMTJServerless, + ) + + # Should have error because SFT_LORA still validates + self.assertEqual(len(errors), 1) + self.assertIn("must be at least 4", errors[0]) + + def test_mtrl_serverless_still_validates_output_s3_path(self): + """MTRL serverless should still validate output_s3_path, data_s3_path, model_type.""" + recipe = {"run": {"output_s3_path": 12345}} # Wrong type for output_s3_path + overrides_template = { + "output_s3_path": {"type": "string", "required": True}, + } + + errors = [] + Validator._validate_recipe( + recipe=recipe, + overrides_template=overrides_template, + instance_type=None, + errors=errors, + method=TrainingMethod.RFT_MULTITURN_LORA, + platform=Platform.SMTJServerless, + ) + + # output_s3_path is NOT skipped, so type validation should still catch this + self.assertEqual(len(errors), 1) + self.assertIn("output_s3_path", errors[0]) + + @patch("amzn_nova_forge.validation.validator.boto3.client") + def test_mtrl_serverless_accepts_agent_core_arn_no_lambda_required(self, mock_boto3_client): + """RFT_MULTITURN_LORA on SMTJServerless does not require rft_lambda_arn.""" + mock_infra = Mock(spec=SMTJRuntimeManager) + mock_infra.instance_type = None + mock_infra.region = "us-east-1" + + # Should NOT raise — MTRL on SMTJServerless uses agent_core_arn, no lambda required + try: + Validator.validate( + platform=Platform.SMTJServerless, + method=TrainingMethod.RFT_MULTITURN_LORA, + infra=mock_infra, + recipe={}, + overrides_template={}, + validation_config=ValidationConfig(iam=False, infra=False, recipe=False), + ) + except ValueError as e: + self.fail(f"MTRL serverless validation should not require lambda ARN, but raised: {e}") + + @patch("amzn_nova_forge.validation.validator.boto3.client") + def test_non_mtrl_method_does_not_skip_validation(self, mock_boto3_client): + """Non-MTRL methods on serverless do NOT get the _is_mtrl_serverless skip.""" + mock_infra = Mock(spec=SMTJRuntimeManager) + mock_infra.instance_type = "ml.p5.48xlarge" + mock_infra.region = "us-east-1" + + # SFT_LORA on serverless should still validate normally (no skip) + # This should not raise because recipe={} with no overrides is valid + Validator.validate( + platform=Platform.SMTJServerless, + method=TrainingMethod.SFT_LORA, + infra=mock_infra, + recipe={}, + overrides_template={}, + validation_config=ValidationConfig(iam=False, infra=False, recipe=False), + ) + + class TestValidationDataS3Path(unittest.TestCase): """Tests for validation_data_s3_path preflight validation."""