Per-job Instance Types and AMIs

CHANGELOG.md

@@ -5,6 +5,13 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [7.8.0]
+
+### Added
+- Allow overriding certain AWS Batch compute environment parameters (including instance types and AMI) within a job spec.
+- Allow job spec tasks to require GPU resource requirements.
+
+
 ## [7.7.2]
 
 ### Change

apps/compute-cf.yml.j2

@@ -29,6 +29,16 @@ Outputs:
   JobQueueArn:
     Value: !Ref BatchJobQueue
 
+  {% for job_type, job_spec in job_types.items()  if 'Default' != job_spec['compute_environment']['name'] %}
+  {% set name = job_spec['compute_environment']['name'] %}
+  {{ name }}ComputeEnvironmentArn:
+    Value: !Ref {{ name }}ComputeEnvironment
+
+  {{ name }}JobQueueArn:
+    Value: !Ref {{ name }}JobQueue
+
+  {% endfor %}
+
   TaskRoleArn:
     Value: !GetAtt TaskRole.Arn
 
@@ -85,17 +95,60 @@ Resources:
         Tags:
           Name: !Ref AWS::StackName
 
+  BatchJobQueue:
+    Type: AWS::Batch::JobQueue
+    Properties:
+      Priority: 1
+      ComputeEnvironmentOrder:
+        - ComputeEnvironment: !Ref ComputeEnvironment
+          Order: 1
+      SchedulingPolicyArn: !Ref SchedulingPolicy
+
   SchedulingPolicy:
     Type: AWS::Batch::SchedulingPolicy
 
-  BatchJobQueue:
+  {% for job_type, job_spec in job_types.items() if 'Default' != job_spec['compute_environment']['name'] %}
+  {% set env = job_spec['compute_environment'] %}
+  {% set name = env['name'] %}
+  {% set instance_types = env['instance_types'] if 'instance_types' in env else ['!Ref InstanceTypes'] %}
+  {% set ami_id = env['ami_id'] if 'ami_id' in env else '!Ref AmiId' %}
+  {% set type = env['allocation_type'] if 'allocation_type' in env else 'SPOT' %}
+  {% set strategy = env['allocation_strategy'] if 'allocation_strategy' in env else 'SPOT_PRICE_CAPACITY_OPTIMIZED' %}
+  {{ name }}ComputeEnvironment:
+    Type: AWS::Batch::ComputeEnvironment
+    Properties:
+      ServiceRole: !GetAtt BatchServiceRole.Arn
+      Type: MANAGED
+      ComputeResources:
+        Type: {{ type }}
+        AllocationStrategy: {{ strategy }}
+        MinvCpus: 0
+        MaxvCpus: !Ref MaxvCpus
+        InstanceTypes:
+        {% for instance_type in instance_types %}
+          - {{ instance_type }}
+        {% endfor %}
+        ImageId: {{ ami_id }}
+        Subnets: !Ref SubnetIds
+        InstanceRole: !Ref InstanceProfile
+        SecurityGroupIds:
+          - !Ref SecurityGroup
+        LaunchTemplate:
+          LaunchTemplateId: !Ref LaunchTemplate
+          Version: !GetAtt LaunchTemplate.LatestVersionNumber
+        Tags:
+          Name: !Ref AWS::StackName
+
+  {{ name }}JobQueue:
     Type: AWS::Batch::JobQueue
     Properties:
       Priority: 1
       ComputeEnvironmentOrder:
-        - ComputeEnvironment: !Ref ComputeEnvironment
+        - ComputeEnvironment: !Ref {{ name }}ComputeEnvironment
           Order: 1
       SchedulingPolicyArn: !Ref SchedulingPolicy
+
+  {% endfor %}
 
   TaskRole:
     Type: {{ 'Custom::JplRole' if security_environment in ('JPL', 'JPL-public') else 'AWS::IAM::Role' }}

apps/handle-batch-event/handle-batch-event-cf.yml.j2

@@ -5,6 +5,11 @@ Parameters:
   JobQueueArn:
     Type: String
 
