[AWS] Add tags on bucket creation time (#3922)

* add tags on creation time * Update comments and docs --------- Co-authored-by: Romil Bhardwaj <[email protected]>
skypilot-org · Sep 9, 2024 · d46fb63 · d46fb63
1 parent 272226b
commit d46fb63
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/docs/source/reference/config.rst b/docs/source/reference/config.rst
@@ -90,7 +90,7 @@ Available fields and semantics:
   # Advanced AWS configurations (optional).
   # Apply to all new instances but not existing ones.
   aws:
-    # Tags to assign to all instances launched by SkyPilot (optional).
+    # Tags to assign to all instances and buckets created by SkyPilot (optional).
     #
     # Example use case: cost tracking by user/team/project.
     #

diff --git a/sky/data/storage.py b/sky/data/storage.py
@@ -1447,6 +1447,20 @@ 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 AWS tags configured in config.yaml to the bucket.
+            # This is useful for cost tracking and external cleanup.
+            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(