Add TTL configuration for AWS instance management

- Introduced a new config module to manage TTL settings for EC2 instances, allowing for auto-termination based on environment variables.
- Updated the AWSProvider and manager to utilize the new TTL settings, including scheduling instance termination via EventBridge Scheduler.
- Added utility functions for resolving the scheduler role ARN and creating termination schedules, ensuring robust error handling and logging.
- Maintained existing code logic while integrating new features for improved instance lifecycle management.

desktop_env/providers/aws/manager.py

@@ -1,9 +1,13 @@
 import os
-from filelock import FileLock
 import boto3
 import logging
 import dotenv
 import signal
+from datetime import datetime, timedelta, timezone
+
+# TTL configuration
+from desktop_env.providers.aws.config import ENABLE_TTL, DEFAULT_TTL_MINUTES, AWS_SCHEDULER_ROLE_ARN
+from desktop_env.providers.aws.scheduler_utils import schedule_instance_termination
 
 
 INSTANCE_TYPE = "t3.medium" 
@@ -92,12 +96,20 @@ def _allocate_vm(region=DEFAULT_REGION, screen_size=(1920, 1080)):
         if not os.getenv('AWS_SUBNET_ID'):
             raise ValueError("AWS_SUBNET_ID is not set in the environment variables.")
 
+        # TTL configuration (cloud-init removed; use cloud-side scheduler only)
+        ttl_enabled = ENABLE_TTL
+        ttl_minutes = DEFAULT_TTL_MINUTES
+        ttl_seconds = max(0, int(ttl_minutes) * 60)
+        eta_utc = datetime.now(timezone.utc) + timedelta(seconds=ttl_seconds)
+        logger.info(f"TTL config: minutes={ttl_minutes}, seconds={ttl_seconds}, ETA(UTC)={eta_utc.isoformat()}")
+
         run_instances_params = {
             "MaxCount": 1,
             "MinCount": 1,
             "ImageId": ami_id,
             "InstanceType": INSTANCE_TYPE,
             "EbsOptimized": True,
+            "InstanceInitiatedShutdownBehavior": "terminate",
             "NetworkInterfaces": [
                 {
                     "SubnetId": os.getenv('AWS_SUBNET_ID'),
@@ -124,12 +136,20 @@ def _allocate_vm(region=DEFAULT_REGION, screen_size=(1920, 1080)):
         
         response = ec2_client.run_instances(**run_instances_params)
         instance_id = response['Instances'][0]['InstanceId']
-        
+
+        # Create TTL schedule immediately after instance is created, to survive early interruptions
+        try:
+            # Always attempt; helper resolves ARN via env or role name
+            if ttl_enabled:
+                schedule_instance_termination(region, instance_id, ttl_seconds, AWS_SCHEDULER_ROLE_ARN, logger)
+        except Exception as e:
+            logger.warning(f"Failed to create EventBridge Scheduler for {instance_id}: {e}")
+
         waiter = ec2_client.get_waiter('instance_running')
         logger.info(f"Waiting for instance {instance_id} to be running...")
         waiter.wait(InstanceIds=[instance_id])
         logger.info(f"Instance {instance_id} is ready.")
-        
+
         try:
             instance_details = ec2_client.describe_instances(InstanceIds=[instance_id])
             instance = instance_details['Reservations'][0]['Instances'][0]