[SOLVED] Override EC2 Image ID for AWS Batch Compute Environment - on AWS CDK Stack

Issue

I'm creating a CloudFormation Stack using AWS CDK (for Python) and one of its main components is a Batch Compute Environment which should run GPU intensive jobs.

Batch jobs would simply not work with GPU unless I manually set up the EC2 configuration, option under Additional configuration when manually creating a Batch Compute Environment (see image below).

Now, I don't want to do that manually, I want to be able to define this in my CDK Stack. Right now, I have something like this:

ecs_machine_image = batch.EcsMachineImage(
    # image="ami-0d625ab7e92ab3a43",
    image_type=batch.EcsMachineImageType.ECS_AL2_NVIDIA
)

compute_env_1 = batch.ManagedEc2EcsComputeEnvironment(
    scope=self,
    id=f"{stack_id}-compute-env",
    vpc=vpc,
    instance_types=[
        ec2.InstanceType("g4dn.2xlarge"),
    ],
    images=[ecs_machine_image],
    maxv_cpus=32,
    minv_cpus=0,
    security_groups=[security_group],
    service_role=batch_service_role,
    instance_role=ecs_instance_role,
)

The comment # image="ami-0d625ab7e92ab3a43" is to illustrate that I thought it would be possible to pass the AMI image id directly to batch.EcsMachineImage, but it isn't.

In summary, my question is: Using AWS CDK configuration, how can I define a specific AMI image to be used by all jobs running in the Batch Compute Environment?

Solution

The image property accepts the IMachineImage interface type. The EC2 MachineImage class has a static method generic_linux that accepts a map of AMIs by region and returns the expected IMachineImage:

image=ec2.MachineImage.generic_linux({'us-east-1': 'ami-12345678'})

Answered By - fedonev

Answer Checked By - Gilberto Lyons (WPSolving Admin)