Workflow communication APIs and Simplified ML Algorithms

nvflare/apis/dxo.py

@@ -29,6 +29,7 @@ class DataKind(object):
     COLLECTION = "COLLECTION"  # Dict or List of DXO objects
     STATISTICS = "STATISTICS"
     PSI = "PSI"
+    RAW = "RAW"
 
 
 class MetaKey(FLMetaKey):

nvflare/apis/wf_controller.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from abc import ABC, abstractmethod
+from typing import Callable, List, Optional
+
+from nvflare.app_common import wf_comm
+
+from .fl_constant import ReturnCode
+
+ABORT_WHEN_IN_ERROR = {
+    ReturnCode.EXECUTION_EXCEPTION: True,
+    ReturnCode.TASK_UNKNOWN: True,
+    ReturnCode.EXECUTION_RESULT_ERROR: False,
+    ReturnCode.TASK_DATA_FILTER_ERROR: True,
+    ReturnCode.TASK_RESULT_FILTER_ERROR: True,
+}
+
+
+class WFController(ABC):
+    def __init__(self):
+        self.communicator = wf_comm.get_wf_comm_api()
+
+    @abstractmethod
+    def run(self):
+        pass
+
+    def broadcast_and_wait(
+        self,
+        task_name: str,
+        min_responses: int,
+        data: any,
+        meta: dict = None,
+        targets: Optional[List[str]] = None,
+        callback: Callable = None,
+    ):
+        return self.communicator.broadcast_and_wait(task_name, min_responses, data, meta, targets, callback)
+
+    def send_and_wait(
+        self,
+        task_name: str,
+        min_responses: int,
+        data: any,
+        meta: dict = None,
+        targets: Optional[List[str]] = None,
+        send_order: str = "sequential",
+        callback: Callable = None,
+    ):
+        return self.communicator.send_and_wait(task_name, min_responses, data, meta, targets, send_order, callback)
+
+    def relay_and_wait(
+        self,
+        task_name: str,
+        min_responses: int,
+        data: any,
+        meta: dict = None,
+        targets: Optional[List[str]] = None,
+        relay_order: str = "sequential",
+        callback: Callable = None,
+    ):
+        return self.communicator.relay_and_wait(task_name, min_responses, data, meta, targets, relay_order, callback)
+
+    def broadcast(self, task_name: str, data: any, meta: dict = None, targets: Optional[List[str]] = None):
+        return self.communicator.broadcast(task_name, data, meta, targets)
+
+    def send(
+        self,
+        task_name: str,
+        data: any,
+        meta: dict = None,
+        targets: Optional[str] = None,
+        send_order: str = "sequential",
+    ):
+        return self.communicator.send(task_name, data, meta, targets, send_order)
+
+    def relay(
+        self,
+        task_name: str,
+        data: any,
+        meta: dict = None,
+        targets: Optional[List[str]] = None,
+        relay_order: str = "sequential",
+    ):
+        return self.communicator.send(task_name, data, meta, targets, relay_order)
+
+    def get_site_names(self) -> List[str]:
+        return self.communicator.get_site_names()

nvflare/app_common/abstract/fl_model.py

@@ -45,6 +45,7 @@ def __init__(
         params: Any = None,
         optimizer_params: Any = None,
         metrics: Optional[Dict] = None,
+        start_round: int = 0,
         current_round: Optional[int] = None,
         total_rounds: Optional[int] = None,
         meta: Optional[Dict] = None,
@@ -79,6 +80,7 @@ def __init__(
         self.params = params
         self.optimizer_params = optimizer_params
         self.metrics = metrics
+        self.start_round = start_round
         self.current_round = current_round
         self.total_rounds = total_rounds
 

nvflare/app_common/utils/fl_model_utils.py

@@ -201,6 +201,9 @@ def get_configs(model: FLModel) -> Optional[dict]:
 
     @staticmethod
     def update_model(model: FLModel, model_update: FLModel, replace_meta: bool = True) -> FLModel:
+
+        model.metrics = model_update.metrics
+
         if model.params_type != ParamsType.FULL:
             raise RuntimeError(f"params_type {model.params_type} of `model` not supported! Expected `ParamsType.FULL`.")
 
@@ -209,8 +212,6 @@ def update_model(model: FLModel, model_update: FLModel, replace_meta: bool = Tru
         else:
             model.meta.update(model_update.meta)
 
-        model.metrics = model_update.metrics
-
         if model_update.params_type == ParamsType.FULL:
             model.params = model_update.params
         elif model_update.params_type == ParamsType.DIFF:

nvflare/app_common/utils/math_utils.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import operator
+from typing import Callable, Optional, Tuple
+
+operator_mapping = {
+    ">=": operator.ge,
+    "<=": operator.le,
+    ">": operator.gt,
+    "<": operator.lt,
+    "=": operator.eq,
+}
+
+
+def parse_compare_criteria(compare_expr: Optional[str] = None) -> Tuple[str, float, Callable]:
+    """
+        Parse the compare expression into individual component
+        compare expression is in the format of string literal : "<key> <op> <value>"
+        such as
+            accuracy >= 0.5
+            loss  > 2.4
+    Args:
+        compare_expr: string literal in the format of  "<key> <op> <value>"
+
+    Returns: Tuple key, value, operator
+
+    """
+    tokens = compare_expr.split(" ")
+    if len(tokens) != 3:
+        raise ValueError(
+            f"Invalid early_stop_condition, expecting form of '<metric> <op> value' but got '{compare_expr}'"
+        )
+
+    key = tokens[0]
+    op = tokens[1]
+    target = tokens[2]
+    op_fn = operator_mapping.get(op, None)
+    if op_fn is None:
+        raise ValueError("Invalid operator symbol: expecting one of <=, =, >=, <, > ")
+    if not target:
+        raise ValueError("Invalid empty or None target value")
+    try:
+        target_value = float(target)
+    except Exception as e:
+        raise ValueError(f"expect a number, but get '{target}' in '{compare_expr}'")
+
+    return key, target_value, op_fn

nvflare/app_common/wf_comm/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nvflare.app_common.wf_comm.wf_comm_api_spec import WFCommAPISpec
+from nvflare.fuel.data_event.data_bus import DataBus
+
+data_bus = DataBus()
+
+
+def get_wf_comm_api() -> WFCommAPISpec:
+    return data_bus.get_data("wf_comm_api")