diff --git a/config/config_template.yaml b/config/config_template.yaml index 19b958d1..fa5e0c49 100644 --- a/config/config_template.yaml +++ b/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: # Mandatory for env_type EC2 / SPOT + #instance_types: ['optimal'] # EC2/SPOT only: list of instance types/families, or ['optimal'] #service_role: #assign_public_ip: True #runtime: diff --git a/docs/source/compute_config/aws_batch.md b/docs/source/compute_config/aws_batch.md index 57c7292b..84384310 100644 --- a/docs/source/compute_config/aws_batch.md +++ b/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 diff --git a/lithops/serverless/backends/aws_batch/aws_batch.py b/lithops/serverless/backends/aws_batch/aws_batch.py index 9c7fa393..14356b00 100644 --- a/lithops/serverless/backends/aws_batch/aws_batch.py +++ b/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, diff --git a/lithops/serverless/backends/aws_batch/config.py b/lithops/serverless/backends/aws_batch/config.py index 02b5a236..dc3a21f1 100644 --- a/lithops/serverless/backends/aws_batch/config.py +++ b/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']