@aws/ml-container-creator 0.6.0 → 0.7.1

package/src/lib/tune-dataset-validator.js

@@ -0,0 +1,279 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Tune Dataset Validator
+ *
+ * Parses dataset arguments (S3 URIs and Hugging Face references) and
+ * validates JSONL dataset lines against catalog-driven schemas.
+ *
+ * Requirements: 3.1, 3.5, 3.6, 3.7, 3.8, 3.10, 3.11, 3.12
+ */
+
+/**
+ * Parse a dataset argument string into a structured object.
+ * Accepts S3 URIs (`s3://bucket/key`) or Hugging Face references
+ * (`hf://org/name` or `hf://org/name/split`).
+ *
+ * @param {string} datasetStr - The dataset argument string
+ * @returns {{ valid: boolean, type?: string, bucket?: string, key?: string, org?: string, name?: string, split?: string, error?: string }}
+ */
+export function parseDatasetArg(datasetStr) {
+    if (!datasetStr || typeof datasetStr !== 'string') {
+        return {
+            valid: false,
+            error: 'Dataset argument is required and must be a non-empty string.'
+        };
+    }
+
+    const trimmed = datasetStr.trim();
+
+    if (trimmed.startsWith('s3://')) {
+        return _parseS3Uri(trimmed);
+    }
+
+    if (trimmed.startsWith('hf://')) {
+        return _parseHfReference(trimmed);
+    }
+
+    return {
+        valid: false,
+        error: `Invalid dataset format: "${trimmed}". Expected s3://bucket/key or hf://org/name[/split].`
+    };
+}
+
+/**
+ * Validate JSONL lines against a dataset schema from the catalog.
+ * Inspects only the first 10 lines per requirement.
+ *
+ * @param {string[]} lines - Array of JSONL line strings
+ * @param {Object} schema - The datasetSchema object from the catalog
+ * @param {string[]} schema.required - Array of required top-level keys
+ * @param {Object} schema.types - Object mapping key to expected type ("string", "array", "object", "number")
+ * @returns {{ valid: boolean, error: string|null, lineNumber: number|null, malformedLine: string|null, expectedFormat: string|null }}
+ */
+export function validateDatasetFormat(lines, schema) {
+    if (!lines || !Array.isArray(lines)) {
+        return {
+            valid: false,
+            error: 'Lines must be provided as an array.',
+            lineNumber: null,
+            malformedLine: null,
+            expectedFormat: _buildExpectedFormat(schema)
+        };
+    }
+
+    if (!schema || !schema.required || !Array.isArray(schema.required)) {
+        return {
+            valid: false,
+            error: 'Schema must include a "required" array of keys.',
+            lineNumber: null,
+            malformedLine: null,
+            expectedFormat: null
+        };
+    }
+
+    const linesToInspect = lines.slice(0, 10);
+
+    for (let i = 0; i < linesToInspect.length; i++) {
+        const line = linesToInspect[i];
+        const lineNumber = i + 1;
+
+        // Skip empty lines
+        if (!line || line.trim() === '') {
+            continue;
+        }
+
+        // Try to parse as JSON
+        let parsed;
+        try {
+            parsed = JSON.parse(line);
+        } catch (e) {
+            return {
+                valid: false,
+                error: `Line ${lineNumber} is not valid JSON: ${e.message}`,
+                lineNumber,
+                malformedLine: line,
+                expectedFormat: _buildExpectedFormat(schema)
+            };
+        }
+
+        // Check that parsed value is an object
+        if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) {
+            return {
+                valid: false,
+                error: `Line ${lineNumber} must be a JSON object.`,
+                lineNumber,
+                malformedLine: line,
+                expectedFormat: _buildExpectedFormat(schema)
+            };
+        }
+
+        // Check required keys
+        for (const key of schema.required) {
+            if (!Object.hasOwn(parsed, key)) {
+                return {
+                    valid: false,
+                    error: `Line ${lineNumber} is missing required key "${key}".`,
+                    lineNumber,
+                    malformedLine: line,
+                    expectedFormat: _buildExpectedFormat(schema)
+                };
+            }
+        }
+
+        // Check types if specified
+        if (schema.types) {
+            for (const [key, expectedType] of Object.entries(schema.types)) {
+                if (!Object.hasOwn(parsed, key)) {
+                    continue;
+                }
+
+                const value = parsed[key];
+                if (!_checkType(value, expectedType)) {
+                    return {
+                        valid: false,
+                        error: `Line ${lineNumber} has key "${key}" with wrong type. Expected "${expectedType}", got "${_getType(value)}".`,
+                        lineNumber,
+                        malformedLine: line,
+                        expectedFormat: _buildExpectedFormat(schema)
+                    };
+                }
+            }
+        }
+    }
+
+    return {
+        valid: true,
+        error: null,
+        lineNumber: null,
+        malformedLine: null,
+        expectedFormat: null
+    };
+}
+
+/**
+ * Parse an S3 URI into bucket and key components.
+ * @param {string} uri - The S3 URI (e.g., "s3://bucket/path/to/file.jsonl")
+ * @returns {Object} Parsed result
+ * @private
+ */
+function _parseS3Uri(uri) {
+    const withoutScheme = uri.slice(5); // Remove "s3://"
+    const slashIndex = withoutScheme.indexOf('/');
+
+    if (slashIndex === -1 || slashIndex === 0) {
+        return {
+            valid: false,
+            error: `Invalid S3 URI: "${uri}". Expected format: s3://bucket/key.`
+        };
+    }
+
+    const bucket = withoutScheme.slice(0, slashIndex);
+    const key = withoutScheme.slice(slashIndex + 1);
+
+    if (!bucket) {
+        return {
+            valid: false,
+            error: `Invalid S3 URI: "${uri}". Bucket name is empty.`
+        };
+    }
+
+    if (!key) {
+        return {
+            valid: false,
+            error: `Invalid S3 URI: "${uri}". Key path is empty.`
+        };
+    }
+
+    return {
+        valid: true,
+        type: 's3',
+        bucket,
+        key
+    };
+}
+
+/**
+ * Parse a Hugging Face dataset reference into org, name, and split.
+ * Defaults to 'train' split if not specified.
+ * @param {string} ref - The HF reference (e.g., "hf://org/name" or "hf://org/name/split")
+ * @returns {Object} Parsed result
+ * @private
+ */
+function _parseHfReference(ref) {
+    const withoutScheme = ref.slice(5); // Remove "hf://"
+    const parts = withoutScheme.split('/');
+
+    if (parts.length < 2 || !parts[0] || !parts[1]) {
+        return {
+            valid: false,
+            error: `Invalid Hugging Face reference: "${ref}". Expected format: hf://org/name[/split].`
+        };
+    }
+
+    const org = parts[0];
+    const name = parts[1];
+    const split = parts.length >= 3 && parts[2] ? parts[2] : 'train';
+
+    return {
+        valid: true,
+        type: 'hf',
+        org,
+        name,
+        split
+    };
+}
+
+/**
+ * Check if a value matches the expected schema type.
+ * @param {*} value - The value to check
+ * @param {string} expectedType - One of "string", "array", "object", "number"
+ * @returns {boolean} True if the value matches the expected type
+ * @private
+ */
+function _checkType(value, expectedType) {
+    switch (expectedType) {
+    case 'string':
+        return typeof value === 'string';
+    case 'number':
+        return typeof value === 'number';
+    case 'array':
+        return Array.isArray(value);
+    case 'object':
+        return typeof value === 'object' && value !== null && !Array.isArray(value);
+    default:
+        return true;
+    }
+}
+
+/**
+ * Get a human-readable type name for a value.
+ * @param {*} value - The value to describe
+ * @returns {string} The type name
+ * @private
+ */
+function _getType(value) {
+    if (value === null) return 'null';
+    if (Array.isArray(value)) return 'array';
+    return typeof value;
+}
+
+/**
+ * Build a human-readable expected format description from a schema.
+ * @param {Object} schema - The dataset schema
+ * @returns {string|null} Description of expected format
+ * @private
+ */
+function _buildExpectedFormat(schema) {
+    if (!schema || !schema.required) {
+        return null;
+    }
+
+    const fields = schema.required.map(key => {
+        const type = schema.types && schema.types[key] ? schema.types[key] : 'any';
+        return `"${key}": <${type}>`;
+    });
+
+    return `Each line must be a JSON object with: {${fields.join(', ')}}`;
+}

