Foreman

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]),
+            )