+  {% for job_type, job_spec in job_types.items() if 'Default' != job_spec['compute_environment']['name'] %}
+  {{ job_spec['compute_environment']['name'] }}JobQueueArn:
+    Type: String
+  {% endfor %}
+
   JobsTable:
     Type: String
 
@@ -95,6 +100,9 @@ Resources:
         detail:
           jobQueue:
             - !Ref JobQueueArn
+            {% for job_type, job_spec in job_types.items() if 'Default' != job_spec['compute_environment']['name'] %}
+            - !Ref {{ job_spec['compute_environment']['name'] }}JobQueueArn
+            {% endfor %}
           status:
             - RUNNING
       Targets:

apps/main-cf.yml.j2

@@ -154,6 +154,10 @@ Resources:
     Properties:
       Parameters:
         ComputeEnvironmentArn: !GetAtt Cluster.Outputs.ComputeEnvironmentArn
+        {% for job_type, job_spec in job_types.items() if 'Default' != job_spec['compute_environment']['name'] %}
+        {% set name = job_spec['compute_environment']['name'] %}
+        {{ name }}ComputeEnvironmentArn: !GetAtt Cluster.Outputs.{{ name }}ComputeEnvironmentArn
+        {% endfor %}
         DefaultMaxvCpus: !Ref DefaultMaxvCpus
         ExpandedMaxvCpus: !Ref ExpandedMaxvCpus
         MonthlyBudget: !Ref MonthlyBudget
@@ -169,6 +173,10 @@ Resources:
     Properties:
       Parameters:
         JobQueueArn: !GetAtt Cluster.Outputs.JobQueueArn
+        {% for job_type, job_spec in job_types.items() if 'Default' != job_spec['compute_environment']['name'] %}
+        {% set name = job_spec['compute_environment']['name'] %}
+        {{ name }}JobQueueArn: !GetAtt Cluster.Outputs.{{ name }}JobQueueArn
+        {% endfor %}
         JobsTable: !Ref JobsTable
         {% if security_environment == 'EDC' %}
         SecurityGroupId: !GetAtt Cluster.Outputs.SecurityGroupId
@@ -181,6 +189,10 @@ Resources:
     Properties:
       Parameters:
         JobQueueArn: !GetAtt Cluster.Outputs.JobQueueArn
+        {% for job_type, job_spec in job_types.items() if 'Default' != job_spec['compute_environment']['name'] %}
+        {% set name = job_spec['compute_environment']['name'] %}
+        {{ name }}JobQueueArn: !GetAtt Cluster.Outputs.{{ name }}JobQueueArn
+        {% endfor %}
         TaskRoleArn: !GetAtt Cluster.Outputs.TaskRoleArn
         JobsTable: !Ref JobsTable
         Bucket: !Ref ContentBucket

apps/scale-cluster/scale-cluster-cf.yml.j2

@@ -5,6 +5,12 @@ Parameters:
   ComputeEnvironmentArn:
     Type: String
 
+  {% for job_type, job_spec in job_types.items() if 'Default' != job_spec['compute_environment']['name'] %}
+  {% set name = job_spec['compute_environment']['name'] %}
+  {{ name }}ComputeEnvironmentArn:
+    Type: String
+  {% endfor %}
+
   DefaultMaxvCpus:
     Type: Number
     MinValue: 0
@@ -79,7 +85,11 @@ Resources:
             Resource: "*"
           - Effect: Allow
             Action: batch:UpdateComputeEnvironment
-            Resource: !Ref ComputeEnvironmentArn
+            Resource: 
+              - !Ref ComputeEnvironmentArn
+              {% for job_type, job_spec in job_types.items() if 'Default' != job_spec['compute_environment']['name'] %}
+              - !Ref {{ job_spec['compute_environment']['name'] }}ComputeEnvironmentArn
+              {% endfor %}
 
   Lambda:
     Type: AWS::Lambda::Function
@@ -118,6 +128,11 @@ Resources:
       Targets:
         - Arn: !GetAtt Lambda.Arn
           Id: lambda
+        {% for job_type, job_spec in job_types.items() if 'Default' != job_spec['compute_environment']['name'] %}
+        {% set name = job_spec['compute_environment']['name'] %}
+        - Arn: !GetAtt {{ name }}Lambda.Arn
+          Id:  {{ name }}lambda
+        {% endfor %}
 
   EventPermission:
     Type: AWS::Lambda::Permission
