@aws/ml-container-creator 1.0.3 → 1.0.4

package/templates/do/training/sft/train.py

@@ -0,0 +1,310 @@
+#!/usr/bin/env python3
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Unmanaged SFT fine-tuning using TRL SFTTrainer + PEFT LoRA.
+
+This script performs LoRA-based supervised fine-tuning on a causal language
+model. It is invoked via `accelerate launch` and works on both single-GPU
+and multi-GPU instances without code changes.
+
+Portable env-var contract (works on SageMaker AI and HyperPod EKS):
+    DATA_DIR / SM_CHANNEL_TRAINING   -> training data path
+    OUTPUT_DIR / SM_MODEL_DIR        -> model artifact output
+    CHECKPOINT_DIR / SM_CHECKPOINT_DIR -> checkpoint path for spot resume
+    HF_MODEL_ID / SM_HP_MODEL_ID     -> base model HuggingFace ID
+    SM_HPS                           -> JSON blob of all hyperparameters
+
+Output:
+    LoRA adapter saved to OUTPUT_DIR (adapter_model.safetensors + adapter_config.json)
+    Metrics logged to stdout in SageMaker-parseable format
+"""
+
+import glob
+import json
+import logging
+import os
+import sys
+
+# ── Logging ───────────────────────────────────────────────────────────────────
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+)
+logger = logging.getLogger("sft-trainer")
+
+# ── Portable Path Resolution ─────────────────────────────────────────────────
+# Fallback chain: generic env var -> SageMaker env var -> default path
+# This allows the same script to run on SageMaker, HyperPod, or locally.
+
+DATA_DIR = (
+    os.environ.get("DATA_DIR")
+    or os.environ.get("SM_CHANNEL_TRAINING")
+    or "/opt/ml/input/data/training"
+)
+
+OUTPUT_DIR = (
+    os.environ.get("OUTPUT_DIR")
+    or os.environ.get("SM_MODEL_DIR")
+    or "/opt/ml/model"
+)
+
+CHECKPOINT_DIR = (
+    os.environ.get("CHECKPOINT_DIR")
+    or os.environ.get("SM_CHECKPOINT_DIR")
+    or "/opt/ml/checkpoints"
+)
+
+MODEL_ID = (
+    os.environ.get("HF_MODEL_ID")
+    or os.environ.get("SM_HP_MODEL_ID")
+    or ""
+)
+
+
+# ── Hyperparameter Loading ────────────────────────────────────────────────────
+
+def load_hyperparameters():
+    """Load hyperparameters from SageMaker SM_HPS env var or individual SM_HP_* vars.
+
+    Returns:
+        dict with typed hyperparameter values.
+    """
+    defaults = {
+        "model_id": MODEL_ID,
+        "lora_r": 16,
+        "lora_alpha": 32,
+        "lora_dropout": 0.05,
+        "learning_rate": 2e-4,
+        "epochs": 3,
+        "batch_size": 4,
+        "max_seq_length": 2048,
+        "gradient_accumulation_steps": 4,
+        "warmup_ratio": 0.03,
+        "dataset_text_field": "text",
+    }
+
+    # Try SM_HPS (JSON blob of all hyperparameters)
+    sm_hps = os.environ.get("SM_HPS")
+    if sm_hps:
+        try:
+            raw = json.loads(sm_hps)
+            # SageMaker passes all values as strings — cast them
+            for key, default_val in defaults.items():
+                if key in raw:
+                    defaults[key] = _cast(raw[key], type(default_val))
+            return defaults
+        except (json.JSONDecodeError, ValueError) as e:
+            logger.warning("Failed to parse SM_HPS: %s", e)
+
+    # Fallback: individual SM_HP_* env vars
+    for key, default_val in defaults.items():
+        env_key = f"SM_HP_{key.upper()}"
+        env_val = os.environ.get(env_key)
+        if env_val is not None:
+            defaults[key] = _cast(env_val, type(default_val))
+
+    return defaults
+
+
+def _cast(value, target_type):
+    """Cast a string value to the target type."""
+    if target_type == bool:
+        return str(value).lower() in ("true", "1", "yes")
+    if target_type == int:
+        return int(float(value))
+    if target_type == float:
+        return float(value)
+    return str(value)
+
+
+# ── Dataset Loading ───────────────────────────────────────────────────────────
+
+def load_dataset(data_dir, text_field):
+    """Load training dataset from data directory.
+
+    Supports .jsonl, .parquet, and .csv files.
+
+    Args:
+        data_dir: Path to directory containing training data files.
+        text_field: Name of the text column in the dataset.
+
+    Returns:
+        A Hugging Face Dataset object.
+    """
+    from datasets import load_dataset as hf_load_dataset
+
+    # Find data files
+    extensions = ["jsonl", "json", "parquet", "csv"]
+    data_files = []
+    for ext in extensions:
+        data_files.extend(glob.glob(os.path.join(data_dir, f"*.{ext}")))
+        data_files.extend(glob.glob(os.path.join(data_dir, f"**/*.{ext}"), recursive=True))
+
+    if not data_files:
+        logger.error("No data files found in %s (searched: %s)", data_dir, extensions)
+        sys.exit(1)
+
+    # Deduplicate and sort
+    data_files = sorted(set(data_files))
+    logger.info("Found %d data file(s) in %s", len(data_files), data_dir)
+
+    # Determine format from first file extension
+    first_ext = data_files[0].rsplit(".", 1)[-1].lower()
+    format_map = {"jsonl": "json", "json": "json", "parquet": "parquet", "csv": "csv"}
+    file_format = format_map.get(first_ext, "json")
+
+    dataset = hf_load_dataset(file_format, data_files=data_files, split="train")
+    logger.info("Loaded dataset: %d rows, columns: %s", len(dataset), dataset.column_names)
+
+    # Verify text field exists
+    if text_field not in dataset.column_names:
+        logger.error(
+            "Text field '%s' not found in dataset. Available columns: %s",
+            text_field, dataset.column_names,
+        )
+        sys.exit(1)
+
+    return dataset
+
+
+# ── Main Training Function ────────────────────────────────────────────────────
+
+def main():
+    """Run SFT training with TRL SFTTrainer + PEFT LoRA."""
+    from accelerate import Accelerator
+    from peft import LoraConfig, TaskType
+    from transformers import (
+        AutoModelForCausalLM,
+        AutoTokenizer,
+        TrainingArguments,
+    )
+    from trl import SFTTrainer
+
+    # Initialize accelerator (handles distributed setup)
+    accelerator = Accelerator()
+
+    # Load hyperparameters
+    hparams = load_hyperparameters()
+    model_id = hparams["model_id"]
+
+    if not model_id:
+        logger.error("No model ID specified. Set HF_MODEL_ID env var or model_id hyperparameter.")
+        sys.exit(1)
+
+    if accelerator.is_main_process:
+        logger.info("=" * 60)
+        logger.info("SFT Training Configuration")
+        logger.info("=" * 60)
+        logger.info("  Model:          %s", model_id)
+        logger.info("  Data dir:       %s", DATA_DIR)
+        logger.info("  Output dir:     %s", OUTPUT_DIR)
+        logger.info("  Checkpoint dir: %s", CHECKPOINT_DIR)
+        logger.info("  LoRA r:         %d", hparams["lora_r"])
+        logger.info("  LoRA alpha:     %d", hparams["lora_alpha"])
+        logger.info("  Learning rate:  %s", hparams["learning_rate"])
+        logger.info("  Epochs:         %d", hparams["epochs"])
+        logger.info("  Batch size:     %d", hparams["batch_size"])
+        logger.info("  Max seq len:    %d", hparams["max_seq_length"])
+        logger.info("  Text field:     %s", hparams["dataset_text_field"])
+        logger.info("=" * 60)
+
+    # ── Load tokenizer and model ──────────────────────────────────────────────
+    if accelerator.is_main_process:
+        logger.info("Loading tokenizer and model: %s", model_id)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        torch_dtype="auto",
+        trust_remote_code=True,
+    )
+
+    # ── Configure LoRA ────────────────────────────────────────────────────────
+    lora_config = LoraConfig(
+        r=hparams["lora_r"],
+        lora_alpha=hparams["lora_alpha"],
+        lora_dropout=hparams["lora_dropout"],
+        target_modules="all-linear",
+        task_type=TaskType.CAUSAL_LM,
+    )
+
+    # ── Load dataset ──────────────────────────────────────────────────────────
+    dataset = load_dataset(DATA_DIR, hparams["dataset_text_field"])
+
+    # ── Training arguments ────────────────────────────────────────────────────
+    training_args = TrainingArguments(
+        output_dir=os.path.join(CHECKPOINT_DIR, "trainer-state"),
+        num_train_epochs=hparams["epochs"],
+        per_device_train_batch_size=hparams["batch_size"],
+        gradient_accumulation_steps=hparams["gradient_accumulation_steps"],
+        learning_rate=hparams["learning_rate"],
+        warmup_ratio=hparams["warmup_ratio"],
+        bf16=True,
+        logging_steps=10,
+        save_strategy="epoch",
+        save_total_limit=2,
+        report_to="none",
+        remove_unused_columns=False,
+    )
+
+    # ── Check for existing checkpoint (spot resume) ───────────────────────────
+    resume_from_checkpoint = None
+    trainer_state_dir = os.path.join(CHECKPOINT_DIR, "trainer-state")
+    if os.path.isdir(trainer_state_dir):
+        checkpoints = sorted(
+            glob.glob(os.path.join(trainer_state_dir, "checkpoint-*")),
+            key=lambda x: int(x.rsplit("-", 1)[-1]) if x.rsplit("-", 1)[-1].isdigit() else 0,
+        )
+        if checkpoints:
+            resume_from_checkpoint = checkpoints[-1]
+            if accelerator.is_main_process:
+                logger.info("Resuming from checkpoint: %s", resume_from_checkpoint)
+
+    # ── Initialize SFTTrainer ─────────────────────────────────────────────────
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=dataset,
+        peft_config=lora_config,
+        args=training_args,
+        max_seq_length=hparams["max_seq_length"],
+        dataset_text_field=hparams["dataset_text_field"],
+    )
+
+    # ── Train ─────────────────────────────────────────────────────────────────
+    if accelerator.is_main_process:
+        logger.info("Starting SFT training...")
+
+    train_result = trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+
+    # ── Save adapter (rank 0 only) ───────────────────────────────────────────
+    if accelerator.is_main_process:
+        logger.info("Saving LoRA adapter to: %s", OUTPUT_DIR)
+        trainer.save_model(OUTPUT_DIR)
+        tokenizer.save_pretrained(OUTPUT_DIR)
+
+        # ── Log final metrics ─────────────────────────────────────────────────
+        # Format: metric_name: value (SageMaker captures via regex in config.yaml)
+        metrics = train_result.metrics
+        print(f"train_loss: {metrics.get('train_loss', 0.0):.4f}")
+        print(f"train_runtime: {metrics.get('train_runtime', 0.0):.1f}")
+        print(f"train_samples_per_second: {metrics.get('train_samples_per_second', 0.0):.2f}")
+        print(f"epochs: {hparams['epochs']}")
+
+        logger.info("Training complete!")
+        logger.info("  Loss:     %.4f", metrics.get("train_loss", 0.0))
+        logger.info("  Runtime:  %.1fs", metrics.get("train_runtime", 0.0))
+        logger.info("  Samples/s: %.2f", metrics.get("train_samples_per_second", 0.0))
+
+    # Wait for all processes before exit
+    accelerator.wait_for_everyone()
+
+
+# ── Entry Point ───────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    main()

package/templates/do/tune

@@ -796,9 +796,16 @@ else:
 _validate_dataset() {
     local dataset="${ARG_DATASET}"
 
-    # ── Parse @v<N> version suffix (AC-2.1, AC-2.3) ──────────────────────────
+    # ── Parse @v<version> suffix (AC-2.1, AC-2.3) ───────────────────────────
     # Syntax: dataset-name@v2 → name="dataset-name", version ordinal=2
-    if [[ "${dataset}" =~ ^(.+)@v([0-9]+)$ ]]; then
+    #         dataset-name@v1.0.0 → name="dataset-name", version="1.0.0" (semver)
+    if [[ "${dataset}" =~ ^(.+)@v([0-9]+\.[0-9]+\.[0-9]+)$ ]]; then
+        # Semver syntax: @v1.0.0
+        ARG_DATASET_NAME="${BASH_REMATCH[1]}"
+        ARG_DATASET_VERSION="${BASH_REMATCH[2]}"
+        dataset=""  # Clear so name-based resolution takes over
+    elif [[ "${dataset}" =~ ^(.+)@v([0-9]+)$ ]]; then
+        # Ordinal syntax: @v1, @v2
         ARG_DATASET_NAME="${BASH_REMATCH[1]}"
         ARG_DATASET_VERSION="${BASH_REMATCH[2]}"
         dataset=""  # Clear so name-based resolution takes over
@@ -824,6 +831,8 @@ _validate_dataset() {
                 echo "   Resolved to: ${resolved_uri}"
                 dataset="${resolved_uri}"
                 ARG_DATASET="${resolved_uri}"
+                RESOLVED_DATASET_S3_URI="${resolved_uri}"
+                return 0
             else
                 echo "❌ Dataset '${ARG_DATASET_NAME}' not found in registry"
                 echo "   Register it first: ./do/register --dataset --dataset-name ${ARG_DATASET_NAME} --dataset-s3-uri s3://..."

package/templates/do/.train_poll_parser.py

@@ -1,135 +0,0 @@
-#!/usr/bin/env python3
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-"""
-Parse DescribeTrainingJob JSON for the polling loop in do/train.
-
-Reads JSON from stdin and outputs structured key=value lines for bash consumption:
-  STATUS=<TrainingJobStatus>
-  SECONDARY=<SecondaryStatus>
-  FAILURE_REASON=<FailureReason or empty>
-  DISPLAY=<formatted single-line status display>
-
-This keeps the bash poll loop simple while handling JSON parsing in Python.
-"""
-
-import json
-import sys
-from datetime import datetime, timezone
-
-
-def format_duration(seconds):
-    """Format seconds into a human-readable duration string."""
-    if seconds is None or seconds < 0:
-        return 'N/A'
-    hours = int(seconds // 3600)
-    minutes = int((seconds % 3600) // 60)
-    secs = int(seconds % 60)
-    if hours > 0:
-        return f'{hours}h {minutes}m {secs}s'
-    elif minutes > 0:
-        return f'{minutes}m {secs}s'
-    else:
-        return f'{secs}s'
-
-
-def parse_iso_time(time_str):
-    """Parse an ISO 8601 timestamp string to a datetime object."""
-    if not time_str:
-        return None
-    try:
-        time_str = time_str.replace('Z', '+00:00')
-        return datetime.fromisoformat(time_str)
-    except (ValueError, TypeError):
-        return None
-
-
-def calculate_elapsed(start_time_str):
-    """Calculate elapsed time from start to now."""
-    start = parse_iso_time(start_time_str)
-    if not start:
-        return None
-    now = datetime.now(timezone.utc)
-    elapsed = (now - start).total_seconds()
-    return max(0, elapsed)
-
-
-def format_metrics(final_metrics):
-    """Format FinalMetricDataList into a compact string."""
-    if not final_metrics:
-        return ''
-    parts = []
-    for metric in final_metrics:
-        name = metric.get('MetricName', 'unknown')
-        value = metric.get('Value', 0)
-        if isinstance(value, float):
-            if abs(value) < 0.001:
-                parts.append(f'{name}={value:.6f}')
-            elif abs(value) < 1:
-                parts.append(f'{name}={value:.4f}')
-            else:
-                parts.append(f'{name}={value:.2f}')
-        else:
-            parts.append(f'{name}={value}')
-    return ', '.join(parts)
-
-
-# Status emoji mapping
-STATUS_EMOJI = {
-    'InProgress': '🔄',
-    'Completed': '✅',
-    'Failed': '❌',
-    'Stopping': '⏸️',
-    'Stopped': '⏹️'
-}
-
-
-def main():
-    """Parse DescribeTrainingJob JSON from stdin and output structured lines."""
-    try:
-        job_data = json.load(sys.stdin)
-    except json.JSONDecodeError as e:
-        print(f'Error parsing JSON: {e}', file=sys.stderr)
-        sys.exit(1)
-
-    status = job_data.get('TrainingJobStatus', 'Unknown')
-    secondary_status = job_data.get('SecondaryStatus', '')
-    failure_reason = job_data.get('FailureReason', '')
-    training_start = job_data.get('TrainingStartTime', '')
-    final_metrics = job_data.get('FinalMetricDataList', [])
-
-    # Calculate elapsed time
-    elapsed_str = ''
-    if training_start:
-        elapsed = calculate_elapsed(training_start)
-        if elapsed is not None:
-            elapsed_str = format_duration(elapsed)
-
-    # Format metrics
-    metrics_str = format_metrics(final_metrics)
-
-    # Build display line
-    emoji = STATUS_EMOJI.get(status, '❓')
-    display_parts = [f'   {emoji} {status}']
-
-    if secondary_status:
-        display_parts.append(f'| {secondary_status}')
-
-    if elapsed_str:
-        display_parts.append(f'| elapsed: {elapsed_str}')
-
-    if metrics_str:
-        display_parts.append(f'| {metrics_str}')
-
-    display_line = ' '.join(display_parts)
-
-    # Output structured lines for bash
-    print(f'STATUS={status}')
-    print(f'SECONDARY={secondary_status}')
-    print(f'FAILURE_REASON={failure_reason}')
-    print(f'DISPLAY={display_line}')
-
-
-if __name__ == '__main__':
-    main()

package/templates/do/.train_status_parser.py

@@ -1,187 +0,0 @@
-#!/usr/bin/env python3
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-"""
-Parse DescribeTrainingJob JSON response and display formatted status.
-
-This helper is called by do/train --status to parse the AWS CLI JSON output
-from DescribeTrainingJob and display a user-friendly status summary.
-"""
-
-import json
-import sys
-import time
-from datetime import datetime, timezone
-
-
-# Status emoji mapping
-STATUS_EMOJI = {
-    'InProgress': '🔄',
-    'Completed': '✅',
-    'Failed': '❌',
-    'Stopping': '⏸️',
-    'Stopped': '⏹️'
-}
-
-# Secondary status descriptions
-SECONDARY_DESCRIPTIONS = {
-    'Starting': 'Preparing training instance',
-    'LaunchingMLInstances': 'Launching ML instances',
-    'PreparingTrainingStack': 'Preparing training stack',
-    'Downloading': 'Downloading training data',
-    'DownloadingTrainingImage': 'Downloading training image',
-    'Training': 'Training in progress',
-    'Uploading': 'Uploading model artifacts',
-    'Completed': 'Training completed',
-    'MaxRuntimeExceeded': 'Max runtime exceeded',
-    'Stopped': 'Training stopped',
-    'MaxWaitTimeExceeded': 'Max wait time exceeded (spot)',
-    'Interrupted': 'Spot instance interrupted'
-}
-
-
-def format_duration(seconds):
-    """Format seconds into a human-readable duration string."""
-    if seconds is None or seconds < 0:
-        return 'N/A'
-    hours = int(seconds // 3600)
-    minutes = int((seconds % 3600) // 60)
-    secs = int(seconds % 60)
-    if hours > 0:
-        return f'{hours}h {minutes}m {secs}s'
-    elif minutes > 0:
-        return f'{minutes}m {secs}s'
-    else:
-        return f'{secs}s'
-
-
-def parse_iso_time(time_str):
-    """Parse an ISO 8601 timestamp string to a datetime object."""
-    if not time_str:
-        return None
-    try:
-        # Handle various AWS timestamp formats
-        # Remove trailing 'Z' and replace with +00:00 for fromisoformat
-        time_str = time_str.replace('Z', '+00:00')
-        return datetime.fromisoformat(time_str)
-    except (ValueError, TypeError):
-        return None
-
-
-def calculate_elapsed(start_time_str):
-    """Calculate elapsed time from start to now."""
-    start = parse_iso_time(start_time_str)
-    if not start:
-        return None
-    now = datetime.now(timezone.utc)
-    elapsed = (now - start).total_seconds()
-    return max(0, elapsed)
-
-
-def display_status(job_data):
-    """Display formatted training job status."""
-    job_name = job_data.get('TrainingJobName', 'Unknown')
-    status = job_data.get('TrainingJobStatus', 'Unknown')
-    secondary_status = job_data.get('SecondaryStatus', '')
-    failure_reason = job_data.get('FailureReason', '')
-    training_start = job_data.get('TrainingStartTime', '')
-    training_end = job_data.get('TrainingEndTime', '')
-    billable_seconds = job_data.get('BillableTimeInSeconds')
-    training_seconds = job_data.get('TrainingTimeInSeconds')
-    final_metrics = job_data.get('FinalMetricDataList', [])
-    output_path = job_data.get('OutputDataConfig', {}).get('S3OutputPath', '')
-    model_artifacts = job_data.get('ModelArtifacts', {}).get('S3ModelArtifacts', '')
-    instance_type = job_data.get('ResourceConfig', {}).get('InstanceType', '')
-    instance_count = job_data.get('ResourceConfig', {}).get('InstanceCount', 1)
-    spot_enabled = job_data.get('EnableManagedSpotTraining', False)
-
-    emoji = STATUS_EMOJI.get(status, '❓')
-
-    print(f'')
-    print(f'   {emoji} Status: {status}')
-
-    # Secondary status with description
-    if secondary_status:
-        desc = SECONDARY_DESCRIPTIONS.get(secondary_status, '')
-        if desc:
-            print(f'   📍 Phase:  {secondary_status} ({desc})')
-        else:
-            print(f'   📍 Phase:  {secondary_status}')
-
-    # Elapsed time
-    if status == 'InProgress' and training_start:
-        elapsed = calculate_elapsed(training_start)
-        if elapsed is not None:
-            print(f'   ⏱️  Elapsed: {format_duration(elapsed)}')
-    elif training_seconds is not None:
-        print(f'   ⏱️  Training time: {format_duration(training_seconds)}')
-
-    # Instance info
-    if instance_type:
-        instance_info = f'{instance_type}'
-        if instance_count and instance_count > 1:
-            instance_info += f' x {instance_count}'
-        if spot_enabled:
-            instance_info += ' (spot)'
-        print(f'   🖥️  Instance: {instance_info}')
-
-    # Billable time and cost savings (for completed spot jobs)
-    if status == 'Completed' and spot_enabled and billable_seconds is not None and training_seconds is not None:
-        savings_seconds = training_seconds - billable_seconds
-        if training_seconds > 0:
-            savings_pct = (savings_seconds / training_seconds) * 100
-            print(f'   💰 Spot savings: {format_duration(savings_seconds)} saved ({savings_pct:.0f}% discount)')
-            print(f'      Billable: {format_duration(billable_seconds)} / Total: {format_duration(training_seconds)}')
-
-    # Training metrics
-    if final_metrics:
-        print(f'   📈 Metrics:')
-        for metric in final_metrics:
-            name = metric.get('MetricName', 'unknown')
-            value = metric.get('Value', 0)
-            # Format value nicely
-            if isinstance(value, float):
-                if abs(value) < 0.001:
-                    print(f'      {name}: {value:.6f}')
-                elif abs(value) < 1:
-                    print(f'      {name}: {value:.4f}')
-                else:
-                    print(f'      {name}: {value:.2f}')
-            else:
-                print(f'      {name}: {value}')
-
-    # Output artifacts (for completed jobs)
-    if status == 'Completed' and model_artifacts:
-        print(f'   📦 Artifacts: {model_artifacts}')
-    elif status == 'Completed' and output_path:
-        print(f'   📦 Output: {output_path}')
-
-    # Failure reason
-    if status == 'Failed' and failure_reason:
-        print(f'   💥 Reason: {failure_reason}')
-        print(f'')
-        print(f'   To start a new job: ./do/train --force')
-
-    # Spot interruption guidance
-    if secondary_status == 'Interrupted':
-        print(f'')
-        print(f'   ℹ️  Spot instance was interrupted. The job will automatically')
-        print(f'      resume from the last checkpoint. Re-run ./do/train to poll.')
-
-    print(f'')
-
-
-def main():
-    """Main entry point — reads JSON from stdin."""
-    try:
-        job_data = json.load(sys.stdin)
-    except json.JSONDecodeError as e:
-        print(f'❌ Failed to parse DescribeTrainingJob response: {e}', file=sys.stderr)
-        sys.exit(1)
-
-    display_status(job_data)
-
-
-if __name__ == '__main__':
-    main()