← back to Exo
periodically update exo_updated_at attribute for tailscale
15a2165d78cbe1d40beed33cfd7fa63a4c1d2afe · 2024-09-23 22:19:15 +0100 · Alex Cheema
Files touched
M exo/networking/tailscale_discovery.pyM exo/networking/tailscale_helpers.py
Diff
commit 15a2165d78cbe1d40beed33cfd7fa63a4c1d2afe
Author: Alex Cheema <alexcheema123@gmail.com>
Date: Mon Sep 23 22:19:15 2024 +0100
periodically update exo_updated_at attribute for tailscale
---
exo/networking/tailscale_discovery.py | 29 ++++++++++++++++++++++++-----
exo/networking/tailscale_helpers.py | 14 +++++++++-----
2 files changed, 33 insertions(+), 10 deletions(-)
diff --git a/exo/networking/tailscale_discovery.py b/exo/networking/tailscale_discovery.py
index f98badf5..b8348d22 100644
--- a/exo/networking/tailscale_discovery.py
+++ b/exo/networking/tailscale_discovery.py
@@ -20,7 +20,7 @@ class TailscaleDiscovery(Discovery):
discovery_timeout: int = 30,
device_capabilities: DeviceCapabilities = UNKNOWN_DEVICE_CAPABILITIES,
tailscale_api_key: str = None,
- tailnet: str = None
+ tailnet: str = None,
):
self.node_id = node_id
self.node_port = node_port
@@ -33,12 +33,25 @@ class TailscaleDiscovery(Discovery):
self.cleanup_task = None
self.tailscale = Tailscale(api_key=tailscale_api_key, tailnet=tailnet)
self._device_id = None
+ self.update_task = None
async def start(self):
self.device_capabilities = device_capabilities()
- await self.update_device_posture_attributes() # Fetch and update device posture attributes
self.discovery_task = asyncio.create_task(self.task_discover_peers())
self.cleanup_task = asyncio.create_task(self.task_cleanup_peers())
+ self.update_task = asyncio.create_task(self.task_update_device_posture_attributes())
+
+ async def task_update_device_posture_attributes(self):
+ while True:
+ try:
+ await self.update_device_posture_attributes()
+ if DEBUG_DISCOVERY >= 2:
+ print(f"Updated device posture attributes")
+ except Exception as e:
+ print(f"Error updating device posture attributes: {e}")
+ print(traceback.format_exc())
+ finally:
+ await asyncio.sleep(self.discovery_interval)
async def get_device_id(self):
if self._device_id:
@@ -67,10 +80,14 @@ class TailscaleDiscovery(Discovery):
for device in active_devices.values():
if device.name != self.node_id:
peer_host = device.addresses[0]
- peer_id, peer_port, device_capabilities = await get_device_attributes(device.device_id, self.tailscale.api_key)
+ peer_id, peer_port, device_capabilities, updated_at = await get_device_attributes(device.device_id, self.tailscale.api_key)
if not peer_id:
if DEBUG_DISCOVERY >= 4: print(f"{device.device_id} does not have exo node attributes. skipping.")
continue
+ if current_time - updated_at > self.discovery_timeout:
+ if DEBUG_DISCOVERY >= 3: print(f"{device.device_id} has outdated exo node attributes. skipping.")
+ continue
+
if peer_id not in self.known_peers or self.known_peers[peer_id][0].addr() != f"{peer_host}:{peer_port}":
if DEBUG >= 1: print(f"Adding {peer_id=} at {peer_host}:{peer_port}. Replace existing peer_id: {peer_id in self.known_peers}")
self.known_peers[peer_id] = (
@@ -92,8 +109,10 @@ class TailscaleDiscovery(Discovery):
self.discovery_task.cancel()
if self.cleanup_task:
self.cleanup_task.cancel()
- if self.discovery_task or self.cleanup_task:
- await asyncio.gather(self.discovery_task, self.cleanup_task, return_exceptions=True)
+ if self.update_task:
+ self.update_task.cancel()
+ if self.discovery_task or self.cleanup_task or self.update_task:
+ await asyncio.gather(self.discovery_task, self.cleanup_task, self.update_task, return_exceptions=True)
async def discover_peers(self, wait_for_peers: int = 0) -> List[PeerHandle]:
if wait_for_peers > 0:
diff --git a/exo/networking/tailscale_helpers.py b/exo/networking/tailscale_helpers.py
index f54e5ad0..a3a576fe 100644
--- a/exo/networking/tailscale_helpers.py
+++ b/exo/networking/tailscale_helpers.py
@@ -2,9 +2,10 @@ import json
import asyncio
import aiohttp
import re
-from typing import Dict, Any, Tuple
+from typing import Dict, Any, Tuple, Optional
from exo.helpers import DEBUG_DISCOVERY
from exo.topology.device_capabilities import DeviceCapabilities, DeviceFlops
+from datetime import datetime, timezone
async def get_device_id() -> str:
try:
@@ -38,7 +39,8 @@ async def update_device_attributes(device_id: str, api_key: str, node_id: str, n
"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),
- "custom:exo_device_capability_flops_int8": str(device_capabilities.flops.int8)
+ "custom:exo_device_capability_flops_int8": str(device_capabilities.flops.int8),
+ "custom:exo_updated_at": str(int(datetime.now(timezone.utc).timestamp()))
}
for attr_name, attr_value in attributes.items():
@@ -50,7 +52,7 @@ async def update_device_attributes(device_id: str, api_key: str, node_id: str, n
else:
print(f"Failed to update device posture attribute {attr_name}: {response.status} {await response.text()}")
-async def get_device_attributes(device_id: str, api_key: str) -> Tuple[str, int, DeviceCapabilities]:
+async def get_device_attributes(device_id: str, api_key: str) -> Tuple[str, int, DeviceCapabilities, int]:
async with aiohttp.ClientSession() as session:
url = f"https://api.tailscale.com/api/v2/device/{device_id}/attributes"
headers = {
@@ -72,10 +74,12 @@ async def get_device_attributes(device_id: str, api_key: str) -> Tuple[str, int,
int8=float(attributes.get("custom:exo_device_capability_flops_int8", 0))
)
)
- return node_id, node_port, device_capabilities
+ updated_at_str = attributes.get("custom:exo_updated_at")
+ updated_at = int(updated_at_str) if updated_at_str else 0
+ return node_id, node_port, device_capabilities, updated_at
else:
print(f"Failed to fetch posture attributes for {device_id}: {response.status}")
- return "", 0, DeviceCapabilities(model="", chip="", memory=0, flops=DeviceFlops(fp16=0, fp32=0, int8=0))
+ return "", 0, DeviceCapabilities(model="", chip="", memory=0, flops=DeviceFlops(fp16=0, fp32=0, int8=0)), 0
def parse_device_attributes(data: Dict[str, str]) -> Dict[str, Any]:
result = {}
← 2e74db8f if it doesnt have exo node attributes then skip
·
back to Exo
·
faster discovery_interval, separate update_interval for tail 69f1fe18 →