diff --git a/pyproject.toml b/pyproject.toml index 409b49a..a3dbe81 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,9 +17,8 @@ dynamic = [ dependencies = [ "aibs-informatics-aws-utils @ git+ssh://git@github.com/AllenInstitute/aibs-informatics-aws-utils.git@main", "aibs-informatics-core @ git+ssh://git@github.com/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 index 806afeb..9fb7267 100644 --- 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,28 +11,28 @@ ) 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, @@ -40,7 +40,7 @@ ) 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 index b816d4d..af4b2b9 100644 --- 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 deleted file mode 100644 index 29181f0..0000000 --- a/src/aibs_informatics_cdk_lib/constructs_/batch/tools.py +++ /dev/null @@ -1,82 +0,0 @@ -from typing import Dict, List, Optional, Tuple - -import constructs -from aibs_informatics_core.env import EnvBase -from aws_cdk import aws_batch_alpha as batch -from aws_cdk import aws_ecr as ecr -from aws_cdk import aws_ecs as ecs - -from aibs_informatics_cdk_lib.constructs_.base import EnvBaseConstruct - -LATEST_TAG = "latest" - - -class BatchTools(EnvBaseConstruct): - def __init__(self, scope: constructs.Construct, id: str, env_base: EnvBase): - super().__init__(scope, id, env_base=env_base) - self._repo_cache: Dict[str, ecr.Repository] = {} - self._image_cache: Dict[Tuple[str, str], ecs.EcrImage] = {} - self._job_def_cache: Dict[str, batch.IJobDefinition] = {} - - def create_job_definition( - self, - container_name: str, - container_tag: str = LATEST_TAG, - command: Optional[List[str]] = None, - memory_limit_mib: int = 1024, - vcpus: int = 1, - mount_points: Optional[List[ecs.MountPoint]] = None, - volumes: Optional[List[ecs.Volume]] = None, - **kwargs, - ) -> batch.IJobDefinition: - job_definition_name = self.get_job_definition_name(container_name) - cache_key = job_definition_name - if cache_key not in self._job_def_cache: - jobdef_container = batch.EcsEc2ContainerDefinition( - image=self.get_ecr_image(container_name, container_tag), - command=command, - memory_limit_mib=memory_limit_mib, - vcpus=vcpus, - mount_points=mount_points, - volumes=volumes, - **kwargs, - ) - # TODO: We should use CfnJobDefinition instead because we have more - # granular control over the configurations (such as defining a more - # complex retry strategy as opposed to specifying an arbitrary number) - self._job_def_cache[cache_key] = batch.EcsJobDefinition( - self, - job_definition_name, - container=jobdef_container, - job_definition_name=job_definition_name, - retry_attempts=3, - ) - return self._job_def_cache[cache_key] - - def get_job_definition(self, container_name: str) -> batch.IJobDefinition: - job_definition_name = self.get_job_definition_name(container_name) - cache_key = job_definition_name - if cache_key not in self._job_def_cache: - self._job_def_cache[cache_key] = batch.EcsJobDefinition.from_job_definition_arn( - self, job_definition_name, job_definition_name=job_definition_name - ) - return self._job_def_cache[cache_key] - - def get_ecr_image(self, container_name: str, container_tag: str = LATEST_TAG) -> ecs.EcrImage: - cache_key = (container_name, container_tag) - if cache_key not in self._image_cache: - repo = self.get_ecr_repo(container_name) - self._image_cache[cache_key] = ecs.EcrImage(repo, container_tag) - return self._image_cache[cache_key] - - def get_ecr_repo(self, container_name: str) -> ecr.Repository: - cache_key = container_name - if cache_key not in self._repo_cache: - repo_name = self.env_base.get_repository_name(container_name) - self._repo_cache[cache_key] = ecr.Repository.from_repository_name( - self, self.env_base.get_construct_id(container_name, "ecr-repo"), repo_name - ) - return self._repo_cache[cache_key] - - def get_job_definition_name(self, container_name: str) -> str: - return self.env_base.prefixed(container_name, "job-definition") diff --git a/src/aibs_informatics_cdk_lib/constructs_/service/compute.py b/src/aibs_informatics_cdk_lib/constructs_/service/compute.py index 8ac7759..2cfcc88 100644 --- 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 index a42b949..e7a1291 100644 --- 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 index dd6a8e3..48b3422 100644 --- 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 index 1e39773..a9b530e 100644 --- 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 index df42f99..1790b73 100644 --- 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 index ca83509..e555378 100644 --- 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 index 120c7b7..6343185 100644 --- 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,