From 71315a9a3dd3fae3349b0cc63b0a026a92d20467 Mon Sep 17 00:00:00 2001
From: Shashank Rammoorthy <shashank@loop.com>
Date: Fri, 6 Sep 2024 12:37:21 -0700
Subject: [PATCH] add tags on creation time

---
 sky/data/storage.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/sky/data/storage.py b/sky/data/storage.py
index b22c4c0da75..9e9de70fc27 100644
--- a/sky/data/storage.py
+++ b/sky/data/storage.py
@@ -1447,6 +1447,19 @@ def _create_s3_bucket(self,
             s3_client.create_bucket(**create_bucket_config)
             logger.info(
                 f'Created S3 bucket {bucket_name!r} in {region or "us-east-1"}')
+
+            # add tags to the bucket
+            bucket_tags = skypilot_config.get_nested(('aws', 'labels'), {})
+            if bucket_tags:
+                s3_client.put_bucket_tagging(
+                    Bucket=bucket_name,
+                    Tagging={
+                        'TagSet': [{
+                            'Key': k,
+                            'Value': v
+                        } for k, v in bucket_tags.items()]
+                    })
+
         except aws.botocore_exceptions().ClientError as e:
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.StorageBucketCreateError(