@@ -126,3 +141,50 @@ Resources:
       Action: lambda:InvokeFunction
       Principal: events.amazonaws.com
       SourceArn: !GetAtt Schedule.Arn
+
+  {% for job_type, job_spec in job_types.items() if 'Default' != job_spec['compute_environment']['name'] %}
+  {% set name = job_spec['compute_environment']['name'] %}
+  {{ name }}LogGroup:
+    Type: AWS::Logs::LogGroup
+    Properties:
+      LogGroupName: !Sub {{ "/aws/lambda/${" + name + "Lambda}" }}
+      RetentionInDays: 90
+
+  {{ name }}Lambda:
+    Type: AWS::Lambda::Function
+    Properties:
+      Environment:
+        Variables:
+          COMPUTE_ENVIRONMENT_ARN: !Ref {{ name }}ComputeEnvironmentArn
+          MONTHLY_BUDGET: !Ref MonthlyBudget
+          DEFAULT_MAX_VCPUS: !Ref DefaultMaxvCpus
+          EXPANDED_MAX_VCPUS: !Ref ExpandedMaxvCpus
+          REQUIRED_SURPLUS: !Ref RequiredSurplus
+      Code: src/
+      Handler: scale_cluster.lambda_handler
+      MemorySize: 128
+      Role: !GetAtt Role.Arn
+      Runtime: python3.9
+      Timeout: 30
+      {% if security_environment == 'EDC' %}
+      VpcConfig:
+        SecurityGroupIds:
+          - !Ref SecurityGroupId
+        SubnetIds: !Ref SubnetIds
+      {% endif %}
+
+  {{ name }}EventInvokeConfig:
+    Type: AWS::Lambda::EventInvokeConfig
+    Properties:
+      FunctionName: !Ref {{ name }}Lambda
+      Qualifier: $LATEST
+      MaximumRetryAttempts: 0
+
+  {{ name }}EventPermission:
+    Type: AWS::Lambda::Permission
+    Properties:
+      FunctionName: !GetAtt  {{ name }}Lambda.Arn
+      Action: lambda:InvokeFunction
+      Principal: events.amazonaws.com
+      SourceArn: !GetAtt Schedule.Arn
+  {% endfor %}

apps/step-function.json.j2

