Foreman

.gitignore

@@ -1,6 +1,7 @@
 .stfolder
 logs
 storage
+state.json
 
 # Created by https://www.toptal.com/developers/gitignore/api/python
 # Edit at https://www.toptal.com/developers/gitignore?templates=python
@@ -165,7 +166,7 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+.idea/
 
 ### Python Patch ###
 # Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration

colony/core/context.py

@@ -1,4 +1,5 @@
 from colony.core.network.nodes import LocalNode
+from colony.role_registry import RoleRegistry
 
 
 class RoleContext:
@@ -20,3 +21,9 @@ class RoleContext:
 
     def set_self_node_roles(self, roles):
         self._role_manager.set_roles(roles)
+
+    def get_role_manager(self):
+        return self._role_manager
+
+    def get_role_registry(self) -> RoleRegistry:
+        return self._role_manager.role_registry

colony/core/network/core.py

@@ -14,8 +14,6 @@ class Network:
     def __init__(self):
         self.logger = get_logger(self.__class__.__name__)
 
-        self.nodes = []
-
         self.self_node = LocalNode(
             hostname=get_hostname(),
             ip=get_local_ip(),
@@ -24,6 +22,8 @@ class Network:
         )
         self.logger.info(f"Self node created: {self.self_node.id}")
 
+        self.nodes = self.load_network_from_snapshot()
+
         if not self.nodes:
             self.nodes.append(self.self_node)
 
@@ -72,12 +72,18 @@ class Network:
             self.logger.warning(f"Snapshot file {file_path} not found. Starting with empty network")
             snapshot = {"nodes": []}
 
-        for node_data in snapshot.get("nodes", []):
+        for node_ip, node_data in snapshot.items():
+            result = node_data.get('result')
+
+            if not isinstance(result, dict):
+                self.logger.warning(f"Skipping node {node_ip}: result is not a dict")
+                continue
+
             node = Node(
-                hostname=node_data["hostname"],
-                ip=node_data["ip"],
-                roles=node_data.get("roles", []),
-                port=node_data["port"]
+                hostname=result["hostname"],
+                ip=result["ip"],
+                port=result["port"],
+                roles=result.get("roles", [])
             )
             self.logger.info(f"Loaded node from snapshot: {node.id}")
             nodes.append(node)

colony/core/network/hardware.py

@@ -41,6 +41,20 @@ class HardwareStatistics:
     def to_json(self):
         return json.dumps(self.to_dict(), indent=2)
 
+    @staticmethod
+    def from_dict(hw_dict: str):
+        hw = HardwareStatistics()
+        hw.cpu = hw_dict.get("cpu")
+        hw.mem = hw_dict.get("memory")
+        hw.gpu = hw_dict.get("gpu")
+        hw.disk = hw_dict.get("disk")
+        hw.os_info = hw_dict.get("os")
+        hw.boot_time = hw_dict.get("boot_time")
+        # Optionally parse boot_time back to datetime
+        if hw.boot_time:
+            hw.boot_time = datetime.datetime.fromisoformat(hw.boot_time)
+        return hw
+
     def _get_cpu_stats(self):
         return {
             "cores": psutil.cpu_count(logical=False),
@@ -98,3 +112,98 @@ class HardwareStatistics:
             "machine": platform.machine(),
             "processor": platform.processor()
         }
