@aws/ml-container-creator 1.0.3 → 1.1.0

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-07-01T20:12:14.996Z
 
 /**
  * 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-07-01T20:12:14.922Z
 
 /**
  * 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})`);
             }
         }

package/src/lib/train-config-builder.js

@@ -0,0 +1,339 @@
+#!/usr/bin/env node
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Interactive Training Job Configuration Builder.
+ *
+ * Guides users through configuring a custom training job by prompting
+ * for technique, model, dataset, instance type, and hyperparameters.
+ * Writes the result to training/config.yaml.
+ *
+ * Invoked from do/train --interactive:
+ *   node -e "import('.../train-config-builder.js').then(m => m.run({...}))"
+ *
+ * Uses @inquirer/prompts via the project's prompt-adapter.js for UX
+ * consistency with the main ml-container-creator generation flow.
+ */
+
+import { select, input, confirm } from '@inquirer/prompts';
+import { readFileSync, writeFileSync, readdirSync, existsSync } from 'node:fs';
+import { join, resolve } from 'node:path';
+import { parseArgs } from 'node:util';
+
+// ── YAML helpers (minimal, no dependency) ────────────────────────────────────
+
+/**
+ * Parse a simple YAML file (flat key-value, no nesting beyond what we need).
+ * Falls back gracefully if format is unexpected.
+ */
+function parseSimpleYaml(content) {
+    const result = {};
+    for (const line of content.split('\n')) {
+        const trimmed = line.trim();
+        if (!trimmed || trimmed.startsWith('#')) continue;
+        const colonIdx = trimmed.indexOf(':');
+        if (colonIdx === -1) continue;
+        const key = trimmed.slice(0, colonIdx).trim();
+        let value = trimmed.slice(colonIdx + 1).trim();
+        // Remove quotes
+        if ((value.startsWith('"') && value.endsWith('"')) ||
+            (value.startsWith('\'') && value.endsWith('\''))) {
+            value = value.slice(1, -1);
+        }
+        // Type coercion
+        if (value === 'true') result[key] = true;
+        else if (value === 'false') result[key] = false;
+        else if (value === '' || value === '""' || value === '\'\'') result[key] = '';
+        else if (!isNaN(value) && value !== '') result[key] = Number(value);
+        else result[key] = value;
+    }
+    return result;
+}
+
+// ── Technique scanning ───────────────────────────────────────────────────────
+
+function scanTechniques(trainingDir) {
+    const techniques = [];
+    try {
+        const entries = readdirSync(trainingDir, { withFileTypes: true });
+        for (const entry of entries) {
+            if (entry.isDirectory()) {
+                const trainScript = join(trainingDir, entry.name, 'train.py');
+                if (existsSync(trainScript)) {
+                    techniques.push(entry.name);
+                }
+            }
+        }
+    } catch {
+        // Directory doesn't exist or not readable
+    }
+    return techniques.length > 0 ? techniques : ['custom'];
+}
+
+// ── Prompts.json loading ─────────────────────────────────────────────────────
+
+function loadTechniquePrompts(trainingDir, technique) {
+    const promptsFile = join(trainingDir, technique, 'prompts.json');
+    if (!existsSync(promptsFile)) return null;
+    try {
+        return JSON.parse(readFileSync(promptsFile, 'utf8'));
+    } catch {
+        return null;
+    }
+}
+
+// ── Defaults loading ─────────────────────────────────────────────────────────
+
+function loadTechniqueDefaults(trainingDir, technique) {
+    const defaultsFile = join(trainingDir, technique, 'defaults.yaml');
+    if (!existsSync(defaultsFile)) return {};
+    try {
+        return parseSimpleYaml(readFileSync(defaultsFile, 'utf8'));
+    } catch {
+        return {};
+    }
+}
+
+// ── Main interactive flow ────────────────────────────────────────────────────
+
+export async function run({ configFile, trainingDir }) {
+    const configPath = resolve(configFile);
+    const trainingPath = resolve(trainingDir);
+
+    // Resolve output_path from bootstrap profile if not already in config
+    let profileOutputPath = '';
+    try {
+        const homedir = process.env.HOME || process.env.USERPROFILE || '';
+        const profilePath = join(homedir, '.ml-container-creator', 'config.json');
+        if (existsSync(profilePath)) {
+            const profileData = JSON.parse(readFileSync(profilePath, 'utf8'));
+            const activeProfile = profileData.profiles?.[profileData.activeProfile] || {};
+            const bucket = activeProfile.benchmarkS3Bucket || '';
+            if (bucket) {
+                // Derive project name from training dir (parent dir name)
+                const projectName = resolve(trainingPath, '..').split('/').pop();
+                profileOutputPath = `s3://${bucket}/${projectName}/training-output/`;
+            }
+        }
+    } catch { /* best-effort */ }
+
+    // Load existing config as defaults
+    let existingConfig = {};
+    if (existsSync(configPath)) {
+        try {
+            existingConfig = parseSimpleYaml(readFileSync(configPath, 'utf8'));
+        } catch {
+            // Ignore parse errors — start fresh
+        }
+    }
+
+    console.log('');
+    console.log('🏋️  Custom Training Job Builder');
+    console.log('━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━');
+    console.log('');
+
+    // ── Technique selection ──────────────────────────────────────────────────
+    const techniques = scanTechniques(trainingPath);
+    const technique = await select({
+        message: 'Training technique?',
+        choices: techniques.map(t => ({ name: t, value: t })),
+        default: existingConfig.technique || 'sft'
+    });
+
+    // ── Common questions ─────────────────────────────────────────────────────
+    const modelId = await input({
+        message: 'Base model (HuggingFace ID)?',
+        default: existingConfig.model_id || process.env.HF_MODEL_ID || 'Qwen/Qwen3-0.6B'
+    });
+
+    const dataset = await input({
+        message: 'Dataset (hf://org/name, s3://..., or registry name)?',
+        default: existingConfig.dataset || ''
+    });
+
+    const instanceType = await input({
+        message: 'Instance type?',
+        default: existingConfig.instance_type || 'ml.g5.xlarge'
+    });
+
+    // ── Load technique defaults for hyperparam questions ─────────────────────
+    const defaults = loadTechniqueDefaults(trainingPath, technique);
+
+    const epochs = await input({
+        message: 'Epochs?',
+        default: String(existingConfig.epochs || defaults.epochs || 3),
+        validate: (v) => !isNaN(v) && Number(v) > 0 ? true : 'Must be a positive number'
+    });
+
+    const learningRate = await input({
+        message: 'Learning rate?',
+        default: String(existingConfig.learning_rate || defaults.learning_rate || '2e-4'),
+        validate: (v) => !isNaN(parseFloat(v)) ? true : 'Must be a number'
+    });
+
+    const loraR = await input({
+        message: 'LoRA rank (r)?',
+        default: String(existingConfig.lora_r || defaults.lora_r || 16),
+        validate: (v) => !isNaN(v) && Number(v) > 0 ? true : 'Must be a positive integer'
+    });
+
+    // ── Technique-specific prompts ───────────────────────────────────────────
+    const techniquePromptsSchema = loadTechniquePrompts(trainingPath, technique);
+    const techniqueAnswers = {};
+
+    if (techniquePromptsSchema && techniquePromptsSchema.prompts) {
+        console.log('');
+        console.log(`─── ${techniquePromptsSchema.section_title || `${technique} settings`} ───`);
+
+        for (const prompt of techniquePromptsSchema.prompts) {
+            const existingVal = existingConfig[prompt.name];
+            const defaultVal = existingVal !== null && existingVal !== undefined ? String(existingVal) :
+                (defaults[prompt.name] !== null && defaults[prompt.name] !== undefined ? String(defaults[prompt.name]) :
+                    (prompt.default || ''));
+
+            const answer = await input({
+                message: `${prompt.message}`,
+                default: defaultVal,
+                validate: (v) => {
+                    if (prompt.validate === 'float') return !isNaN(parseFloat(v)) ? true : 'Must be a number';
+                    if (prompt.validate === 'int') return !isNaN(parseInt(v)) ? true : 'Must be an integer';
+                    return true;
+                }
+            });
+            techniqueAnswers[prompt.name] = answer;
+        }
+    }
+
+    // ── Build config ─────────────────────────────────────────────────────────
+    const hyperparameters = {
+        epochs,
+        learning_rate: learningRate,
+        lora_r: loraR,
+        ...techniqueAnswers
+    };
+
+    // ── Write config ─────────────────────────────────────────────────────────
+    // Build YAML output (preserving the original file structure where possible)
+    const yamlLines = [
+        '# do/training/config.yaml — Generated by interactive builder',
+        `# Technique: ${technique}`,
+        `# Generated: ${new Date().toISOString()}`,
+        '',
+        `technique: "${technique}"`,
+        '',
+        '# Base model',
+        `model_id: "${modelId}"`,
+        '',
+        '# Dataset',
+        `dataset: "${dataset}"`,
+        '',
+        '# Instance',
+        `instance_type: "${instanceType}"`,
+        `instance_count: ${existingConfig.instance_count || 1}`,
+        '',
+        '# Container image',
+        `image: "${existingConfig.image || ''}"`,
+        '',
+        '# Script (auto-selected from technique)',
+        `script: "do/training/${technique}/train.py"`,
+        '',
+        '# Output',
+        `output_path: "${existingConfig.output_path || profileOutputPath}"`,
+        '',
+        '# Hyperparameters',
+        'hyperparameters:'
+    ];
+
+    for (const [key, val] of Object.entries(hyperparameters)) {
+        yamlLines.push(`  ${key}: "${val}"`);
+    }
+
+    // Preserve other existing fields
+    if (existingConfig.max_runtime_seconds) {
+        yamlLines.push('', `max_runtime_seconds: ${existingConfig.max_runtime_seconds}`);
+    }
+    if (existingConfig.volume_size_gb) {
+        yamlLines.push(`volume_size_gb: ${existingConfig.volume_size_gb}`);
+    }
+    if (existingConfig.enable_spot) {
+        yamlLines.push(`enable_spot: ${existingConfig.enable_spot}`);
+    }
+
+    yamlLines.push('');
+    writeFileSync(configPath, yamlLines.join('\n'), 'utf8');
+
+    // ── Summary ──────────────────────────────────────────────────────────────
+    console.log('');
+    console.log('✅ Configuration written to training/config.yaml');
+    console.log('');
+    console.log(`  technique: ${technique}`);
+    console.log(`  model: ${modelId}`);
+    console.log(`  dataset: ${dataset || '(none)'}`);
+    console.log(`  instance_type: ${instanceType}`);
+    console.log(`  epochs: ${epochs}`);
+    console.log(`  learning_rate: ${learningRate}`);
+    console.log(`  lora_r: ${loraR}`);
+    if (Object.keys(techniqueAnswers).length > 0) {
+        for (const [k, v] of Object.entries(techniqueAnswers)) {
+            console.log(`  ${k}: ${v}`);
+        }
+    }
+    console.log('');
+
+    // ── Run now? ─────────────────────────────────────────────────────────────
+    const runNow = await confirm({
+        message: 'Run training job now?',
+        default: false
+    });
+
+    // Output JSON for bash consumption
+    const resultObj = {
+        config_written: true,
+        technique,
+        run_now: runNow
+    };
+
+    // Print to stdout (for CLI entry point / backward compat)
+    console.log(JSON.stringify(resultObj));
+
+    // Return for programmatic callers (do/train writes to temp file)
+    return resultObj;
+}
+
+// ── CLI entry point ──────────────────────────────────────────────────────────
+
+async function main() {
+    const { values } = parseArgs({
+        options: {
+            'config-file': { type: 'string' },
+            'training-dir': { type: 'string' }
+        }
+    });
+
+    const configFile = values['config-file'];
+    const trainingDir = values['training-dir'];
+
+    if (!configFile || !trainingDir) {
+        console.error('Usage: train-config-builder --config-file <path> --training-dir <path>');
+        process.exit(1);
+    }
+
+    try {
+        await run({ configFile, trainingDir });
+    } catch (err) {
+        if (err.name === 'ExitPromptError') {
+            // User pressed Ctrl+C
+            console.log('\n⚠️  Cancelled.');
+            process.exit(130);
+        }
+        console.error(`❌ Error: ${err.message}`);
+        process.exit(1);
+    }
+}
+
+// Run if invoked directly
+const isMainModule = process.argv[1] && resolve(process.argv[1]) === resolve(new URL(import.meta.url).pathname);
+if (isMainModule) {
+    main();
+}