package/src/lib/tune-output-resolver.js

@@ -0,0 +1,66 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Tune Output Resolver
+ *
+ * Detects output type from training type and generates context-aware
+ * next-step commands for deploying tune job artifacts.
+ *
+ * Requirements: 8.3, 8.11
+ */
+
+/**
+ * Detect the output type based on the training type used for the job.
+ * LoRA training produces adapter weights; full-rank produces a full model.
+ *
+ * @param {string} trainingType - The training type ('lora' or 'full-rank')
+ * @returns {string} The output type: 'adapter' for lora, 'full-model' for full-rank
+ */
+export function detectOutputType(trainingType) {
+    if (trainingType === 'lora') {
+        return 'adapter';
+    }
+    if (trainingType === 'full-rank') {
+        return 'full-model';
+    }
+    return 'adapter';
+}
+
+/**
+ * Generate context-aware next-step commands based on the output type.
+ *
+ * For adapter output:
+ *   - Quick path: ./do/adapter add tuned-${technique} --from-tune
+ *   - Technique-specific: ./do/adapter add tuned-${technique} --from-tune ${technique}
+ *   - Explicit path: ./do/adapter add tuned-${technique} --weights ${artifactPath}
+ *
+ * For full-model output:
+ *   - Deploy as new IC: ./do/add-ic tuned-v1 --from-tune
+ *   - Explicit path: ./do/add-ic tuned-v1 --model-data ${artifactPath}
+ *   - Replace current base: ./do/deploy --force-ic --model-data ${artifactPath}
+ *
+ * @param {string} outputType - The output type ('adapter' or 'full-model')
+ * @param {string} technique - The technique used (e.g., 'sft', 'dpo')
+ * @param {string} artifactPath - The S3 path to the output artifact
+ * @returns {string[]} Array of suggested next-step commands
+ */
+export function generateNextStepCommands(outputType, technique, artifactPath) {
+    if (outputType === 'adapter') {
+        return [
+            `./do/adapter add tuned-${technique} --from-tune`,
+            `./do/adapter add tuned-${technique} --from-tune ${technique}`,
+            `./do/adapter add tuned-${technique} --weights ${artifactPath}`
+        ];
+    }
+
+    if (outputType === 'full-model') {
+        return [
+            './do/add-ic tuned-v1 --from-tune',
+            `./do/add-ic tuned-v1 --model-data ${artifactPath}`,
+            `./do/deploy --force-ic --model-data ${artifactPath}`
+        ];
+    }
+
+    return [];
+}

