@aws/ml-container-creator 1.0.3 → 1.0.4

package/servers/lib/schemas/image-catalog.schema.json

@@ -159,6 +159,18 @@
                     "items": {
                         "type": "string"
                     }
+                },
+                "min_driver_version": {
+                    "type": "string",
+                    "description": "Minimum GPU driver version required (e.g., '550.54')"
+                },
+                "cuda_toolkit": {
+                    "type": "string",
+                    "description": "CUDA toolkit version bundled in the image (e.g., '12.4')"
+                },
+                "transformers_version": {
+                    "type": "string",
+                    "description": "Bundled transformers library version (e.g., '4.52.0')"
                 }
             },
             "additionalProperties": false

package/src/app.js

@@ -366,10 +366,11 @@ export async function writeProject(templateDir, destDir, answers, registryConfig
     const trainIncluded = answers.deploymentTarget !== 'batch-transform';
     if (!trainIncluded) {
         ignorePatterns.push('**/do/train');
+        ignorePatterns.push('**/do/.train_helper.py');
         ignorePatterns.push('**/do/.train_build_request.py');
-        ignorePatterns.push('**/do/.train_status_parser.py');
-        ignorePatterns.push('**/do/.train_poll_parser.py');
         ignorePatterns.push('**/do/training/**');
+        ignorePatterns.push('**/do/evaluate');
+        ignorePatterns.push('**/do/.eval_helper.py');
     }
 
     // Exclude feedback.sh when neither tune nor train is included
@@ -404,10 +405,11 @@ export async function writeProject(templateDir, destDir, answers, registryConfig
         ignorePatterns.push('**/do/.adapter_helper.py');
         ignorePatterns.push('**/do/.register_helper.py');
         ignorePatterns.push('**/do/train');
+        ignorePatterns.push('**/do/.train_helper.py');
         ignorePatterns.push('**/do/.train_build_request.py');
-        ignorePatterns.push('**/do/.train_status_parser.py');
-        ignorePatterns.push('**/do/.train_poll_parser.py');
         ignorePatterns.push('**/do/training/**');
+        ignorePatterns.push('**/do/evaluate');
+        ignorePatterns.push('**/do/.eval_helper.py');
         ignorePatterns.push('**/do/add-ic');
         ignorePatterns.push('**/do/run');
         ignorePatterns.push('**/sample_model/**');

package/src/lib/generated/cli-options.js

@@ -1,6 +1,6 @@
 // AUTO-GENERATED by scripts/codegen-cli.js — DO NOT EDIT
 // Source: config/parameter-schema-v2.json
-// Generated: 2026-06-29T13:37:06.271Z
+// Generated: 2026-06-30T16:45:56.916Z
 
 /**
  * CLI option definitions derived from parameter-schema-v2.json.

package/src/lib/generated/parameter-matrix.js

@@ -1,6 +1,6 @@
 // AUTO-GENERATED by scripts/codegen-parameter-matrix.js — DO NOT EDIT
 // Source: config/parameter-schema-v2.json
-// Generated: 2026-06-29T13:37:06.375Z
+// Generated: 2026-06-30T16:45:57.021Z
 
 /**
  * Parameter matrix defining how each parameter is loaded from various sources.

package/src/lib/generated/validation-rules.js

@@ -1,6 +1,6 @@
 // AUTO-GENERATED by scripts/codegen-validator.js — DO NOT EDIT
 // Source: config/parameter-schema-v2.json
-// Generated: 2026-06-29T13:37:06.303Z
+// Generated: 2026-06-30T16:45:56.949Z
 
 /**
  * Validation rules derived from parameter-schema-v2.json.

package/src/lib/mcp-query-runner.js

@@ -384,6 +384,9 @@ export default class McpQueryRunner {
                 const endpointNames = result.choices.endpointName;
                 const metadata = result.metadata || {};
 
+                // Store endpoint metadata for later instance type resolution (US-1)
+                this.runner._endpointPickerMetadata = metadata;
+
                 // Build choices with metadata annotations
                 this.runner._mcpEndpointChoices = endpointNames.map(name => {
                     const meta = metadata[name];
@@ -412,12 +415,15 @@ export default class McpQueryRunner {
     }
 
     /**
-     * Query MCP base-image-picker server after deployment config is selected.
+     * Query MCP base-image-picker server after deployment config and instance type are known.
      * Populates _mcpBaseImageChoices for the base image selection prompt.
-     * Requirements: 5.1, 5.2, 5.3, 5.4, 9.1, 9.2, 9.3
+     * Requirements: 5.1, 5.2, 5.3, 5.4, 9.1, 9.2, 9.3, US-1 (ordering constraint)
+     * @param {Object} frameworkAnswers - Framework/architecture answers
+     * @param {Object} _explicitConfig - Explicit CLI/config values
+     * @param {Object} [infraContext] - Infrastructure context (instanceType, tensorParallelSize, modelId)
      * @private
      */
-    async _queryMcpForBaseImage(frameworkAnswers, _explicitConfig) {
+    async _queryMcpForBaseImage(frameworkAnswers, _explicitConfig, infraContext = {}) {
         // Skip if base image provided via CLI --base-image flag
         if (this.runner.options['base-image']) return;
 
@@ -454,6 +460,17 @@ export default class McpQueryRunner {
             context.searchCriteria = searchCriteria.trim();
         }
 
+        // Pass infrastructure context for driver-aware filtering (US-1 ordering constraint)
+        if (infraContext.instanceType) {
+            context.instanceType = infraContext.instanceType;
+        }
+        if (infraContext.tensorParallelSize !== null && infraContext.tensorParallelSize !== undefined) {
+            context.tensorParallelSize = infraContext.tensorParallelSize;
+        }
+        if (infraContext.modelId) {
+            context.modelId = infraContext.modelId;
+        }
+
         const result = await cm.queryMcpServer('base-image-picker', context);
 
         if (result && result.metadata?.baseImage?.length > 0) {
@@ -716,6 +733,96 @@ export default class McpQueryRunner {
         }
     }
 
+    /**
+     * Resolve instance type from an existing endpoint.
+     * Priority:
+     *   1. Endpoint-picker metadata (already fetched, no network call)
+     *   2. Direct AWS SDK call: DescribeEndpoint → DescribeEndpointConfig
+     *
+     * Reuses the resolution pattern from do/lib/resolve-instance.sh:
+     *   - Check ProductionVariants[0].CurrentInstanceType or InstanceType
+     *   - Fallback: DescribeEndpointConfig → ProductionVariants[0].InstanceType
+     *   - Final fallback: InstancePools[0] (highest priority)
+     *
+     * Requirements: US-1 (ordering constraint — resolve instance type before base image picker)
+     * @param {string} endpointName - The existing endpoint name
+     * @param {string} awsRegion - AWS region for API calls
+     * @returns {Promise<string|null>} Resolved instance type or null on failure
+     * @private
+     */
+    async _resolveEndpointInstanceType(endpointName, awsRegion) {
+        // Strategy 1: Use endpoint-picker metadata (already fetched, no network call)
+        if (this.runner._endpointPickerMetadata) {
+            const meta = this.runner._endpointPickerMetadata[endpointName];
+            if (meta?.instanceType) {
+                // Strip pool annotation if present: "ml.g5.12xlarge (pool: 3 types)" → "ml.g5.12xlarge"
+                const rawInstanceType = meta.instanceType.includes(' (pool:')
+                    ? meta.instanceType.split(' (pool:')[0]
+                    : meta.instanceType;
+                if (rawInstanceType && rawInstanceType !== 'unknown') {
+                    console.log(`   ✓ Resolved instance type from endpoint metadata: ${rawInstanceType}`);
+                    return rawInstanceType;
+                }
+            }
+        }
+
+        // Strategy 2: Direct AWS SDK call (for custom endpoint names not in picker results)
+        console.log('   🔍 Resolving instance type from existing endpoint...');
+        try {
+            const { SageMakerClient, DescribeEndpointCommand, DescribeEndpointConfigCommand } = await import('@aws-sdk/client-sagemaker');
+
+            const region = awsRegion || process.env.AWS_REGION || 'us-east-1';
+            const clientOptions = { region };
+
+            // Use AWS profile if available
+            const awsProfile = this.runner.configManager?.config?.awsProfile
+                || this.runner.options?.profile || process.env.AWS_PROFILE || null;
+            if (awsProfile) {
+                try {
+                    const { fromIni } = await import('@aws-sdk/credential-providers');
+                    clientOptions.credentials = fromIni({ profile: awsProfile });
+                } catch {
+                    // credential-providers not available, use default chain
+                }
+            }
+
+            const client = new SageMakerClient(clientOptions);
+
+            // DescribeEndpoint — check ProductionVariants for instance type
+            const epResponse = await client.send(new DescribeEndpointCommand({ EndpointName: endpointName }));
+
+            const primaryVariant = (epResponse.ProductionVariants || [])[0] || {};
+            let instanceType = primaryVariant.CurrentInstanceType || primaryVariant.InstanceType || null;
+
+            // Fallback: DescribeEndpointConfig
+            if (!instanceType && epResponse.EndpointConfigName) {
+                const ecResponse = await client.send(
+                    new DescribeEndpointConfigCommand({ EndpointConfigName: epResponse.EndpointConfigName })
+                );
+                const ecVariant = (ecResponse.ProductionVariants || [])[0];
+                if (ecVariant?.InstanceType) {
+                    instanceType = ecVariant.InstanceType;
+                } else if (ecVariant?.InstancePools?.length > 0) {
+                    // Use highest-priority pool entry (lowest Priority number)
+                    const sorted = [...ecVariant.InstancePools].sort((a, b) => (a.Priority || 99) - (b.Priority || 99));
+                    instanceType = sorted[0].InstanceType || null;
+                }
+            }
+
+            if (instanceType) {
+                console.log(`   ✓ Resolved instance type from endpoint: ${instanceType}`);
+                return instanceType;
+            }
+
+            console.log('   ↳ Could not determine instance type from endpoint');
+            return null;
+        } catch (err) {
+            // Graceful fallback: if AWS call fails, skip filtering (no driver-aware filter)
+            console.log(`   ⚠️  Could not resolve instance type from endpoint: ${err.message}`);
+            return null;
+        }
+    }
+
     /**
      * Validate and display instance type compatibility
      * Requirements: 4.1, 4.2, 4.3, 4.4, 4.5, 4.6

package/src/lib/prompt-runner.js

@@ -68,6 +68,7 @@ export default class PromptRunner {
     _queryMcpForInstance(...args) { return this.mcpQueryRunner._queryMcpForInstance(...args); }
     _queryMcpForInstanceSizing(...args) { return this.mcpQueryRunner._queryMcpForInstanceSizing(...args); }
     _queryMcpForEndpoints(...args) { return this.mcpQueryRunner._queryMcpForEndpoints(...args); }
+    _resolveEndpointInstanceType(...args) { return this.mcpQueryRunner._resolveEndpointInstanceType(...args); }
     _queryMcpForHyperPod(...args) { return this.mcpQueryRunner._queryMcpForHyperPod(...args); }
     _fetchAndDisplayModelInfo(...args) { return this.mcpQueryRunner._fetchAndDisplayModelInfo(...args); }
     _validateAndDisplayInstanceType(...args) { return this.mcpQueryRunner._validateAndDisplayInstanceType(...args); }
@@ -182,8 +183,8 @@ export default class PromptRunner {
         }
 
         // ══════════════════════════════════════════════════════════════════════
-        // Phase 2 — How (deployment target + serving profile + base image)
-        // Requirements: 4.3 — instance prompt appears AFTER base image is known
+        // Phase 2 — How (deployment target + serving profile)
+        // Requirements: US-1 — base image selection moved AFTER instance resolution
         // ══════════════════════════════════════════════════════════════════════
         console.log('\n💪 Infrastructure & Deployment');
 
@@ -192,25 +193,8 @@ export default class PromptRunner {
         const regionPreviousAnswers = bootstrapRegion ? { _bootstrapRegion: bootstrapRegion } : {};
         const regionAndTargetAnswers = await this._runPhase(infraRegionAndTargetPrompts, { ...frameworkAnswers, ...regionPreviousAnswers }, explicitConfig, existingConfig);
 
-        // 2b. Query base-image-picker MCP server for base image choices
-        await this.mcpQueryRunner._queryMcpForBaseImage(frameworkAnswers, explicitConfig);
-        const baseImagePreviousAnswers = {
-            ...frameworkAnswers,
-            ...engineAnswers,
-            ...(this._mcpBaseImageChoices ? { _mcpBaseImageChoices: this._mcpBaseImageChoices } : {})
-        };
-        const baseImageAnswers = await this._runPhase(
-            baseImagePrompts,
-            baseImagePreviousAnswers,
-            explicitConfig,
-            existingConfig
-        );
-
-        // Requirements: 4.2-4.5 — Check model architecture compatibility after base image selection
-        this._checkModelArchitectureCompatibility(baseImageAnswers, frameworkAnswers);
-
-        // Extract CUDA version from selected base image for instance-sizer context
-        const selectedBaseImageCuda = this._extractCudaFromBaseImage(baseImageAnswers);
+        // NOTE: Base image selection moved to Phase 3 (after instance type resolution)
+        // to enable driver-aware filtering. See US-1 ordering constraint in requirements.
 
         // ══════════════════════════════════════════════════════════════════════
         // Phase 3 — Where (region + instance [derived] + CUDA/AMI + HyperPod + build target)
@@ -283,7 +267,7 @@ export default class PromptRunner {
                 } else if (phase1ModelId && phase1ModelId !== 'Custom (enter manually)') {
                     // Query instance-sizer with full context
                     await this.mcpQueryRunner._queryMcpForInstanceSizing(frameworkAnswers, modelFormatAnswers, explicitConfig, {
-                        cudaVersion: selectedBaseImageCuda,
+                        cudaVersion: null, // base image not yet selected (moved after instance resolution)
                         profileEnvVars: this._selectedProfileEnvVars || {}
                     });
                 } else {
@@ -422,6 +406,66 @@ export default class PromptRunner {
             }
         }
 
+        // 3b2. Base image selection — AFTER instance type resolved (US-1 ordering constraint)
+        // Pass resolved instanceType and tensorParallelSize for driver-aware filtering
+        let resolvedInstanceType = instanceAnswers.customInstanceType || instanceAnswers.instanceType;
+        let resolvedTensorParallelSize = this._autoTensorParallelism || 1;
+
+        // For existing endpoints: resolve instance type from the endpoint (US-1 ordering constraint)
+        // The instance type is needed for driver-aware base image filtering even though the user
+        // doesn't select it manually. Pattern reused from do/lib/resolve-instance.sh.
+        if (!resolvedInstanceType && existingEndpointAnswers.existingEndpointName) {
+            const resolvedRegion = regionAndTargetAnswers.customAwsRegion || regionAndTargetAnswers.awsRegion;
+            resolvedInstanceType = await this.mcpQueryRunner._resolveEndpointInstanceType(
+                existingEndpointAnswers.existingEndpointName,
+                resolvedRegion
+            );
+            // Store resolved instance type for downstream use (IC config, GPU count derivation)
+            if (resolvedInstanceType) {
+                existingEndpointAnswers._resolvedEndpointInstanceType = resolvedInstanceType;
+                // Propagate as instanceType so template-variable-resolver derives
+                // icGpuCount and tensorParallelSize from the instance catalog.
+                // Without this, IC_GPU_COUNT defaults to 1 even for multi-GPU instances.
+                existingEndpointAnswers.instanceType = resolvedInstanceType;
+
+                // Derive GPU count from instance catalog for immediate use (TP for base image filtering)
+                const endpointInstanceEntry = instanceCatalogRaw[resolvedInstanceType];
+                if (endpointInstanceEntry?.gpus && endpointInstanceEntry.gpus > 1) {
+                    existingEndpointAnswers.gpuCount = endpointInstanceEntry.gpus;
+                    existingEndpointAnswers.tensorParallelSize = endpointInstanceEntry.gpus;
+                    this._autoTensorParallelism = endpointInstanceEntry.gpus;
+                    this._autoGpuCount = endpointInstanceEntry.gpus;
+                    console.log(`   ✓ Endpoint instance ${resolvedInstanceType}: ${endpointInstanceEntry.gpus} GPUs (TP=${endpointInstanceEntry.gpus})`);
+                }
+            }
+        }
+
+        // Re-read tensor parallel size after potential endpoint resolution update
+        resolvedTensorParallelSize = this._autoTensorParallelism || 1;
+
+        await this.mcpQueryRunner._queryMcpForBaseImage(frameworkAnswers, explicitConfig, {
+            instanceType: resolvedInstanceType,
+            tensorParallelSize: resolvedTensorParallelSize,
+            modelId: phase1ModelId || undefined
+        });
+        const baseImagePreviousAnswers = {
+            ...frameworkAnswers,
+            ...engineAnswers,
+            ...(this._mcpBaseImageChoices ? { _mcpBaseImageChoices: this._mcpBaseImageChoices } : {})
+        };
+        const baseImageAnswers = await this._runPhase(
+            baseImagePrompts,
+            baseImagePreviousAnswers,
+            explicitConfig,
+            existingConfig
+        );
+
+        // Requirements: 4.2-4.5 — Check model architecture compatibility after base image selection
+        this._checkModelArchitectureCompatibility(baseImageAnswers, frameworkAnswers);
+
+        // Extract CUDA version from selected base image for CUDA/AMI auto-resolution
+        const selectedBaseImageCuda = this._extractCudaFromBaseImage(baseImageAnswers);
+
         // 3c. Async-specific prompts (only when deploymentTarget === 'async-inference')
         let asyncAnswers = {};
         if (regionAndTargetAnswers.deploymentTarget === 'async-inference') {

package/src/lib/template-variable-resolver.js

@@ -454,6 +454,14 @@ export async function _ensureTemplateVariables(answers, registryConfigManager =
                 answers.tensorParallelSize = instanceGpuCount;
                 answers._tpAutoResolved = true;
                 answers._tpAutoResolvedFrom = answers.instanceType;
+
+                // Also propagate to icEnvVars so IC_ENV_VLLM_TENSOR_PARALLEL_SIZE
+                // (or equivalent) is written in do/config for deploy-time IC creation.
+                if (!answers.icEnvVars) {
+                    answers.icEnvVars = {};
+                }
+                answers.icEnvVars[tpEnvKey] = String(instanceGpuCount);
+
                 console.log(`    ℹ️  TP degree: ${instanceGpuCount} (auto-detected from ${answers.instanceType})`);
             }
         }