+
+
+class HardwareScorer:
+    @staticmethod
+    def cpu_score(hw_obj) -> float:
+        """
+        Compute available CPU units:
+        (cores + threads) adjusted for current usage.
+        Returns a float representing "effective compute units" currently free.
+        """
+        cpu = hw_obj.cpu
+        if not cpu:
+            return 0.0
+        cores = cpu.get("cores", 1)
+        threads = cpu.get("threads", 1)
+        usage = cpu.get("usage_percent", 0)  # current CPU utilization in percent
+        available = (cores + threads) * (1 - usage / 100)
+        return round(available, 2)
+
+    @staticmethod
+    def memory_score(hw_obj) -> float:
+        """
+        Compute memory available in GB.
+        Converts total bytes to GB and multiplies by free percentage.
+        """
+        mem = hw_obj.mem
+        if not mem:
+            return 0.0
+        total = mem.get("total", 1)
+        percent_used = mem.get("percent", 100)
+        free_gb = (total / (1024**3)) * (1 - percent_used / 100)
+        return round(free_gb, 2)
+
+    @staticmethod
+    def gpu_score(hw_obj) -> float:
+        """
+        Compute total GPU memory free in GB across all GPUs.
+        Accounts for current GPU load.
+        """
+        gpus = hw_obj.gpu
+        if not gpus:
+            return 0.0
+        total_free_gb = 0
+        for gpu in gpus:
+            memory_total = gpu.get("memory_total", 1)
+            load = gpu.get("load", 100)  # GPU utilization in percent
+            free_gb = (memory_total / 1024) * (1 - load / 100)
+            total_free_gb += free_gb
+        return round(total_free_gb, 2)
+
+    @staticmethod
+    def disk_score(hw_obj) -> float:
+        """
+        Compute total disk space free in GB across all disks.
+        Accounts for current usage percentage.
+        """
+        disks = hw_obj.disk
+        if not disks:
+            return 0.0
+        total_free_gb = 0
+        for d in disks:
+            total_gb = d.get("total", 1) / (1024**3)
+            percent_used = d.get("percent", 100)
+            free_gb = total_gb * (1 - percent_used / 100)
+            total_free_gb += free_gb
+        return round(total_free_gb, 2)
+
+    @staticmethod
+    def longevity_score(hw_obj) -> float:
+        """
+        Compute uptime in hours.
+        Converts boot_time from ISO format to datetime, then calculates
+        the difference from the current time.
+        """
+        boot_time_iso = hw_obj.boot_time
+        try:
+            boot_time = datetime.datetime.fromisoformat(boot_time_iso)
+        except Exception:
+            return 0.0
+        uptime_hours = (datetime.datetime.now() - boot_time).total_seconds() / 3600
+        return round(uptime_hours, 2)
+
+    @classmethod
+    def from_object(cls, hw_obj) -> dict:
+        """
+        Convenience method to get all scores in one dictionary.
+        Keys: 'cpu', 'memory', 'gpu', 'disk', 'longevity'.
+        """
+        return {
+            "cpu": cls.cpu_score(hw_obj),
+            "memory": cls.memory_score(hw_obj),
+            "gpu": cls.gpu_score(hw_obj),
+            "disk": cls.disk_score(hw_obj),
+            "longevity": cls.longevity_score(hw_obj)
+        }

colony/core/roles/explorer/role.py

@@ -64,7 +64,7 @@ class Explorer(Role):
             try:
                 port = Discover.request(str(ip))
                 peers[str(ip)] = port
-            except (TimeoutError, ConnectionRefusedError):
+            except (TimeoutError, ConnectionRefusedError, OSError):
                 continue
         return peers
 

colony/core/roles/foreman/role.py

@@ -1,6 +1,17 @@
 from enum import Enum
+
+from colony.core.network.hardware import HardwareStatistics, HardwareScorer
 from colony.core.roles.role import Role
 from colony.core.roles.unit.role import Unit
+from colony.role_registry import HardwareType, RolePriority
+
+_HARDWARE_SCORE_KEY = {
+    HardwareType.DISK_SPACE: "disk",
+    HardwareType.CPU: "cpu",
+    HardwareType.GPU: "gpu",
+    HardwareType.RAM: "memory",
+    HardwareType.LONGEVITY: "longevity",
+}
 
 
 class Foreman(Role):
@@ -17,5 +28,113 @@ class Foreman(Role):
 
         super().__init__(commands, context)
 
