From 94003ecb056e56623aa6621a2d013c1a7e3dcebe Mon Sep 17 00:00:00 2001
From: "k.goto" <24818752+go-to-k@users.noreply.github.com>
Date: Sat, 30 Dec 2023 22:24:18 +0900
Subject: [PATCH] feat(stepfunctions-tasks): additional allocation strategies
 for spot instance fleets in EmrCreateCluster (#28525)

This PR adds new allocation strategies for spot instance fleets in EmrCreateCluster.

- price-capacity-optimized
  - recommended
- lowest-price
- diversified

https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-emr-instancefleetconfig-spotprovisioningspecification.html

https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-fleet-allocation-strategy.html

----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
---
 .../lib/emr/emr-create-cluster.ts                  | 14 ++++++++++++++
 .../test/emr/emr-create-cluster.test.ts            | 13 +++++++++----
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/emr/emr-create-cluster.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/emr/emr-create-cluster.ts
index 78d504884b629..46d6f79a06704 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/emr/emr-create-cluster.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/emr/emr-create-cluster.ts
@@ -684,6 +684,20 @@ export namespace EmrCreateCluster {
      * Capacity-optimized, which launches instances from Spot Instance pools with optimal capacity for the number of instances that are launching.
      */
     CAPACITY_OPTIMIZED = 'capacity-optimized',
+    /**
+     * Price-capacity-optimized, which launches instances from Spot Instance pools with the highest capacity availability for the number of instances that are launching.
+     *
+     * Recommended.
+     */
+    PRICE_CAPACITY_OPTIMIZED = 'price-capacity-optimized',
+    /**
+     * Lowest-price, which launches instances from the lowest priced pool that has available capacity.
+     */
+    LOWEST_PRICE = 'lowest-price',
+    /**
+     * Diversified, which launches instances across all Spot capacity pools.
+     */
+    DIVERSIFIED = 'diversified',
   }
 
   /**
diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/emr/emr-create-cluster.test.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/emr/emr-create-cluster.test.ts
index 2038d779c12f6..69f9044056ca0 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/emr/emr-create-cluster.test.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/emr/emr-create-cluster.test.ts
@@ -1,4 +1,4 @@
-import { Template, Match } from '../../../assertions';
+import { Template } from '../../../assertions';
 import * as iam from '../../../aws-iam';
 import * as sfn from '../../../aws-stepfunctions';
 import * as cdk from '../../../core';
@@ -883,7 +883,12 @@ test('Create Cluster with Instances configuration', () => {
   });
 });
 
-test('Create Cluster with InstanceFleet with allocation strategy=capacity-optimized for Spot instances', () => {
+test.each([
+  [EmrCreateCluster.SpotAllocationStrategy.CAPACITY_OPTIMIZED, 'capacity-optimized'],
+  [EmrCreateCluster.SpotAllocationStrategy.PRICE_CAPACITY_OPTIMIZED, 'price-capacity-optimized'],
+  [EmrCreateCluster.SpotAllocationStrategy.LOWEST_PRICE, 'lowest-price'],
+  [EmrCreateCluster.SpotAllocationStrategy.DIVERSIFIED, 'diversified'],
+])('Create Cluster with InstanceFleet with allocation strategy %s for Spot instances', (strategy, expected) => {
   // WHEN
   const task = new EmrCreateCluster(stack, 'Task', {
     instances: {
@@ -913,7 +918,7 @@ test('Create Cluster with InstanceFleet with allocation strategy=capacity-optimi
         }],
         launchSpecifications: {
           spotSpecification: {
-            allocationStrategy: EmrCreateCluster.SpotAllocationStrategy.CAPACITY_OPTIMIZED,
+            allocationStrategy: strategy,
             blockDurationMinutes: 1,
             timeoutAction: EmrCreateCluster.SpotTimeoutAction.TERMINATE_CLUSTER,
             timeoutDurationMinutes: 5,
@@ -975,7 +980,7 @@ test('Create Cluster with InstanceFleet with allocation strategy=capacity-optimi
           }],
           LaunchSpecifications: {
             SpotSpecification: {
-              AllocationStrategy: 'capacity-optimized',
+              AllocationStrategy: expected,
               BlockDurationMinutes: 1,
               TimeoutAction: 'TERMINATE_CLUSTER',
               TimeoutDurationMinutes: 5,