[AWS Batch] Add instance_types config option for AWS Batch EC2/SPOT environments

config/config_template.yaml

@@ -90,6 +90,7 @@
     #env_type: FARGATE_SPOT               # One of: EC2, SPOT, FARGATE, FARGATE_SPOT
     #env_max_cpus: 10                     # Max total vCPUs of the compute environment
     #instance_role: <INSTANCE_ROLE_ARN>   # Mandatory for env_type EC2 / SPOT
+    #instance_types: ['optimal']          # EC2/SPOT only: list of instance types/families, or ['optimal']
     #service_role: <SERVICE_ROLE_ARN>
     #assign_public_ip: True
     #runtime: <RUNTIME_NAME>

docs/source/compute_config/aws_batch.md

@@ -133,6 +133,7 @@ In summary, you can use one of the following settings:
 | aws_batch  | service_role     | | no | Service role for AWS Batch. Leave empty to use a service-linked execution role. More info [here](https://docs.aws.amazon.com/batch/latest/userguide/using-service-linked-roles.html) |
 | aws_batch  | env_max_cpus     | 10 | no | Maximum total CPUs of the compute environment  |
 | aws_batch  | env_type         | FARGATE_SPOT | no | Compute environment type, one of: `["EC2", "SPOT", "FARGATE", "FARGATE_SPOT"]` |
+| aws_batch  | instance_types   | `['optimal']` | no | EC2/SPOT only. List of EC2 instance types passed to AWS Batch `instanceTypes` (ignored for Fargate). Accepts instance types (`m5.large`, `c5.xlarge`), family prefixes (`m5`, `c5`), or `['optimal']` (AWS picks based on vCPU/mem demand). |
 
 
 ## Test Lithops

lithops/serverless/backends/aws_batch/aws_batch.py

@@ -141,7 +141,7 @@ def _create_compute_env(self):
             if self.env_type in {'EC2', 'SPOT'}:
                 compute_resources_spec['instanceRole'] = self.aws_batch_config['instance_role']
                 compute_resources_spec['minvCpus'] = 0
-                compute_resources_spec['instanceTypes'] = ['optimal']
+                compute_resources_spec['instanceTypes'] = self.aws_batch_config.get('instance_types') or ['optimal']
 
             res = self.batch_client.create_compute_environment(
                 computeEnvironmentName=self._compute_env_name,

lithops/serverless/backends/aws_batch/config.py

@@ -129,6 +129,13 @@ def load_config(config_data):
     if config_data['aws_batch']['env_type'] in {'EC2', 'SPOT'}:
         if 'instance_role' not in config_data['aws_batch']:
             raise Exception("'instance_role' mandatory for EC2 or SPOT environments")
+        if 'instance_types' in config_data['aws_batch']:
+            it = config_data['aws_batch']['instance_types']
+            if not isinstance(it, list) or not all(isinstance(x, str) and x for x in it):
+                raise Exception(
+                    "'instance_types' must be a list of strings, e.g. "
+                    "['m5.large','c5.xlarge'], ['m5'], or ['optimal']"
+                )
 
     config_data['aws_batch']['max_workers'] = config_data['aws_batch']['env_max_cpus'] \
         // config_data['aws_batch']['runtime_cpu']