package/templates/Dockerfile

@@ -290,6 +290,7 @@ RUN chmod +x /usr/bin/serve_trtllm
 
 # Copy startup script
 COPY code/cuda_compat.sh /usr/bin/cuda_compat.sh
+COPY code/cw_log_forwarder.py /usr/bin/cw_log_forwarder.py
 COPY code/start_server.sh /usr/bin/start_server.sh
 RUN chmod +x /usr/bin/start_server.sh /usr/bin/cuda_compat.sh
 
@@ -307,6 +308,7 @@ COPY code/serving.properties /opt/ml/model/serving.properties
 # The container will automatically start DJL Serving with the configuration
 <% } else { %>
 COPY code/cuda_compat.sh /usr/bin/cuda_compat.sh
+COPY code/cw_log_forwarder.py /usr/bin/cw_log_forwarder.py
 COPY code/serve /usr/bin/serve
 RUN chmod 777 /usr/bin/serve /usr/bin/cuda_compat.sh
 

package/templates/code/cw_log_forwarder.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+"""CloudWatch log forwarder — workaround for IC platform log routing gap.
+Pipes stdin to a CW log stream while passing through to stderr.
+Usage: exec > >(python3 /usr/bin/cw_log_forwarder.py) 2>&1
+"""
+import sys, os, time, threading
+import boto3
+from botocore.config import Config
+
+LOG_GROUP = os.environ.get("CW_LOG_GROUP",
+    f"/aws/sagemaker/InferenceComponents/{os.environ.get('INFERENCE_COMPONENT_NAME', os.environ.get('HOSTNAME', 'unknown'))}")
+LOG_STREAM = f"AllTraffic/{os.environ.get('HOSTNAME', 'container')}"
+REGION = os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION", "us-west-2"))
+
+def main():
+    client = boto3.client("logs", region_name=REGION, config=Config(retries={"max_attempts": 2}))
+    try:
+        client.create_log_group(logGroupName=LOG_GROUP)
+    except Exception:
+        pass
+    try:
+        client.create_log_stream(logGroupName=LOG_GROUP, logStreamName=LOG_STREAM)
+    except Exception as e:
+        # Can't create stream — just passthrough
+        for line in sys.stdin:
+            sys.stderr.write(line)
+        return
+
+    buf, lock, seq = [], threading.Lock(), [None]
+
+    def flush():
+        with lock:
+            if not buf:
+                return
+            batch = buf[:50]
+            del buf[:50]
+        events = [{"timestamp": int(t * 1000), "message": m} for t, m in batch]
+        kw = {"logGroupName": LOG_GROUP, "logStreamName": LOG_STREAM, "logEvents": events}
+        if seq[0]:
+            kw["sequenceToken"] = seq[0]
+        try:
+            r = client.put_log_events(**kw)
+            seq[0] = r.get("nextSequenceToken")
+        except Exception:
+            pass
+
+    def loop():
+        while True:
+            time.sleep(2)
+            flush()
+
+    threading.Thread(target=loop, daemon=True).start()
+    try:
+        for line in sys.stdin:
+            sys.stderr.write(line)
+            with lock:
+                buf.append((time.time(), line.rstrip("\n")))
+    except (KeyboardInterrupt, BrokenPipeError):
+        pass
+    finally:
+        flush()
+
+if __name__ == "__main__":
+    main()

package/templates/code/serve

@@ -2,6 +2,11 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+# CloudWatch log forwarder — workaround for IC platform log routing gap
+exec > >(python3 /usr/bin/cw_log_forwarder.py) 2>&1
+
+echo "$(date -u '+%Y-%m-%dT%H:%M:%SZ') [serve] Container started — PID $$"
+
 # CUDA compatibility setup (required for newer SageMaker inference AMIs)
 source /usr/bin/cuda_compat.sh 2>/dev/null || true
 
@@ -270,8 +275,14 @@ for var in "${env_vars[@]}"; do
     
     # Remove prefix, convert to lowercase, and replace underscores with dashes
     arg_name=$(echo "${key#"${PREFIX}"}" | tr '[:upper:]' '[:lower:]' | tr '_' '-')
+
+    # Boolean handling: true = flag only, false = skip entirely
+    if [ "$value" = "false" ]; then
+        continue
+    fi
+
     SERVER_ARGS+=("${ARG_PREFIX}${arg_name}")
-    if [ -n "$value" ]; then
+    if [ -n "$value" ] && [ "$value" != "true" ]; then
         SERVER_ARGS+=("$value")
     fi
 done