diff --git a/docs/spec/index.md b/docs/spec/index.md index 9f731fc..332b4c2 100644 --- a/docs/spec/index.md +++ b/docs/spec/index.md @@ -11,6 +11,7 @@ - [Enums and Configuration](enums.md) — Platform, Model, TrainingMethod, etc. - [RFT Multiturn](rft-multiturn.md) — Multi-turn RFT infrastructure - [Notifications](notifications.md) — Job notifications and error handling +- [InspectLens Evaluation Guide](../user-guides/inspect_eval.md) — Prerequisites, benchmark authoring, inference modes, MLflow tracking --- @@ -23,10 +24,13 @@ 6. **Use descriptive job names** to help track and organize your experiments 7. **Save results incrementally** during long-running jobs 8. **Test with small datasets** before scaling up to full training + --- + ## Additional Resources - AWS Documentation: [Amazon Bedrock](https://docs.aws.amazon.com/bedrock/) - AWS Documentation: [Amazon SageMaker](https://docs.aws.amazon.com/sagemaker/) - SDK GitHub Repository: Check for updates and examples - Support: Use AWS Support for technical assistance + --- diff --git a/docs/spec/service-classes.md b/docs/spec/service-classes.md index 39e3fd6..8b0319f 100644 --- a/docs/spec/service-classes.md +++ b/docs/spec/service-classes.md @@ -28,7 +28,7 @@ class ForgeConfig: - `output_s3_path` (Optional[str]): S3 path for output artifacts. Auto-generated if not provided - `generated_recipe_dir` (Optional[str]): Local path to save generated recipe files - `validation_config` (Optional[ValidationConfig]): Controls pre-flight validation. Fields: `iam` (bool), `infra` (bool), `recipe` (bool) — all default to True -- `image_uri` (Optional[str]): Custom container image URI override +- `image_uri` (Optional[str]): Custom container image URI override. For InspectLens, overrides the default orchestrator container image. - `mlflow_monitor` (Optional[MLflowMonitor]): MLflow monitoring configuration (SageMaker only) - `enable_job_caching` (bool): Enable caching of completed job results for reuse. Default: False - `job_cache_dir` (str): Directory for cached job results. Default: `~/.nova-forge/cache` @@ -429,6 +429,131 @@ evaluator.get_logs(job_result=eval_result, limit=50) ``` --- +##### InspectLens Evaluation (`EvaluationTask.INSPECT_LENS`) + +InspectLens runs [Inspect AI](https://inspect.ai-safety-institute.org.uk/) benchmarks as a SageMaker Training Job. The job acts as a CPU-only orchestrator — inference is delegated to a Bedrock endpoint or a SageMaker endpoint. No GPU is needed for the orchestrator instance. + +Pass `eval_task=EvaluationTask.INSPECT_LENS` and provide an `InspectLensConfig` via the `inspect_lens_config` parameter. + +**`InspectLensConfig`** + +```python +@dataclass +class InspectLensConfig: + benchmarks_path: Optional[str] = None + tasks: List[Dict[str, Any]] = field(default_factory=list) + output_s3_path: Optional[str] = None + output_format: Optional[str] = None + bedrock_model_id: Optional[str] = None + endpoint_name: Optional[str] = None + model_s3_uri: Optional[str] = None + inference_image_uri: Optional[str] = None + endpoint_instance_type: Optional[str] = None + endpoint_instance_count: int = 1 + endpoint_execution_role_arn: Optional[str] = None + context_length: Optional[str] = None + max_concurrency: Optional[str] = None + enable_rai: bool = True + cleanup_endpoint: bool = True + endpoint_prefix: str = "inspectlens" + endpoint_environment: Optional[Dict[str, str]] = None + extra_args: Optional[List[str]] = None + environment: Optional[Dict[str, str]] = None +``` + +**Parameters:** +- `benchmarks_path` (str): S3 URI to benchmark `.py` files — required at runtime. Use `evaluator.upload_benchmarks(local_dir, s3_path)` to upload a local directory first. Validated: must be a non-empty `s3://` URI; task names are checked against `@task`-decorated functions at job submission time. +- `tasks` (List[Dict]): List of task dicts with a required `"name"` key and optional `"limit"` and `"epochs"`. Leave empty to run all `@task` functions in `benchmarks_path`. Validated: each entry must be a dict with a `"name"` key. +- `output_s3_path` (Optional[str]): S3 prefix for eval result JSON logs. Defaults to the `ForgeEvaluator` output path. Validated: must be an `s3://` URI if provided. +- `output_format` (Optional[str]): Output format for eval results. One of `"eval"`, `"csv"`, `"jsonl"`, `"json"`. Defaults to `"eval"`. +- `bedrock_model_id` (Optional[str]): Bedrock model ID or ARN. Falls back to the `ForgeEvaluator.model` cross-region inference profile. +- `endpoint_name` (Optional[str]): Existing SageMaker endpoint name. Mutually exclusive with `model_s3_uri`. +- `model_s3_uri` (Optional[str]): S3 URI of model artifacts for creating a new endpoint. Requires `inference_image_uri`. Validated: must be an `s3://` URI if provided. +- `inference_image_uri` (Optional[str]): ECR image URI for the new endpoint container. Requires `model_s3_uri`. Validated: must match ECR URI format if provided. +- `endpoint_instance_type` (Optional[str]): Instance type for the new endpoint. +- `endpoint_instance_count` (int): Number of instances for the new endpoint. Default: 1. +- `endpoint_execution_role_arn` (Optional[str]): IAM role ARN for the new endpoint. Validated: must match IAM ARN format if provided. +- `context_length` (Optional[str]): Context length for the new endpoint. +- `max_concurrency` (Optional[str]): Max concurrency for the new endpoint. +- `enable_rai` (bool): Enable RAI guardrails on the endpoint. Default: True. +- `cleanup_endpoint` (bool): Delete the endpoint after evaluation. Default: True. +- `endpoint_prefix` (str): Prefix for auto-created endpoint names. Default: `"inspectlens"`. +- `endpoint_environment` (Optional[Dict[str, str]]): Extra environment variables for the inference endpoint container. Example: `{"HF_TOKEN": "hf_xxx"}`. +- `extra_args` (Optional[List[str]]): Additional CLI args forwarded to `inspect eval`. +- `environment` (Optional[Dict[str, str]]): Arbitrary environment variables passed to the SageMaker Training Job container. Example: `{"HF_TOKEN": "hf_xxx"}`. + +**Inference provider modes:** + +| Mode | When used | +|---------------------|-----------------------------------------------------------------------------------------------------------| +| Bedrock | No `endpoint_name` or `model_s3_uri` set. Uses `bedrock_model_id` or the evaluator's model enum. | +| Existing endpoint | `endpoint_name` is set. | +| Create new endpoint | `model_s3_uri` + `inference_image_uri` are set. Endpoint is created, evaluated, then optionally deleted. | + +**Decoding overrides:** (`temperature`, `top_p`, `max_tokens`, `max_connections`, `max_retries`, `timeout`) are passed via the `overrides` parameter on `evaluate()`. + +**MLflow tracking:** Pass an `MLflowMonitor` via `ForgeConfig.mlflow_monitor`. The SDK automatically injects the tracking section into the InspectLens YAML config. + +**Example:** + +```python +from amzn_nova_forge.evaluator import ForgeEvaluator, InspectLensConfig +from amzn_nova_forge.core.enums import EvaluationTask, Model +from amzn_nova_forge.core.types import ForgeConfig +from amzn_nova_forge.manager import SMTJRuntimeManager + +infra = SMTJRuntimeManager( + instance_type="ml.m5.large", # CPU only — no GPU needed + instance_count=1, + execution_role="arn:aws:iam::123456789012:role/InspectLensEvalRole", +) +evaluator = ForgeEvaluator( + model=Model.NOVA_LITE_2, + infra=infra, + config=ForgeConfig(output_s3_path="s3://my-bucket/inspectlens/"), +) + +# Upload local benchmarks first +benchmarks_s3_uri = evaluator.upload_benchmarks( + "./my_benchmarks/", + "s3://my-bucket/inspectlens/benchmarks/my_benchmarks/", +) + +# Run evaluation +eval_result = evaluator.evaluate( + job_name="inspectlens-eval", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=InspectLensConfig( + benchmarks_path=benchmarks_s3_uri, + tasks=[{"name": "boolq_pt", "limit": 100}], + ), + overrides={"temperature": 0.0, "max_tokens": 4096}, +) +``` + +See `samples/inspectlens_quickstart.ipynb` for a full walkthrough including all three inference modes (Bedrock, existing endpoint, new endpoint). + +For a conceptual guide covering prerequisites, IAM setup, benchmark file format, inference modes, MLflow tracking, and job caching, see [docs/user-guides/inspect_eval.md](../user-guides/inspect_eval.md). + +--- + +##### `upload_benchmarks()` +Uploads a local benchmarks directory to S3 before starting an InspectLens job. + +**Signature:** +```python +def upload_benchmarks(self, local_dir: str, s3_path: str) -> str +``` + +**Parameters:** +- `local_dir` (str): Path to the local directory containing benchmark `.py` files with `@task` decorators. +- `s3_path` (str): S3 URI (`s3://bucket/prefix/`) where the benchmark files will be uploaded. + +**Returns:** +- The S3 URI where benchmarks were uploaded (normalized with a trailing slash). + +--- + ### ForgeDeployer Handles model deployment to Amazon Bedrock and SageMaker endpoints. diff --git a/docs/user-guides/inspect_eval.md b/docs/user-guides/inspect_eval.md new file mode 100644 index 0000000..373f79d --- /dev/null +++ b/docs/user-guides/inspect_eval.md @@ -0,0 +1,315 @@ +# InspectLens Evaluation Guide + +InspectLens is a job-based evaluation framework in the Nova Forge SDK. It runs [Inspect AI](https://inspect.ai-safety-institute.org.uk/) benchmarks against your Nova model via a SageMaker Training Job — no GPU required, since inference is delegated to a Bedrock endpoint or a SageMaker endpoint. + +## Overview + +| Concept | Description | +|---------|-------------| +| **Orchestrator** | CPU-only SageMaker Training Job (e.g. `ml.m5.large`) | +| **Inference** | Bedrock, existing SageMaker endpoint, or SDK-created endpoint | +| **Benchmarks** | Python files with `@task` decorators uploaded to S3 | +| **Results** | JSON logs written to an S3 output path | + +## Prerequisites + +- AWS credentials configured +- S3 bucket for evaluation artifacts and results +- IAM execution role with SageMaker and S3 permissions +- Nova Forge SDK installed (`pip install .` from the repo root) + +## Quick Start + +```python +from amzn_nova_forge import * +from amzn_nova_forge.core.enums import EvaluationTask, Model +from amzn_nova_forge.evaluator import ForgeEvaluator, InspectLensConfig + +# 1. Configure infrastructure (CPU-only — no GPU needed) +eval_infra = SMTJRuntimeManager( + instance_type="ml.m5.large", + instance_count=1, + execution_role="arn:aws:iam::123456789012:role/YourSageMakerRole", +) + +# 2. Initialize the evaluator +evaluator = ForgeEvaluator( + model=Model.NOVA_LITE_2, + infra=eval_infra, + config=ForgeConfig(output_s3_path="s3://your-bucket/inspectlens/output"), +) + +# 3. Upload benchmarks to S3 +benchmarks_s3_uri = evaluator.upload_benchmarks( + "./my_benchmarks", + "s3://your-bucket/inspectlens/benchmarks/my_benchmarks/", +) + +# 4. Configure and run evaluation +inspect_config = InspectLensConfig( + benchmarks_path=benchmarks_s3_uri, + tasks=[ + {"name": "boolq_pt", "limit": 100}, + {"name": "mmlu_pro_pt", "limit": 50}, + ], + output_s3_path="s3://your-bucket/inspectlens/results/", +) + +result = evaluator.evaluate( + job_name="my-inspectlens-eval", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=inspect_config, +) + +print(f"Job ID: {result.job_id}") +print(f"Results: {result.eval_output_path}") +``` + +## Writing Benchmarks + +Benchmarks are Python files with `@task` decorators. They can wrap built-in `inspect-evals` tasks or define custom evaluation logic. + +```python +# my_benchmarks/boolq_pt.py +from inspect_ai import task +from inspect_evals.boolq import boolq + +@task +def boolq_pt(): + return boolq() +``` + +```python +# my_benchmarks/mmlu_pro_pt.py +from inspect_ai import task +from inspect_evals.mmlu_pro import mmlu_pro + +@task +def mmlu_pro_pt(): + return mmlu_pro() +``` + +Upload benchmarks to S3 before starting a job: + +```python +benchmarks_s3_uri = evaluator.upload_benchmarks("./my_benchmarks", "s3://bucket/benchmarks/") +``` + +Only `.py`, `pyproject.toml`, and `requirements.txt` files are uploaded. + +## Inference Providers + +InspectLens supports three inference modes configured via `InspectLensConfig`. + +### Option A: Bedrock (default) + +The simplest path — no endpoint management needed. If `bedrock_model_id` is omitted, the SDK uses the cross-region inference profile for the `model` passed to `ForgeEvaluator`. + +```python +config = InspectLensConfig( + benchmarks_path=benchmarks_s3_uri, + # bedrock_model_id="us.amazon.nova-2-lite-v1:0", # optional override + tasks=[{"name": "boolq_pt", "limit": 100}], + output_s3_path="s3://bucket/results/bedrock-eval/", +) +``` + +### Option B: Existing SageMaker Endpoint + +Use when you already have a deployed SageMaker endpoint. + +```python +config = InspectLensConfig( + benchmarks_path=benchmarks_s3_uri, + endpoint_name="my-deployed-nova-endpoint", + tasks=[{"name": "boolq_pt", "limit": 100}], + output_s3_path="s3://bucket/results/endpoint-eval/", +) +``` + +### Option C: Create a New SageMaker Endpoint + +Use when you have model artifacts in S3 and want the SDK to spin up an endpoint, run the evaluation, and clean up afterwards. + +```python +config = InspectLensConfig( + benchmarks_path=benchmarks_s3_uri, + model_s3_uri="s3://bucket/model-artifacts/", + inference_image_uri="123456789012.dkr.ecr.us-east-1.amazonaws.com/your-image:latest", + endpoint_instance_type="ml.g5.12xlarge", + endpoint_execution_role_arn="arn:aws:iam::123456789012:role/YourRole", + cleanup_endpoint=True, # delete endpoint after eval (default) + tasks=[{"name": "boolq_pt", "limit": 100}], + output_s3_path="s3://bucket/results/new-endpoint-eval/", +) +``` + +## InspectLensConfig Reference + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `benchmarks_path` | `str` | Required | S3 URI containing benchmark `.py` files | +| `tasks` | `list[dict]` | `[]` | Task dicts with `name`, optional `limit` and `epochs`. Empty = run all tasks | +| `output_s3_path` | `str` | Auto | S3 prefix for result JSON logs | +| `output_format` | `str` | `"eval"` | Output format: `"eval"`, `"csv"`, `"jsonl"`, `"json"` | +| `bedrock_model_id` | `str` | Auto | Bedrock model ID/ARN (Bedrock mode) | +| `endpoint_name` | `str` | None | Existing SageMaker endpoint (endpoint mode) | +| `model_s3_uri` | `str` | None | Model artifacts S3 URI (create-endpoint mode) | +| `inference_image_uri` | `str` | None | ECR image for new endpoint (requires `model_s3_uri`) | +| `endpoint_instance_type` | `str` | None | Instance type for new endpoint | +| `endpoint_instance_count` | `int` | 1 | Instance count for new endpoint | +| `endpoint_execution_role_arn` | `str` | None | IAM role for new endpoint | +| `context_length` | `str` | None | Context length for new endpoint | +| `max_concurrency` | `str` | None | Max concurrency for new endpoint | +| `enable_rai` | `bool` | True | Enable RAI guardrails on the endpoint | +| `cleanup_endpoint` | `bool` | True | Delete endpoint after evaluation | +| `endpoint_prefix` | `str` | `"inspectlens"` | Prefix for auto-created endpoint names | +| `endpoint_environment` | `dict[str, str]` | None | Env vars for the inference endpoint container | +| `extra_args` | `list[str]` | None | Additional CLI args forwarded to `inspect eval` | +| `environment` | `dict[str, str]` | None | Env vars for the container (e.g. `{"HF_TOKEN": "..."}`) | + +To override the InspectLens orchestrator container image, set `image_uri` on `ForgeConfig` (not on `InspectLensConfig`): + +```python +config = ForgeConfig( + output_s3_path="s3://bucket/output/", + image_uri="763104351884.dkr.ecr.us-east-1.amazonaws.com/sagemaker-inspect-ai", +) +``` + +The default image is `763104351884.dkr.ecr.{region}.amazonaws.com/sagemaker-inspect-ai`, resolved automatically from the evaluator's region. + +## Decoding Overrides + +Decoding parameters are passed via the `overrides` dict on `evaluator.evaluate()`: + +```python +result = evaluator.evaluate( + job_name="my-eval", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=config, + overrides={ + "temperature": 0.0, + "top_p": 1.0, + "max_tokens": 8192, + "max_connections": 16, + "max_retries": 100, + "timeout": 600, + }, +) +``` + +| Override | Default | Description | +|---------|---------|-------------| +| `temperature` | 0.0 | Sampling temperature | +| `top_p` | 1.0 | Nucleus sampling | +| `max_tokens` | 8192 | Max output tokens | +| `max_connections` | 16 | Parallel inference connections | +| `max_retries` | 100 | Retry count for failed requests | +| `timeout` | 600 | Request timeout (seconds) | + +## Dry Run Validation + +Use `dry_run=True` to validate your configuration locally without submitting a job: + +```python +evaluator.evaluate( + job_name="my-eval", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=config, + dry_run=True, +) +``` + +This validates config fields, resolves the inference provider, and saves the generated YAML locally. + +## Monitoring and Results + +### Check job status + +```python +status, message = result.get_job_status() +print(f"Status: {status} — {message}") +``` + +### Stream CloudWatch logs + +```python +evaluator.get_logs(job_result=result, limit=50, start_from_head=False) +``` + +### Retrieve results + +Results are written as JSON logs to the configured `output_s3_path` after the job completes: + +```python +print(f"Results location: {result.eval_output_path}") +``` + +## MLflow Tracking + +Pass an `MLflowMonitor` to `ForgeConfig` and the SDK automatically populates the MLflow tracking section in the InspectLens config. + +Both tracking server ARNs (`mlflow-tracking-server/...`) and app ARNs (`mlflow-app/...`) are supported. + +```python +from amzn_nova_forge.monitor import MLflowMonitor + +mlflow_monitor = MLflowMonitor( + tracking_uri="arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/my-app", + experiment_name="nova-inspectlens-evals", + run_name="inspectlens-run-1", +) + +evaluator = ForgeEvaluator( + model=Model.NOVA_LITE_2, + infra=eval_infra, + config=ForgeConfig( + output_s3_path="s3://bucket/output", + mlflow_monitor=mlflow_monitor, + ), +) + +result = evaluator.evaluate( + job_name="inspectlens-mlflow-eval", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=InspectLensConfig( + benchmarks_path=benchmarks_s3_uri, + tasks=[{"name": "boolq_pt", "limit": 100}], + ), +) +``` + +## Evaluating a Fine-Tuned Model + +You can pass a `TrainingResult` from a previous training job to evaluate the trained checkpoint: + +```python +# After training +training_result = trainer.train(job_name="my-sft-job") + +# Evaluate the trained model via Bedrock or endpoint +eval_result = evaluator.evaluate( + job_name="post-training-eval", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=config, + job_result=training_result, +) +``` + +If the checkpoint is an S3 URI and Bedrock mode is selected, the SDK falls back to the base model and logs a warning. To evaluate an S3 checkpoint, use the endpoint modes (`endpoint_name` or `model_s3_uri` + `inference_image_uri`). + +## Job Caching + +When `enable_job_caching=True` in `ForgeConfig`, the SDK caches completed InspectLens results. Subsequent calls with identical parameters return the cached result without submitting a new job. + +Cache keys include: `job_name`, `benchmarks_path`, `tasks`, inference scenario, endpoint name, Bedrock model ID, and overrides. + +## Summary + +- `InspectLensConfig` controls benchmarks, inference provider, and output location +- `benchmarks_path` must be an S3 URI — use `evaluator.upload_benchmarks()` to upload first +- Pass `eval_task=EvaluationTask.INSPECT_LENS` to `ForgeEvaluator.evaluate()` +- Use `dry_run=True` to validate without submitting +- Decoding params go in `overrides`, not `InspectLensConfig` +- Pass `MLflowMonitor` via `ForgeConfig` for experiment tracking diff --git a/samples/inspectlens_quickstart.ipynb b/samples/inspectlens_quickstart.ipynb new file mode 100644 index 0000000..124667f --- /dev/null +++ b/samples/inspectlens_quickstart.ipynb @@ -0,0 +1,583 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Amazon Nova Forge SDK - InspectLens Evaluation Quick Start\n", + "\n", + "This notebook walks through running **InspectLens** evaluations using the Nova Forge SDK.\n", + "InspectLens is a job-based evaluation framework that runs benchmarks against your Nova model\n", + "via a SageMaker Training Job — no GPU required, since inference is delegated to a Bedrock\n", + "endpoint or an existing/new SageMaker endpoint.\n", + "\n", + "## What You'll Learn\n", + "\n", + "1. Setting up AWS credentials and S3 paths\n", + "2. Configuring `InspectLensConfig` for different inference providers\n", + "3. Creating benchmark tasks and running evaluations\n", + "4. Monitoring progress and retrieving results\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 ForgeEvaluator](#step-4-initialize-forgeevaluator)\n", + "- [Step 5: Create Sample Benchmarks](#step-5-create-sample-benchmarks)\n", + "- [Step 6: Run InspectLens Evaluation](#step-6-run-inspectlens-evaluation)\n", + " - [Option A: Bedrock Inference (default)](#option-a-bedrock-inference-default)\n", + " - [Option B: Existing SageMaker Endpoint](#option-b-existing-sagemaker-endpoint)\n", + " - [Option C: Create a New SageMaker Endpoint](#option-c-create-a-new-sagemaker-endpoint)\n", + "- [Step 7: Monitor and Retrieve Results](#step-7-monitor-and-retrieve-results)\n", + "- [Advanced: MLflow Tracking](#advanced-mlflow-tracking)\n", + "- [Summary](#summary)\n", + "\n", + "## Prerequisites\n", + "\n", + "- AWS credentials configured\n", + "- S3 bucket for evaluation artifacts and results\n", + "- IAM permissions for SageMaker and Bedrock\n", + "- Nova Forge SDK installed per [its README](https://github.com/aws/nova-forge-sdk/blob/main/README.md#installation)\n", + "- An `execution_role` with SageMaker and S3 permissions (required for InspectLens)\n", + "\n", + "## Helpful Links\n", + "- See `docs/spec.md` for the full parameter reference.\n", + "- See `README.md` for a high-level overview of the Nova Forge SDK." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 1: Import Required Modules" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!cd ../ && pip install ." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "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", + " \"\"\"\n", + " Load AWS credentials with fallback behavior.\n", + "\n", + " Args:\n", + " profile (str, optional): AWS profile name. If provided, loads from credentials file.\n", + " If None, uses current authenticated AWS session.\n", + "\n", + " Returns:\n", + " dict: Dictionary containing AWS credentials and region\n", + "\n", + " Raises:\n", + " RuntimeError: If credential loading fails\n", + " \"\"\"\n", + " if profile:\n", + " # Try loading from credentials file\n", + " try:\n", + " session = boto3.Session(profile_name=profile)\n", + " credentials = session.get_credentials()\n", + "\n", + " if not credentials:\n", + " raise RuntimeError(f\"No credentials found for profile '{profile}'\")\n", + "\n", + " except ProfileNotFound:\n", + " raise RuntimeError(f\"Profile '{profile}' not found in credentials file\")\n", + " except Exception as e:\n", + " raise RuntimeError(f\"Failed to load credentials from file: {e}\")\n", + "\n", + " else:\n", + " # Try loading from current authenticated session\n", + " try:\n", + " session = boto3.Session()\n", + " credentials = session.get_credentials()\n", + "\n", + " if not credentials:\n", + " raise RuntimeError(\"No credentials found in current AWS session\")\n", + "\n", + " except NoCredentialsError:\n", + " raise RuntimeError(\"No AWS credentials configured\")\n", + " except Exception as e:\n", + " raise RuntimeError(f\"Failed to load credentials from current session: {e}\")\n", + "\n", + " # Validate credentials by making a test call\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", + " except Exception as e:\n", + " raise RuntimeError(f\"Failed to validate 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, + "metadata": {}, + "outputs": [], + "source": [ + "creds = load_credentials()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from amzn_nova_forge import *\n", + "from amzn_nova_forge.core.enums import EvaluationTask, Model\n", + "from amzn_nova_forge.evaluator import ForgeEvaluator, InspectLensConfig\n", + "\n", + "print(\"✅ SDK imported successfully!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Configure Your AWS Resources" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Replace the values below with your own S3 bucket and IAM execution role\n", + "S3_BUCKET = \"your-bucket-name\" # Replace with your S3 bucket name, e.g. \"my-nova-eval-bucket\"\n", + "S3_OUTPUT_PATH = f\"s3://{S3_BUCKET}/inspectlens/output\"\n", + "\n", + "# IAM role ARN with SageMaker + S3 permissions — required for InspectLens\n", + "# e.g. \"arn:aws:iam::123456789012:role/YourSageMakerExecutionRole\"\n", + "EXECUTION_ROLE = \"arn:aws:iam::123456789012:role/YourSageMakerRole\"\n", + "\n", + "print(f\"Output Path: {S3_OUTPUT_PATH}\")\n", + "print(f\"Execution Role: {EXECUTION_ROLE}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3: Configure Runtime Infrastructure\n", + "\n", + "InspectLens runs as a CPU-only SageMaker Training Job (the container acts as an orchestrator).\n", + "Use `SMTJRuntimeManager` with a CPU instance — no GPU needed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "eval_infra = SMTJRuntimeManager(\n", + " instance_type=\"ml.m5.large\", # CPU instance — no GPU required\n", + " instance_count=1,\n", + " execution_role=EXECUTION_ROLE, # Required for InspectLens\n", + ")\n", + "\n", + "print(\"✅ Runtime configured\")\n", + "print(f\" Instance Type: {eval_infra.instance_type}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4: Initialize ForgeEvaluator" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "evaluator = ForgeEvaluator(\n", + " model=Model.NOVA_LITE_2,\n", + " infra=eval_infra,\n", + " config=ForgeConfig(\n", + " output_s3_path=S3_OUTPUT_PATH,\n", + " ),\n", + ")\n", + "\n", + "print(\"✅ ForgeEvaluator initialized — Model: Nova Lite 2.0\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 5: Create Sample Benchmarks\n", + "\n", + "InspectLens requires benchmark `.py` files with `@task` decorators.\n", + "The cell below creates a local `my_benchmarks/` directory with two ready-to-use tasks\n", + "that wrap built-in `inspect-evals` tasks.\n", + "\n", + "**Important:** You must upload your benchmarks to S3 before starting an evaluation job.\n", + "Use `evaluator.upload_benchmarks(local_dir, s3_path)` to upload, then pass the returned\n", + "S3 URI as `benchmarks_path` in `InspectLensConfig`.\n", + "\n", + "### Pre-installed packages\n", + "\n", + "The InspectLens container comes with the following packages pre-installed:\n", + "`inspect-ai`, `boto3`, `aioboto3`, `openai`, `mlflow`, `pyyaml`.\n", + "\n", + "Any additional dependencies your benchmarks require (e.g. `inspect-evals`) must be declared\n", + "in a `pyproject.toml` placed alongside your benchmark files. The container installs it automatically at job start." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "BENCHMARKS_DIR = \"./my_benchmarks\"\n", + "os.makedirs(BENCHMARKS_DIR, exist_ok=True)\n", + "\n", + "# BoolQ — reading comprehension (yes/no questions)\n", + "with open(f\"{BENCHMARKS_DIR}/boolq_pt.py\", \"w\") as f:\n", + " f.write(\"\"\"\\\n", + "from inspect_ai import task\n", + "from inspect_evals.boolq import boolq\n", + "\n", + "\n", + "@task\n", + "def boolq_pt():\n", + " return boolq()\n", + "\"\"\")\n", + "\n", + "# MMLU Pro — multi-domain multiple choice\n", + "with open(f\"{BENCHMARKS_DIR}/mmlu_pro_pt.py\", \"w\") as f:\n", + " f.write(\"\"\"\\\n", + "from inspect_ai import task\n", + "from inspect_evals.mmlu_pro import mmlu_pro\n", + "\n", + "\n", + "@task\n", + "def mmlu_pro_pt():\n", + " return mmlu_pro()\n", + "\"\"\")\n", + "\n", + "print(f\"✅ Sample benchmarks created in {BENCHMARKS_DIR}/\")\n", + "for f in os.listdir(BENCHMARKS_DIR):\n", + " print(f\" {f}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Create a pyproject.toml in the benchmarks directory so inspect-evals is installed\n", + "# when the SageMaker Training Job container starts up.\n", + "with open(f\"{BENCHMARKS_DIR}/pyproject.toml\", \"w\") as f:\n", + " f.write(\"\"\"\\\n", + "[build-system]\n", + "requires = [\"setuptools\"]\n", + "build-backend = \"setuptools.backends.legacy:build\"\n", + "\n", + "[project]\n", + "name = \"my-benchmarks\"\n", + "version = \"0.1.0\"\n", + "dependencies = [\n", + " \"inspect-evals\",\n", + "]\n", + "\"\"\")\n", + "\n", + "print(f\"✅ pyproject.toml created in {BENCHMARKS_DIR}/\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Upload benchmarks to S3 — this is a separate step from starting the eval job.\n", + "# You choose where to store them and pass the S3 URI when configuring InspectLensConfig.\n", + "BENCHMARKS_S3_PATH = f\"s3://{S3_BUCKET}/inspectlens/benchmarks/my_benchmarks/\"\n", + "\n", + "benchmarks_s3_uri = evaluator.upload_benchmarks(BENCHMARKS_DIR, BENCHMARKS_S3_PATH)\n", + "print(f\"Benchmarks uploaded to: {benchmarks_s3_uri}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 6: Run InspectLens Evaluation\n", + "\n", + "InspectLens supports three inference provider modes:\n", + "- **Bedrock** (default) — evaluates against a Bedrock model ID or ARN\n", + "- **Existing SageMaker endpoint** — evaluates against a deployed endpoint\n", + "- **New SageMaker endpoint** — creates an endpoint from model artifacts, evaluates, then optionally cleans up\n", + "\n", + "Use `dry_run=True` first to validate your config locally without submitting a job." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Option A: Bedrock Inference (default)\n", + "\n", + "The simplest path — no endpoint management needed. Inference is routed through Bedrock.\n", + "If `bedrock_model_id` is omitted, the SDK uses the cross-region inference profile for\n", + "the `model` passed to `ForgeEvaluator`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "inspect_config_bedrock = InspectLensConfig(\n", + " benchmarks_path=benchmarks_s3_uri, # S3 URI from upload_benchmarks()\n", + " # bedrock_model_id=\"us.amazon.nova-2-lite-v1:0\", # defaults to ForgeEvaluator model\n", + " tasks=[\n", + " {\"name\": \"boolq_pt\", \"limit\": 100},\n", + " {\"name\": \"mmlu_pro_pt\", \"limit\": 50},\n", + " ],\n", + " output_s3_path=f\"{S3_OUTPUT_PATH}/bedrock-eval/\",\n", + ")\n", + "\n", + "# Validate config locally before submitting\n", + "evaluator.evaluate(\n", + " job_name=\"inspectlens-bedrock-eval\",\n", + " eval_task=EvaluationTask.INSPECT_LENS,\n", + " inspect_lens_config=inspect_config_bedrock,\n", + " dry_run=True,\n", + ")\n", + "print(\"Config validated — ready to submit\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "eval_result_bedrock = evaluator.evaluate(\n", + " job_name=\"inspectlens-bedrock-eval\",\n", + " eval_task=EvaluationTask.INSPECT_LENS,\n", + " inspect_lens_config=inspect_config_bedrock,\n", + " # overrides={\"temperature\": 0.0, \"max_tokens\": 4096, \"max_connections\": 8},\n", + ")\n", + "\n", + "print(\"🚀 InspectLens evaluation started (Bedrock)!\")\n", + "print(f\" Job ID: {eval_result_bedrock.job_id}\")\n", + "print(f\" Results: {eval_result_bedrock.eval_output_path}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Option B: Existing SageMaker Endpoint\n", + "\n", + "Use this when you already have a deployed SageMaker endpoint." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# inspect_config_endpoint = InspectLensConfig(\n", + "# benchmarks_path=benchmarks_s3_uri, # S3 URI from upload_benchmarks()\n", + "# endpoint_name=\"my-custom-nova-model-sagemaker\", # TODO: Replace\n", + "# tasks=[{\"name\": \"boolq_pt\", \"limit\": 100}],\n", + "# output_s3_path=f\"{S3_OUTPUT_PATH}/endpoint-eval/\",\n", + "# )\n", + "#\n", + "# eval_result_endpoint = evaluator.evaluate(\n", + "# job_name=\"inspectlens-endpoint-eval\",\n", + "# eval_task=EvaluationTask.INSPECT_LENS,\n", + "# inspect_lens_config=inspect_config_endpoint,\n", + "# )\n", + "# print(f\" Job ID: {eval_result_endpoint.job_id}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Option C: Create a New SageMaker Endpoint\n", + "\n", + "Use this when you have model artifacts in S3 and want the SDK to spin up an endpoint,\n", + "run the evaluation, and optionally clean up the endpoint afterwards." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# inspect_config_new_endpoint = InspectLensConfig(\n", + "# benchmarks_path=benchmarks_s3_uri, # S3 URI from upload_benchmarks()\n", + "# model_s3_uri=\"s3://your-bucket/model-artifacts/\", # Replace with your model artifacts S3 path\n", + "# inference_image_uri=\"123456789012.dkr.ecr.us-east-1.amazonaws.com/your-image:latest\", # Replace with your ECR image URI: .dkr.ecr..amazonaws.com/:\n", + "# endpoint_instance_type=\"ml.g5.12xlarge\",\n", + "# endpoint_execution_role_arn=EXECUTION_ROLE,\n", + "# cleanup_endpoint=True, # delete endpoint after eval (default: True)\n", + "# tasks=[{\"name\": \"boolq_pt\", \"limit\": 100}],\n", + "# output_s3_path=f\"{S3_OUTPUT_PATH}/new-endpoint-eval/\",\n", + "# )\n", + "#\n", + "# eval_result_new_ep = evaluator.evaluate(\n", + "# job_name=\"inspectlens-new-endpoint-eval\",\n", + "# eval_task=EvaluationTask.INSPECT_LENS,\n", + "# inspect_lens_config=inspect_config_new_endpoint,\n", + "# )\n", + "# print(f\" Job ID: {eval_result_new_ep.job_id}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 7: Monitor and Retrieve Results" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Check job status\n", + "status, message = eval_result_bedrock.get_job_status()\n", + "print(f\"Status: {status} — {message}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Stream CloudWatch logs\n", + "evaluator.get_logs(job_result=eval_result_bedrock, limit=50, start_from_head=False)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Results are written as JSON logs to the output S3 path after the job completes\n", + "print(f\"Results: {eval_result_bedrock.eval_output_path}\")\n", + "print(f\"Job ID: {eval_result_bedrock.job_id}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Advanced: MLflow Tracking\n", + "\n", + "Pass an `MLflowMonitor` to `ForgeConfig` and the SDK automatically populates the MLflow\n", + "tracking section in the InspectLens config.\n", + "\n", + "Both tracking server ARNs (`mlflow-tracking-server/...`) and app ARNs (`mlflow-app/...`) are supported." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# mlflow_monitor = MLflowMonitor(\n", + "# tracking_uri=\"\", # TODO: Replace\n", + "# # arn:aws:sagemaker:::mlflow-tracking-server/\n", + "# # arn:aws:sagemaker:::mlflow-app/\n", + "# experiment_name=\"nova-inspectlens-evals\",\n", + "# run_name=\"inspectlens-run-1\",\n", + "# )\n", + "#\n", + "# evaluator_with_mlflow = ForgeEvaluator(\n", + "# model=Model.NOVA_LITE_2,\n", + "# infra=eval_infra,\n", + "# config=ForgeConfig(\n", + "# output_s3_path=S3_OUTPUT_PATH,\n", + "# mlflow_monitor=mlflow_monitor, # SDK auto-populates MLflow in InspectLens config\n", + "# ),\n", + "# )\n", + "#\n", + "# eval_result_mlflow = evaluator_with_mlflow.evaluate(\n", + "# job_name=\"inspectlens-mlflow-eval\",\n", + "# eval_task=EvaluationTask.INSPECT_LENS,\n", + "# inspect_lens_config=InspectLensConfig(\n", + "# benchmarks_path=benchmarks_s3_uri, # S3 URI from upload_benchmarks()\n", + "# tasks=[{\"name\": \"boolq_pt\", \"limit\": 100}],\n", + "# output_s3_path=f\"{S3_OUTPUT_PATH}/mlflow-eval/\",\n", + "# ),\n", + "# )\n", + "# print(f\" Job ID: {eval_result_mlflow.job_id}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "- `InspectLensConfig` controls benchmarks, inference provider, decoding, and output\n", + "- `benchmarks_path` must be an S3 URI — use `evaluator.upload_benchmarks(local_dir, s3_path)` to upload a local directory first\n", + "- Pass `eval_task=EvaluationTask.INSPECT_LENS` to `ForgeEvaluator.evaluate()`\n", + "- Use `dry_run=True` to validate config without submitting a job\n", + "- Pass `MLflowMonitor` via `ForgeConfig` to enable MLflow tracking\n", + "\n", + "For more details, see `docs/spec.md` or the SDK README." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/src/amzn_nova_forge/__init__.py b/src/amzn_nova_forge/__init__.py index 679b318..99ae5ec 100644 --- a/src/amzn_nova_forge/__init__.py +++ b/src/amzn_nova_forge/__init__.py @@ -50,7 +50,7 @@ from .dataset.operations.transform_operation import TransformMethod from .dataset.operations.validate_operation import ValidateMethod from .deployer import ForgeDeployer -from .evaluator import EvalTaskConfig, ForgeEvaluator +from .evaluator import EvalTaskConfig, ForgeEvaluator, InspectLensConfig from .inference import ForgeInference from .manager import ( BedrockRuntimeManager, diff --git a/src/amzn_nova_forge/__version__.py b/src/amzn_nova_forge/__version__.py index 9f90d55..699e2a5 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.7" # pragma: no cover +VERSION = "1.4.8" # pragma: no cover diff --git a/src/amzn_nova_forge/core/constants.py b/src/amzn_nova_forge/core/constants.py index 0af06d8..03a72ab 100644 --- a/src/amzn_nova_forge/core/constants.py +++ b/src/amzn_nova_forge/core/constants.py @@ -294,6 +294,13 @@ def get_available_subtasks(task: EvaluationTask) -> List[str]: ESCROW_URI_TAG_KEY = "sagemaker.amazonaws.com/forge/escrow-uri" +INSPECT_LENS_DEFAULT_IMAGE_REPO = "763104351884.dkr.ecr.{region}.amazonaws.com/sagemaker-inspect-ai" + + +def get_inspect_lens_default_image_uri(region: str) -> str: + """Return the default InspectLens image URI for the given region.""" + return INSPECT_LENS_DEFAULT_IMAGE_REPO.format(region=region) + def _escrow_tag_value(escrow_uri: str) -> str: """Normalize escrow URI for use as a tag value (max 256 chars).""" diff --git a/src/amzn_nova_forge/core/enums.py b/src/amzn_nova_forge/core/enums.py index b74a19a..7b5dd34 100644 --- a/src/amzn_nova_forge/core/enums.py +++ b/src/amzn_nova_forge/core/enums.py @@ -81,6 +81,16 @@ def from_model_name(cls, model_name: str) -> "Model": return model raise ValueError(f"Unknown model name: {model_name}") + @property + def bedrock_model_id(self) -> str: + """Return the Bedrock cross-region inference profile ID. + + Strips the context-length suffix from model_type and prepends 'us.' + e.g. 'amazon.nova-micro-v1:0:128k' -> 'us.amazon.nova-micro-v1:0' + """ + base = self.model_type.rsplit(":", 1)[0] # strip ':128k', ':300k', ':256k' + return f"us.{base}" + NOVA_MICRO = ( "nova_micro", Version.ONE, @@ -179,6 +189,7 @@ class EvaluationTask(enum.Enum): RUBRIC_LLM_JUDGE = "rubric_llm_judge" RFT_EVAL = "rft_eval" RFT_MULTITURN_EVAL = "rft_multiturn_eval" # Maps to "rft_eval" in recipes + INSPECT_LENS = "inspect_lens" # InspectLens job-based evaluation via SageMaker Training Job def get_recipe_value(self) -> str: """Get the value to use in recipe templates.""" diff --git a/src/amzn_nova_forge/core/types.py b/src/amzn_nova_forge/core/types.py index 05a1b8f..bb1c300 100644 --- a/src/amzn_nova_forge/core/types.py +++ b/src/amzn_nova_forge/core/types.py @@ -19,7 +19,7 @@ from dataclasses import dataclass, field from datetime import datetime -from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, Optional, TypedDict +from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, List, Optional, TypedDict from pydantic import BaseModel @@ -156,6 +156,7 @@ class JobConfig: mlflow_run_name: Optional[str] = None method: Optional[TrainingMethod] = None # Training method (required for Bedrock) data_mixing_config: Optional[Dict[str, Any]] = None # Datamix percent fields (SMTJServerless) + environment: Optional[Dict[str, str]] = None # Environment variables for the training container # 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/evaluator/__init__.py b/src/amzn_nova_forge/evaluator/__init__.py index 14a6053..c00866b 100644 --- a/src/amzn_nova_forge/evaluator/__init__.py +++ b/src/amzn_nova_forge/evaluator/__init__.py @@ -12,5 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. from amzn_nova_forge.evaluator.forge_evaluator import EvalTaskConfig, ForgeEvaluator +from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig -__all__ = ["ForgeEvaluator", "EvalTaskConfig"] +__all__ = ["ForgeEvaluator", "EvalTaskConfig", "InspectLensConfig"] diff --git a/src/amzn_nova_forge/evaluator/forge_evaluator.py b/src/amzn_nova_forge/evaluator/forge_evaluator.py index 124b356..aa7a4e3 100644 --- a/src/amzn_nova_forge/evaluator/forge_evaluator.py +++ b/src/amzn_nova_forge/evaluator/forge_evaluator.py @@ -15,14 +15,18 @@ from __future__ import annotations +import io +import os import uuid from dataclasses import dataclass from datetime import datetime, timezone +from tempfile import mkdtemp from typing import TYPE_CHECKING, Any, Dict, Optional, cast import boto3 +import yaml -from amzn_nova_forge.core.constants import DEFAULT_REGION +from amzn_nova_forge.core.constants import DEFAULT_REGION, get_inspect_lens_default_image_uri from amzn_nova_forge.core.enums import EvaluationTask, Model, Platform, TrainingMethod from amzn_nova_forge.core.job_cache import ( build_cache_context, @@ -41,6 +45,7 @@ JobConfig, 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.model.nova_model_customizer_util import ( requires_custom_eval_data, @@ -55,6 +60,8 @@ detect_platform_from_path, validate_platform_compatibility, ) +from amzn_nova_forge.util.s3_utils import parse_s3_uri +from amzn_nova_forge.validation.inspect_lens_validator import validate_task_names_against_benchmarks from amzn_nova_forge.validation.validator import validate_rft_lambda_name if TYPE_CHECKING: @@ -136,6 +143,7 @@ def evaluate( dry_run: bool = False, job_result: Optional[TrainingResult] = None, rft_multiturn_infra: Optional[RFTMultiturnInfrastructure] = None, + inspect_lens_config: Optional[InspectLensConfig] = None, ) -> Optional[EvaluationResult]: """Launch an evaluation job. @@ -150,10 +158,64 @@ def evaluate( job_result: Optional TrainingResult to extract checkpoint path from. rft_multiturn_infra: Optional RFT multiturn infrastructure (passed through to RecipeBuilder for multi-turn RFT evaluation). + inspect_lens_config: Required when eval_task is INSPECT_LENS. + Configures benchmarks, inference provider, decoding, and output. Returns: EvaluationResult on success, None if dry_run is True. """ + if eval_task == EvaluationTask.INSPECT_LENS: + if overrides: + # Warn only if overrides contains keys that are not InspectLens decoding params + _INSPECT_LENS_VALID_OVERRIDES = { + "temperature", + "top_p", + "top_k", + "max_tokens", + "max_connections", + "max_retries", + "timeout", + } + unknown = set(overrides.keys()) - _INSPECT_LENS_VALID_OVERRIDES + if unknown: + logger.warning( + "'overrides' keys %s are not applied for EvaluationTask.INSPECT_LENS " + "and will be ignored. Use InspectLensConfig for benchmark, endpoint, " + "and output settings.", + sorted(unknown), + ) + + # Check job cache before doing any work + cached = load_existing_result( + self._cache_context, + job_name=job_name, + job_type="inspect_lens", + model_path=model_path, + benchmarks_path=inspect_lens_config.benchmarks_path + if inspect_lens_config + else None, + tasks=str(inspect_lens_config.tasks) if inspect_lens_config else None, + inference_scenario=inspect_lens_config._infer_scenario() + if inspect_lens_config + else None, + endpoint_name=inspect_lens_config.endpoint_name if inspect_lens_config else None, + bedrock_model_id=inspect_lens_config.bedrock_model_id + if inspect_lens_config + else None, + overrides=overrides or {}, + ) + if cached: + logger.info("Returning cached InspectLens result for '%s'.", job_name) + return cached # type: ignore[return-value] + + return self._evaluate_inspect_lens( + job_name=job_name, + inspect_lens_config=inspect_lens_config, + model_path=model_path, + job_result=job_result, + dry_run=dry_run, + overrides=overrides, + ) if self._platform == Platform.BEDROCK: raise NotImplementedError( "Evaluation is not supported on the Bedrock platform. " @@ -336,6 +398,382 @@ def evaluate( return evaluation_result + def _evaluate_inspect_lens( + self, + job_name: str, + inspect_lens_config: Optional[InspectLensConfig], + model_path: Optional[str], + job_result: Optional[TrainingResult], + dry_run: bool, + overrides: Optional[Dict[str, Any]] = None, + ) -> Optional[EvaluationResult]: + """Submit an InspectLens evaluation as a SageMaker Training Job. + + The training job runs the InspectLens container as an orchestrator only + (no GPU needed). Inference is delegated to a Bedrock endpoint or an + existing SageMaker endpoint as configured in ``inspect_lens_config``. + """ + if inspect_lens_config is None: + raise ValueError( + "inspect_lens_config is required when eval_task=EvaluationTask.INSPECT_LENS." + ) + + # Resolve the inference provider from model_path / job_result. + # Priority: explicit model_path > job_result checkpoint > job_result model ARN. + resolved_model_path = model_path + if resolved_model_path is None and job_result is not None: + resolved_model_path = ( + job_result.model_artifacts.checkpoint_s3_path + or job_result.model_artifacts.output_model_arn + ) + + # Guard: an S3 checkpoint URI is not a valid Bedrock model ID. + # If the resolved path is an S3 URI and no endpoint is configured, fall back to + # the evaluator's model enum so the job doesn't fail at runtime. + if ( + resolved_model_path is not None + and resolved_model_path.startswith("s3://") + and inspect_lens_config._infer_scenario() == "bedrock" + ): + logger.warning( + "model_path '%s' is an S3 URI and cannot be used as a Bedrock model ID. " + "Falling back to the evaluator model (%s) for Bedrock inference. " + "To evaluate an S3 checkpoint, set endpoint_name or model_s3_uri + " + "inference_image_uri on InspectLensConfig.", + resolved_model_path, + self.model.bedrock_model_id, + ) + resolved_model_path = None + + inference_provider = self._build_inspect_lens_inference_provider( + inspect_lens_config, resolved_model_path + ) + + # Resolve output S3 path: config override > evaluator default + results_s3_path = inspect_lens_config.output_s3_path or ( + f"{self.output_s3_path.rstrip('/')}/inspectlens-results/" + ) + + # Generate run ID here — used for benchmarks, config, and job S3 paths + run_id = str(uuid.uuid4()) + + # Always use output_s3_path as the base for config and output. + # benchmarks_path is just a reference in the YAML — no need to co-locate. + base_s3_path = self.output_s3_path + + # Validate benchmarks_path is already an S3 URI. + # Users must call upload_benchmarks() separately before starting the job. + if ( + inspect_lens_config.benchmarks_path + and not inspect_lens_config.benchmarks_path.startswith("s3://") + ): + raise ValueError( + f"benchmarks_path must be an s3:// URI, got: '{inspect_lens_config.benchmarks_path}'. " + f"Upload your local benchmarks first using evaluator.upload_benchmarks(local_dir, s3_path), " + f"then pass the returned S3 URI as benchmarks_path in InspectLensConfig." + ) + + # Validate task names against @task functions in the benchmarks S3 path. + # Done here (not in __post_init__) to avoid S3 calls during config construction. + task_name_errors = validate_task_names_against_benchmarks( + tasks=inspect_lens_config.tasks, + benchmarks_path=inspect_lens_config.benchmarks_path, + region=self.region, + ) + if task_name_errors: + bullet_list = "\n".join(f" - {e}" for e in task_name_errors) + raise ValueError( + f"InspectLensConfig task validation failed with {len(task_name_errors)} error(s):\n{bullet_list}" + ) + + # Serialize config to YAML (needed for both dry_run local save and actual upload) + # Populate MLflow tracking section from ForgeConfig.mlflow_monitor if provided. + mlflow_monitor = self._config.mlflow_monitor + mlflow_tracking_section: Optional[Dict[str, Any]] = None + if mlflow_monitor and mlflow_monitor.tracking_uri: + mlflow_tracking_section = { + "mlflow_tracking_arn": mlflow_monitor.tracking_uri, + "mlflow_experiment_name": mlflow_monitor.experiment_name or "nova-evals", + "mlflow_tracing": True, + "mlflow_log_artifacts": True, + } + logger.info( + f"Populated InspectLens MLflow tracking from ForgeConfig.mlflow_monitor: " + f"{mlflow_monitor.tracking_uri}" + ) + + config_dict = inspect_lens_config.to_yaml_dict( + inference_provider=inference_provider, + output_s3_path=results_s3_path, + overrides=overrides, + ) + if mlflow_tracking_section: + config_dict["tracking"] = mlflow_tracking_section + config_yaml = yaml.dump(config_dict, default_flow_style=False, sort_keys=False) + + # Save config YAML locally — always write to a temp dir (mirrors RecipeBuilder default), + # or to generated_recipe_dir if explicitly set. + if self._config.generated_recipe_dir: + local_dir = os.path.expanduser(self._config.generated_recipe_dir) + else: + local_dir = mkdtemp() + + os.makedirs(local_dir, exist_ok=True) + local_config_path = os.path.join(local_dir, f"{job_name}_inspect_config.yaml") + with open(local_config_path, "w") as fh: + fh.write(config_yaml) + logger.info(f"InspectLens config saved locally to {local_config_path}") + + if dry_run: + logger.info( + f"[dry_run] InspectLens job '{job_name}' validated. " + f"Inference provider: {list(inference_provider.keys())[0]}. " + f"Results would be written to: {results_s3_path}" + ) + return None + + # Truncate job_name first to guarantee the full UUID is always preserved + config_s3_prefix = f"{base_s3_path.rstrip('/')}/{run_id}/config/" + config_s3_uri = f"{config_s3_prefix}inspect_config.yaml" + self._upload_config_to_s3(config_yaml, config_s3_uri) + + # Resolve container image: ForgeConfig > default (region-specific DLC image) + image_uri = self._config.image_uri or get_inspect_lens_default_image_uri(self.region) + + start_time = datetime.now(timezone.utc) + + # Build environment vars for the SageMaker Training Job + env_vars: Dict[str, Any] = dict(inspect_lens_config.environment or {}) + environment: Optional[Dict[str, Any]] = env_vars or None + + # InspectLens uses ml.m5.large (CPU orchestrator only — no GPU needed). + # SMTJRuntimeManager.execute() uses ModelTrainer.from_recipe() which rejects + # CPU instances, so we call create_training_job directly here. + # TODO: extract sagemaker_client/execution_role resolution into a helper + sm_client = getattr(self.infra, "sagemaker_client", None) or boto3.client( + "sagemaker", region_name=self.region + ) + execution_role = getattr(self.infra, "execution_role", None) + if not execution_role: + raise ValueError( + "InspectLens requires an execution_role on the RuntimeManager. " + "Pass execution_role to SMTJRuntimeManager." + ) + + job_id = self._submit_inspect_lens_training_job( + job_name=job_name, + run_id=run_id, + base_s3_path=base_s3_path, + config_s3_prefix=config_s3_prefix, + image_uri=image_uri, + environment=environment, + sm_client=sm_client, + execution_role=execution_role, + ) + logger.info( + f"InspectLens job '{job_id}' submitted. Results will be written to: {results_s3_path}" + ) + + evaluation_result = SMTJEvaluationResult( + job_id=job_id, + eval_task=EvaluationTask.INSPECT_LENS, + started_time=start_time, + eval_output_path=results_s3_path, + sagemaker_client=sm_client, + ) + persist_result( + self._cache_context, + evaluation_result, + job_name=job_name, + job_type="inspect_lens", + model_path=model_path, + benchmarks_path=inspect_lens_config.benchmarks_path, + tasks=str(inspect_lens_config.tasks), + inference_scenario=inspect_lens_config._infer_scenario(), + endpoint_name=inspect_lens_config.endpoint_name, + bedrock_model_id=inspect_lens_config.bedrock_model_id, + overrides=overrides or {}, + ) + return evaluation_result + + def _submit_inspect_lens_training_job( + self, + job_name: str, + run_id: str, + base_s3_path: str, + config_s3_prefix: str, + image_uri: str, + environment: Optional[Dict[str, Any]], + sm_client: Any, + execution_role: str, + ) -> str: + """Create the SageMaker Training Job for an InspectLens run and return the job name.""" + # Truncate job_name first to guarantee the full UUID is always preserved + # TODO move to a separate helper function + max_name_len = 63 - len(run_id) - 1 # 63 - 36 - 1 = 26 + unique_job_name = f"{job_name[:max_name_len]}-{run_id}" + + algorithm_spec: Dict[str, Any] = { + "TrainingImage": image_uri, + "TrainingInputMode": "File", + } + + create_kwargs: Dict[str, Any] = { + "TrainingJobName": unique_job_name, + "AlgorithmSpecification": algorithm_spec, + "RoleArn": execution_role, + "InputDataConfig": [ + { + "ChannelName": "config", + "DataSource": { + "S3DataSource": { + "S3DataType": "S3Prefix", + "S3Uri": config_s3_prefix, + "S3DataDistributionType": "FullyReplicated", + } + }, + "ContentType": "application/x-yaml", + } + ], + "OutputDataConfig": { + "S3OutputPath": f"{base_s3_path.rstrip('/')}/{run_id}/output/", + }, + "ResourceConfig": { + "InstanceType": self.infra.instance_type, + "InstanceCount": 1, + "VolumeSizeInGB": 30, + }, + "StoppingCondition": { + "MaxRuntimeInSeconds": getattr(self.infra, "max_job_runtime", 86400), + }, + } + if environment: + create_kwargs["Environment"] = environment + if self.infra.kms_key_id: + create_kwargs["OutputDataConfig"]["KmsKeyId"] = self.infra.kms_key_id + + sm_client.create_training_job(**create_kwargs) + return unique_job_name + + def _build_inspect_lens_inference_provider( + self, + inspect_lens_config: InspectLensConfig, + model_path: Optional[str], + ) -> Dict[str, Any]: + """Build the ``inference_provider`` block for inspect_config.yaml. + + Priority: + 1. InspectLensConfig endpoint fields (endpoint_name / model_s3_uri) + 2. model_path / job_result passed to evaluate() + 3. Fallback: Bedrock with the evaluator's model enum value + """ + scenario = inspect_lens_config._infer_scenario() + + if scenario == "existing_endpoint": + ep_existing: Dict[str, Any] = { + "endpoint_name": inspect_lens_config.endpoint_name, + "region": self.region, + } + if inspect_lens_config.context_length: + ep_existing["context_length"] = inspect_lens_config.context_length + if inspect_lens_config.max_concurrency: + ep_existing["max_concurrency"] = inspect_lens_config.max_concurrency + ep_existing["enable_rai"] = inspect_lens_config.enable_rai + return {"sagemaker_endpoint": ep_existing} + + if scenario == "create_endpoint": + ep: Dict[str, Any] = { + "endpoint_name": None, + "region": self.region, + "model_s3_uri": (inspect_lens_config.model_s3_uri or "").rstrip("/") + "/", + "inference_image_uri": inspect_lens_config.inference_image_uri, + "cleanup_endpoint": inspect_lens_config.cleanup_endpoint, + "instance_count": inspect_lens_config.endpoint_instance_count, + "enable_rai": inspect_lens_config.enable_rai, + "endpoint_prefix": inspect_lens_config.endpoint_prefix, + } + if inspect_lens_config.endpoint_instance_type: + ep["instance_type"] = inspect_lens_config.endpoint_instance_type + if inspect_lens_config.endpoint_execution_role_arn: + ep["execution_role_arn"] = inspect_lens_config.endpoint_execution_role_arn + if inspect_lens_config.context_length: + ep["context_length"] = inspect_lens_config.context_length + if inspect_lens_config.max_concurrency: + ep["max_concurrency"] = inspect_lens_config.max_concurrency + if inspect_lens_config.endpoint_environment: + ep["environment"] = inspect_lens_config.endpoint_environment + return {"sagemaker_endpoint": ep} + + # Bedrock — resolve model_id from model_path, InspectLensConfig.bedrock_model_id, + # or fall back to the evaluator's model enum (cross-region inference profile ID). + if model_path is not None: + bedrock_model_id = model_path + elif inspect_lens_config.bedrock_model_id is not None: + bedrock_model_id = inspect_lens_config.bedrock_model_id + else: + bedrock_model_id = self.model.bedrock_model_id + return { + "bedrock": { + "model_id": bedrock_model_id, + "region": self.region, + } + } + + @_telemetry_emitter(Feature.EVAL, "upload_benchmarks") + def upload_benchmarks(self, local_dir: str, s3_path: str) -> str: + """Upload a local benchmarks directory to S3. + + Call this before starting an InspectLens evaluation job. Pass the + returned S3 URI as ``benchmarks_path`` in ``InspectLensConfig``. + + Args: + local_dir: Path to the local directory containing benchmark ``.py`` + files with ``@task`` decorators. + s3_path: S3 URI (``s3://bucket/prefix/``) where the benchmark files + will be uploaded. + + Returns: + The S3 URI where benchmarks were uploaded (same as ``s3_path``, + normalized with a trailing slash). + + Raises: + ValueError: If ``local_dir`` is not an existing directory or + ``s3_path`` is not an ``s3://`` URI. + """ + expanded = os.path.expanduser(local_dir) + if not os.path.isdir(expanded): + raise ValueError(f"local_dir '{local_dir}' is not an existing directory.") + if not s3_path.startswith("s3://"): + raise ValueError(f"s3_path must be an s3:// URI, got: '{s3_path}'") + + s3_path = s3_path.rstrip("/") + "/" + s3_client = boto3.client("s3", region_name=self.region) + bucket, prefix = parse_s3_uri(s3_path) + uploaded = 0 + for fname in os.listdir(expanded): + local_file = os.path.join(expanded, fname) + if os.path.isfile(local_file) and ( + fname.endswith(".py") or fname in ("pyproject.toml", "requirements.txt") + ): + s3_key = f"{prefix}{fname}" + s3_client.upload_file(local_file, bucket, s3_key) + uploaded += 1 + logger.info(f"Uploaded {uploaded} benchmark file(s) from '{expanded}' to {s3_path}") + return s3_path + + def _upload_config_to_s3(self, config_yaml: str, s3_uri: str) -> None: + """Upload YAML string to an s3:// URI.""" + assert s3_uri.startswith("s3://"), f"Expected s3:// URI, got: {s3_uri}" + bucket, key = parse_s3_uri(s3_uri) + s3_client = boto3.client("s3", region_name=self.region) + s3_client.put_object( + Bucket=bucket, + Key=key, + Body=config_yaml.encode("utf-8"), + ContentType="application/x-yaml", + ) + logger.info(f"Uploaded InspectLens config to s3://{bucket}/{key}") + @_telemetry_emitter(Feature.EVAL, "get_logs") def get_logs( self, diff --git a/src/amzn_nova_forge/evaluator/inspect_lens_config.py b/src/amzn_nova_forge/evaluator/inspect_lens_config.py new file mode 100644 index 0000000..df75fe0 --- /dev/null +++ b/src/amzn_nova_forge/evaluator/inspect_lens_config.py @@ -0,0 +1,179 @@ +# 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. +"""InspectLens evaluation configuration.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from amzn_nova_forge.validation.inspect_lens_validator import validate_inspect_lens_config + + +@dataclass +class InspectLensConfig: + """Configuration for an InspectLens job-based evaluation. + + Serialized to ``inspect_config.yaml`` and uploaded to S3 before the + SageMaker Training Job is submitted. + + Decoding parameters (``temperature``, ``top_p``, ``top_k``, ``max_tokens``, + ``max_connections``, ``max_retries``, ``timeout``) are passed via the + ``overrides`` parameter on ``ForgeEvaluator.evaluate()``, consistent with + how other evaluation tasks handle inference configuration. + + Args: + benchmarks_path: ``s3://`` URI containing benchmark ``.py`` files with + ``@task`` decorators. Required — the container always needs a + benchmarks location. Use ``ForgeEvaluator.upload_benchmarks()`` + to upload a local directory to S3 first. + tasks: List of task dicts, each with a required ``"name"`` key and + optional ``"path"``, ``"limit"``, ``"epochs"``, and ``"task_args"`` + keys. Empty list runs all ``@task`` functions in + ``benchmarks_path`` (eval-set mode). + Example: ``[{"name": "mmlu", "path": "mmlu/mmlu_sft.py", "limit": 100, "epochs": 1, "task_args": {"subject": "anatomy"}}]`` + output_s3_path: S3 prefix where InspectLens writes eval result JSON + logs. Defaults to the ``ForgeEvaluator`` output path. + output_format: Output format for eval results. One of ``"eval"``, + ``"csv"``, ``"jsonl"``, ``"json"``. Defaults to ``"eval"``. + bedrock_model_id: Bedrock model ID or ARN for Bedrock inference mode. + e.g. ``"us.amazon.nova-micro-v1:0"``. Falls back to the + ``ForgeEvaluator.model`` enum value when not set. + endpoint_name: Existing SageMaker endpoint name to evaluate against. + Mutually exclusive with ``model_s3_uri`` / ``inference_image_uri``. + model_s3_uri: S3 URI of model artifacts for creating a new endpoint. + Requires ``inference_image_uri``. + inference_image_uri: ECR image URI for the new endpoint container. + Requires ``model_s3_uri``. + endpoint_instance_type: Instance type for the new endpoint. + endpoint_instance_count: Number of instances for the new endpoint. + endpoint_execution_role_arn: IAM role ARN for the new endpoint. + context_length: Context length for the new endpoint (as string). + max_concurrency: Max concurrency for the new endpoint (as string). + enable_rai: Enable RAI guardrails on the endpoint. Default: True. + cleanup_endpoint: Delete the endpoint after evaluation. Default: True. + endpoint_prefix: Prefix for auto-created endpoint names. Default: ``"inspectlens"``. + endpoint_environment: Extra environment variables for the inference + endpoint container (e.g. ``{"HF_TOKEN": "hf_xxx"}``). + extra_args: Additional CLI args forwarded to ``inspect eval``. + environment: Arbitrary environment variables passed to the SageMaker Training Job + container (e.g. ``{"HF_TOKEN": "hf_xxx", "HF_HUB_DOWNLOAD_TIMEOUT": "300"}``). + Any key-value pairs are accepted. + """ + + benchmarks_path: Optional[str] = None + tasks: List[Dict[str, Any]] = field(default_factory=list) + output_s3_path: Optional[str] = None + output_format: Optional[str] = None + bedrock_model_id: Optional[str] = None + endpoint_name: Optional[str] = None + model_s3_uri: Optional[str] = None + inference_image_uri: Optional[str] = None + endpoint_instance_type: Optional[str] = None + endpoint_instance_count: int = 1 + endpoint_execution_role_arn: Optional[str] = None + context_length: Optional[str] = None + max_concurrency: Optional[str] = None + enable_rai: bool = True + cleanup_endpoint: bool = True + endpoint_prefix: str = "inspectlens" + endpoint_environment: Optional[Dict[str, str]] = None + extra_args: Optional[List[str]] = None + environment: Optional[Dict[str, str]] = None + + def __post_init__(self) -> None: + new_endpoint_fields_set = any( + getattr(self, f) is not None + for f in ("model_s3_uri", "inference_image_uri", "endpoint_instance_type") + ) + if self.endpoint_name is not None and new_endpoint_fields_set: + raise ValueError( + "Cannot combine endpoint_name with model_s3_uri / " + "inference_image_uri / endpoint_instance_type. " + "Set endpoint_name=None to create a new endpoint." + ) + if self.endpoint_name is None and ( + bool(self.model_s3_uri) != bool(self.inference_image_uri) + ): + raise ValueError( + "Creating a new endpoint requires both model_s3_uri and " + "inference_image_uri. Provide both or neither." + ) + self._validate() + + def _validate(self) -> None: + """Validate InspectLensConfig fields and raise ValueError with actionable messages.""" + validate_inspect_lens_config( + benchmarks_path=self.benchmarks_path, + tasks=self.tasks, + output_s3_path=self.output_s3_path, + output_format=self.output_format, + inference_image_uri=self.inference_image_uri, + model_s3_uri=self.model_s3_uri, + endpoint_execution_role_arn=self.endpoint_execution_role_arn, + endpoint_instance_type=self.endpoint_instance_type, + endpoint_instance_count=self.endpoint_instance_count, + context_length=self.context_length, + max_concurrency=self.max_concurrency, + ) + + def _infer_scenario(self) -> str: + """Return the inference provider mode based on which fields are set.""" + if self.endpoint_name is not None: + return "existing_endpoint" + if self.model_s3_uri is not None: + return "create_endpoint" + return "bedrock" + + def to_yaml_dict( + self, + inference_provider: Dict[str, Any], + output_s3_path: str, + overrides: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Serialize to the dict structure expected by ``inspect_config.yaml``. + + Decoding and eval parameters are read from ``overrides`` with sensible + defaults, consistent with how other SDK eval tasks handle configuration. + """ + ov = overrides or {} + benchmarks_section: Dict[str, Any] = {"tasks": self.tasks} + if self.benchmarks_path is not None: + benchmarks_section["s3_path"] = self.benchmarks_path + + decoding: Dict[str, Any] = { + "temperature": ov.get("temperature", 0.0), + "top_p": ov.get("top_p", 1.0), + "top_k": ov.get("top_k", -1), + "max_tokens": ov.get("max_tokens", 8192), + } + + config: Dict[str, Any] = { + "inference_provider": inference_provider, + "benchmarks": benchmarks_section, + "eval": { + "max_connections": ov.get("max_connections", 16), + "max_retries": ov.get("max_retries", 100), + "timeout": ov.get("timeout", 600), + "decoding": decoding, + }, + "output": { + "s3_path": output_s3_path, + }, + } + if self.output_format: + config["output"]["output_format"] = self.output_format + if self.extra_args: + config["eval"]["extra_args"] = self.extra_args + return config diff --git a/src/amzn_nova_forge/manager/runtime_manager.py b/src/amzn_nova_forge/manager/runtime_manager.py index ac1b314..aedf8f9 100644 --- a/src/amzn_nova_forge/manager/runtime_manager.py +++ b/src/amzn_nova_forge/manager/runtime_manager.py @@ -21,6 +21,7 @@ import textwrap import time import zipfile +from urllib.parse import urlparse from dataclasses import dataclass, field from typing import Any, Dict, List, Optional @@ -710,6 +711,10 @@ def execute(self, job_config: JobConfig) -> str: "sagemaker_session": self.sagemaker_session, "training_image": job_config.image_uri, } + if job_config.environment: + # Pass arbitrary environment variables to the training container + # (used by InspectLens to inject HF_TOKEN, timeouts, etc.) + trainer_config["environment"] = job_config.environment # For eval job, the input could be none # https://docs.aws.amazon.com/sagemaker/latest/dg/nova-model-evaluation.html#nova-model-evaluation-notebook if job_config.data_s3_path: @@ -1372,6 +1377,10 @@ def _build_create_training_job_params( vpc_config = self._build_vpc_config() if vpc_config: create_params["VpcConfig"] = vpc_config + if not self._is_standard_ecr_image(self.image_uri): + create_params["AlgorithmSpecification"]["TrainingImageConfig"] = { + "TrainingRepositoryAccessMode": "Vpc", + } return create_params @@ -1384,6 +1393,22 @@ def _build_vpc_config(self) -> Dict[str, Any]: vpc_config["SecurityGroupIds"] = self.security_group_ids return vpc_config + @staticmethod + 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 + if not host: + host = parsed.path.split("/", 1)[0] + if not host: + return False + return bool( + re.fullmatch( + r"\d{12}\.dkr\.ecr\.[a-z0-9-]+\.amazonaws\.com(?:\.cn)?", + host, + ) + ) + def _raise_if_not_completed(self, job_name: str, status: str) -> None: """Raise ``RuntimeError`` with the FailureReason if the job didn't succeed.""" if status == "Completed": diff --git a/src/amzn_nova_forge/util/sagemaker.py b/src/amzn_nova_forge/util/sagemaker.py index c0b12a5..a5a4135 100644 --- a/src/amzn_nova_forge/util/sagemaker.py +++ b/src/amzn_nova_forge/util/sagemaker.py @@ -507,8 +507,19 @@ def find_sagemaker_model_by_tag(escrow_uri: str, sagemaker_client: BaseClient) - ResourceTypeFilters=["sagemaker:model"], ) results = response.get("ResourceTagMappingList", []) - if results: + if len(results) == 1: return results[0]["ResourceARN"] + if len(results) > 1: + # Multiple models share the same tag (e.g. from repeated test runs). + # Return the most recently created one to avoid stale dedup hits. + def _creation_time(arn: str) -> Any: + try: + model_name = arn.split("/")[-1] + return sagemaker_client.describe_model(ModelName=model_name)["CreationTime"] + except Exception: + return datetime.min.replace(tzinfo=timezone.utc) + + return max((r["ResourceARN"] for r in results), key=_creation_time) except ClientError as e: logger.warning( f"Could not search SageMaker models by tag (may lack tag:GetResources permission): {e}" diff --git a/src/amzn_nova_forge/validation/inspect_lens_validator.py b/src/amzn_nova_forge/validation/inspect_lens_validator.py new file mode 100644 index 0000000..59d7e63 --- /dev/null +++ b/src/amzn_nova_forge/validation/inspect_lens_validator.py @@ -0,0 +1,296 @@ +# 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. +"""Validation for InspectLensConfig.""" + +from __future__ import annotations + +import os +import re +from typing import Any, Dict, List, Optional, Set + +import boto3 + +from amzn_nova_forge.util.logging import logger +from amzn_nova_forge.util.s3_utils import parse_s3_uri + +_ECR_URI_REGEX = re.compile(r"^\d{12}\.dkr\.ecr\.[a-z0-9-]+\.amazonaws\.com/[^:]+:[^:]+$") +_IAM_ROLE_ARN_REGEX = re.compile(r"^arn:aws[^:]*:iam::\d{12}:role/.+$") +_TASK_DECORATOR_REGEX = re.compile(r"@task\b[^\n]*\n(?:@[^\n]*\n)*def\s+(\w+)") + + +def _extract_task_names_from_dir(local_dir: str) -> Set[str]: + """Return all @task-decorated function names found in .py files under local_dir.""" + task_names: Set[str] = set() + expanded = os.path.expanduser(local_dir) + if not os.path.isdir(expanded): + return task_names + for fname in os.listdir(expanded): + if not fname.endswith(".py"): + continue + fpath = os.path.join(expanded, fname) + try: + with open(fpath) as f: + source = f.read() + except OSError: + continue + for match in _TASK_DECORATOR_REGEX.finditer(source): + task_names.add(match.group(1)) + return task_names + + +def _extract_task_names_from_s3(s3_uri: str, region: Optional[str] = None) -> Set[str]: + """Return all @task-decorated function names found in .py files at an S3 prefix. + + Returns an empty set silently if S3 access fails (permissions, network) — + we don't block the user with a validation error they can't fix at config time. + """ + task_names: Set[str] = set() + bucket, prefix = parse_s3_uri(s3_uri) + + try: + s3 = boto3.client("s3", region_name=region) + paginator = s3.get_paginator("list_objects_v2") + for page in paginator.paginate(Bucket=bucket, Prefix=prefix): + for obj in page.get("Contents", []): + key = obj["Key"] + if not key.endswith(".py"): + continue + try: + response = s3.get_object(Bucket=bucket, Key=key) + source = response["Body"].read().decode("utf-8") + for match in _TASK_DECORATOR_REGEX.finditer(source): + task_names.add(match.group(1)) + except Exception as e: + logger.debug("Skipping S3 key '%s' during task extraction: %s", key, e) + continue + except Exception as e: + logger.debug("Could not list S3 prefix '%s' for task validation: %s", s3_uri, e) + return task_names + + +def _validate_benchmarks_path(benchmarks_path: Optional[str], errors: List[str]) -> None: + """Validate benchmarks_path value and append errors to the list.""" + if not benchmarks_path: + errors.append( + "benchmarks_path is required. Provide an s3:// URI containing benchmark .py files " + "with @task decorators. Use evaluator.upload_benchmarks(local_dir, s3_path) to " + "upload a local directory first.\n" + " Example: benchmarks_path='s3://your-bucket/benchmarks/my_benchmarks/'" + ) + elif not benchmarks_path.strip(): + errors.append("benchmarks_path must not be an empty string.") + elif not benchmarks_path.strip().startswith("s3://"): + errors.append( + f"benchmarks_path must be an s3:// URI, got: '{benchmarks_path}'. " + f"Use evaluator.upload_benchmarks(local_dir, s3_path) to upload a local " + f"directory first, then pass the returned S3 URI as benchmarks_path." + ) + + +_VALID_OUTPUT_FORMATS = ("eval", "csv", "jsonl", "json") + + +def _validate_tasks( + tasks: Any, + errors: List[str], +) -> None: + """Validate tasks structure (format only — no S3 calls). + + S3-based task name matching is deferred to job submission time via + validate_task_names_against_benchmarks() to avoid network calls in __post_init__. + """ + if not isinstance(tasks, list): + errors.append("tasks must be a list of dicts, e.g. [{'name': 'boolq_pt', 'limit': 100}].") + return + + for i, t in enumerate(tasks): + if not isinstance(t, dict): + errors.append(f"tasks[{i}] must be a dict with a 'name' key, got: {t!r}") + continue + if "name" not in t: + errors.append( + f"tasks[{i}] is missing required key 'name'. " + f"Each task must have at least {{'name': ''}}." + ) + if "path" in t and t["path"] is not None: + path_val = t["path"] + if not isinstance(path_val, str) or not path_val.endswith(".py"): + errors.append( + f"tasks[{i}].path must be a relative path ending in .py, got: {path_val!r}" + ) + if "limit" in t and t["limit"] is not None: + limit_val = t["limit"] + if not isinstance(limit_val, int) or limit_val < 1: + errors.append(f"tasks[{i}].limit must be an int >= 1, got: {limit_val!r}") + if "epochs" in t and t["epochs"] is not None: + epochs_val = t["epochs"] + if not isinstance(epochs_val, int) or epochs_val < 1: + errors.append(f"tasks[{i}].epochs must be an int >= 1, got: {epochs_val!r}") + if "task_args" in t and t["task_args"] is not None: + task_args_val = t["task_args"] + if not isinstance(task_args_val, dict): + errors.append( + f"tasks[{i}].task_args must be a dict, got: {type(task_args_val).__name__}" + ) + + +def validate_task_names_against_benchmarks( + tasks: Any, + benchmarks_path: Optional[str], + region: Optional[str] = None, +) -> List[str]: + """Validate task names against @task-decorated functions in the benchmarks S3 path. + + This is intentionally separate from validate_inspect_lens_config() to avoid + making S3 API calls during InspectLensConfig construction (__post_init__). + Call this at job submission time when region and credentials are available. + + Returns a list of error strings (empty if all tasks match). + """ + errors: List[str] = [] + + if ( + not benchmarks_path + or not tasks + or not isinstance(tasks, list) + or any(not isinstance(t, dict) or "name" not in t for t in tasks) + ): + return errors + + path = benchmarks_path.strip() + if not path.startswith("s3://"): + return errors + + known_tasks = _extract_task_names_from_s3(path, region=region) + if not known_tasks: + return errors + + for t in tasks: + if not isinstance(t, dict) or "name" not in t: + continue + task_name = t["name"] + if not any(task_name in k for k in known_tasks): + errors.append( + f"Task '{task_name}' was not found in any @task-decorated function " + f"in '{path}'.\n" + f" Available @task functions: {sorted(known_tasks)}\n" + f" The container matches tasks by substring — make sure the task " + f"name matches (or is a substring of) a @task function name in your " + f"benchmark .py files." + ) + return errors + + +def _validate_s3_and_arn_fields( + output_s3_path: Optional[str], + inference_image_uri: Optional[str], + model_s3_uri: Optional[str], + endpoint_execution_role_arn: Optional[str], + errors: List[str], +) -> None: + """Validate S3 URIs and ARN fields, appending errors to the list.""" + if output_s3_path is not None and not output_s3_path.startswith("s3://"): + errors.append(f"output_s3_path must be an s3:// URI, got: '{output_s3_path}'") + + if inference_image_uri is not None and not _ECR_URI_REGEX.match(inference_image_uri): + errors.append( + f"inference_image_uri must be a valid ECR URI in the format " + f".dkr.ecr..amazonaws.com/:. " + f"Got: '{inference_image_uri}'" + ) + + if model_s3_uri is not None and not model_s3_uri.startswith("s3://"): + errors.append(f"model_s3_uri must be an s3:// URI, got: '{model_s3_uri}'") + + if endpoint_execution_role_arn is not None and not _IAM_ROLE_ARN_REGEX.match( + endpoint_execution_role_arn + ): + errors.append( + f"endpoint_execution_role_arn must be a valid IAM role ARN " + f"(arn:aws:iam:::role/). " + f"Got: '{endpoint_execution_role_arn}'" + ) + + +def _validate_output_format(output_format: Optional[str], errors: List[str]) -> None: + """Validate output_format value if provided.""" + if output_format is not None and output_format not in _VALID_OUTPUT_FORMATS: + errors.append( + f"output_format must be one of {_VALID_OUTPUT_FORMATS}, got: '{output_format}'" + ) + + +def _validate_endpoint_fields( + endpoint_instance_type: Optional[str], + endpoint_instance_count: int, + context_length: Optional[str], + max_concurrency: Optional[str], + errors: List[str], +) -> None: + """Validate endpoint-related numeric and string fields.""" + if endpoint_instance_type is not None and not endpoint_instance_type.startswith("ml."): + errors.append( + f"endpoint_instance_type must start with 'ml.', got: '{endpoint_instance_type}'" + ) + if endpoint_instance_count < 1: + errors.append(f"endpoint_instance_count must be >= 1, got: {endpoint_instance_count}") + if context_length is not None: + if not context_length.isdigit() or int(context_length) <= 0: + errors.append( + f"context_length must be a positive integer as string, got: '{context_length}'" + ) + if max_concurrency is not None: + if not max_concurrency.isdigit() or int(max_concurrency) <= 0: + errors.append( + f"max_concurrency must be a positive integer as string, got: '{max_concurrency}'" + ) + + +def validate_inspect_lens_config( + benchmarks_path: Optional[str], + tasks: Any, + output_s3_path: Optional[str], + output_format: Optional[str], + inference_image_uri: Optional[str], + model_s3_uri: Optional[str], + endpoint_execution_role_arn: Optional[str], + endpoint_instance_type: Optional[str] = None, + endpoint_instance_count: int = 1, + context_length: Optional[str] = None, + max_concurrency: Optional[str] = None, +) -> None: + """Validate InspectLensConfig fields and raise ValueError with actionable messages. + + Takes individual field values rather than the config object to avoid + circular imports between the config and validator modules. + + Raises: + ValueError: If any required fields are missing or malformed. + """ + errors: List[str] = [] + _validate_benchmarks_path(benchmarks_path, errors) + _validate_tasks(tasks, errors) + _validate_output_format(output_format, errors) + _validate_s3_and_arn_fields( + output_s3_path, inference_image_uri, model_s3_uri, endpoint_execution_role_arn, errors + ) + _validate_endpoint_fields( + endpoint_instance_type, endpoint_instance_count, context_length, max_concurrency, errors + ) + + if errors: + bullet_list = "\n".join(f" - {e}" for e in errors) + raise ValueError( + f"InspectLensConfig validation failed with {len(errors)} error(s):\n{bullet_list}" + ) diff --git a/tests/unit/evaluator/test_forge_evaluator.py b/tests/unit/evaluator/test_forge_evaluator.py index 4418fe0..e7cef14 100644 --- a/tests/unit/evaluator/test_forge_evaluator.py +++ b/tests/unit/evaluator/test_forge_evaluator.py @@ -763,5 +763,516 @@ def test_get_logs_smhp_includes_cluster_namespace(self, mock_monitor_cls): ) +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" + mock_infra.instance_count = 1 + mock_infra.platform = Platform.SMTJ + mock_infra.execution_role = "arn:aws:iam::123:role/MyRole" + mock_infra.max_job_runtime = 7200 + + config = ForgeConfig( + output_s3_path="s3://bucket/output/", + image_uri=image_uri + or "123456789012.dkr.ecr.us-east-1.amazonaws.com/inspect-lens:beta-latest", + ) + + with ( + patch( + "amzn_nova_forge.evaluator.forge_evaluator.set_output_s3_path", + return_value="s3://bucket/output/", + ), + patch("boto3.session.Session") as mock_session, + patch("boto3.client"), + ): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + evaluator = ForgeEvaluator( + model=Model.NOVA_MICRO, + infra=mock_infra, + config=config, + ) + 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/", + tasks=[{"name": "boolq_pt", "limit": 50}], + output_s3_path="s3://bucket/results/", + ) + with ( + patch("boto3.client"), + patch("amzn_nova_forge.evaluator.forge_evaluator.yaml") as mock_yaml, + ): + mock_yaml.dump.return_value = "inference_provider: {}\n" + result = evaluator.evaluate( + job_name="test-job", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=config, + dry_run=True, + ) + self.assertIsNone(result) + + def test_missing_inspect_lens_config_raises(self): + evaluator = self._make_evaluator() + with self.assertRaises(ValueError) as ctx: + evaluator.evaluate( + job_name="test-job", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=None, + ) + 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/", + output_s3_path="s3://bucket/results/", + ) + with ( + patch("boto3.client"), + patch("amzn_nova_forge.evaluator.forge_evaluator.yaml") as mock_yaml, + patch("amzn_nova_forge.evaluator.forge_evaluator.logger") as mock_logger, + ): + mock_yaml.dump.return_value = "inference_provider: {}\n" + evaluator.evaluate( + job_name="test-job", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=config, + overrides={"benchmarks_path": "s3://other/"}, + dry_run=True, + ) + # Should warn about unknown key + warning_calls = [str(c) for c in mock_logger.warning.call_args_list] + 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/", + output_s3_path="s3://bucket/results/", + ) + with ( + patch("boto3.client"), + patch("amzn_nova_forge.evaluator.forge_evaluator.yaml") as mock_yaml, + patch("amzn_nova_forge.evaluator.forge_evaluator.logger") as mock_logger, + ): + mock_yaml.dump.return_value = "inference_provider: {}\n" + evaluator.evaluate( + job_name="test-job", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=config, + overrides={"temperature": 0.0, "max_tokens": 512, "max_connections": 4}, + dry_run=True, + ) + # No warning for valid decoding keys + warning_calls = [str(c) for c in mock_logger.warning.call_args_list] + 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/", + output_s3_path="s3://bucket/results/", + ) + provider = evaluator._build_inspect_lens_inference_provider(config, model_path=None) + self.assertIn("bedrock", provider) + 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/", + endpoint_name="my-endpoint", + ) + provider = evaluator._build_inspect_lens_inference_provider(config, model_path=None) + self.assertIn("sagemaker_endpoint", provider) + 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/", + output_s3_path="s3://bucket/results/", + ) + provider = evaluator._build_inspect_lens_inference_provider( + config, model_path="us.amazon.nova-lite-v1:0" + ) + self.assertIn("bedrock", provider) + self.assertEqual(provider["bedrock"]["model_id"], "us.amazon.nova-lite-v1:0") + + 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/", + tasks=[{"name": "boolq_pt", "limit": 50}], + output_s3_path="s3://bucket/results/", + ) + mock_cached = MagicMock() + + with patch( + "amzn_nova_forge.evaluator.forge_evaluator.load_existing_result", + return_value=mock_cached, + ) as mock_load: + result = evaluator.evaluate( + job_name="test-job", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=config, + ) + + self.assertIs(result, mock_cached) + mock_load.assert_called_once_with( + evaluator._cache_context, + job_name="test-job", + job_type="inspect_lens", + model_path=None, + benchmarks_path="s3://bucket/benchmarks/boolq/", + tasks=str([{"name": "boolq_pt", "limit": 50}]), + inference_scenario="bedrock", + endpoint_name=None, + bedrock_model_id=None, + overrides={}, + ) + + 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" + mock_infra.instance_count = 1 + mock_infra.platform = Platform.SMTJ + mock_infra.execution_role = "arn:aws:iam::123:role/MyRole" + mock_infra.max_job_runtime = 7200 + + # Patch validate_mlflow_overrides to avoid real AWS calls during MLflowMonitor init + with patch( + "amzn_nova_forge.monitor.mlflow_monitor.validate_mlflow_overrides", + return_value=[], + ): + mlflow_monitor = MLflowMonitor( + tracking_uri="arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", + experiment_name="nova-evals", + ) + + forge_config = ForgeConfig( + output_s3_path="s3://bucket/output/", + image_uri="123456789012.dkr.ecr.us-east-1.amazonaws.com/inspect-lens:beta-latest", + mlflow_monitor=mlflow_monitor, + ) + + with ( + patch( + "amzn_nova_forge.evaluator.forge_evaluator.set_output_s3_path", + return_value="s3://bucket/output/", + ), + patch("boto3.session.Session") as mock_session, + patch("boto3.client"), + ): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + evaluator = ForgeEvaluator( + model=Model.NOVA_MICRO, + infra=mock_infra, + config=forge_config, + ) + + inspect_config = InspectLensConfig( + benchmarks_path="s3://bucket/benchmarks/", + output_s3_path="s3://bucket/results/", + ) + + captured_dict = {} + + def capture_yaml_dump(d, **kwargs): + captured_dict.update(d) + return "mocked_yaml\n" + + with ( + patch("boto3.client"), + patch("amzn_nova_forge.evaluator.forge_evaluator.yaml") as mock_yaml, + ): + mock_yaml.dump.side_effect = capture_yaml_dump + evaluator.evaluate( + job_name="test-mlflow-job", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=inspect_config, + dry_run=True, + ) + + self.assertIn("tracking", captured_dict) + tracking = captured_dict["tracking"] + self.assertEqual( + tracking["mlflow_tracking_arn"], + "arn:aws:sagemaker:us-east-1:123456789012:mlflow-app/app-xxx", + ) + self.assertEqual(tracking["mlflow_experiment_name"], "nova-evals") + self.assertTrue(tracking["mlflow_tracing"]) + self.assertTrue(tracking["mlflow_log_artifacts"]) + + +class TestInspectLensS3Paths(unittest.TestCase): + """Tests for run_id-based S3 path layout in _evaluate_inspect_lens.""" + + 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" + mock_infra.instance_count = 1 + mock_infra.platform = Platform.SMTJ + mock_infra.execution_role = "arn:aws:iam::123:role/MyRole" + mock_infra.max_job_runtime = 7200 + + config = ForgeConfig( + output_s3_path=output_s3_path, + image_uri="123456789012.dkr.ecr.us-east-1.amazonaws.com/inspect-lens:beta-latest", + ) + + with ( + patch( + "amzn_nova_forge.evaluator.forge_evaluator.set_output_s3_path", + return_value=output_s3_path, + ), + patch("boto3.session.Session") as mock_session, + patch("boto3.client"), + ): + type(mock_session.return_value).region_name = PropertyMock(return_value="us-east-1") + evaluator = ForgeEvaluator( + model=Model.NOVA_MICRO, + infra=mock_infra, + config=config, + ) + return evaluator + + def _run_evaluate(self, evaluator, inspect_lens_config, mock_sm_client): + """Helper: run evaluate() with all external calls mocked.""" + mock_sm_client.create_training_job.return_value = {} + with ( + patch("amzn_nova_forge.evaluator.forge_evaluator.uuid") as mock_uuid, + patch("amzn_nova_forge.evaluator.forge_evaluator.yaml") as mock_yaml, + patch("amzn_nova_forge.evaluator.forge_evaluator.boto3") as mock_boto3, + ): + mock_uuid.uuid4.return_value = self.FIXED_RUN_ID + mock_yaml.dump.return_value = "inference_provider: {}\n" + mock_boto3.client.return_value = mock_sm_client + result = evaluator.evaluate( + job_name="test-job", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=inspect_lens_config, + ) + return result, 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: + with open(os.path.join(tmpdir, "boolq_pt.py"), "w") as f: + f.write("@task\ndef boolq_pt():\n pass\n") + + # Bypass InspectLensConfig validation to test the evaluator-level check + cfg = object.__new__(InspectLensConfig) + for field, default in InspectLensConfig.__dataclass_fields__.items(): + try: + setattr(cfg, field, default.default) + except Exception: + setattr(cfg, field, default.default_factory()) + cfg.benchmarks_path = tmpdir + cfg.tasks = [{"name": "boolq_pt"}] + + with ( + patch("amzn_nova_forge.evaluator.forge_evaluator.boto3") as mock_boto3, + ): + mock_boto3.client.return_value = MagicMock() + with self.assertRaises(ValueError) as ctx: + evaluator.evaluate( + job_name="test-job", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=cfg, + ) + self.assertIn("upload_benchmarks", str(ctx.exception)) + + 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: + with open(os.path.join(tmpdir, "boolq_pt.py"), "w") as f: + f.write("@task\ndef boolq_pt():\n pass\n") + with open(os.path.join(tmpdir, "README.md"), "w") as f: + f.write("# not uploaded\n") + + with patch("amzn_nova_forge.evaluator.forge_evaluator.boto3") as mock_boto3: + mock_s3 = MagicMock() + mock_boto3.client.return_value = mock_s3 + + result = evaluator.upload_benchmarks(tmpdir, "s3://bucket/my-benchmarks/") + + self.assertEqual(result, "s3://bucket/my-benchmarks/") + mock_s3.upload_file.assert_called_once() + call_args = mock_s3.upload_file.call_args + self.assertIn("boolq_pt.py", call_args[0][0]) + self.assertEqual(call_args[0][1], "bucket") + self.assertEqual(call_args[0][2], "my-benchmarks/boolq_pt.py") + + def test_upload_benchmarks_invalid_local_dir_raises(self): + """upload_benchmarks() raises ValueError for non-existent directory.""" + evaluator = self._make_evaluator() + with self.assertRaises(ValueError) as ctx: + evaluator.upload_benchmarks("/nonexistent/path", "s3://bucket/benchmarks/") + self.assertIn("not an existing directory", str(ctx.exception)) + + 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: + evaluator.upload_benchmarks(tmpdir, "/local/path/") + self.assertIn("s3://", str(ctx.exception)) + + 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( + benchmarks_path="s3://separate-benchmarks-bucket/my-project/benchmarks/my_benchmarks/", + tasks=[{"name": "boolq_pt"}], + ) + + with ( + patch("amzn_nova_forge.evaluator.forge_evaluator.uuid") as mock_uuid, + patch("amzn_nova_forge.evaluator.forge_evaluator.yaml") as mock_yaml, + patch("amzn_nova_forge.evaluator.forge_evaluator.boto3") as mock_boto3, + ): + mock_uuid.uuid4.return_value = self.FIXED_RUN_ID + mock_yaml.dump.return_value = "inference_provider: {}\n" + mock_s3 = MagicMock() + mock_boto3.client.return_value = mock_s3 + mock_s3.create_training_job = MagicMock(return_value={}) + + evaluator.evaluate( + job_name="test-job", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=cfg, + ) + + # Config always goes under output_s3_path//config/ — not the separate benchmarks bucket + put_calls = mock_s3.put_object.call_args_list + config_call = next(c for c in put_calls if "inspect_config.yaml" in str(c)) + kwargs = config_call[1] + self.assertEqual(kwargs["Bucket"], "bucket") + self.assertEqual(kwargs["Key"], f"output/{self.FIXED_RUN_ID}/config/inspect_config.yaml") + + # SageMaker output also under output_s3_path//output/ + create_call = mock_s3.create_training_job.call_args + output_path = create_call[1]["OutputDataConfig"]["S3OutputPath"] + self.assertEqual(output_path, f"s3://bucket/output/{self.FIXED_RUN_ID}/output/") + + 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( + benchmarks_path="s3://bucket/benchmarks/", + tasks=[{"name": "boolq_pt"}], + ) + + with ( + patch("amzn_nova_forge.evaluator.forge_evaluator.uuid") as mock_uuid, + patch("amzn_nova_forge.evaluator.forge_evaluator.yaml") as mock_yaml, + patch("amzn_nova_forge.evaluator.forge_evaluator.boto3") as mock_boto3, + ): + mock_uuid.uuid4.return_value = self.FIXED_RUN_ID + mock_yaml.dump.return_value = "inference_provider: {}\n" + mock_s3 = MagicMock() + mock_boto3.client.return_value = mock_s3 + mock_s3.create_training_job = MagicMock(return_value={}) + + result = evaluator.evaluate( + job_name="my-eval", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=cfg, + ) + + self.assertIn(self.FIXED_RUN_ID, result.job_id) + self.assertTrue(result.job_id.startswith("my-eval-")) + + 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( + benchmarks_path="s3://separate-benchmarks-bucket/benchmarks/", + tasks=[{"name": "boolq_pt"}], + ) + + with ( + patch("amzn_nova_forge.evaluator.forge_evaluator.uuid") as mock_uuid, + patch("amzn_nova_forge.evaluator.forge_evaluator.yaml") as mock_yaml, + patch("amzn_nova_forge.evaluator.forge_evaluator.boto3") as mock_boto3, + ): + mock_uuid.uuid4.return_value = self.FIXED_RUN_ID + mock_yaml.dump.return_value = "inference_provider: {}\n" + mock_s3 = MagicMock() + mock_boto3.client.return_value = mock_s3 + mock_s3.create_training_job = MagicMock(return_value={}) + + evaluator.evaluate( + job_name="test-job", + eval_task=EvaluationTask.INSPECT_LENS, + inspect_lens_config=cfg, + ) + + # parent of s3://separate-benchmarks-bucket/benchmarks/ → s3://separate-benchmarks-bucket/ + create_call = mock_s3.create_training_job.call_args + output_path = create_call[1]["OutputDataConfig"]["S3OutputPath"] + # Output always goes under output_s3_path, not the benchmarks bucket + self.assertEqual(output_path, f"s3://bucket/output/{self.FIXED_RUN_ID}/output/") + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/evaluator/test_inspect_lens.py b/tests/unit/evaluator/test_inspect_lens.py new file mode 100644 index 0000000..af90276 --- /dev/null +++ b/tests/unit/evaluator/test_inspect_lens.py @@ -0,0 +1,592 @@ +# Copyright Amazon.com, Inc. or its affiliates + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for InspectLensConfig and inspect_lens_validator.""" + +import os +import tempfile +import unittest +from io import BytesIO +from unittest.mock import MagicMock, patch + +from amzn_nova_forge.core.enums import Model +from amzn_nova_forge.evaluator.inspect_lens_config import InspectLensConfig +from amzn_nova_forge.validation.inspect_lens_validator import ( + _extract_task_names_from_dir, + _extract_task_names_from_s3, +) + +VALID_BENCHMARKS_PATH = "s3://bucket/benchmarks/" +VALID_TASKS = [{"name": "boolq_pt", "limit": 100}] +VALID_OUTPUT = "s3://bucket/output/" +VALID_ECR = "123456789012.dkr.ecr.us-east-1.amazonaws.com/my-repo:latest" +VALID_ROLE = "arn:aws:iam::123456789012:role/MyRole" + + +def _make_valid_config(**kwargs): + """Return a minimal valid InspectLensConfig, bypassing __post_init__ validation.""" + defaults = dict( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=list(VALID_TASKS), + ) + defaults.update(kwargs) + # Bypass __post_init__ to test validate_inspect_lens_config directly + obj = object.__new__(InspectLensConfig) + for field, default in InspectLensConfig.__dataclass_fields__.items(): + setattr( + obj, + field, + defaults.get( + field, + default.default + if default.default is not default.default_factory + else default.default_factory(), + ), + ) + for k, v in defaults.items(): + setattr(obj, k, v) + return obj + + +class TestInspectLensConfigPostInit(unittest.TestCase): + """Tests for field-level mutual exclusion checks in __post_init__.""" + + def test_valid_bedrock_config(self): + """Minimal valid config with S3 benchmarks_path should not raise.""" + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + ) + self.assertEqual(cfg.benchmarks_path, VALID_BENCHMARKS_PATH) + + def test_endpoint_name_with_model_s3_uri_raises(self): + with self.assertRaises(ValueError) as ctx: + InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + endpoint_name="my-endpoint", + model_s3_uri="s3://bucket/model/", + inference_image_uri=VALID_ECR, + ) + self.assertIn("endpoint_name", str(ctx.exception)) + + def test_model_s3_uri_without_inference_image_raises(self): + with self.assertRaises(ValueError) as ctx: + InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + model_s3_uri="s3://bucket/model/", + ) + self.assertIn("inference_image_uri", str(ctx.exception)) + + def test_inference_image_without_model_s3_uri_raises(self): + with self.assertRaises(ValueError) as ctx: + InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + inference_image_uri=VALID_ECR, + ) + self.assertIn("model_s3_uri", str(ctx.exception)) + + def test_both_new_endpoint_fields_valid(self): + """Providing both model_s3_uri and inference_image_uri should not raise.""" + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + model_s3_uri="s3://bucket/model/", + inference_image_uri=VALID_ECR, + ) + self.assertEqual(cfg.model_s3_uri, "s3://bucket/model/") + + +class TestValidateInspectLensConfig(unittest.TestCase): + """Tests for validate_inspect_lens_config().""" + + def _assert_error(self, config, fragment: str): + with self.assertRaises(ValueError) as ctx: + config._validate() + self.assertIn(fragment, str(ctx.exception)) + + def test_missing_benchmarks_path_raises(self): + cfg = _make_valid_config(benchmarks_path=None) + self._assert_error(cfg, "benchmarks_path is required") + + def test_local_benchmarks_path_raises(self): + cfg = _make_valid_config(benchmarks_path="./my_benchmarks/") + self._assert_error(cfg, "must be an s3:// URI") + + def test_empty_benchmarks_path_raises(self): + cfg = _make_valid_config(benchmarks_path=" ") + self._assert_error(cfg, "must not be an empty string") + + def test_valid_s3_benchmarks_path_passes(self): + cfg = _make_valid_config(benchmarks_path="s3://bucket/benchmarks/") + cfg._validate() # should not raise + + def test_tasks_not_a_list_raises(self): + cfg = _make_valid_config(tasks="boolq_pt") + self._assert_error(cfg, "tasks must be a list") + + def test_task_missing_name_raises(self): + cfg = _make_valid_config(tasks=[{"limit": 10}]) + self._assert_error(cfg, "missing required key 'name'") + + def test_task_not_a_dict_raises(self): + cfg = _make_valid_config(tasks=["boolq_pt"]) + self._assert_error(cfg, "must be a dict") + + def test_valid_tasks_pass(self): + cfg = _make_valid_config(tasks=[{"name": "boolq_pt"}, {"name": "mmlu_pro_pt", "limit": 50}]) + cfg._validate() + + def test_empty_tasks_list_passes(self): + cfg = _make_valid_config(tasks=[]) + cfg._validate() + + def test_invalid_output_s3_path_raises(self): + cfg = _make_valid_config(output_s3_path="not-an-s3-path") + self._assert_error(cfg, "output_s3_path must be an s3:// URI") + + def test_valid_output_s3_path_passes(self): + cfg = _make_valid_config(output_s3_path=VALID_OUTPUT) + cfg._validate() + + def test_none_output_s3_path_passes(self): + cfg = _make_valid_config(output_s3_path=None) + cfg._validate() + + def test_invalid_ecr_uri_raises(self): + cfg = _make_valid_config(inference_image_uri="not-an-ecr-uri") + self._assert_error(cfg, "inference_image_uri must be a valid ECR URI") + + def test_valid_ecr_uri_passes(self): + cfg = _make_valid_config(inference_image_uri=VALID_ECR) + cfg._validate() + + def test_invalid_model_s3_uri_raises(self): + cfg = _make_valid_config(model_s3_uri="not-s3") + self._assert_error(cfg, "model_s3_uri must be an s3:// URI") + + def test_valid_model_s3_uri_passes(self): + cfg = _make_valid_config(model_s3_uri="s3://bucket/model/") + cfg._validate() + + def test_invalid_role_arn_raises(self): + cfg = _make_valid_config(endpoint_execution_role_arn="not-an-arn") + self._assert_error(cfg, "endpoint_execution_role_arn must be a valid IAM role ARN") + + def test_valid_role_arn_passes(self): + cfg = _make_valid_config(endpoint_execution_role_arn=VALID_ROLE) + cfg._validate() + + def test_mlflow_fields_not_on_config(self): + """mlflow and hf_token fields should not exist on InspectLensConfig.""" + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + ) + self.assertFalse(hasattr(cfg, "mlflow_tracking_arn")) + self.assertFalse(hasattr(cfg, "mlflow_experiment_name")) + self.assertFalse(hasattr(cfg, "mlflow_tracing")) + self.assertFalse(hasattr(cfg, "mlflow_log_artifacts")) + self.assertFalse(hasattr(cfg, "hf_token")) + + def test_invalid_output_format_raises(self): + cfg = _make_valid_config(output_format="xml") + self._assert_error(cfg, "output_format must be one of") + + def test_valid_output_formats_pass(self): + for fmt in ("eval", "csv", "jsonl", "json"): + cfg = _make_valid_config(output_format=fmt) + cfg._validate() + + def test_none_output_format_passes(self): + cfg = _make_valid_config(output_format=None) + cfg._validate() + + def test_invalid_endpoint_instance_type_raises(self): + cfg = _make_valid_config(endpoint_instance_type="p5.48xlarge") + self._assert_error(cfg, "must start with 'ml.'") + + def test_valid_endpoint_instance_type_passes(self): + cfg = _make_valid_config(endpoint_instance_type="ml.p5.48xlarge") + cfg._validate() + + def test_invalid_endpoint_instance_count_raises(self): + cfg = _make_valid_config(endpoint_instance_count=0) + self._assert_error(cfg, "endpoint_instance_count must be >= 1") + + def test_invalid_context_length_raises(self): + cfg = _make_valid_config(context_length="abc") + self._assert_error(cfg, "context_length must be a positive integer") + + def test_zero_context_length_raises(self): + cfg = _make_valid_config(context_length="0") + self._assert_error(cfg, "context_length must be a positive integer") + + def test_valid_context_length_passes(self): + cfg = _make_valid_config(context_length="12000") + cfg._validate() + + def test_invalid_max_concurrency_raises(self): + cfg = _make_valid_config(max_concurrency="foo") + self._assert_error(cfg, "max_concurrency must be a positive integer") + + def test_valid_max_concurrency_passes(self): + cfg = _make_valid_config(max_concurrency="16") + cfg._validate() + + def test_task_with_invalid_path_raises(self): + cfg = _make_valid_config(tasks=[{"name": "mmlu", "path": "mmlu/mmlu_sft.txt"}]) + self._assert_error(cfg, "path must be a relative path ending in .py") + + def test_task_with_valid_path_passes(self): + cfg = _make_valid_config(tasks=[{"name": "mmlu", "path": "mmlu/mmlu_sft.py"}]) + cfg._validate() + + def test_task_with_invalid_epochs_raises(self): + cfg = _make_valid_config(tasks=[{"name": "mmlu", "epochs": 0}]) + self._assert_error(cfg, "epochs must be an int >= 1") + + def test_task_with_valid_epochs_passes(self): + cfg = _make_valid_config(tasks=[{"name": "mmlu", "epochs": 3}]) + cfg._validate() + + def test_task_with_invalid_task_args_raises(self): + cfg = _make_valid_config(tasks=[{"name": "mmlu", "task_args": "invalid"}]) + self._assert_error(cfg, "task_args must be a dict") + + def test_task_with_valid_task_args_passes(self): + cfg = _make_valid_config(tasks=[{"name": "mmlu", "task_args": {"subject": "anatomy"}}]) + cfg._validate() + + def test_multiple_errors_raised_together(self): + cfg = _make_valid_config( + benchmarks_path=None, + output_s3_path="bad-path", + inference_image_uri="not-ecr", + ) + with self.assertRaises(ValueError) as ctx: + cfg._validate() + msg = str(ctx.exception) + self.assertIn("3 error(s)", msg) + self.assertIn("benchmarks_path is required", msg) + self.assertIn("output_s3_path", msg) + self.assertIn("inference_image_uri", msg) + + +class TestBenchmarksPathMustBeS3(unittest.TestCase): + """Tests that benchmarks_path must be an S3 URI.""" + + def test_local_path_raises(self): + cfg = _make_valid_config(benchmarks_path="/local/path/benchmarks") + with self.assertRaises(ValueError) as ctx: + cfg._validate() + self.assertIn("must be an s3:// URI", str(ctx.exception)) + self.assertIn("upload_benchmarks", str(ctx.exception)) + + def test_relative_path_raises(self): + cfg = _make_valid_config(benchmarks_path="./my_benchmarks/") + with self.assertRaises(ValueError) as ctx: + cfg._validate() + self.assertIn("must be an s3:// URI", str(ctx.exception)) + + def test_s3_path_passes(self): + cfg = _make_valid_config(benchmarks_path="s3://bucket/benchmarks/") + cfg._validate() # should not raise + + +class TestExtractTaskNamesFromDir(unittest.TestCase): + def test_extracts_task_decorated_functions(self): + with tempfile.TemporaryDirectory() as tmpdir: + source = "@task\ndef boolq_pt():\n pass\n\n@task\ndef mmlu_pro_pt():\n pass\n" + with open(os.path.join(tmpdir, "tasks.py"), "w") as f: + f.write(source) + result = _extract_task_names_from_dir(tmpdir) + self.assertEqual(result, {"boolq_pt", "mmlu_pro_pt"}) + + def test_ignores_non_task_functions(self): + with tempfile.TemporaryDirectory() as tmpdir: + source = "def helper():\n pass\n\n@task\ndef real_task():\n pass\n" + with open(os.path.join(tmpdir, "tasks.py"), "w") as f: + f.write(source) + result = _extract_task_names_from_dir(tmpdir) + self.assertEqual(result, {"real_task"}) + + def test_ignores_non_py_files(self): + with tempfile.TemporaryDirectory() as tmpdir: + with open(os.path.join(tmpdir, "README.md"), "w") as f: + f.write("@task\ndef fake_task():\n pass\n") + result = _extract_task_names_from_dir(tmpdir) + self.assertEqual(result, set()) + + def test_nonexistent_dir_returns_empty(self): + result = _extract_task_names_from_dir("/nonexistent/path") + self.assertEqual(result, set()) + + def test_multiple_files(self): + with tempfile.TemporaryDirectory() as tmpdir: + for name in ("boolq_pt", "mmlu_pro_pt", "arc_c_pt"): + with open(os.path.join(tmpdir, f"{name}.py"), "w") as f: + f.write(f"@task\ndef {name}():\n pass\n") + result = _extract_task_names_from_dir(tmpdir) + self.assertEqual(result, {"boolq_pt", "mmlu_pro_pt", "arc_c_pt"}) + + +class TestExtractTaskNamesFromS3(unittest.TestCase): + @patch("boto3.client") + def test_extracts_task_names_from_s3(self, mock_client): + source = b"@task\ndef boolq_pt():\n pass\n\n@task\ndef mmlu_pro_pt():\n pass\n" + mock_s3 = MagicMock() + mock_client.return_value = mock_s3 + mock_paginator = MagicMock() + mock_s3.get_paginator.return_value = mock_paginator + mock_paginator.paginate.return_value = [{"Contents": [{"Key": "benchmarks/boolq_pt.py"}]}] + mock_s3.get_object.return_value = {"Body": BytesIO(source)} + + result = _extract_task_names_from_s3("s3://bucket/benchmarks/") + self.assertEqual(result, {"boolq_pt", "mmlu_pro_pt"}) + + @patch("boto3.client") + def test_s3_access_failure_returns_empty(self, mock_client): + mock_client.side_effect = Exception("no credentials") + result = _extract_task_names_from_s3("s3://bucket/benchmarks/") + self.assertEqual(result, set()) + + @patch("boto3.client") + def test_skips_non_py_keys(self, mock_client): + mock_s3 = MagicMock() + mock_client.return_value = mock_s3 + mock_paginator = MagicMock() + mock_s3.get_paginator.return_value = mock_paginator + mock_paginator.paginate.return_value = [ + {"Contents": [{"Key": "benchmarks/README.md"}, {"Key": "benchmarks/config.yaml"}]} + ] + result = _extract_task_names_from_s3("s3://bucket/benchmarks/") + self.assertEqual(result, set()) + mock_s3.get_object.assert_not_called() + + @patch("amzn_nova_forge.validation.inspect_lens_validator.boto3") + def test_task_name_mismatch_via_s3_raises(self, mock_boto3): + from amzn_nova_forge.validation.inspect_lens_validator import ( + validate_task_names_against_benchmarks, + ) + + source = b"@task\ndef boolq_pt():\n pass\n" + mock_s3 = MagicMock() + mock_boto3.client.return_value = mock_s3 + mock_paginator = MagicMock() + mock_s3.get_paginator.return_value = mock_paginator + mock_paginator.paginate.return_value = [{"Contents": [{"Key": "benchmarks/boolq_pt.py"}]}] + mock_s3.get_object.return_value = {"Body": BytesIO(source)} + + errors = validate_task_names_against_benchmarks( + tasks=[{"name": "mmlu_0_shot"}], + benchmarks_path="s3://bucket/benchmarks/", + ) + self.assertTrue(len(errors) > 0) + self.assertTrue(any("mmlu_0_shot" in e for e in errors)) + self.assertTrue(any("boolq_pt" in e for e in errors)) + + +class TestInferScenario(unittest.TestCase): + def test_bedrock_scenario(self): + cfg = InspectLensConfig(benchmarks_path=VALID_BENCHMARKS_PATH, tasks=VALID_TASKS) + self.assertEqual(cfg._infer_scenario(), "bedrock") + + def test_existing_endpoint_scenario(self): + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + endpoint_name="my-endpoint", + ) + self.assertEqual(cfg._infer_scenario(), "existing_endpoint") + + def test_create_endpoint_scenario(self): + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + model_s3_uri="s3://bucket/model/", + inference_image_uri=VALID_ECR, + ) + self.assertEqual(cfg._infer_scenario(), "create_endpoint") + + +class TestToYamlDict(unittest.TestCase): + def _bedrock_provider(self): + return {"bedrock": {"model_id": "us.amazon.nova-2-lite-v1:0", "region": "us-east-1"}} + + def test_benchmarks_s3_path_included_when_set(self): + cfg = InspectLensConfig( + benchmarks_path="s3://bucket/benchmarks/", + tasks=[{"name": "boolq_pt"}], + ) + result = cfg.to_yaml_dict(self._bedrock_provider(), "s3://bucket/output/") + self.assertEqual(result["benchmarks"]["s3_path"], "s3://bucket/benchmarks/") + + def test_benchmarks_s3_path_omitted_when_none(self): + cfg = object.__new__(InspectLensConfig) + for field, default in InspectLensConfig.__dataclass_fields__.items(): + try: + setattr(cfg, field, default.default) + except Exception: + setattr(cfg, field, default.default_factory()) + cfg.benchmarks_path = None + cfg.tasks = [{"name": "boolq_pt"}] + result = cfg.to_yaml_dict(self._bedrock_provider(), "s3://bucket/output/") + self.assertNotIn("s3_path", result["benchmarks"]) + + def test_decoding_overrides_applied(self): + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + ) + result = cfg.to_yaml_dict( + self._bedrock_provider(), + "s3://bucket/output/", + overrides={"temperature": 0.5, "max_tokens": 2048}, + ) + self.assertEqual(result["eval"]["decoding"]["temperature"], 0.5) + self.assertEqual(result["eval"]["decoding"]["max_tokens"], 2048) + + def test_mlflow_section_not_in_to_yaml_dict(self): + """MLflow tracking is injected by ForgeEvaluator, not to_yaml_dict.""" + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + ) + result = cfg.to_yaml_dict(self._bedrock_provider(), "s3://bucket/output/") + self.assertNotIn("tracking", result) + + def test_mlflow_section_omitted_when_not_set(self): + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + ) + result = cfg.to_yaml_dict(self._bedrock_provider(), "s3://bucket/output/") + self.assertNotIn("tracking", result) + + def test_output_s3_path_in_result(self): + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + ) + result = cfg.to_yaml_dict(self._bedrock_provider(), "s3://bucket/my-output/") + self.assertEqual(result["output"]["s3_path"], "s3://bucket/my-output/") + + def test_extra_args_included_when_set(self): + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + extra_args=["-M", "completion_mode=True"], + ) + result = cfg.to_yaml_dict(self._bedrock_provider(), "s3://bucket/output/") + self.assertEqual(result["eval"]["extra_args"], ["-M", "completion_mode=True"]) + + def test_extra_args_omitted_when_not_set(self): + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + ) + result = cfg.to_yaml_dict(self._bedrock_provider(), "s3://bucket/output/") + self.assertNotIn("extra_args", result["eval"]) + + def test_top_k_in_decoding(self): + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + ) + result = cfg.to_yaml_dict(self._bedrock_provider(), "s3://bucket/output/") + self.assertEqual(result["eval"]["decoding"]["top_k"], -1) + + def test_top_k_override(self): + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + ) + result = cfg.to_yaml_dict( + self._bedrock_provider(), "s3://bucket/output/", overrides={"top_k": 50} + ) + self.assertEqual(result["eval"]["decoding"]["top_k"], 50) + + def test_output_format_included_when_set(self): + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + output_format="csv", + ) + result = cfg.to_yaml_dict(self._bedrock_provider(), "s3://bucket/output/") + self.assertEqual(result["output"]["output_format"], "csv") + + def test_output_format_omitted_when_none(self): + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + ) + result = cfg.to_yaml_dict(self._bedrock_provider(), "s3://bucket/output/") + self.assertNotIn("output_format", result["output"]) + + def test_default_eval_values_match_spec(self): + cfg = InspectLensConfig( + benchmarks_path=VALID_BENCHMARKS_PATH, + tasks=VALID_TASKS, + ) + result = cfg.to_yaml_dict(self._bedrock_provider(), "s3://bucket/output/") + self.assertEqual(result["eval"]["max_connections"], 16) + self.assertEqual(result["eval"]["max_retries"], 100) + self.assertEqual(result["eval"]["timeout"], 600) + self.assertEqual(result["eval"]["decoding"]["temperature"], 0.0) + self.assertEqual(result["eval"]["decoding"]["top_p"], 1.0) + self.assertEqual(result["eval"]["decoding"]["top_k"], -1) + self.assertEqual(result["eval"]["decoding"]["max_tokens"], 8192) + + +class TestModelBedrockModelId(unittest.TestCase): + """Tests for Model.bedrock_model_id — context-length suffix stripping.""" + + def test_nova_micro_strips_suffix(self): + # amazon.nova-micro-v1:0:128k → us.amazon.nova-micro-v1:0 + self.assertEqual(Model.NOVA_MICRO.bedrock_model_id, "us.amazon.nova-micro-v1:0") + + def test_nova_lite_strips_suffix(self): + # amazon.nova-lite-v1:0:300k → us.amazon.nova-lite-v1:0 + self.assertEqual(Model.NOVA_LITE.bedrock_model_id, "us.amazon.nova-lite-v1:0") + + def test_nova_lite_2_strips_suffix(self): + # amazon.nova-2-lite-v1:0:256k → us.amazon.nova-2-lite-v1:0 + self.assertEqual(Model.NOVA_LITE_2.bedrock_model_id, "us.amazon.nova-2-lite-v1:0") + + def test_nova_pro_strips_suffix(self): + # amazon.nova-pro-v1:0:300k → us.amazon.nova-pro-v1:0 + self.assertEqual(Model.NOVA_PRO.bedrock_model_id, "us.amazon.nova-pro-v1:0") + + def test_all_models_start_with_us_prefix(self): + for model in Model: + self.assertTrue( + model.bedrock_model_id.startswith("us."), + f"{model.name}.bedrock_model_id should start with 'us.', got: {model.bedrock_model_id}", + ) + + def test_no_model_has_context_length_suffix(self): + for model in Model: + bid = model.bedrock_model_id + # Should not end with :128k, :256k, :300k etc. + self.assertNotRegex( + bid, + r":\d+k$", + f"{model.name}.bedrock_model_id should not have context-length suffix, got: {bid}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/manager/test_runtime_manager.py b/tests/unit/manager/test_runtime_manager.py index d29ea03..278df75 100644 --- a/tests/unit/manager/test_runtime_manager.py +++ b/tests/unit/manager/test_runtime_manager.py @@ -515,6 +515,38 @@ def test_validate_lambda_local_raises_on_s3_read_failure( finally: os.unlink(src) + @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_passes_environment_to_trainer( + self, mock_setup, mock_model_trainer_cls, mock_boto_client + ): + """When job_config.environment is set, it must be passed to ModelTrainer.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": "test-job-env"}] + } + + env = {"HF_TOKEN": "hf_xxx", "HF_HUB_DOWNLOAD_TIMEOUT": "300"} + job_config = JobConfig( + job_name="test-job", + image_uri="123456789012.dkr.ecr.us-east-1.amazonaws.com/my-image:latest", + recipe_path="/path/to/recipe", + output_s3_path="s3://output-bucket/output", + data_s3_path="s3://input-bucket/data", + input_s3_data_type="S3Prefix", + environment=env, + ) + + manager.execute(job_config) + + call_kwargs = mock_model_trainer_cls.from_recipe.call_args.kwargs + self.assertEqual(call_kwargs["environment"], env) + class TestSMTJValidationDataset(unittest.TestCase): """Tests for SFT validation dataset support in SMTJRuntimeManager."""