@@ -207,7 +207,9 @@
       "Parameters": {
         "JobDefinition": "{{ '${'+ snake_to_pascal_case(task['name']) + '}' }}",
         "JobName.$": "$.job_id",
-        "JobQueue": "${JobQueueArn}",
+        {% set name = job_spec['compute_environment']['name'] %}
+        {% set job_queue = name + 'JobQueueArn' if 'Default' != name else 'JobQueueArn' %}
+        "JobQueue": "{{ '${' + job_queue + '}' }}",
         "ShareIdentifier": "default",
         "SchedulingPriorityOverride.$": "$.priority",
         "Parameters.$": "$.job_parameters",

apps/workflow-cf.yml.j2

@@ -5,6 +5,11 @@ Parameters:
   JobQueueArn:
     Type: String
 
+  {% for job_type, job_spec in job_types.items() if 'Default' != job_spec['compute_environment']['name'] %}
+  {{ job_spec['compute_environment']['name'] }}JobQueueArn:
+    Type: String
+  {% endfor %}
+
   JobsTable:
     Type: String
 
@@ -36,6 +41,7 @@ Outputs:
   StepFunctionArn:
     Value: !Ref StepFunction
 
+
 Resources:
   {% for job_type, job_spec in job_types.items() %}
   {% for task in job_spec['tasks'] %}
@@ -60,6 +66,10 @@ Resources:
             Value: "{{ task['vcpu'] }}"
           - Type: MEMORY
             Value: "{{ task['memory'] }}"
+          {% if 'gpu' in task %}
+          - Type: GPU
+            Value: "{{ task['gpu'] }}"
+          {% endif %}
         Command:
           {% for command in task['command'] %}
           - {{ command }}
@@ -83,6 +93,10 @@ Resources:
       DefinitionS3Location: step-function.json
       DefinitionSubstitutions:
         JobQueueArn: !Ref JobQueueArn
+        {% for job_type, job_spec in job_types.items() if 'Default' != job_spec['compute_environment']['name'] %}
+        {% set name = job_spec['compute_environment']['name'] %}
+        {{ name }}JobQueueArn: !Ref {{ name }}JobQueueArn
+        {% endfor %}
         {% for job_type, job_spec in job_types.items() %}
         {% for task in job_spec['tasks'] %}
         {{ snake_to_pascal_case(task['name']) }}: !Ref {{ snake_to_pascal_case(task['name']) }}
@@ -124,6 +138,9 @@ Resources:
             Action: batch:SubmitJob
             Resource:
               - !Ref JobQueueArn
+              {% for job_type, job_spec in job_types.items() if 'Default' != job_spec['compute_environment']['name'] %}
+              - !Ref {{ job_spec['compute_environment']['name'] }}JobQueueArn
+              {% endfor %}
               {% for job_type, job_spec in job_types.items() %}
               {% for task in job_spec['tasks'] %}
               - !Ref {{ snake_to_pascal_case(task['name']) }}

job_spec/ARIA_AUTORIFT.yml

@@ -42,6 +42,8 @@ AUTORIFT:
     DEFAULT:
       cost: 1.0
   validators: []
+  compute_environment:
+    name: 'Default'
   tasks:
     - name: ''
       image: ghcr.io/asfhyp3/hyp3-autorift

job_spec/ARIA_RAIDER.yml

@@ -25,6 +25,8 @@ ARIA_RAIDER:
     DEFAULT:
       cost: 1.0
   validators: []
+  compute_environment:
+    name: 'Default'
   tasks:
     - name: ''
       image: ghcr.io/dbekaert/raider

job_spec/AUTORIFT.yml

@@ -39,6 +39,8 @@ AUTORIFT:
     DEFAULT:
       cost: 1.0
   validators: []
+  compute_environment:
+    name: 'Default'
   tasks:
     - name: ''
       image: ghcr.io/asfhyp3/hyp3-autorift

job_spec/AUTORIFT_ITS_LIVE.yml

@@ -51,6 +51,8 @@ AUTORIFT:
     DEFAULT:
       cost: 1.0
   validators: []
+  compute_environment:
+    name: 'Default'
   tasks:
     - name: ''
       image: ghcr.io/asfhyp3/hyp3-autorift

job_spec/INSAR_GAMMA.yml

@@ -82,6 +82,8 @@ INSAR_GAMMA:
       cost: 1.0
   validators:
     - check_dem_coverage
+  compute_environment:
+    name: 'Default'
   tasks:
     - name: ''
       image: 845172464411.dkr.ecr.us-west-2.amazonaws.com/hyp3-gamma

job_spec/INSAR_ISCE.yml

@@ -92,6 +92,8 @@ INSAR_ISCE:
     DEFAULT:
       cost: 1.0
   validators: []
+  compute_environment:
+    name: 'Default'
   tasks:
     - name: ''
       image: ghcr.io/access-cloud-based-insar/dockerizedtopsapp

job_spec/INSAR_ISCE_BURST.yml

@@ -44,6 +44,8 @@ INSAR_ISCE_BURST:
     - check_valid_polarizations
     - check_same_burst_ids
     - check_not_antimeridian 
+  compute_environment:
+    name: 'Default'
   tasks:
     - name: ''
       image: ghcr.io/asfhyp3/hyp3-isce2

job_spec/RTC_GAMMA.yml

@@ -103,6 +103,8 @@ RTC_GAMMA:
       cost: 1.0
   validators:
     - check_dem_coverage
+  compute_environment:
+    name: 'Default'
   tasks:
     - name: ''
       image: 845172464411.dkr.ecr.us-west-2.amazonaws.com/hyp3-gamma

job_spec/S1_CORRECTION_ITS_LIVE.yml

@@ -29,6 +29,8 @@ S1_CORRECTION_TEST:
     DEFAULT:
       cost: 1.0
   validators: []
+  compute_environment:
+    name: 'Default'
   tasks:
     - name: ''
       image: ghcr.io/asfhyp3/hyp3-autorift