Add guaranteed instance termination on all exit scenarios

[tomtao57]

Problem

Previously, EC2/GCE instances would only terminate if the launch script completed successfully. If any error occurred during setup, upload, training, or post-processing, the script would exit early and never reach the termination commands, leaving instances running indefinitely and incurring unnecessary costs.

Solution

Implemented bash trap handlers that guarantee instance termination regardless of how the script exits.

How It Works

Added a cleanup() function containing termination commands and set a trap to call it on EXIT, ERR, INT, and TERM signals:

# Define cleanup function for guaranteed termination
cleanup() {
    EXIT_CODE=$?
    echo "Cleanup triggered (exit code: $EXIT_CODE)"
    # Termination commands (with configured delay)
    aws ec2 terminate-instances --instance-ids $EC2_INSTANCE_ID --region $REGION
}
# Set trap to call cleanup on EXIT, ERR, INT, and TERM
trap cleanup EXIT ERR INT TERM

Termination Guaranteed On

• ✅ Normal completion (exit 0)
• ✅ Command failures (exit 1)
• ✅ Command not found errors (exit 127)
• ✅ Python/training script errors
• ✅ Missing dependencies
• ✅ File not found errors
• ✅ Script interruptions (Ctrl+C)
• ✅ Kill signals (SIGTERM)
• ✅ Any other exit scenario

Changes

jaynes/launchers/base_launcher.py

• Modified make_launch_script() to generate cleanup trap when terminate_after=True
• Trap is only added when termination is configured
• Supports both EC2 and GCE termination
• Cleanup includes configured delay before termination

Testing

Unit Tests: test_termination_trap.py

• ✅ EC2 termination trap generation
• ✅ GCE termination trap generation
• ✅ No trap when terminate_after=False

Integration Tests: test_error_scenarios.sh

Tests various failure scenarios to verify termination:

• ✅ Normal success termination
• ✅ Setup failure termination
• ✅ Command not found termination
• ✅ Python error termination
• ✅ Missing dependency termination
• ✅ File not found termination

All tests pass! 🎉

Benefits

Cost Safety

Prevents runaway costs from instances stuck due to:

• Configuration errors
• Missing dependencies
• Training script failures
• Network issues during setup
• User interruptions

Reliability

Ensures instances always terminate, even when things go wrong, providing peace of mind for long-running jobs and large-scale deployments.

Example Generated Script

Click to see example EC2 launch script with trap

#!/bin/bash
set +o posix

mkdir -p /tmp/jaynes-launch
JAYNES_LAUNCH_DIR=/tmp/jaynes-launch

# Define cleanup function for guaranteed termination
cleanup() {
    EXIT_CODE=$?
    echo "Cleanup triggered (exit code: $EXIT_CODE)"
    sleep 60
    die() { status=$1; shift; echo "FATAL: $*"; exit $status; }
    echo "Now terminate this instance"
    export REGION="$(wget -q -O - http://169.254.169.254/latest/meta-data/placement/region)"
    export EC2_INSTANCE_ID="`wget -q -O - http://169.254.169.254/latest/meta-data/instance-id || die "wget instance-id has failed: $?"`"
    aws ec2 terminate-instances --instance-ids $EC2_INSTANCE_ID --region $REGION
}
# Set trap to call cleanup on EXIT, ERR, INT, and TERM
trap cleanup EXIT ERR INT TERM

{
    # launch.setup script
    # ... (user scripts here)
    # run script
    python train.py
    # post script
    # ... (post-processing here)
} > >(tee -a /tmp/jaynes-launch/jaynes-launch.log) 2> >(tee -a /tmp/jaynes-launch/jaynes-launch.err.log >&2)

Backward Compatibility

✅ Fully backward compatible - Only affects instances with terminate_after: true configured in .jaynes.yml. Existing configurations without auto-termination are unchanged.

🤖 Generated with Claude Code