From c4303729f9e6726881df6830a3a1dfb30ccb61dd Mon Sep 17 00:00:00 2001
From: Sebastien Poirier <sebastien@h2o.ai>
Date: Tue, 28 Jun 2022 15:54:38 +0200
Subject: [PATCH] tag all AWS entities

---
 amlb/runners/aws.py   | 10 ++++------
 resources/config.yaml |  2 ++
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/amlb/runners/aws.py b/amlb/runners/aws.py
index 3bddb8a27..980905bb1 100644
--- a/amlb/runners/aws.py
+++ b/amlb/runners/aws.py
@@ -568,6 +568,8 @@ def _start_instance(self, instance_def, script_params="", instance_key=None, tim
             if instance_def.volume_size:
                 ebs['VolumeSize'] = instance_def.volume_size
 
+            instance_tags = ec2_config.instance_tags | ns(Name=f"amlb_{inst_key}")
+            volume_tags = (ec2_config.volume_tags or instance_tags) | ns(Name=f"amlb_{inst_key}")
             instance_params = dict(
                 BlockDeviceMappings=[dict(
                     DeviceName=ec2_config.root_device_name,
@@ -582,15 +584,11 @@ def _start_instance(self, instance_def, script_params="", instance_key=None, tim
                 TagSpecifications=[
                     dict(
                         ResourceType='instance',
-                        Tags=[
-                            dict(Key='Name', Value=f"benchmark_{inst_key}")
-                        ]
+                        Tags=[dict(Key=k, Value=v) for k, v in instance_tags]
                     ),
                     dict(
                         ResourceType='volume',
-                        Tags=[
-                            dict(Key='Name', Value=f"benchmark_{inst_key}")
-                        ]
+                        Tags=[dict(Key=k, Value=v) for k, v in volume_tags]
                     ),
                 ],
                 UserData=self._ec2_startup_script(inst_key, script_params=script_params, timeout_secs=timeout_secs)
diff --git a/resources/config.yaml b/resources/config.yaml
index c03a742eb..ba3a9f930 100644
--- a/resources/config.yaml
+++ b/resources/config.yaml
@@ -157,7 +157,9 @@ aws:                    # configuration namespace for AWS mode.
         '4': xlarge
         '8': 2xlarge
         '16': 4xlarge
+    instance_tags: {}                    # specify custom tags for the EC2 instances here
     volume_type: standard                # one of gp2, io1, st1, sc1, or standard (default).
+    volume_tags: {}                      # specify custom tags for the volume tags here (if empty, will apply the same tags as the corresponding instance)
     root_device_name: '/dev/sda1'        #
     availability_zone:                   # the availability zone where the instances will be created (if not set, aws will pick a default one).
     subnet_id: ''                        #