Implement custom per-resource tags (#2533)

* Implement custom per-resource tags

* Handle client backward compatibility

Author: r4victor

src/dstack/_internal/core/backends/aws/compute.py

@@ -169,14 +169,19 @@ def create_instance(
             raise NoCapacityError("No eligible availability zones")
 
         instance_name = generate_unique_instance_name(instance_config)
-        tags = {
+        base_tags = {
             "Name": instance_name,
             "owner": "dstack",
             "dstack_project": project_name,
             "dstack_name": instance_config.instance_name,
             "dstack_user": instance_config.user,
         }
-        tags = merge_tags(tags=tags, backend_tags=self.config.tags)
+        tags = merge_tags(
+            base_tags=base_tags,
+            backend_tags=self.config.tags,
+            resource_tags=instance_config.tags,
+        )
+        tags = aws_resources.filter_invalid_tags(tags)
 
         disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024)
         max_efa_interfaces = _get_maximum_efa_interfaces(
@@ -326,15 +331,20 @@ def create_gateway(
         ec2_client = self.session.client("ec2", region_name=configuration.region)
 
         instance_name = generate_unique_gateway_instance_name(configuration)
-        tags = {
+        base_tags = {
             "Name": instance_name,
             "owner": "dstack",
             "dstack_project": configuration.project_name,
             "dstack_name": configuration.instance_name,
         }
         if settings.DSTACK_VERSION is not None:
-            tags["dstack_version"] = settings.DSTACK_VERSION
-        tags = merge_tags(tags=tags, backend_tags=self.config.tags)
+            base_tags["dstack_version"] = settings.DSTACK_VERSION
+        tags = merge_tags(
+            base_tags=base_tags,
+            backend_tags=self.config.tags,
+            resource_tags=configuration.tags,
+        )
+        tags = aws_resources.filter_invalid_tags(tags)
         tags = aws_resources.make_tags(tags)
 
         vpc_id, subnets_ids = get_vpc_id_subnet_id_or_error(
@@ -522,14 +532,19 @@ def create_volume(self, volume: Volume) -> VolumeProvisioningData:
         ec2_client = self.session.client("ec2", region_name=volume.configuration.region)
 
         volume_name = generate_unique_volume_name(volume)
-        tags = {
+        base_tags = {
             "Name": volume_name,
             "owner": "dstack",
             "dstack_project": volume.project_name,
             "dstack_name": volume.name,
             "dstack_user": volume.user,
         }
-        tags = merge_tags(tags=tags, backend_tags=self.config.tags)
+        tags = merge_tags(
+            base_tags=base_tags,
+            backend_tags=self.config.tags,
+            resource_tags=volume.configuration.tags,
+        )
+        tags = aws_resources.filter_invalid_tags(tags)
 
         zones = aws_resources.get_availability_zones(
             ec2_client=ec2_client, region=volume.configuration.region

src/dstack/_internal/core/backends/aws/resources.py

@@ -448,6 +448,16 @@ def make_tags(tags: Dict[str, str]) -> List[Dict[str, str]]:
     return tags_list
 
 
+def filter_invalid_tags(tags: Dict[str, str]) -> Dict[str, str]:
+    filtered_tags = {}
+    for k, v in tags.items():
+        if not _is_valid_tag(k, v):
+            logger.warning("Skipping invalid tag '%s: %s'", k, v)
+            continue
+        filtered_tags[k] = v
+    return filtered_tags
+
+
 def validate_tags(tags: Dict[str, str]):
     for k, v in tags.items():
         if not _is_valid_tag(k, v):

src/dstack/_internal/core/backends/azure/compute.py

@@ -136,13 +136,18 @@ def create_instance(
             location=location,
         )
 
-        tags = {
+        base_tags = {
             "owner": "dstack",
             "dstack_project": instance_config.project_name,
             "dstack_name": instance_config.instance_name,
             "dstack_user": instance_config.user,
         }
-        tags = merge_tags(tags=tags, backend_tags=self.config.tags)
+        tags = merge_tags(
+            base_tags=base_tags,
+            backend_tags=self.config.tags,
+            resource_tags=instance_config.tags,
+        )
+        tags = azure_resources.filter_invalid_tags(tags)
 
         # TODO: Support custom availability_zones.
         # Currently, VMs are regional, which means they don't have zone info.
@@ -228,14 +233,19 @@ def create_gateway(
             location=configuration.region,
         )
 
-        tags = {
+        base_tags = {
             "owner": "dstack",
             "dstack_project": configuration.project_name,
             "dstack_name": configuration.instance_name,
         }
         if settings.DSTACK_VERSION is not None:
-            tags["dstack_version"] = settings.DSTACK_VERSION
-        tags = merge_tags(tags=tags, backend_tags=self.config.tags)
+            base_tags["dstack_version"] = settings.DSTACK_VERSION
+        tags = merge_tags(
+            base_tags=base_tags,
+            backend_tags=self.config.tags,
+            resource_tags=configuration.tags,
+        )
+        tags = azure_resources.filter_invalid_tags(tags)
 
         vm = _launch_instance(
             compute_client=self._compute_client,

src/dstack/_internal/core/backends/azure/resources.py

@@ -5,6 +5,10 @@
 from azure.mgmt.network.models import Subnet
 
 from dstack._internal.core.errors import BackendError
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
 
 MAX_RESOURCE_NAME_LEN = 64
 
@@ -77,6 +81,16 @@ def _is_eligible_private_subnet(
     return False
 
 
+def filter_invalid_tags(tags: Dict[str, str]) -> Dict[str, str]:
+    filtered_tags = {}
+    for k, v in tags.items():
+        if not _is_valid_tag(k, v):
+            logger.warning("Skipping invalid tag '%s: %s'", k, v)
+            continue
+        filtered_tags[k] = v
+    return filtered_tags
+
+
 def validate_tags(tags: Dict[str, str]):
     for k, v in tags.items():
         if not _is_valid_tag(k, v):

src/dstack/_internal/core/backends/base/compute.py

@@ -173,6 +173,7 @@ def run_job(
             ssh_keys=[SSHKey(public=project_ssh_public_key.strip())],
             volumes=volumes,
             reservation=run.run_spec.configuration.reservation,
+            tags=run.run_spec.merged_profile.tags,
         )
         instance_offer = instance_offer.copy()
         self._restrict_instance_offer_az_to_volumes_az(instance_offer, volumes)
@@ -692,9 +693,18 @@ def get_dstack_gateway_commands() -> List[str]:
     ]
 
 
-def merge_tags(tags: Dict[str, str], backend_tags: Optional[Dict[str, str]]) -> Dict[str, str]:
-    res = tags.copy()
+def merge_tags(
+    base_tags: Dict[str, str],
+    backend_tags: Optional[Dict[str, str]] = None,
+    resource_tags: Optional[Dict[str, str]] = None,
+) -> Dict[str, str]:
+    res = base_tags.copy()
+    # backend_tags have priority over resource_tags
+    # so that regular users do not override the tags set by admins
     if backend_tags is not None:
         for k, v in backend_tags.items():
             res.setdefault(k, v)
+    if resource_tags is not None:
+        for k, v in resource_tags.items():
+            res.setdefault(k, v)
     return res

src/dstack/_internal/core/backends/gcp/compute.py

@@ -211,8 +211,12 @@ def create_instance(
             "dstack_name": instance_config.instance_name,
             "dstack_user": instance_config.user.lower(),
         }
-        labels = {k: v for k, v in labels.items() if gcp_resources.is_valid_label_value(v)}
-        labels = merge_tags(tags=labels, backend_tags=self.config.tags)
+        labels = merge_tags(
+            base_tags=labels,
+            backend_tags=self.config.tags,
+            resource_tags=instance_config.tags,
+        )
+        labels = gcp_resources.filter_invalid_labels(labels)
         is_tpu = (
             _is_tpu(instance_offer.instance.resources.gpus[0].name)
             if instance_offer.instance.resources.gpus
@@ -471,8 +475,12 @@ def create_gateway(
             "dstack_project": configuration.project_name.lower(),
             "dstack_name": configuration.instance_name,
         }
-        labels = {k: v for k, v in labels.items() if gcp_resources.is_valid_label_value(v)}
-        labels = merge_tags(tags=labels, backend_tags=self.config.tags)
+        labels = merge_tags(
+            base_tags=labels,
+            backend_tags=self.config.tags,
+            resource_tags=configuration.tags,
+        )
+        labels = gcp_resources.filter_invalid_labels(labels)
 
         request = compute_v1.InsertInstanceRequest()
         request.zone = zone
@@ -573,8 +581,12 @@ def create_volume(self, volume: Volume) -> VolumeProvisioningData:
             "dstack_name": volume.name,
             "dstack_user": volume.user,
         }
-        labels = {k: v for k, v in labels.items() if gcp_resources.is_valid_label_value(v)}
-        labels = merge_tags(tags=labels, backend_tags=self.config.tags)
+        labels = merge_tags(
+            base_tags=labels,
+            backend_tags=self.config.tags,
+            resource_tags=volume.configuration.tags,
+        )
+        labels = gcp_resources.filter_invalid_labels(labels)
 
         disk = compute_v1.Disk()
         disk.name = disk_name

src/dstack/_internal/core/backends/gcp/resources.py

@@ -332,6 +332,16 @@ def get_accelerators(
     return [accelerator_config]
 
 
+def filter_invalid_labels(labels: Dict[str, str]) -> Dict[str, str]:
+    filtered_labels = {}
+    for k, v in labels.items():
+        if not _is_valid_label(k, v):
+            logger.warning("Skipping invalid label '%s: %s'", k, v)
+            continue
+        filtered_labels[k] = v
+    return filtered_labels
+
+
 def validate_labels(labels: Dict[str, str]):
     for k, v in labels.items():
         if not _is_valid_label(k, v):

src/dstack/_internal/core/models/fleets.py

@@ -21,6 +21,7 @@
 )
 from dstack._internal.core.models.resources import Range, ResourcesSpec
 from dstack._internal.utils.json_schema import add_extra_schema_types
+from dstack._internal.utils.tags import tags_validator
 
 
 class FleetStatus(str, Enum):
@@ -249,7 +250,18 @@ class FleetProps(CoreModel):
 
 
 class FleetConfiguration(InstanceGroupParams, FleetProps):
-    pass
+    tags: Annotated[
+        Optional[Dict[str, str]],
+        Field(
+            description=(
+                "The custom tags to associate with the resource."
+                " The tags are also propagated to the underlying backend resources."
+                " If there is a conflict with backend-level tags, does not override them"
+            )
+        ),
+    ] = None
+
+    _validate_tags = validator("tags", pre=True, allow_reuse=True)(tags_validator)
 
 
 class FleetSpec(CoreModel):

src/dstack/_internal/core/models/gateways.py

@@ -1,12 +1,13 @@
 import datetime
 from enum import Enum
-from typing import Optional, Union
+from typing import Dict, Optional, Union
 
-from pydantic import Field
+from pydantic import Field, validator
 from typing_extensions import Annotated, Literal
 
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import CoreModel
+from dstack._internal.utils.tags import tags_validator
 
 
 class GatewayStatus(str, Enum):
@@ -57,6 +58,18 @@ class GatewayConfiguration(CoreModel):
         Optional[AnyGatewayCertificate],
         Field(description="The SSL certificate configuration. Defaults to `type: lets-encrypt`"),
     ] = LetsEncryptGatewayCertificate()
+    tags: Annotated[
+        Optional[Dict[str, str]],
+        Field(
+            description=(
+                "The custom tags to associate with the gateway."
+                " The tags are also propagated to the underlying backend resources."
+                " If there is a conflict with backend-level tags, does not override them"
+            )
+        ),
+    ] = None
+
+    _validate_tags = validator("tags", pre=True, allow_reuse=True)(tags_validator)
 
 
 class GatewaySpec(CoreModel):
@@ -88,7 +101,7 @@ class GatewayPlan(CoreModel):
     project_name: str
     user: str
     spec: GatewaySpec
-    current_resource: Optional[Gateway]
+    current_resource: Optional[Gateway] = None
 
 
 class GatewayComputeConfiguration(CoreModel):
@@ -98,7 +111,8 @@ class GatewayComputeConfiguration(CoreModel):
     region: str
     public_ip: bool
     ssh_key_pub: str
-    certificate: Optional[AnyGatewayCertificate]
+    certificate: Optional[AnyGatewayCertificate] = None
+    tags: Optional[Dict[str, str]] = None
 
 
 class GatewayProvisioningData(CoreModel):

src/dstack/_internal/core/models/instances.py

@@ -1,6 +1,6 @@
 import datetime
 from enum import Enum
-from typing import List, Optional
+from typing import Dict, List, Optional
 from uuid import UUID
 
 import gpuhunt
@@ -108,6 +108,7 @@ class InstanceConfiguration(CoreModel):
     placement_group_name: Optional[str] = None
     reservation: Optional[str] = None
     volumes: Optional[List[Volume]] = None
+    tags: Optional[Dict[str, str]] = None
 
     def get_public_keys(self) -> List[str]:
         return [ssh_key.public.strip() for ssh_key in self.ssh_keys]