more robust sanitization for chip and model

e6808a5cf18cc4f177d39377654ea1b37d2d7249 · 2024-09-23 15:25:22 +0100 · Alex Cheema

Files touched

M exo/networking/tailscale_helpers.py

Diff

commit e6808a5cf18cc4f177d39377654ea1b37d2d7249
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Mon Sep 23 15:25:22 2024 +0100

    more robust sanitization for chip and model
---
 exo/networking/tailscale_helpers.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/exo/networking/tailscale_helpers.py b/exo/networking/tailscale_helpers.py
index 69f9c211..f54e5ad0 100644
--- a/exo/networking/tailscale_helpers.py
+++ b/exo/networking/tailscale_helpers.py
@@ -1,6 +1,7 @@
 import json
 import asyncio
 import aiohttp
+import re
 from typing import Dict, Any, Tuple
 from exo.helpers import DEBUG_DISCOVERY
 from exo.topology.device_capabilities import DeviceCapabilities, DeviceFlops
@@ -32,8 +33,8 @@ async def update_device_attributes(device_id: str, api_key: str, node_id: str, n
     attributes = {
       "custom:exo_node_id": node_id.replace('-', '_'),
       "custom:exo_node_port": node_port,
-      "custom:exo_device_capability_chip": device_capabilities.chip.replace(' ', '_'),
-      "custom:exo_device_capability_model": device_capabilities.model.replace(' ', '_'),
+      "custom:exo_device_capability_chip": sanitize_attribute(device_capabilities.chip),
+      "custom:exo_device_capability_model": sanitize_attribute(device_capabilities.model),
       "custom:exo_device_capability_memory": str(device_capabilities.memory),
       "custom:exo_device_capability_flops_fp16": str(device_capabilities.flops.fp16),
       "custom:exo_device_capability_flops_fp32": str(device_capabilities.flops.fp32),
@@ -87,3 +88,9 @@ def parse_device_attributes(data: Dict[str, str]) -> Dict[str, Any]:
       elif attr_name in ["device_capability_memory", "device_capability_flops_fp16", "device_capability_flops_fp32", "device_capability_flops_int8"]:
         result[attr_name] = float(value)
   return result
+
+def sanitize_attribute(value: str) -> str:
+  # Replace invalid characters with underscores
+  sanitized_value = re.sub(r'[^a-zA-Z0-9_.]', '_', value)
+  # Truncate to 50 characters
+  return sanitized_value[:50]

← 1798fc07 support node_id, node_port and device_capabilities with tail · back to Exo · if it doesnt have exo node attributes then skip 2e74db8f →