MEDfl 2.0.10__tar.gz → 2.0.10.dev1__tar.gz

{medfl-2.0.10 → medfl-2.0.10.dev1}/MEDfl/rw/xgboost/strategy.py

@@ -1,6 +1,8 @@
 import json
 import os
 from typing import Dict, Any, Optional
+import tempfile
+import xgboost as xgb
 
 import flwr as fl
 from flwr.common import (
@@ -13,7 +15,7 @@ from flwr.common import GetPropertiesIns
 from MEDfl.rw.xgboost.config import XGBoostConfig
 from MEDfl.rw.xgboost.metrics import aggregate_weighted_metrics
 from MEDfl.rw.xgboost.serialization import (
-    parameters_to_booster,
+    parameters_to_booster,booster_to_parameters,
     save_booster,
     XGBOOST_TENSOR_TYPE,
 )
@@ -108,6 +110,8 @@ class XGBoostStrategy:
         self.total_rounds = total_rounds
 
         self.strategy_object = None
+        self.current_booster: Optional[xgb.Booster] = None
+        self.current_num_trees: int = 0
 
     def create_strategy(self) -> None:
         """
@@ -278,19 +282,14 @@ class XGBoostStrategy:
                     f"Client:{client.cid}: {fit_res.metrics}"
                 )
 
+            if failures:
+                print(f"[Server/XGBoost] ⚠️ Round {server_round} failures: {failures}")
+
             if not results:
                 print("[Server/XGBoost] No fit results received.")
                 return None, {}
 
-            # Select booster from largest client
-            best_client, best_fit_res = max(
-                results,
-                key=lambda item: item[1].num_examples,
-            )
-
-            agg_parameters = best_fit_res.parameters
-
-            # Manual weighted metric aggregation
+            # Aggregate metrics manually, same as before.
             metric_results = []
 
             for _, fit_res in results:
@@ -303,18 +302,95 @@ class XGBoostStrategy:
 
             metrics = aggregate_weighted_metrics(metric_results)
 
+            previous_num_trees = int(self.current_num_trees)
+            new_local_boosters = []
+
             print(
-                f"[Server/XGBoost] ✅ Round {server_round} - "
-                f"Selected booster from client {best_client.cid} "
-                f"with {best_fit_res.num_examples} examples."
+                f"[Server/XGBoost] Previous global tree count before round "
+                f"{server_round}: {previous_num_trees}"
             )
 
+            for client, fit_res in results:
+                client_booster = parameters_to_booster(fit_res.parameters)
+
+                if client_booster is None:
+                    print(
+                        f"[Server/XGBoost] Client {client.cid} returned no booster. Skipping."
+                    )
+                    continue
+
+                client_total_trees = self._num_boosted_rounds(client_booster)
+
+                print(
+                    f"[Server/XGBoost] Client {client.cid} returned booster with "
+                    f"{client_total_trees} trees."
+                )
+
+                # In round 1, previous_num_trees is 0, so this keeps all local trees.
+                # In later rounds, this keeps only the new trees added by the client.
+                new_trees_booster = self._slice_new_trees(
+                    client_booster=client_booster,
+                    previous_num_trees=previous_num_trees,
+                )
+
+                if new_trees_booster is None:
+                    continue
+
+                new_tree_count = self._num_boosted_rounds(new_trees_booster)
+
+                print(
+                    f"[Server/XGBoost] Client {client.cid} contributed "
+                    f"{new_tree_count} new trees."
+                )
+
+                new_local_boosters.append(new_trees_booster)
+
+            if not new_local_boosters:
+                print(
+                    "[Server/XGBoost] No new local trees received. "
+                    "Keeping previous global booster."
+                )
+
+                if self.current_booster is None:
+                    return None, metrics
+
+                agg_parameters = booster_to_parameters(self.current_booster)
+                return agg_parameters, metrics
+
+            boosters_to_combine = []
+
+            if self.current_booster is not None:
+                boosters_to_combine.append(self.current_booster)
+
+            boosters_to_combine.extend(new_local_boosters)
+
+            combined_booster = self._combine_boosters_by_json(boosters_to_combine)
+
+            if combined_booster is None:
+                print("[Server/XGBoost] Failed to combine boosters.")
+                return None, metrics
+
+            self.current_booster = combined_booster
+            self.current_num_trees = self._num_boosted_rounds(combined_booster)
+
+            print(
+                f"[Server/XGBoost] ✅ Round {server_round} - Flower-style bagging aggregation done."
+            )
+            print(
+                f"[Server/XGBoost] Global booster now has {self.current_num_trees} trees."
+            )
             print(f"[Server/XGBoost] Aggregated metrics: {metrics}\n")
 
+            # Optional: print importance directly, no saving required.
+            if server_round == self.total_rounds:
+                self._print_feature_importance(combined_booster, server_round)
+
+            agg_parameters = booster_to_parameters(combined_booster)
+
+            # Optional: keep your existing save behavior if savingPath is provided.
             self._save_if_needed(server_round, agg_parameters)
 
             return agg_parameters, metrics
-
         strategy.aggregate_fit = wrapped_aggregate_fit
 
 
@@ -377,4 +453,216 @@ class XGBoostStrategy:
 
         save_booster(booster, filepath)
 
-        print(f"[Server/XGBoost] 💾 Saved booster to: {filepath}")
+        print(f"[Server/XGBoost] 💾 Saved booster to: {filepath}")
+    def _num_boosted_rounds(self, booster: xgb.Booster) -> int:
+        """
+        Return the number of boosted rounds/trees in the booster.
+        """
+        try:
+            return int(booster.num_boosted_rounds())
+        except Exception:
+            try:
+                return len(booster.get_dump())
+            except Exception:
+                return 0
+
+
+    def _slice_new_trees(
+        self,
+        client_booster: xgb.Booster,
+        previous_num_trees: int,
+    ) -> Optional[xgb.Booster]:
+        """
+        Extract only the trees added by the client in this round.
+
+        Client booster = previous global booster + new local trees.
+        We keep only new local trees.
+        """
+
+        total_trees = self._num_boosted_rounds(client_booster)
+
+        if total_trees <= previous_num_trees:
+            print(
+                "[Server/XGBoost] Client booster has no new trees. "
+                f"total_trees={total_trees}, previous_num_trees={previous_num_trees}"
+            )
+            return None
+
+        try:
+            return client_booster[previous_num_trees:total_trees]
+        except Exception as exc:
+            print(f"[Server/XGBoost] Failed to slice new trees: {exc}")
+            return None
+
+
+    def _booster_to_raw_bytes(self, booster: xgb.Booster) -> bytes:
+        """
+        Serialize booster to raw UBJ bytes.
+        """
+        with tempfile.NamedTemporaryFile(suffix=".ubj", delete=False) as tmp:
+            tmp_path = tmp.name
+
+        try:
+            booster.save_model(tmp_path)
+            with open(tmp_path, "rb") as f:
+                return f.read()
+        finally:
+            try:
+                os.remove(tmp_path)
+            except Exception:
+                pass
+
+
+    def _booster_from_raw_bytes(self, raw: bytes) -> Optional[xgb.Booster]:
+        """
+        Deserialize booster from raw UBJ bytes.
+        """
+        if not raw:
+            return None
+
+        with tempfile.NamedTemporaryFile(suffix=".ubj", delete=False) as tmp:
+            tmp_path = tmp.name
+            tmp.write(raw)
+
+        try:
+            booster = xgb.Booster()
+            booster.load_model(tmp_path)
+            return booster
+        except Exception as exc:
+            print(f"[Server/XGBoost] Failed to load booster from bytes: {exc}")
+            return None
+        finally:
+            try:
+                os.remove(tmp_path)
+            except Exception:
+                pass
+
+    def _combine_boosters_by_json(
+        self,
+        boosters: list[xgb.Booster],
+    ) -> Optional[xgb.Booster]:
+        """
+        Combine multiple XGBoost boosters by concatenating their trees.
+
+        Fixes:
+        - updates trees
+        - updates tree_info
+        - updates gbtree_model_param["num_trees"]
+        - updates iteration_indptr
+
+        This is the missing part that caused:
+        Check failed: model.iteration_indptr.back() == model.param.num_trees
+        """
+
+        boosters = [b for b in boosters if b is not None]
+
+        if not boosters:
+            return None
+
+        if len(boosters) == 1:
+            return boosters[0]
+
+        model_jsons = []
+
+        for booster in boosters:
+            raw_json = booster.save_raw(raw_format="json")
+            if isinstance(raw_json, bytes):
+                raw_json = raw_json.decode("utf-8")
+            model_jsons.append(json.loads(raw_json))
+
+        base = model_jsons[0]
+
+        base_model = base["learner"]["gradient_booster"]["model"]
+        base_trees = base_model["trees"]
+        base_tree_info = base_model["tree_info"]
+
+        # Start from the base model iteration structure
+        base_iteration_indptr = base_model.get("iteration_indptr", None)
+
+        if base_iteration_indptr is None:
+            # Safe fallback for binary/regression: one tree per iteration
+            combined_iteration_indptr = list(range(0, len(base_trees) + 1))
+        else:
+            combined_iteration_indptr = list(base_iteration_indptr)
+
+        next_tree_id = len(base_trees)
+
+        for extra in model_jsons[1:]:
+            extra_model = extra["learner"]["gradient_booster"]["model"]
+
+            extra_trees = extra_model["trees"]
+            extra_tree_info = extra_model["tree_info"]
+
+            # Append trees and reassign IDs
+            for tree in extra_trees:
+                tree["id"] = next_tree_id
+                base_trees.append(tree)
+                next_tree_id += 1
+
+            # Append tree_info
+            base_tree_info.extend(extra_tree_info)
+
+            # Append iteration_indptr correctly
+            extra_iteration_indptr = extra_model.get("iteration_indptr", None)
+
+            if extra_iteration_indptr is None:
+                # Fallback: one tree per boosting iteration
+                extra_iteration_indptr = list(range(0, len(extra_trees) + 1))
+
+            offset = combined_iteration_indptr[-1]
+
+            # Skip the first 0 from the extra booster
+            for ptr in extra_iteration_indptr[1:]:
+                combined_iteration_indptr.append(offset + ptr)
+
+        # Write back merged model fields
+        base_model["trees"] = base_trees
+        base_model["tree_info"] = base_tree_info
+        base_model["iteration_indptr"] = combined_iteration_indptr
+
+        # Update num_trees
+        if "gbtree_model_param" in base_model:
+            base_model["gbtree_model_param"]["num_trees"] = str(len(base_trees))
+
+        # Save/load to validate the merged model
+        with tempfile.NamedTemporaryFile(
+            mode="w",
+            suffix=".json",
+            delete=False,
+            encoding="utf-8",
+        ) as tmp:
+            tmp_path = tmp.name
+            json.dump(base, tmp)
+
+        try:
+            combined = xgb.Booster()
+            combined.load_model(tmp_path)
+            return combined
+        finally:
+            try:
+                os.remove(tmp_path)
+            except Exception:
+                pass
+    def _print_feature_importance(
+        self,
+        booster: xgb.Booster,
+        server_round: int,
+    ) -> None:
+        """
+        Print feature importance for the current global booster.
+        """
+        if booster is None:
+            return
+
+        importance_gain = booster.get_score(importance_type="gain")
+        importance_weight = booster.get_score(importance_type="weight")
+        importance_cover = booster.get_score(importance_type="cover")
+
+        print(f"\n[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Gain:")
+        print(importance_gain)
+
+        print(f"[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Weight:")
+        print(importance_weight)
+
+        print(f"[Server/XGBoost] 📊 Round {server_round} Global Feature Importance - Cover:")
+        print(importance_cover)

{medfl-2.0.10 → medfl-2.0.10.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: medfl
-Version: 2.0.10
+Version: 2.0.10.dev1
 Summary: Python Open-source package for simulating federated learning and differential privacy
 Home-page: https://github.com/MEDomics-UdeS/MEDfl
 Author: MEDomics consortium

{medfl-2.0.10 → medfl-2.0.10.dev1}/medfl.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: medfl
-Version: 2.0.10
+Version: 2.0.10.dev1
 Summary: Python Open-source package for simulating federated learning and differential privacy
 Home-page: https://github.com/MEDomics-UdeS/MEDfl
 Author: MEDomics consortium

{medfl-2.0.10 → medfl-2.0.10.dev1}/medfl.egg-info/requires.txt

@@ -14,3 +14,4 @@ PyYAML<7.0,>=6.0
 Sphinx<6.0,>=5.3
 sphinx-jsonschema==1.19.1
 sphinx-rtd-dark-mode==1.2.4
+xgboost

{medfl-2.0.10 → medfl-2.0.10.dev1}/setup.py

@@ -14,7 +14,7 @@ with open('requirements.txt') as f:
 
 setup(
     name="medfl",
-    version="2.0.10",
+    version="2.0.10.dev1",
     author="MEDomics consortium",
     author_email="medomics.info@gmail.com",
     description="Python Open-source package for simulating federated learning and differential privacy",