@aws/ml-container-creator 0.13.4 → 0.15.0

package/src/lib/prompt-runner.js

@@ -710,6 +710,11 @@ export default class PromptRunner {
             delete combinedAnswers.customHyperPodCluster;
         }
 
+        // Propagate max_model_len from instance-sizer context capping (AC-1.7)
+        if (this._sizerMaxModelLen) {
+            combinedAnswers.sizerMaxModelLen = this._sizerMaxModelLen;
+        }
+
         // Apply CUDA version selection → inference AMI override
         if (combinedAnswers._resolvedInferenceAmiVersion) {
             combinedAnswers.inferenceAmiVersion = combinedAnswers._resolvedInferenceAmiVersion;

package/src/lib/prompts/feature-prompts.js

@@ -90,7 +90,7 @@ const loraPrompts = [
         type: 'confirm',
         name: 'enableLora',
         message: 'Enable LoRA adapter serving?',
-        default: false,
+        default: true,
         when: (answers) => {
             const architecture = answers.architecture || answers.deploymentConfig?.split('-')[0];
             const backend = answers.backend || answers.deploymentConfig?.split('-').slice(1).join('-');

package/src/lib/template-manager.js

@@ -314,13 +314,6 @@ export default class TemplateManager {
     _validateBenchmarkConfig() {
         if (!this.answers.includeBenchmark) return;
 
-        // Gate to supported architectures
-        const dc = this.answers.deploymentConfig;
-        const arch = dc ? dc.split('-')[0] : this.answers.architecture;
-        if (arch !== 'transformers' && arch !== 'diffusors') {
-            throw new Error('⚠️  Benchmarking is only supported with transformers and diffusors architectures.');
-        }
-
         // Gate to supported deployment targets
         if (this.answers.deploymentTarget === 'hyperpod-eks') {
             throw new Error('⚠️  Benchmarking is only supported with managed-inference, async-inference, and batch-transform deployment targets');

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

@@ -232,7 +232,7 @@ export async function _ensureTemplateVariables(answers, registryConfigManager =
         artifactUri: '',
         modelLoadStrategy: 'runtime',
         existingEndpointName: null,
-        enableLora: false,
+        enableLora: true,
         maxLoras: 30,
         maxLoraRank: 64
     };
@@ -261,6 +261,20 @@ export async function _ensureTemplateVariables(answers, registryConfigManager =
         }
     }
 
+    // Always include benchmarking by default (AC-2.3 — enabled for all architectures).
+    // Only set when not explicitly provided by user (AC-2.4, AC-2.7 — respect explicit opt-out).
+    if (answers.includeBenchmark === undefined) {
+        answers.includeBenchmark = true;
+    }
+
+    // Enforce enableLora scoping: only LoRA-capable servers get enableLora=true
+    // (AC-2.1, NFR-2). All incompatible backends are forced to false.
+    const loraCapableServers = ['vllm', 'sglang', 'djl-lmi', 'lmi', 'djl'];
+    const resolvedBackend = answers.backend || answers.modelServer;
+    if (!loraCapableServers.includes(resolvedBackend)) {
+        answers.enableLora = false;
+    }
+
     // Merge catalog env vars into answers.envVars with correct precedence
     await _mergeEnvVarsWithPrecedence(answers, registryConfigManager);
 
@@ -445,6 +459,35 @@ export async function _ensureTemplateVariables(answers, registryConfigManager =
         }
     }
 
+    // Propagate max_model_len from instance-sizer context capping to env vars (AC-1.7).
+    // The instance-sizer sets sizerMaxModelLen when the model's full context doesn't fit
+    // on the recommended instance. Write as VLLM_MAX_MODEL_LEN or SGLANG_MAX_MODEL_LEN.
+    const _MAX_MODEL_LEN_ENGINE_MAP = {
+        'vllm': 'VLLM_MAX_MODEL_LEN',
+        'vllm-omni': 'VLLM_MAX_MODEL_LEN',
+        'sglang': 'SGLANG_MAX_MODEL_LEN'
+    };
+
+    if (answers.sizerMaxModelLen) {
+        const maxLenEngine = answers.backend || answers.modelServer;
+        const maxLenEnvKey = maxLenEngine ? _MAX_MODEL_LEN_ENGINE_MAP[maxLenEngine] : null;
+        if (maxLenEnvKey) {
+            // Only set if user hasn't explicitly provided this env var
+            const userServerEnvVars = answers.serverEnvVars || {};
+            const userExplicitlySetMaxLen = (
+                userServerEnvVars['MAX_MODEL_LEN'] !== undefined ||
+                userServerEnvVars[maxLenEnvKey] !== undefined
+            );
+            if (!userExplicitlySetMaxLen && (!answers.envVars || !answers.envVars[maxLenEnvKey])) {
+                if (!answers.envVars) {
+                    answers.envVars = {};
+                }
+                answers.envVars[maxLenEnvKey] = String(answers.sizerMaxModelLen);
+                console.log(`    ℹ️  max_model_len: ${answers.sizerMaxModelLen} (context capped by instance-sizer)`);
+            }
+        }
+    }
+
     // Determine tune support based on model presence in the tune catalog.
     // Used by the do/config template to write TUNE_SUPPORTED=true|false.
     if (answers.tuneSupported === undefined) {
@@ -481,4 +524,11 @@ export async function _ensureTemplateVariables(answers, registryConfigManager =
             answers.tuneModelId = null;
         }
     }
+
+    // Propagate --ic-env KEY=VALUE pairs to icEnvVars for do/config template rendering.
+    // These are rendered as IC_ENV_* exports in do/config, which inference-component.sh
+    // reads at deploy time and passes as the Environment field in InferenceComponent.create().
+    if (!answers.icEnvVars) {
+        answers.icEnvVars = {};
+    }
 }

package/src/lib/tune-config-state.js

@@ -74,22 +74,35 @@ export function persistSubmissionState(configPath, { technique, trainingType, da
  * Simulate the config writes that happen after a job completes successfully.
  * This mirrors the behavior in do/tune's _handle_completion() function.
  *
+ * Writes three levels of tracking (AC-4.1, AC-4.2):
+ * - Level 1: TUNE_OUTPUT_PATH_LATEST (always the last run, any technique)
+ * - Level 2: TUNE_ADAPTER_PATH_<TECHNIQUE> (last run per technique)
+ * - Level 3: TUNE_ADAPTER_PATH_<TECHNIQUE>_<SLUG> (per technique + dataset slug)
+ *
  * @param {string} configPath - Path to the config file
  * @param {object} params - Completion parameters
  * @param {string} params.technique - Technique (sft, dpo, rlaif, rlvr)
  * @param {string} params.trainingType - Training type (lora, full-rank)
  * @param {string} params.artifactPath - S3 path to the output artifact
  * @param {string} params.outputType - Output type (adapter, full-model)
+ * @param {string} [params.datasetSlug] - Optional dataset slug for per-technique-per-dataset tracking
  */
-export function persistCompletionState(configPath, { technique, trainingType, artifactPath, outputType }) {
+export function persistCompletionState(configPath, { technique, trainingType, artifactPath, outputType, datasetSlug }) {
     const techniqueUpper = technique.toUpperCase();
 
     if (trainingType === 'lora') {
+        // Level 2: per-technique
         updateConfigVar(configPath, `TUNE_ADAPTER_PATH_${techniqueUpper}`, artifactPath);
+        // Level 3: per-technique + per-dataset (if slug available)
+        if (datasetSlug) {
+            const slugUpper = datasetSlug.toUpperCase().replace(/-/g, '_');
+            updateConfigVar(configPath, `TUNE_ADAPTER_PATH_${techniqueUpper}_${slugUpper}`, artifactPath);
+        }
     } else if (trainingType === 'full-rank') {
         updateConfigVar(configPath, `TUNE_MODEL_PATH_${techniqueUpper}`, artifactPath);
     }
 
+    // Level 1: latest
     updateConfigVar(configPath, 'TUNE_OUTPUT_PATH_LATEST', artifactPath);
     updateConfigVar(configPath, 'TUNE_OUTPUT_TYPE_LATEST', outputType);
 }

package/templates/do/.adapter_helper.py

@@ -0,0 +1,451 @@
+#!/usr/bin/env python3
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""SageMaker Processing Job helper for adapter staging.
+
+Subcommands:
+    stage-from-tune  - Submit Processing Job to copy adapter from training output to S3
+    status           - Check Processing Job status
+
+All output is JSON on stdout for bash consumption.
+
+Uses sagemaker-core ProcessingJob.create() / ProcessingJob.get() per SDK v3 policy.
+"""
+
+import argparse
+import logging
+import json
+import os
+import sys
+import time
+import warnings
+
+# Suppress noisy dependency version warnings
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+warnings.filterwarnings("ignore", message=".*urllib3.*")
+
+# Suppress sagemaker-core INFO/WARNING logging that pollutes stdout
+logging.getLogger("sagemaker.config").setLevel(logging.ERROR)
+logging.getLogger("sagemaker.core").setLevel(logging.ERROR)
+logging.getLogger("sagemaker").setLevel(logging.ERROR)
+
+# ── Constants ─────────────────────────────────────────────────────────────────
+POLL_INTERVAL_SECONDS = 30
+MAX_RUNTIME_SECONDS = 3600  # 1 hour timeout for adapter staging
+INSTANCE_TYPE = "ml.m5.large"
+VOLUME_SIZE_GB = 100
+
+# ── Utility functions ─────────────────────────────────────────────────────────
+
+
+def _error_exit(message, exit_code=1):
+    """Print error to stderr and exit."""
+    print(f"Error: {message}", file=sys.stderr)
+    sys.exit(exit_code)
+
+
+def _output(data):
+    """Print JSON result to stdout."""
+    print(json.dumps(data))
+    sys.exit(0)
+
+
+# ── Dependency checks ─────────────────────────────────────────────────────────
+
+
+def _check_sagemaker_core():
+    """Verify sagemaker-core is installed."""
+    try:
+        from sagemaker.core.resources import ProcessingJob  # noqa: F401
+    except ImportError:
+        _error_exit(
+            "sagemaker-core is not installed. "
+            "Please install: pip install 'sagemaker>=3.0.0' (includes sagemaker-core)"
+        )
+
+
+def _check_boto3():
+    """Verify boto3 is installed (needed for S3 entrypoint upload)."""
+    try:
+        import boto3  # noqa: F401
+    except ImportError:
+        _error_exit(
+            "boto3 is not installed. "
+            "Please install: pip install boto3"
+        )
+
+
+# ── Processing Job helpers ────────────────────────────────────────────────────
+
+
+def _generate_job_name(project_name, adapter_name):
+    """Generate a unique Processing Job name."""
+    timestamp = time.strftime("%Y%m%d-%H%M%S")
+    # Job names must be <= 63 chars, start with alphanumeric
+    base = f"mlcc-adapter-{project_name}-{adapter_name}"
+    # Truncate base to leave room for timestamp
+    max_base = 63 - len(timestamp) - 1
+    if len(base) > max_base:
+        base = base[:max_base]
+    return f"{base}-{timestamp}"
+
+
+def _upload_entrypoint(bucket, job_name, region):
+    """Upload the processing job entrypoint script to S3.
+
+    The entrypoint simply copies files from the Processing input path
+    to the Processing output path (SageMaker handles S3 download/upload).
+
+    Returns the S3 URI of the uploaded entrypoint.
+    """
+    import boto3
+
+    entrypoint_content = """#!/bin/bash
+set -e
+echo "Adapter staging: copying input to output..."
+echo "Input contents:"
+ls -la /opt/ml/processing/input/adapter/ || echo "No input files found"
+echo ""
+echo "Copying adapter files..."
+cp -r /opt/ml/processing/input/adapter/* /opt/ml/processing/output/ 2>/dev/null || \
+cp -r /opt/ml/processing/input/adapter/. /opt/ml/processing/output/
+echo "Output contents:"
+ls -la /opt/ml/processing/output/
+echo ""
+echo "Adapter staging complete."
+"""
+
+    s3_key = f"staging-jobs/{job_name}/entrypoint.sh"
+    s3_uri = f"s3://{bucket}/{s3_key}"
+
+    s3_client = boto3.client("s3", region_name=region)
+    try:
+        s3_client.put_object(
+            Bucket=bucket,
+            Key=s3_key,
+            Body=entrypoint_content.encode("utf-8"),
+            ContentType="text/x-shellscript",
+        )
+    except Exception as e:
+        _error_exit(f"Failed to upload entrypoint to S3: {e}")
+
+    return s3_uri
+
+
+def _resolve_container_image(region):
+    """Resolve the SageMaker-managed PyTorch CPU image URI for the region.
+
+    Uses the standard SageMaker DLC (Deep Learning Container) PyTorch CPU image
+    which includes AWS CLI and Python 3.10.
+    """
+    # SageMaker DLC account IDs per region
+    # https://docs.aws.amazon.com/sagemaker/latest/dg/ecr-us-east-1.html
+    dlc_accounts = {
+        "us-east-1": "763104351884",
+        "us-east-2": "763104351884",
+        "us-west-1": "763104351884",
+        "us-west-2": "763104351884",
+        "eu-west-1": "763104351884",
+        "eu-west-2": "763104351884",
+        "eu-central-1": "763104351884",
+        "ap-northeast-1": "763104351884",
+        "ap-southeast-1": "763104351884",
+        "ap-southeast-2": "763104351884",
+        "ap-south-1": "763104351884",
+        "ca-central-1": "763104351884",
+    }
+    account_id = dlc_accounts.get(region, "763104351884")
+    # Use PyTorch CPU processing image
+    return f"{account_id}.dkr.ecr.{region}.amazonaws.com/pytorch-training:2.2.0-cpu-py310-ubuntu20.04-sagemaker"
+
+
+# ── Subcommand: stage-from-tune ───────────────────────────────────────────────
+
+
+def cmd_stage_from_tune(args):
+    """Submit a Processing Job to copy adapter from training output to S3 adapter location.
+
+    Returns: {"job_name": str, "status": str, "adapter_s3_uri": str}
+    """
+    _check_sagemaker_core()
+    _check_boto3()
+
+    from sagemaker.core.resources import ProcessingJob
+
+    # Validate required arguments
+    if not args.training_output_s3_uri:
+        _error_exit("--training-output-s3-uri is required")
+    if not args.adapter_name:
+        _error_exit("--adapter-name is required")
+    if not args.bucket:
+        _error_exit("--bucket is required")
+    if not args.project:
+        _error_exit("--project is required")
+    if not args.role_arn:
+        _error_exit("--role-arn is required")
+
+    region = args.region or os.environ.get("AWS_DEFAULT_REGION") or os.environ.get("AWS_REGION", "us-west-2")
+    # Ensure region is set in env for sagemaker-core
+    os.environ["AWS_DEFAULT_REGION"] = region
+    os.environ.setdefault("AWS_REGION", region)
+
+    # Generate job name
+    job_name = _generate_job_name(args.project, args.adapter_name)
+
+    # Build adapter output S3 URI
+    adapter_s3_uri = f"s3://{args.bucket}/{args.project}/adapters/{args.adapter_name}/"
+
+    # Resolve container image
+    container_image = args.container_image or _resolve_container_image(region)
+
+    # Upload entrypoint script to S3
+    entrypoint_s3_uri = _upload_entrypoint(args.bucket, job_name, region)
+
+    # Build entrypoint command — download script from S3 then execute
+    entrypoint_cmd = (
+        f"aws s3 cp {entrypoint_s3_uri} /tmp/entrypoint.sh && "
+        "chmod +x /tmp/entrypoint.sh && /tmp/entrypoint.sh"
+    )
+
+    # Normalize training output S3 URI (ensure trailing slash for S3Prefix)
+    training_output_s3_uri = args.training_output_s3_uri
+    if not training_output_s3_uri.endswith("/"):
+        training_output_s3_uri += "/"
+
+    # Submit Processing Job via sagemaker-core
+    try:
+        job = ProcessingJob.create(
+            processing_job_name=job_name,
+            processing_resources={
+                "cluster_config": {
+                    "instance_count": 1,
+                    "instance_type": INSTANCE_TYPE,
+                    "volume_size_in_gb": VOLUME_SIZE_GB,
+                }
+            },
+            processing_inputs=[{
+                "input_name": "adapter",
+                "s3_input": {
+                    "s3_uri": training_output_s3_uri,
+                    "s3_data_type": "S3Prefix",
+                    "s3_input_mode": "File",
+                    "local_path": "/opt/ml/processing/input/adapter",
+                }
+            }],
+            processing_output_config={
+                "outputs": [{
+                    "output_name": "staged-adapter",
+                    "s3_output": {
+                        "s3_uri": adapter_s3_uri,
+                        "s3_upload_mode": "EndOfJob",
+                        "local_path": "/opt/ml/processing/output",
+                    }
+                }]
+            },
+            app_specification={
+                "image_uri": container_image,
+                "container_entrypoint": ["bash", "-c", entrypoint_cmd],
+            },
+            role_arn=args.role_arn,
+            stopping_condition={"max_runtime_in_seconds": MAX_RUNTIME_SECONDS},
+        )
+    except Exception as e:
+        error_msg = str(e)
+        if "AccessDeniedException" in error_msg or "AccessDenied" in error_msg:
+            _error_exit(
+                f"Access denied when creating Processing Job. "
+                f"Ensure the role has sagemaker:CreateProcessingJob permission. "
+                f"Details: {error_msg}"
+            )
+        elif "ResourceLimitExceeded" in error_msg:
+            _error_exit(
+                f"Resource limit exceeded. You may need to request a quota increase. "
+                f"Details: {error_msg}"
+            )
+        else:
+            _error_exit(f"Failed to create Processing Job: {error_msg}")
+
+    print(f"Processing Job submitted: {job_name}", file=sys.stderr)
+    print(f"Adapter output: {adapter_s3_uri}", file=sys.stderr)
+
+    # If --no-wait, return immediately
+    if args.no_wait:
+        _output({
+            "job_name": job_name,
+            "status": "InProgress",
+            "adapter_s3_uri": adapter_s3_uri,
+        })
+
+    # Poll until completion
+    print(f"Polling every {POLL_INTERVAL_SECONDS}s...", file=sys.stderr)
+    while True:
+        try:
+            job_desc = ProcessingJob.get(processing_job_name=job_name)
+            status = job_desc.processing_job_status
+        except Exception as e:
+            print(f"Warning: failed to get job status: {e}", file=sys.stderr)
+            time.sleep(POLL_INTERVAL_SECONDS)
+            continue
+
+        print(
+            f"  [{time.strftime('%H:%M:%S')}] Status: {status}",
+            file=sys.stderr,
+        )
+
+        if status in ("Completed", "Failed", "Stopped"):
+            break
+
+        time.sleep(POLL_INTERVAL_SECONDS)
+
+    # Handle terminal states
+    if status == "Failed":
+        failure_reason = getattr(job_desc, "failure_reason", None) or "Unknown failure"
+        print(f"Processing Job failed: {failure_reason}", file=sys.stderr)
+        sys.exit(1)
+
+    if status == "Stopped":
+        print("Processing Job was stopped.", file=sys.stderr)
+        sys.exit(1)
+
+    # Success
+    _output({
+        "job_name": job_name,
+        "status": "Completed",
+        "adapter_s3_uri": adapter_s3_uri,
+    })
+
+
+# ── Subcommand: status ────────────────────────────────────────────────────────
+
+
+def cmd_status(args):
+    """Check Processing Job status.
+
+    Returns: {"job_name": str, "status": str, "failure_reason": str|None}
+    """
+    _check_sagemaker_core()
+
+    from sagemaker.core.resources import ProcessingJob
+
+    if not args.job_name:
+        _error_exit("--job-name is required")
+
+    region = args.region or os.environ.get("AWS_DEFAULT_REGION") or os.environ.get("AWS_REGION", "us-west-2")
+    os.environ["AWS_DEFAULT_REGION"] = region
+    os.environ.setdefault("AWS_REGION", region)
+
+    try:
+        job_desc = ProcessingJob.get(processing_job_name=args.job_name)
+    except Exception as e:
+        error_msg = str(e)
+        if "does not exist" in error_msg or "ValidationException" in error_msg:
+            _error_exit(f"Processing Job not found: {args.job_name}")
+        else:
+            _error_exit(f"Failed to get Processing Job status: {error_msg}")
+
+    status = job_desc.processing_job_status
+    failure_reason = None
+
+    if status == "Failed":
+        failure_reason = getattr(job_desc, "failure_reason", None) or "Unknown failure"
+        print(f"Processing Job failed: {failure_reason}", file=sys.stderr)
+
+    _output({
+        "job_name": args.job_name,
+        "status": status,
+        "failure_reason": failure_reason,
+    })
+
+
+# ── Argument parsing ──────────────────────────────────────────────────────────
+
+
+def main():
+    """Parse arguments and dispatch to subcommand."""
+    parser = argparse.ArgumentParser(
+        description="SageMaker Processing Job helper for adapter staging",
+        prog=".adapter_helper.py",
+    )
+    subparsers = parser.add_subparsers(dest="subcommand", help="Subcommand")
+
+    # ── stage-from-tune ───────────────────────────────────────────────────
+    stage_parser = subparsers.add_parser(
+        "stage-from-tune",
+        help="Submit Processing Job to stage adapter from training output",
+    )
+    stage_parser.add_argument(
+        "--training-output-s3-uri",
+        required=True,
+        help="S3 URI of training output (adapter artifacts)",
+    )
+    stage_parser.add_argument(
+        "--adapter-name",
+        required=True,
+        help="Name of the adapter (used in output S3 path)",
+    )
+    stage_parser.add_argument(
+        "--bucket",
+        required=True,
+        help="S3 bucket for adapter output",
+    )
+    stage_parser.add_argument(
+        "--project",
+        required=True,
+        help="Project name (used in S3 path prefix)",
+    )
+    stage_parser.add_argument(
+        "--role-arn",
+        required=True,
+        help="SageMaker execution role ARN",
+    )
+    stage_parser.add_argument(
+        "--region",
+        default=None,
+        help="AWS region (default: from environment)",
+    )
+    stage_parser.add_argument(
+        "--container-image",
+        default=None,
+        help="Override container image URI (default: SageMaker PyTorch CPU image)",
+    )
+    stage_parser.add_argument(
+        "--no-wait",
+        action="store_true",
+        default=False,
+        help="Return immediately after submitting the job",
+    )
+
+    # ── status ────────────────────────────────────────────────────────────
+    status_parser = subparsers.add_parser(
+        "status",
+        help="Check Processing Job status",
+    )
+    status_parser.add_argument(
+        "--job-name",
+        required=True,
+        help="Processing Job name to check",
+    )
+    status_parser.add_argument(
+        "--region",
+        default=None,
+        help="AWS region (default: from environment)",
+    )
+
+    # ── Parse and dispatch ────────────────────────────────────────────────
+    args = parser.parse_args()
+
+    if not args.subcommand:
+        parser.print_help()
+        sys.exit(1)
+
+    if args.subcommand == "stage-from-tune":
+        cmd_stage_from_tune(args)
+    elif args.subcommand == "status":
+        cmd_status(args)
+    else:
+        _error_exit(f"Unknown subcommand: {args.subcommand}")
+
+
+if __name__ == "__main__":
+    main()

package/templates/do/.benchmark_writer.py

@@ -487,6 +487,7 @@ def enrich_records(config, results, run_timestamp=None):
             'mcc_version': mcc_version,
             'run_timestamp': run_timestamp.isoformat(),
             'region': region,
+            'adapter_name': config.get('adapter_name', ''),
         }
         records.append(record)
 
@@ -859,6 +860,7 @@ def get_parquet_schema():
         pa.field("mcc_version", pa.string()),
         pa.field("run_timestamp", pa.string()),
         pa.field("region", pa.string()),
+        pa.field("adapter_name", pa.string()),
     ])
 
 
@@ -1177,6 +1179,8 @@ def cmd_write(args):
         input_data['workload'] = args.workload
     if args.region:
         input_data['region'] = args.region
+    if args.adapter_name:
+        input_data['adapter_name'] = args.adapter_name
 
     # ── Validate before any S3 interaction ────────────────────────────────
     errors = validate_benchmark_input(input_data)
@@ -1385,6 +1389,7 @@ def _load_config_file(config_path):
                 shell_map = {
                     'PROJECT_NAME': 'project_name',
                     'MODEL_NAME': 'model_name',
+                    'HF_MODEL_ID': 'hf_model_id',
                     'INSTANCE_TYPE': 'instance_type',
                     'DEPLOYMENT_CONFIG': 'deployment_config',
                     'DEPLOYMENT_TARGET': 'deployment_target',
@@ -1402,6 +1407,18 @@ def _load_config_file(config_path):
     except Exception:
         pass
 
+    # Prefer HF_MODEL_ID over MODEL_NAME for the model_name field.
+    # After do/stage runs, MODEL_NAME is rewritten to an S3 URI which is
+    # unsuitable for S3 result paths (nested s3:// in path) and model family
+    # derivation.  HF_MODEL_ID preserves the original HuggingFace repo ID.
+    if context.get('hf_model_id'):
+        context['model_name'] = context.pop('hf_model_id')
+    elif context.get('model_name', '').startswith('s3://'):
+        # Fallback: if no HF_MODEL_ID but MODEL_NAME is an S3 URI, extract
+        # the model slug from the S3 path (last non-empty segment)
+        parts = context['model_name'].rstrip('/').split('/')
+        context['model_name'] = parts[-1] if parts else context['model_name']
+
     return context
 
 
@@ -1449,6 +1466,11 @@ def main():
         '--region',
         help='AWS region'
     )
+    write_parser.add_argument(
+        '--adapter-name', dest='adapter_name', default=None,
+        help='LoRA adapter name (differentiates adapter benchmarks from base model in Athena)'
+    )
+
     write_parser.add_argument(
         '--dry-run', dest='dry_run', action='store_true',
         help='Output enriched records as JSON without writing to S3'