@aws/ml-container-creator 0.12.1 → 0.13.4

package/src/lib/prompt-runner.js

@@ -18,8 +18,6 @@ import {
     modelLoadStrategyPrompts,
     modelProfilePrompts,
     modulePrompts,
-    loraPrompts,
-    benchmarkPrompts,
     infraRegionAndTargetPrompts,
     infraExistingEndpointPrompts,
     infraInstancePrompts,
@@ -521,38 +519,23 @@ export default class PromptRunner {
         const ngcApiKeyAnswers = { ngcApiKey: secretAnswers.ngcApiKey, ngcTokenArn: secretAnswers.ngcTokenArn };
 
         // Module selection
-        const moduleAnswers = await this._runPhase(modulePrompts, { ...frameworkAnswers, ...engineAnswers }, explicitConfig, existingConfig);
-        
-        // Ensure transformers, diffusors, and ineligible Triton backends don't get sample model
-        if (frameworkAnswers.architecture === 'transformers' ||
-            frameworkAnswers.architecture === 'diffusors' ||
-            (frameworkAnswers.architecture === 'triton' && 
-             !this._tritonBackends[frameworkAnswers.backend]?.supportsSampleModel)) {
+        // Only ask about sample model for non-transformers/diffusors (Triton etc.)
+        const moduleAnswers = {};
+        if (frameworkAnswers.architecture !== 'transformers' &&
+            frameworkAnswers.architecture !== 'diffusors') {
+            const sampleModelAnswers = await this._runPhase(
+                modulePrompts.filter(p => p.name === 'includeSampleModel'),
+                { ...frameworkAnswers, ...engineAnswers }, explicitConfig, existingConfig
+            );
+            Object.assign(moduleAnswers, sampleModelAnswers);
+        } else {
             moduleAnswers.includeSampleModel = false;
         }
 
-        // Benchmark prompts — derive includeBenchmark from testTypes selection or CLI flag
-        // Requirements: 1.1, 1.2
-        let benchmarkAnswers = {};
-        if (frameworkAnswers.architecture === 'transformers' || frameworkAnswers.architecture === 'diffusors') {
-            const testTypes = moduleAnswers.testTypes || [];
-            const includeBenchmark = testTypes.includes('sagemaker-ai-automated-benchmarking') ||
-                explicitConfig.includeBenchmark === true ||
-                explicitConfig.includeBenchmark === 'true';
-            benchmarkAnswers.includeBenchmark = includeBenchmark;
-            if (includeBenchmark) {
-                const subAnswers = await this._runPhase(benchmarkPrompts, { ...frameworkAnswers, ...moduleAnswers, includeBenchmark }, explicitConfig, existingConfig);
-                benchmarkAnswers = { ...benchmarkAnswers, ...subAnswers };
-            }
-        }
-
-        // LoRA adapter prompts — only for transformers with vllm/sglang/djl-lmi
-        // Requirements: 1.1, 1.2, 1.4
-        let loraAnswers = {};
-        const loraSubAnswers = await this._runPhase(loraPrompts, { ...frameworkAnswers, ...engineAnswers }, explicitConfig, existingConfig);
-        if (loraSubAnswers.enableLora !== undefined) {
-            loraAnswers = loraSubAnswers;
-        }
+        // Test types, benchmark, and LoRA are always-on (BL-122)
+        moduleAnswers.testTypes = ['hosted-model-endpoint', 'sagemaker-ai-automated-benchmarking'];
+        const benchmarkAnswers = { includeBenchmark: true };
+        const loraAnswers = { enableLora: true };
 
         // Validate instance type against framework requirements (now that framework version is known)
         const finalInstanceType = infraAnswers.customInstanceType || infraAnswers.instanceType;

package/src/lib/prove-pipeline-executor.js

@@ -0,0 +1,294 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Prove Pipeline Executor
+ *
+ * Executes lifecycle stages for validation targets in the `mcc prove` workflow.
+ * Handles stage-specific logic including idempotency checks, status tracking,
+ * and fail-fast behavior.
+ *
+ * Feature: s3-model-loading
+ * Requirements: 5.1, 5.2, 5.3, 5.4, 5.5
+ */
+
+import { execFile } from 'node:child_process';
+import { promisify } from 'node:util';
+import { existsSync, readFileSync } from 'node:fs';
+import path from 'node:path';
+
+const execFileAsync = promisify(execFile);
+
+// ── Valid Lifecycle Stages ────────────────────────────────────────────────────
+
+/**
+ * All recognized lifecycle stages for the prove pipeline.
+ * The "stage" step pre-stages model weights from HuggingFace to S3.
+ */
+export const VALID_LIFECYCLE_STAGES = [
+    'generate',
+    'stage',
+    'build',
+    'push',
+    'deploy',
+    'test',
+    'tune',
+    'adapter',
+    'test-adapter',
+    'benchmark',
+    'register',
+    'clean'
+];
+
+/**
+ * Possible staging states for status output.
+ */
+export const STAGING_STATES = {
+    STAGED: 'staged',
+    NOT_STAGED: 'not-staged',
+    STAGE_FAILED: 'stage-failed'
+};
+
+// ── Stage Lifecycle Step ─────────────────────────────────────────────────────
+
+/**
+ * Check if a model has already been staged by looking for `.mlcc/staged-assets.json`.
+ *
+ * @param {string} projectDir - Path to the generated project directory
+ * @returns {boolean} True if the model has already been staged
+ */
+export function isAlreadyStaged(projectDir) {
+    const stagedAssetsPath = path.join(projectDir, '.mlcc', 'staged-assets.json');
+    if (!existsSync(stagedAssetsPath)) {
+        return false;
+    }
+
+    try {
+        const content = readFileSync(stagedAssetsPath, 'utf8');
+        const data = JSON.parse(content);
+        // Check that there's a valid staged URI
+        return !!(data?.models?.default?.staged_uri);
+    } catch {
+        return false;
+    }
+}
+
+/**
+ * Get the current staging state for a project.
+ *
+ * @param {string} projectDir - Path to the generated project directory
+ * @param {object} [stepResults] - Previous step results (to check for stage-failed)
+ * @returns {string} One of: 'staged', 'not-staged', 'stage-failed'
+ */
+export function getStagingState(projectDir, stepResults = null) {
+    // Check if stage previously failed
+    if (stepResults?.stage?.status === 'fail') {
+        return STAGING_STATES.STAGE_FAILED;
+    }
+
+    if (isAlreadyStaged(projectDir)) {
+        return STAGING_STATES.STAGED;
+    }
+
+    return STAGING_STATES.NOT_STAGED;
+}
+
+/**
+ * Execute the stage lifecycle step with idempotency support.
+ *
+ * If the model is already staged (`.mlcc/staged-assets.json` exists with a valid URI),
+ * the step is skipped and marked as passed.
+ *
+ * If `do/stage` exits non-zero, the model is marked as stage-failed.
+ *
+ * @param {string} projectDir - Path to the generated project directory
+ * @param {object} [options] - Execution options
+ * @param {number} [options.timeout=1800] - Timeout in seconds (default: 30 minutes)
+ * @param {boolean} [options.verbose=false] - Stream stdout/stderr in real time
+ * @returns {Promise<object>} StepResult with name, status, duration, stagingState, and optional error
+ */
+export async function executeStageStep(projectDir, options = {}) {
+    const { timeout = 1800, verbose = false } = options;
+    const startTime = Date.now();
+
+    // Idempotency check: skip if already staged (Requirement 5.4)
+    if (isAlreadyStaged(projectDir)) {
+        return {
+            name: 'stage',
+            status: 'pass',
+            duration: Date.now() - startTime,
+            stagingState: STAGING_STATES.STAGED,
+            skipped: true,
+            message: '✓ Model already staged — skipping'
+        };
+    }
+
+    // Execute do/stage and verify exit code 0 (Requirement 5.2)
+    const command = './do/stage';
+
+    try {
+        if (verbose) {
+            // Verbose: stream output in real time
+            const { spawn } = await import('node:child_process');
+            const result = await new Promise((resolve) => {
+                const child = spawn('bash', ['-c', command], {
+                    cwd: projectDir,
+                    stdio: ['pipe', 'inherit', 'inherit']
+                });
+
+                let killed = false;
+                const timer = setTimeout(() => {
+                    killed = true;
+                    child.kill('SIGTERM');
+                }, timeout * 1000);
+
+                child.on('close', (code) => {
+                    clearTimeout(timer);
+                    if (code === 0) {
+                        resolve({
+                            name: 'stage',
+                            status: 'pass',
+                            duration: Date.now() - startTime,
+                            stagingState: STAGING_STATES.STAGED
+                        });
+                    } else {
+                        const error = killed
+                            ? `Timeout after ${timeout}s`
+                            : `do/stage exited with code ${code}`;
+                        resolve({
+                            name: 'stage',
+                            status: 'fail',
+                            duration: Date.now() - startTime,
+                            stagingState: STAGING_STATES.STAGE_FAILED,
+                            error
+                        });
+                    }
+                });
+
+                child.on('error', (err) => {
+                    clearTimeout(timer);
+                    resolve({
+                        name: 'stage',
+                        status: 'fail',
+                        duration: Date.now() - startTime,
+                        stagingState: STAGING_STATES.STAGE_FAILED,
+                        error: err.message.slice(-500)
+                    });
+                });
+            });
+            return result;
+        }
+
+        // Non-verbose: buffer output
+        await execFileAsync('bash', ['-c', command], {
+            cwd: projectDir,
+            timeout: timeout * 1000,
+            maxBuffer: 10 * 1024 * 1024
+        });
+
+        return {
+            name: 'stage',
+            status: 'pass',
+            duration: Date.now() - startTime,
+            stagingState: STAGING_STATES.STAGED
+        };
+    } catch (err) {
+        // Mark model as failed if staging fails (Requirement 5.3)
+        const error = err.killed
+            ? `Timeout after ${timeout}s`
+            : (err.stderr || err.message).slice(-500);
+
+        return {
+            name: 'stage',
+            status: 'fail',
+            duration: Date.now() - startTime,
+            stagingState: STAGING_STATES.STAGE_FAILED,
+            error
+        };
+    }
+}
+
+// ── Stage Validation ─────────────────────────────────────────────────────────
+
+/**
+ * Validate that a lifecycle stage name is recognized by the prove pipeline.
+ *
+ * @param {string} stageName - The stage name to validate
+ * @returns {boolean} True if the stage is valid
+ */
+export function isValidLifecycleStage(stageName) {
+    return VALID_LIFECYCLE_STAGES.includes(stageName);
+}
+
+/**
+ * Validate a stages array from validation-targets configuration.
+ *
+ * @param {string[]} stages - Array of stage names
+ * @returns {object} Validation result: { valid: boolean, errors: string[] }
+ */
+export function validateStagesArray(stages) {
+    const errors = [];
+
+    if (!Array.isArray(stages)) {
+        return { valid: false, errors: ['stages must be an array'] };
+    }
+
+    if (stages.length === 0) {
+        return { valid: false, errors: ['stages array must not be empty'] };
+    }
+
+    for (const stage of stages) {
+        if (typeof stage !== 'string') {
+            errors.push(`Invalid stage type: expected string, got ${typeof stage}`);
+            continue;
+        }
+        if (!isValidLifecycleStage(stage)) {
+            errors.push(`Unrecognized lifecycle stage: "${stage}"`);
+        }
+    }
+
+    return { valid: errors.length === 0, errors };
+}
+
+// ── Status Output ────────────────────────────────────────────────────────────
+
+/**
+ * Format the staging state for status output display.
+ *
+ * @param {string} state - One of STAGING_STATES values
+ * @returns {string} Formatted status string with emoji
+ */
+export function formatStagingStatus(state) {
+    switch (state) {
+    case STAGING_STATES.STAGED:
+        return '✓ staged';
+    case STAGING_STATES.NOT_STAGED:
+        return '○ not-staged';
+    case STAGING_STATES.STAGE_FAILED:
+        return '✗ stage-failed';
+    default:
+        return '? unknown';
+    }
+}
+
+/**
+ * Build a status summary for a prove target including staging state.
+ *
+ * @param {object} target - The validation target
+ * @param {string} target.model_name - Model name
+ * @param {string} projectDir - Path to the project directory
+ * @param {object} [stepResults] - Results of executed steps
+ * @returns {object} Status summary including stagingState
+ */
+export function buildTargetStatus(target, projectDir, stepResults = null) {
+    const stagingState = getStagingState(projectDir, stepResults);
+    const stages = target.stages || [];
+    const includesStage = stages.includes('stage');
+
+    return {
+        model_name: target.model_name,
+        stagingState,
+        stagingStatus: formatStagingStatus(stagingState),
+        includesStageStep: includesStage
+    };
+}

package/templates/IAM_PERMISSIONS.md

@@ -10,14 +10,47 @@ This project uses three sets of IAM permissions:
 
 ## SageMaker Execution Role
 
-The bootstrap command creates an IAM role (`mlcc-sagemaker-execution-role`) with permissions for:
+The bootstrap command creates an IAM role (`mlcc-sagemaker-execution-role`) with these permission groups:
 
-- **SageMaker**: Create, update, delete, and invoke endpoints, endpoint configs, models, and inference components
-- **ECR**: Pull images from the `ml-container-creator` repository
-- **CloudWatch Logs**: Write container logs
-- **S3**: Read model artifacts from `ml-container-creator-*` buckets
+### Endpoint Management
+Create, update, delete, describe, and invoke endpoints, endpoint configs, models, and inference components.
 
-The role is defined in the CloudFormation stack template (`config/bootstrap-stack.json`) and updated automatically when you re-run bootstrap after upgrading.
+### AI Benchmarking
+Create, describe, list, stop, and delete AI benchmark jobs, AI recommendation jobs, and AI workload configs.
+
+### Training & Model Customization
+Create/describe/stop training jobs, model packages, model package groups. Access SageMaker Hub contents. Manage training plans.
+
+### MLflow Integration
+List/describe MLflow tracking servers and apps. Create presigned URLs. Call MLflow app APIs.
+
+### ECR
+Pull container images (GetAuthorizationToken, BatchGetImage, GetDownloadUrlForLayer, BatchCheckLayerAvailability).
+
+### S3
+Read and write model artifacts, adapters, benchmark results:
+- `s3:GetObject`, `s3:PutObject`, `s3:AbortMultipartUpload`, `s3:ListBucket`
+- Scoped to `mlcc-*` and `ml-container-creator-*` buckets
+
+### CloudWatch Logs
+Create log groups/streams and put log events.
+
+### Secrets Manager
+Read and write secrets prefixed with `mlcc/` or `ml-container-creator/` (used for HF tokens, API keys).
+
+### SNS
+Publish notifications to `mlcc-*` and `ml-container-creator-*` topics (benchmark completion alerts).
+
+### Service Quotas & Capacity
+Query service quotas and training plan availability for instance selection.
+
+### Lambda
+Invoke functions (reward model evaluation during training/tuning).
+
+### PassRole
+Self-pass to SageMaker service, scoped to `mlcc-sagemaker-execution-role`.
+
+The role is defined in `config/bootstrap-stack.json` and updated automatically when you re-run bootstrap after upgrading.
 
 If you use a custom role (`--role-arn`), ensure it has at minimum:
 
@@ -25,12 +58,15 @@ If you use a custom role (`--role-arn`), ensure it has at minimum:
 |-----------|---------|
 | `sagemaker:CreateEndpoint`, `CreateEndpointConfig`, `CreateModel`, `CreateInferenceComponent` | Deploy |
 | `sagemaker:DeleteEndpoint`, `DeleteEndpointConfig`, `DeleteModel`, `DeleteInferenceComponent` | Clean up |
-| `sagemaker:DescribeEndpoint`, `DescribeEndpointConfig`, `DescribeModel`, `DescribeInferenceComponent` | Status checks |
+| `sagemaker:DescribeEndpoint`, `DescribeEndpointConfig`, `DescribeModel`, `DescribeInferenceComponent`, `ListInferenceComponents` | Status |
 | `sagemaker:InvokeEndpoint`, `InvokeEndpointAsync` | Inference |
 | `sagemaker:UpdateEndpoint`, `UpdateEndpointWeightsAndCapacities`, `UpdateInferenceComponent` | Updates |
-| `ecr:GetAuthorizationToken`, `BatchGetImage`, `GetDownloadUrlForLayer`, `BatchCheckLayerAvailability` | Pull container image |
-| `logs:CreateLogGroup`, `CreateLogStream`, `PutLogEvents` | Container logging |
-| `s3:GetObject`, `s3:ListBucket` on `ml-container-creator-*` | Model artifact access |
+| `sagemaker:CreateAIBenchmarkJob`, `DescribeAIBenchmarkJob`, `ListAIBenchmarkJobs` | Benchmark |
+| `sagemaker:CreateTrainingJob`, `DescribeTrainingJob`, `StopTrainingJob` | Training/tuning |
+| `ecr:GetAuthorizationToken`, `BatchGetImage`, `GetDownloadUrlForLayer`, `BatchCheckLayerAvailability` | Pull image |
+| `logs:CreateLogGroup`, `CreateLogStream`, `PutLogEvents` | Logging |
+| `s3:GetObject`, `s3:PutObject`, `s3:ListBucket` on project buckets | Artifacts |
+| `iam:PassRole` (to sagemaker.amazonaws.com) | Role delegation |
 
 Trust policy must allow `sagemaker.amazonaws.com` to assume the role.
 
@@ -48,12 +84,27 @@ Your AWS user or CI system needs these permissions to run the do-scripts:
 
 | Script | Permissions Needed |
 |--------|-------------------|
+| `./do/build` | Local only — no AWS permissions |
+| `./do/run` | Local only — no AWS permissions |
 | `./do/push` | `ecr:GetAuthorizationToken`, `ecr:PutImage`, `ecr:InitiateLayerUpload`, `ecr:UploadLayerPart`, `ecr:CompleteLayerUpload`, `ecr:BatchCheckLayerAvailability` |
 | `./do/submit` | `codebuild:CreateProject`, `codebuild:StartBuild`, `codebuild:BatchGetBuilds`, `iam:CreateRole`, `iam:PutRolePolicy`, `iam:PassRole`, `s3:PutObject`, `s3:CreateBucket` |
-| `./do/deploy` | `sagemaker:CreateEndpointConfig`, `sagemaker:CreateEndpoint`, `sagemaker:CreateInferenceComponent`, `sagemaker:DescribeEndpoint`, `iam:PassRole` |
-| `./do/clean` | `sagemaker:DeleteEndpoint`, `sagemaker:DeleteEndpointConfig`, `sagemaker:DeleteInferenceComponent`, `codebuild:DeleteProject`, `iam:DeleteRole`, `iam:DeleteRolePolicy` |
+| `./do/stage` | `s3:PutObject`, `s3:GetObject`, `s3:ListBucket` on mlcc-* buckets |
+| `./do/deploy` | `sagemaker:CreateEndpointConfig`, `sagemaker:CreateEndpoint`, `sagemaker:CreateModel`, `sagemaker:CreateInferenceComponent`, `sagemaker:DescribeEndpoint`, `iam:PassRole` |
+| `./do/add-ic` | `sagemaker:CreateInferenceComponent`, `sagemaker:DescribeEndpoint`, `sagemaker:ListInferenceComponents`, `iam:PassRole` |
 | `./do/test` | `sagemaker-runtime:InvokeEndpoint` |
-| `bootstrap` | `cloudformation:*`, `iam:CreateRole`, `iam:PutRolePolicy`, `iam:TagRole`, `ecr:CreateRepository`, `s3:CreateBucket` (and `sts:GetCallerIdentity`) |
+| `./do/benchmark` | `sagemaker:CreateAIBenchmarkJob`, `sagemaker:DescribeAIBenchmarkJob`, `sagemaker:ListAIBenchmarkJobs`, `sagemaker:CreateAIWorkloadConfig`, `iam:PassRole`, `s3:GetObject` |
+| `./do/train` | `sagemaker:CreateTrainingJob`, `sagemaker:DescribeTrainingJob`, `iam:PassRole`, `s3:GetObject`, `s3:PutObject` |
+| `./do/tune` | `sagemaker:CreateTrainingJob`, `sagemaker:DescribeTrainingJob`, `iam:PassRole`, `s3:GetObject`, `s3:PutObject` |
+| `./do/adapter` | `sagemaker:CreateInferenceComponent`, `sagemaker:UpdateInferenceComponent`, `sagemaker:DescribeInferenceComponent`, `s3:GetObject` |
+| `./do/optimize` | `sagemaker:CreateModel`, `sagemaker:DescribeModel`, `s3:GetObject`, `s3:PutObject` |
+| `./do/register` | `sagemaker:CreateModelPackage`, `sagemaker:CreateModelPackageGroup`, `sagemaker:DescribeModelPackage` |
+| `./do/logs` | `logs:GetLogEvents`, `logs:FilterLogEvents`, `logs:DescribeLogStreams` |
+| `./do/status` | `sagemaker:DescribeEndpoint`, `sagemaker:DescribeInferenceComponent`, `sagemaker:ListInferenceComponents` |
+| `./do/clean` | `sagemaker:DeleteEndpoint`, `sagemaker:DeleteEndpointConfig`, `sagemaker:DeleteModel`, `sagemaker:DeleteInferenceComponent`, `codebuild:DeleteProject`, `iam:DeleteRole`, `iam:DeleteRolePolicy` |
+| `./do/export` | Local only — reads config files |
+| `./do/validate` | Local only — validates project structure |
+| `./do/manifest` | Local only — generates deployment manifest |
+| `bootstrap` | `cloudformation:*`, `iam:CreateRole`, `iam:PutRolePolicy`, `iam:TagRole`, `ecr:CreateRepository`, `s3:CreateBucket`, `sts:GetCallerIdentity` |
 
 <% if (framework === 'transformers' && hfToken) { %>
 ## HuggingFace Token Security

package/templates/do/.tune_helper.py

@@ -1510,8 +1510,6 @@ def cmd_discover(args):
 
     Returns: {"models": [str], "count": int}
     """
-    import boto3
-
     region = args.region or os.environ.get('AWS_REGION', 'us-east-1')
 
     family = args.family or ""
@@ -1528,6 +1526,11 @@ def cmd_discover(args):
     if not prefix:
         _error_exit("No family or filter provided for discovery")
 
+    try:
+        import boto3
+    except ImportError:
+        _error_exit("Hub discovery failed: boto3 is not installed. Install with: pip install boto3")
+
     try:
         client = boto3.client("sagemaker", region_name=region)
         models = []