+    def _peers_for_assignment(self):
+        self_node = self.context.self()
+        by_key = {}
+        for p in self.context.peers():
+            by_key[(p.ip, p.port)] = p
+        k = (self_node.ip, self_node.port)
+        if k not in by_key:
+            by_key[k] = self_node
+        return list(by_key.values())
+
+    def _fetch_live_roles(self, peers):
+        live = {}
+        for peer in peers:
+            resp = peer.execute(Unit.COMMANDS.SHOW_STATS)
+            if resp.get("status") != "ok":
+                self.logger.warning(
+                    "SHOW_STATS failed for %s:%s, assuming no roles",
+                    peer.ip,
+                    peer.port,
+                )
+                live[peer] = []
+                continue
+            live[peer] = list(resp["result"].get("roles", []))
+        return live
+
+    def _assign_critical_holders(self, peers, critical_roles, live_roles):
+        peer_list = list(peers)
+        n = len(peer_list)
+        holders = {}
+        for i, critical_role in enumerate(critical_roles):
+            current = next(
+                (p for p in peer_list if critical_role in live_roles.get(p, ())),
+                None,
+            )
+            if current is not None:
+                holders[critical_role] = current
+            else:
+                holders[critical_role] = peer_list[i % n]
+        return holders
+
+    def _best_peer_for_optional(self, peers, role_stat):
+        best_score = float("-inf")
+        best_peer = None
+        for peer in peers:
+            hw_resp = peer.execute(Unit.COMMANDS.SHOW_HARDWARE)
+            if hw_resp.get("status") != "ok":
+                continue
+            hw_obj = HardwareStatistics.from_dict(hw_resp["result"])
+            hs = HardwareScorer.from_object(hw_obj)
+            score = 0.0
+            if role_stat.hardware_preferences:
+                for hw_type, weight in role_stat.hardware_preferences.priorities.items():
+                    key = _HARDWARE_SCORE_KEY.get(hw_type)
+                    if key is not None:
+                        score += hs.get(key, 0) * weight
+            if score > best_score:
+                best_score = score
+                best_peer = peer
+        return best_peer
+
+    @staticmethod
+    def _dedupe_preserve_order(role_names):
+        seen = set()
+        out = []
+        for r in role_names:
+            if r not in seen:
+                seen.add(r)
+                out.append(r)
+        return out
+
     def loop(self):
-        self.context.self().execute(Unit.COMMANDS.SET_ROLES, ['Unit', 'Foreman', 'Explorer', 'VaultWarden'])
+        role_registry = self.context.get_role_registry()
+        roles = role_registry.roles
+
+        general_roles = sorted(
+            name for name, stats in roles.items() if stats.priority == RolePriority.GENERAL
+        )
+        critical_roles = sorted(
+            name for name, stats in roles.items() if stats.priority == RolePriority.CRITICAL
+        )
+        optional_roles = sorted(
+            name
+            for name, stats in roles.items()
+            if stats.priority not in (RolePriority.GENERAL, RolePriority.CRITICAL)
+        )
+
+        peers = self._peers_for_assignment()
+        if not peers:
+            return
+
+        live_roles = self._fetch_live_roles(peers)
+        critical_holders = self._assign_critical_holders(
+            peers, critical_roles, live_roles
+        )
+
+        desired = {peer: list(general_roles) for peer in peers}
+        for critical_role, holder in critical_holders.items():
+            desired[holder].append(critical_role)
+
+        for role_name in optional_roles:
+            role_stat = roles[role_name]
+            best_peer = self._best_peer_for_optional(peers, role_stat)
+            if best_peer is not None:
+                desired[best_peer].append(role_name)
+
+        for peer in peers:
+            peer.execute(
+                Unit.COMMANDS.SET_ROLES,
+                self._dedupe_preserve_order(desired[peer]),
+            )

colony/core/state_manager.py

@@ -17,7 +17,11 @@ class StateManager:
     def _load(self):
         if self.STATE_FILE.exists():
             with open(self.STATE_FILE, "r", encoding="utf-8") as f:
+                try:
                     self._data = json.load(f)
+                except json.decoder.JSONDecodeError:
+                    self._data = {}
+
         else:
             self._data = {}
             self._save()

main.py

@@ -13,6 +13,10 @@ from colony.settings import configuration
 
 def main():
     state_manager = StateManager()
+
+    if 'Unit' not in state_manager.get(StateManager.Data.Roles):
+        state_manager.set(StateManager.Data.Roles, ['Unit', 'Foreman'])
+
     network = Network()
     role_registry = RoleRegistry(role_list=available_roles)
     dispatcher = Dispatcher(network=network, role_registry=role_registry)