Merge pull request #22 from AllenInstitute/update-aws-cdk-lib-dependency

Update aibs-informatics-cdk-lib dependencies
AllenInstitute · Sep 12, 2024 · 6f69f88 · 6f69f88
2 parents 65da335 + e45030a
commit 6f69f88
Show file tree

Hide file tree

Showing 11 changed files with 27 additions and 141 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -17,9 +17,8 @@ dynamic = [
 dependencies = [
     "aibs-informatics-aws-utils @ git+ssh://[email protected]/AllenInstitute/aibs-informatics-aws-utils.git@main",
     "aibs-informatics-core @ git+ssh://[email protected]/AllenInstitute/aibs-informatics-core.git@main",
-    "aws-cdk-lib>=2.90.0",
-    "aws-cdk.aws-batch-alpha>=2.90.0a",
-    "constructs>=10.0.0",
+    "aws-cdk-lib~=2.96",
+    "constructs~=10.0",
     "pydantic~=2.0",
 ]
 

diff --git a/src/aibs_informatics_cdk_lib/constructs_/batch/defaults.py b/src/aibs_informatics_cdk_lib/constructs_/batch/defaults.py
@@ -1,4 +1,4 @@
-from aws_cdk import aws_batch_alpha as batch
+from aws_cdk import aws_batch as batch
 
 from aibs_informatics_cdk_lib.constructs_.batch.infrastructure import BatchEnvironmentConfig
 from aibs_informatics_cdk_lib.constructs_.batch.instance_types import (
@@ -11,36 +11,36 @@
 )
 
 LOW_PRIORITY_BATCH_ENV_CONFIG = BatchEnvironmentConfig(
-    allocation_strategy=batch.AllocationStrategy.SPOT_CAPACITY_OPTIMIZED,
+    allocation_strategy=batch.AllocationStrategy.SPOT_PRICE_CAPACITY_OPTIMIZED,
     instance_types=[*SPOT_INSTANCE_TYPES],
     use_spot=True,
     use_fargate=False,
     use_public_subnets=False,
 )
 NORMAL_PRIORITY_BATCH_ENV_CONFIG = BatchEnvironmentConfig(
-    allocation_strategy=batch.AllocationStrategy.BEST_FIT_PROGRESSIVE,
+    allocation_strategy=batch.AllocationStrategy.SPOT_PRICE_CAPACITY_OPTIMIZED,
     instance_types=[*SPOT_INSTANCE_TYPES],
     use_spot=True,
     use_fargate=False,
     use_public_subnets=False,
 )
 HIGH_PRIORITY_BATCH_ENV_CONFIG = BatchEnvironmentConfig(
-    allocation_strategy=batch.AllocationStrategy.BEST_FIT_PROGRESSIVE,
+    allocation_strategy=batch.AllocationStrategy.BEST_FIT,
     instance_types=[*ON_DEMAND_INSTANCE_TYPES],
     use_spot=False,
     use_fargate=False,
     use_public_subnets=False,
 )
 PUBLIC_SUBNET_BATCH_ENV_CONFIG = BatchEnvironmentConfig(
-    allocation_strategy=batch.AllocationStrategy.BEST_FIT_PROGRESSIVE,
+    allocation_strategy=batch.AllocationStrategy.BEST_FIT,
     instance_types=[*TRANSFER_INSTANCE_TYPES],
     use_spot=False,
     use_fargate=False,
     use_public_subnets=True,
 )
 
 LAMBDA_BATCH_ENV_CONFIG = BatchEnvironmentConfig(
-    allocation_strategy=batch.AllocationStrategy.BEST_FIT_PROGRESSIVE,
+    allocation_strategy=batch.AllocationStrategy.BEST_FIT,
     instance_types=[
         *LAMBDA_SMALL_INSTANCE_TYPES,
         *LAMBDA_MEDIUM_INSTANCE_TYPES,
@@ -51,21 +51,21 @@
     use_public_subnets=False,
 )
 LAMBDA_SMALL_BATCH_ENV_CONFIG = BatchEnvironmentConfig(
-    allocation_strategy=batch.AllocationStrategy.BEST_FIT_PROGRESSIVE,
+    allocation_strategy=batch.AllocationStrategy.BEST_FIT,
     instance_types=[*LAMBDA_SMALL_INSTANCE_TYPES],
     use_spot=False,
     use_fargate=False,
     use_public_subnets=False,
 )
 LAMBDA_MEDIUM_BATCH_ENV_CONFIG = BatchEnvironmentConfig(
-    allocation_strategy=batch.AllocationStrategy.BEST_FIT_PROGRESSIVE,
+    allocation_strategy=batch.AllocationStrategy.BEST_FIT,
     instance_types=[*LAMBDA_MEDIUM_INSTANCE_TYPES],
     use_spot=False,
     use_fargate=False,
     use_public_subnets=False,
 )
 LAMBDA_LARGE_BATCH_ENV_CONFIG = BatchEnvironmentConfig(
-    allocation_strategy=batch.AllocationStrategy.BEST_FIT_PROGRESSIVE,
+    allocation_strategy=batch.AllocationStrategy.BEST_FIT,
     instance_types=[*LAMBDA_LARGE_INSTANCE_TYPES],
     use_spot=False,
     use_fargate=False,

diff --git a/src/aibs_informatics_cdk_lib/constructs_/batch/infrastructure.py b/src/aibs_informatics_cdk_lib/constructs_/batch/infrastructure.py
@@ -16,7 +16,7 @@
 from aibs_informatics_core.utils.decorators import cached_property
 from aibs_informatics_core.utils.hashing import sha256_hexdigest
 from aibs_informatics_core.utils.tools.dicttools import remove_null_values
-from aws_cdk import aws_batch_alpha as batch
+from aws_cdk import aws_batch as batch
 from aws_cdk import aws_ec2 as ec2
 from aws_cdk import aws_efs as efs
 from aws_cdk import aws_iam as iam

diff --git a/src/aibs_informatics_cdk_lib/constructs_/batch/tools.py b/src/aibs_informatics_cdk_lib/constructs_/batch/tools.py
diff --git a/src/aibs_informatics_cdk_lib/constructs_/service/compute.py b/src/aibs_informatics_cdk_lib/constructs_/service/compute.py
@@ -1,27 +1,14 @@
 from abc import abstractmethod
-from typing import Any, Iterable, List, Optional, Union
+from typing import Iterable, List, Optional, Union
 
-import aws_cdk as cdk
 from aibs_informatics_core.env import EnvBase
-from aws_cdk import aws_batch_alpha as batch
+from aws_cdk import aws_batch as batch
 from aws_cdk import aws_ec2 as ec2
 from aws_cdk import aws_efs as efs
 from aws_cdk import aws_iam as iam
-from aws_cdk import aws_lambda as lambda_
 from aws_cdk import aws_s3 as s3
-from aws_cdk import aws_stepfunctions as sfn
-from aws_cdk import aws_stepfunctions_tasks as sfn_tasks
 from constructs import Construct
 
-from aibs_informatics_cdk_lib.common.aws.iam_utils import (
-    LAMBDA_READ_ONLY_ACTIONS,
-    batch_policy_statement,
-    lambda_policy_statement,
-    s3_policy_statement,
-)
-from aibs_informatics_cdk_lib.constructs_.assets.code_asset_definitions import (
-    AIBSInformaticsCodeAssets,
-)
 from aibs_informatics_cdk_lib.constructs_.base import EnvBaseConstruct
 from aibs_informatics_cdk_lib.constructs_.batch.infrastructure import (
     Batch,
@@ -38,14 +25,6 @@
 from aibs_informatics_cdk_lib.constructs_.batch.launch_template import BatchLaunchTemplateBuilder
 from aibs_informatics_cdk_lib.constructs_.batch.types import BatchEnvironmentDescriptor
 from aibs_informatics_cdk_lib.constructs_.efs.file_system import MountPointConfiguration
-from aibs_informatics_cdk_lib.constructs_.sfn.fragments.base import create_state_machine
-from aibs_informatics_cdk_lib.constructs_.sfn.fragments.batch import (
-    AWSBatchMixins,
-    SubmitJobWithDefaultsFragment,
-)
-from aibs_informatics_cdk_lib.constructs_.sfn.fragments.informatics import (
-    BatchInvokedLambdaFunction,
-)
 
 
 class BaseBatchComputeConstruct(EnvBaseConstruct):
@@ -167,7 +146,7 @@ def create_batch_environments(self):
         self.on_demand_batch_environment = self.batch.setup_batch_environment(
             descriptor=BatchEnvironmentDescriptor(f"{self.name}-on-demand"),
             config=BatchEnvironmentConfig(
-                allocation_strategy=batch.AllocationStrategy.BEST_FIT_PROGRESSIVE,
+                allocation_strategy=batch.AllocationStrategy.BEST_FIT,
                 instance_types=[ec2.InstanceType(_) for _ in ON_DEMAND_INSTANCE_TYPES],
                 use_spot=False,
                 use_fargate=False,
@@ -179,7 +158,7 @@ def create_batch_environments(self):
         self.spot_batch_environment = self.batch.setup_batch_environment(
             descriptor=BatchEnvironmentDescriptor(f"{self.name}-spot"),
             config=BatchEnvironmentConfig(
-                allocation_strategy=batch.AllocationStrategy.BEST_FIT_PROGRESSIVE,
+                allocation_strategy=batch.AllocationStrategy.SPOT_PRICE_CAPACITY_OPTIMIZED,
                 instance_types=[ec2.InstanceType(_) for _ in SPOT_INSTANCE_TYPES],
                 use_spot=True,
                 use_fargate=False,
@@ -213,7 +192,7 @@ def create_batch_environments(self):
         self.lambda_batch_environment = self.batch.setup_batch_environment(
             descriptor=BatchEnvironmentDescriptor(f"{self.name}-lambda"),
             config=BatchEnvironmentConfig(
-                allocation_strategy=batch.AllocationStrategy.BEST_FIT_PROGRESSIVE,
+                allocation_strategy=batch.AllocationStrategy.BEST_FIT,
                 instance_types=[
                     *LAMBDA_SMALL_INSTANCE_TYPES,
                     *LAMBDA_MEDIUM_INSTANCE_TYPES,
@@ -230,7 +209,7 @@ def create_batch_environments(self):
         self.lambda_small_batch_environment = self.batch.setup_batch_environment(
             descriptor=BatchEnvironmentDescriptor(f"{self.name}-lambda-small"),
             config=BatchEnvironmentConfig(
-                allocation_strategy=batch.AllocationStrategy.BEST_FIT_PROGRESSIVE,
+                allocation_strategy=batch.AllocationStrategy.BEST_FIT,
                 instance_types=[*LAMBDA_SMALL_INSTANCE_TYPES],
                 use_spot=False,
                 use_fargate=False,
@@ -242,7 +221,7 @@ def create_batch_environments(self):
         self.lambda_medium_batch_environment = self.batch.setup_batch_environment(
             descriptor=BatchEnvironmentDescriptor(f"{self.name}-lambda-medium"),
             config=BatchEnvironmentConfig(
-                allocation_strategy=batch.AllocationStrategy.BEST_FIT_PROGRESSIVE,
+                allocation_strategy=batch.AllocationStrategy.BEST_FIT,
                 instance_types=[*LAMBDA_MEDIUM_INSTANCE_TYPES],
                 use_spot=False,
                 use_fargate=False,
@@ -255,7 +234,7 @@ def create_batch_environments(self):
         self.lambda_large_batch_environment = self.batch.setup_batch_environment(
             descriptor=BatchEnvironmentDescriptor(f"{self.name}-lambda-large"),
             config=BatchEnvironmentConfig(
-                allocation_strategy=batch.AllocationStrategy.BEST_FIT_PROGRESSIVE,
+                allocation_strategy=batch.AllocationStrategy.BEST_FIT,
                 instance_types=[*LAMBDA_LARGE_INSTANCE_TYPES],
                 use_spot=False,
                 use_fargate=False,

diff --git a/src/aibs_informatics_cdk_lib/constructs_/sfn/fragments/informatics/batch.py b/src/aibs_informatics_cdk_lib/constructs_/sfn/fragments/informatics/batch.py
@@ -1,5 +1,4 @@
-from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Mapping, Optional, Sequence, Union
-from unittest import result
+from typing import TYPE_CHECKING, Any, List, Literal, Mapping, Optional, Sequence, Union
 
 import constructs
 from aibs_informatics_aws_utils.constants.lambda_ import (
@@ -10,19 +9,12 @@
 )
 from aibs_informatics_aws_utils.constants.s3 import S3_SCRATCH_KEY_PREFIX
 from aibs_informatics_core.env import EnvBase
-from aws_cdk import aws_batch_alpha as batch
-from aws_cdk import aws_ecr_assets as ecr_assets
 from aws_cdk import aws_iam as iam
-from aws_cdk import aws_s3 as s3
 from aws_cdk import aws_stepfunctions as sfn
-from aws_cdk import aws_stepfunctions_tasks as sfn_tasks
 
 from aibs_informatics_cdk_lib.common.aws.iam_utils import (
-    SFN_STATES_EXECUTION_ACTIONS,
-    SFN_STATES_READ_ACCESS_ACTIONS,
     batch_policy_statement,
     s3_policy_statement,
-    sfn_policy_statement,
 )
 from aibs_informatics_cdk_lib.constructs_.base import EnvBaseConstructMixins
 from aibs_informatics_cdk_lib.constructs_.efs.file_system import MountPointConfiguration

diff --git a/src/aibs_informatics_cdk_lib/constructs_/sfn/fragments/informatics/data_sync.py b/src/aibs_informatics_cdk_lib/constructs_/sfn/fragments/informatics/data_sync.py
@@ -2,7 +2,7 @@
 
 import constructs
 from aibs_informatics_core.env import EnvBase
-from aws_cdk import aws_batch_alpha as batch
+from aws_cdk import aws_batch as batch
 from aws_cdk import aws_ecr_assets as ecr_assets
 from aws_cdk import aws_iam as iam
 from aws_cdk import aws_s3 as s3
@@ -15,7 +15,6 @@
 )
 from aibs_informatics_cdk_lib.constructs_.base import EnvBaseConstructMixins
 from aibs_informatics_cdk_lib.constructs_.efs.file_system import MountPointConfiguration
-from aibs_informatics_cdk_lib.constructs_.sfn.fragments.base import EnvBaseStateMachineFragment
 from aibs_informatics_cdk_lib.constructs_.sfn.fragments.informatics.batch import (
     BatchInvokedBaseFragment,
     BatchInvokedLambdaFunction,

diff --git a/src/aibs_informatics_cdk_lib/constructs_/sfn/fragments/informatics/demand_execution.py b/src/aibs_informatics_cdk_lib/constructs_/sfn/fragments/informatics/demand_execution.py
@@ -2,8 +2,7 @@
 
 import constructs
 from aibs_informatics_core.env import EnvBase
-from aibs_informatics_core.utils.tools.dicttools import remove_null_values
-from aws_cdk import aws_batch_alpha as batch
+from aws_cdk import aws_batch as batch
 from aws_cdk import aws_ecr_assets as ecr_assets
 from aws_cdk import aws_iam as iam
 from aws_cdk import aws_s3 as s3

diff --git a/src/aibs_informatics_cdk_lib/constructs_/sfn/fragments/informatics/efs.py b/src/aibs_informatics_cdk_lib/constructs_/sfn/fragments/informatics/efs.py
@@ -3,7 +3,7 @@
 
 import constructs
 from aibs_informatics_core.env import EnvBase
-from aws_cdk import aws_batch_alpha as batch
+from aws_cdk import aws_batch as batch
 from aws_cdk import aws_ecr_assets as ecr_assets
 from aws_cdk import aws_efs as efs
 from aws_cdk import aws_events as events

diff --git a/src/aibs_informatics_core_app/stacks/demand_execution.py b/src/aibs_informatics_core_app/stacks/demand_execution.py
@@ -7,7 +7,7 @@
     EFS_TMP_PATH,
 )
 from aibs_informatics_core.env import EnvBase
-from aws_cdk import aws_batch_alpha as batch
+from aws_cdk import aws_batch as batch
 from aws_cdk import aws_ec2 as ec2
 from aws_cdk import aws_s3 as s3
 

diff --git a/test/aibs_informatics_cdk_lib/constructs_/batch/test_infrastructure.py b/test/aibs_informatics_cdk_lib/constructs_/batch/test_infrastructure.py
@@ -2,7 +2,7 @@
 from test.aibs_informatics_cdk_lib.base import CdkBaseTest
 
 from aibs_informatics_core.env import EnvBase
-from aws_cdk import aws_batch_alpha
+from aws_cdk import aws_batch as batch
 from aws_cdk.aws_ec2 import Vpc
 
 from aibs_informatics_cdk_lib.constructs_.batch.infrastructure import (
@@ -32,7 +32,7 @@ def test__init__simple(self):
         stack = self.get_dummy_stack("test")
         vpc = Vpc(stack, "vpc")
         config = BatchEnvironmentConfig(
-            allocation_strategy=aws_batch_alpha.AllocationStrategy.SPOT_CAPACITY_OPTIMIZED,
+            allocation_strategy=batch.AllocationStrategy.SPOT_PRICE_CAPACITY_OPTIMIZED,
             instance_types=["t2.micro"],
             use_public_subnets=False,
             use_spot=True,