Enh(aws): Disk usage management

Author: klieret

aws/setup/batch/environment.json

@@ -31,6 +31,9 @@
       "sg-014723b4566a8214d"
     ],
     "instanceRole": "arn:aws:iam::039984708918:instance-profile/kilian-codeclash-ecsInstanceRole",
+    "launchTemplate": {
+      "launchTemplateName": "kilian-codeclash-launch-template"
+    },
     "ec2Configuration": [
       {
         "imageType": "ECS_AL2"

aws/setup/batch/iam-ebs.json

@@ -0,0 +1,30 @@
+{
+  "Role": {
+    "Path": "/",
+    "RoleName": "kilian-codeclash-ecs-role-for-ebs-volumes",
+    "AssumeRolePolicyDocument": {
+      "Version": "2012-10-17",
+      "Statement": [
+        {
+          "Effect": "Allow",
+          "Principal": {
+            "Service": "ecs.amazonaws.com"
+          },
+          "Action": "sts:AssumeRole"
+        }
+      ]
+    },
+    "Description": "ECS Infrastructure role for managing EBS volumes",
+    "MaxSessionDuration": 3600,
+    "Tags": [
+      {
+        "Key": "project",
+        "Value": "codeclash"
+      },
+      {
+        "Key": "author",
+        "Value": "kilian"
+      }
+    ]
+  }
+}

aws/setup/batch/launch_template.json

@@ -0,0 +1,29 @@
+{
+  "LaunchTemplateName": "kilian-codeclash-launch-template",
+  "LaunchTemplateData": {
+    "BlockDeviceMappings": [
+      {
+        "DeviceName": "/dev/xvda",
+        "Ebs": {
+          "VolumeSize": 50,
+          "VolumeType": "gp3"
+        }
+      }
+    ]
+  },
+  "TagSpecifications": [
+    {
+      "ResourceType": "launch-template",
+      "Tags": [
+        {
+          "Key": "project",
+          "Value": "codeclash"
+        },
+        {
+          "Key": "author",
+          "Value": "kilian"
+        }
+      ]
+    }
+  ]
+}

aws/setup/batch/pull.py

@@ -46,6 +46,36 @@ def pull_job_definition(job_name: str, batch_client) -> dict:
     return job_data
 
 
+def pull_launch_template(template_name: str, ec2_client) -> dict:
+    """Pull EC2 launch template from AWS."""
+    response = ec2_client.describe_launch_templates(LaunchTemplateNames=[template_name])
+    template_data = response["LaunchTemplates"][0]
+
+    # Get the latest version of the template
+    version_response = ec2_client.describe_launch_template_versions(
+        LaunchTemplateId=template_data["LaunchTemplateId"], Versions=["$Latest"]
+    )
+    template_version = version_response["LaunchTemplateVersions"][0]
+
+    # Combine template metadata with version data
+    result = {
+        "LaunchTemplateName": template_data["LaunchTemplateName"],
+        "LaunchTemplateData": template_version["LaunchTemplateData"],
+    }
+
+    # Add tags if they exist
+    if template_data.get("Tags"):
+        result["TagSpecifications"] = [
+            {
+                "ResourceType": "launch-template",
+                "Tags": [{"Key": tag["Key"], "Value": tag["Value"]} for tag in template_data["Tags"]],
+            }
+        ]
+
+    logger.info(f"✅ Successfully pulled launch template: {template_name}")
+    return result
+
+
 def clean_response_data(data: dict, is_iam_role: bool) -> dict:
     """Clean AWS response data to remove metadata and format for JSON storage."""
     if is_iam_role:
@@ -55,6 +85,11 @@ def clean_response_data(data: dict, is_iam_role: bool) -> dict:
     cleaned.pop("ResponseMetadata", None)
     cleaned.pop("createdAt", None)
     cleaned.pop("schedulingPriority", None)
+    # Remove AWS-specific metadata fields for job definitions
+    cleaned.pop("jobDefinitionArn", None)
+    cleaned.pop("revision", None)
+    cleaned.pop("status", None)
+    cleaned.pop("containerOrchestrationType", None)
     return cleaned
 
 
@@ -72,15 +107,21 @@ def pull_file(json_path: Path, region: str) -> None:
         case "iam-job-role.json":
             data = pull_iam_role("kilian-codeclash-job-role", boto3.client("iam", region_name=region))
             cleaned_data = clean_response_data(data, is_iam_role=True)
+        case "iam-ebs.json":
+            data = pull_iam_role("kilian-codeclash-ecs-role-for-ebs-volumes", boto3.client("iam", region_name=region))
+            cleaned_data = clean_response_data(data, is_iam_role=True)
         case "environment.json":
             data = pull_compute_environment("codeclash-batch", boto3.client("batch", region_name=region))
             cleaned_data = clean_response_data(data, is_iam_role=False)
         case "job_queue.json":
             data = pull_job_queue("codeclash-queue", boto3.client("batch", region_name=region))
             cleaned_data = clean_response_data(data, is_iam_role=False)
         case "job_definition.json":
-            data = pull_job_definition("codeclash-default-job", boto3.client("batch", region_name=region))
+            data = pull_job_definition("codeclash-ebs-20g", boto3.client("batch", region_name=region))
             cleaned_data = clean_response_data(data, is_iam_role=False)
+        case "launch_template.json":
+            data = pull_launch_template("kilian-codeclash-launch-template", boto3.client("ec2", region_name=region))
+            cleaned_data = data  # No cleaning needed for launch templates
         case _:
             logger.error(f"Unknown filename: {filename}")
             raise ValueError(f"Unknown filename: {filename}")

aws/setup/batch/push.py

@@ -8,6 +8,7 @@
 from typing import Any
 
 import boto3
+from botocore.exceptions import ClientError
 
 logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
 logger = logging.getLogger(__name__)
@@ -173,10 +174,15 @@ def push_job_definition(job_data: dict, batch_client) -> None:
     register_params = {
         "jobDefinitionName": job_name,
         "type": job_data.get("type", "container"),
-        "containerProperties": job_data["containerProperties"],
         "platformCapabilities": job_data.get("platformCapabilities", ["EC2"]),
     }
 
+    # Handle both old containerProperties and new ecsProperties formats
+    if "ecsProperties" in job_data:
+        register_params["ecsProperties"] = job_data["ecsProperties"]
+    elif "containerProperties" in job_data:
+        register_params["containerProperties"] = job_data["containerProperties"]
+
     if "parameters" in job_data:
         register_params["parameters"] = job_data["parameters"]
     if "tags" in job_data:
@@ -188,20 +194,70 @@ def push_job_definition(job_data: dict, batch_client) -> None:
     logger.info(f"✅ Successfully pushed job definition: {job_name}:{revision}")
 
 
+def _update_launch_template(template_data: dict, existing_template: dict, ec2_client) -> None:
+    """Update existing launch template by creating a new version."""
+    template_name = template_data["LaunchTemplateName"]
+    logger.info(f"Launch template {template_name} already exists, creating new version")
+
+    create_params = {
+        "LaunchTemplateId": existing_template["LaunchTemplateId"],
+        "LaunchTemplateData": template_data["LaunchTemplateData"],
+    }
+
+    response = ec2_client.create_launch_template_version(**create_params)
+    version = response["LaunchTemplateVersion"]["VersionNumber"]
+
+    logger.info(f"✅ Successfully created new version {version} for launch template: {template_name}")
+
+
+def _create_launch_template(template_data: dict, ec2_client) -> None:
+    """Create new launch template."""
+    template_name = template_data["LaunchTemplateName"]
+    logger.info(f"Creating new launch template: {template_name}")
+
+    create_params = {
+        "LaunchTemplateName": template_name,
+        "LaunchTemplateData": template_data["LaunchTemplateData"],
+    }
+
+    if "TagSpecifications" in template_data:
+        create_params["TagSpecifications"] = template_data["TagSpecifications"]
+
+    ec2_client.create_launch_template(**create_params)
+    logger.info(f"✅ Successfully created launch template: {template_name}")
+
+
+def push_launch_template(template_data: dict, ec2_client) -> None:
+    """Push EC2 launch template to AWS."""
+    template_name = template_data["LaunchTemplateName"]
+
+    try:
+        response = ec2_client.describe_launch_templates(LaunchTemplateNames=[template_name])
+        existing_template = response["LaunchTemplates"][0]
+        _update_launch_template(template_data, existing_template, ec2_client)
+    except ClientError as e:
+        if e.response["Error"]["Code"] == "InvalidLaunchTemplateName.NotFoundException":
+            _create_launch_template(template_data, ec2_client)
+        else:
+            raise
+
+
 def push_file(json_path: Path, region: str) -> None:
     """Push a single JSON file to AWS."""
     data = json.loads(json_path.read_text())
     filename = json_path.name
 
     match filename:
-        case "iam-environment-role.json" | "iam-execution-role.json" | "iam-job-role.json":
+        case "iam-environment-role.json" | "iam-execution-role.json" | "iam-job-role.json" | "iam-ebs.json":
             push_iam_role(data, boto3.client("iam", region_name=region))
         case "environment.json":
             push_compute_environment(data, boto3.client("batch", region_name=region))
         case "job_queue.json":
             push_job_queue(data, boto3.client("batch", region_name=region))
         case "job_definition.json":
             push_job_definition(data, boto3.client("batch", region_name=region))
+        case "launch_template.json":
+            push_launch_template(data, boto3.client("ec2", region_name=region))
         case _:
             logger.error(f"Unknown filename: {filename}")
             raise ValueError(f"Unknown filename: {filename}")

aws/setup/docker/docker_entrypoint.sh

@@ -28,6 +28,9 @@ echo "  AWS_BATCH_JOB_ATTEMPT: ${AWS_BATCH_JOB_ATTEMPT:-<not set>}"
 echo "  AWS_BATCH_JOB_ID: ${AWS_BATCH_JOB_ID:-<not set>}"
 echo "  AWS_BATCH_JQ_NAME: ${AWS_BATCH_JQ_NAME:-<not set>}"
 
+echo "💾 Disk usage before job:"
+df -h
+
 source .venv/bin/activate
 
 echo "═══════════════════════════════════════════════════════════════════════════════"
@@ -36,6 +39,13 @@ echo "$@"
 echo "═══════════════════════════════════════════════════════════════════════════════"
 # Execute the remaining command arguments
 exec "$@"
+echo "═══════════════════════════════════════════════════════════════════════════════"
+echo "✅ User provided command finished"
+echo "═══════════════════════════════════════════════════════════════════════════════"
+
+echo "💾 Disk usage after job:"
+df -h
+
 echo "═══════════════════════════════════════════════════════════════════════════════"
 echo "✅ docker_entrypoint.sh finished"
 echo "═══════════════════════════════════════════════════════════════════════════════"