Repository: NVIDIA/NVFlare

Issue:
Errors in streaming.py
@yanchengnv notice some issues in nvflare/app_common/widgets/streaming.py:

- Line 47 to Line 52, the checking of the args and error messages are wrong.
- All these write_xxx() methods, should check the tag and data arg and make sure they are what we expect (str, dict, …)
- Line 257, in the call self.log_xxx(), we should set send_event=False; otherwise it may cause recursive events
- Since fed events are handled by a separate thread, there is a potential racing condition that a fed event could be fired after END_RUN event. In the Receiver code, we need to make sure to discard other events after END_RUN (and hence finalize) is done.

Comments:

Patch:
diff --git a/nvflare/apis/analytix.py b/nvflare/apis/analytix.py
index b99c10f3..26a6423b 100644
--- a/nvflare/apis/analytix.py
+++ b/nvflare/apis/analytix.py
@@ -17,8 +17,8 @@ from typing import Optional
 
 from nvflare.apis.dxo import DXO, DataKind
 
-DATA_TYPE_KEY = "analytics_data_type"
-KWARGS_KEY = "analytics_kwargs"
+_DATA_TYPE_KEY = "analytics_data_type"
+_KWARGS_KEY = "analytics_kwargs"
 
 
 class AnalyticsDataType(Enum):
@@ -29,7 +29,7 @@ class AnalyticsDataType(Enum):
 
 
 class AnalyticsData:
-    def __init__(self, tag: str, value, data_type: AnalyticsDataType, kwargs: Optional[dict]):
+    def __init__(self, tag: str, value, data_type: AnalyticsDataType, kwargs: Optional[dict] = None):
         """This class defines AnalyticsData format.
 
             It is a wrapper to provide from / to DXO conversion.
@@ -39,6 +39,18 @@ class AnalyticsData:
             data_type (AnalyticDataType): type of the analytic data.
             kwargs (optional, dict): additional arguments to be passed.
         """
+        if not isinstance(tag, str):
+            raise TypeError(f"expect tag to be an instance of str, but got {type(tag)}.")
+        if not isinstance(data_type, AnalyticsDataType):
+            raise TypeError(f"expect data_type to be an instance of AnalyticsDataType, but got {type(data_type)}.")
+        if kwargs and not isinstance(kwargs, dict):
+            raise TypeError(f"expect kwargs to be an instance of dict, but got {type(kwargs)}.")
+        if data_type == AnalyticsDataType.SCALAR and not isinstance(value, float):
+            raise TypeError(f"expect value to be an instance of float, but got {type(value)}")
+        elif data_type == AnalyticsDataType.SCALARS and not isinstance(value, dict):
+            raise TypeError(f"expect value to be an instance of dict, but got {type(value)}")
+        elif data_type == AnalyticsDataType.TEXT and not isinstance(value, str):
+            raise TypeError(f"expect value to be an instance of str, but got {type(value)}")
         self.tag = tag
         self.value = value
         self.data_type = data_type
@@ -47,8 +59,8 @@ class AnalyticsData:
     def to_dxo(self):
         """Converts the AnalyticsData to DXO object."""
         dxo = DXO(data_kind=DataKind.ANALYTIC, data={self.tag: self.value})
-        dxo.set_meta_prop(DATA_TYPE_KEY, self.data_type)
-        dxo.set_meta_prop(KWARGS_KEY, self.kwargs)
+        dxo.set_meta_prop(_DATA_TYPE_KEY, self.data_type)
+        dxo.set_meta_prop(_KWARGS_KEY, self.kwargs)
         return dxo
 
     @classmethod
@@ -59,16 +71,14 @@ class AnalyticsData:
             dxo (DXO): The DXO object to convert.
         """
         if not isinstance(dxo, DXO):
-            raise TypeError(f"dxo is not of type DXO, instead it has type {type(dxo)}.")
+            raise TypeError(f"expect dxo to be an instance of DXO, but got {type(dxo)}.")
 
         if len(dxo.data) != 1:
             raise ValueError("dxo does not have the correct format for AnalyticsData.")
 
         tag, value = list(dxo.data.items())[0]
 
-        data_type = dxo.get_meta_prop(DATA_TYPE_KEY)
-        kwargs = dxo.get_meta_prop(KWARGS_KEY)
-        if not isinstance(data_type, AnalyticsDataType):
-            raise ValueError(f"data_type {data_type} is not supported.")
+        data_type = dxo.get_meta_prop(_DATA_TYPE_KEY)
+        kwargs = dxo.get_meta_prop(_KWARGS_KEY)
 
         return cls(tag, value, data_type, kwargs)
diff --git a/nvflare/app_common/widgets/streaming.py b/nvflare/app_common/widgets/streaming.py
index 588f381e..dcd82aaa 100644
--- a/nvflare/app_common/widgets/streaming.py
+++ b/nvflare/app_common/widgets/streaming.py
@@ -45,11 +45,11 @@ def send_analytic_dxo(comp: FLComponent, dxo: DXO, fl_ctx: FLContext, event_type
         event_type (str): Event type.
     """
     if not isinstance(comp, FLComponent):
-        raise TypeError("expect comp to be FLComponent, but got {}".format(type(fl_ctx)))
+        raise TypeError(f"expect comp to be an instance of FLComponent, but got {type(comp)}")
     if not isinstance(dxo, DXO):
-        raise TypeError("expect fl_ctx to be FLContext, but got {}".format(type(fl_ctx)))
+        raise TypeError(f"expect dxo to be an instance of DXO, but got {type(dxo)}")
     if not isinstance(fl_ctx, FLContext):
-        raise TypeError("expect fl_ctx to be FLContext, but got {}".format(type(fl_ctx)))
+        raise TypeError(f"expect fl_ctx to be an instance of FLContext, but got {type(fl_ctx)}")
 
     fl_ctx.set_prop(key=FLContextKey.EVENT_DATA, value=dxo.to_shareable(), private=True, sticky=False)
     comp.fire_event(event_type=event_type, fl_ctx=fl_ctx)
@@ -117,7 +117,7 @@ class AnalyticsSender(Widget):
         """Sends analytics data.
 
         This class implements some common methods follows signatures from PyTorch SummaryWriter and Python logger.
-        It provides a convenient way for LearnerService to use.
+        It provides a convenient way for Learner to use.
         """
         super().__init__()
         self.engine = None
@@ -126,50 +126,66 @@ class AnalyticsSender(Widget):
         if event_type == EventType.START_RUN:
             self.engine = fl_ctx.get_engine()
 
-    def _add(self, tag: str, value, data_type: AnalyticsDataType, kwargs: Optional[dict] = None):
+    def _add(
+        self,
+        tag: str,
+        value,
+        data_type: AnalyticsDataType,
+        global_step: Optional[int] = None,
+        kwargs: Optional[dict] = None,
+    ):
+        kwargs = kwargs if kwargs else {}
+        if global_step:
+            if not isinstance(global_step, int):
+                raise TypeError(f"Expect global step to be an instance of int, but got {type(global_step)}")
+            kwargs["global_step"] = global_step
         dxo = _write(tag=tag, value=value, data_type=data_type, kwargs=kwargs)
         with self.engine.new_context() as fl_ctx:
             send_analytic_dxo(self, dxo=dxo, fl_ctx=fl_ctx)
 
-    def add_scalar(self, tag: str, scalar: float, **kwargs):
+    def add_scalar(self, tag: str, scalar: float, global_step: Optional[int] = None, **kwargs):
         """Sends a scalar.
 
         Args:
             tag (str): Data identifier.
             scalar (float): Value to send.
+            global_step (optional, int): Global step value.
             **kwargs: Additional arguments to pass to the receiver side.
         """
-        self._add(tag=tag, value=scalar, data_type=AnalyticsDataType.SCALAR, kwargs=kwargs)
+        self._add(tag=tag, value=scalar, data_type=AnalyticsDataType.SCALAR, global_step=global_step, kwargs=kwargs)
 
-    def add_scalars(self, tag: str, scalars: dict, **kwargs):
+    def add_scalars(self, tag: str, scalars: dict, global_step: Optional[int] = None, **kwargs):
         """Sends scalars.
 
         Args:
             tag (str): The parent name for the tags.
             scalars (dict): Key-value pair storing the tag and corresponding values.
+            global_step (optional, int): Global step value.
             **kwargs: Additional arguments to pass to the receiver side.
         """
-        self._add(tag=tag, value=scalars, data_type=AnalyticsDataType.SCALARS, kwargs=kwargs)
+        self._add(tag=tag, value=scalars, data_type=AnalyticsDataType.SCALARS, global_step=global_step, kwargs=kwargs)
 
-    def add_text(self, tag: str, text: str, **kwargs):
+    def add_text(self, tag: str, text: str, global_step: Optional[int] = None, **kwargs):
         """Sends a text.
 
         Args:
             tag (str): Data identifier.
             text (str): String to send.
+            global_step (optional, int): Global step value.
             **kwargs: Additional arguments to pass to the receiver side.
         """
-        self._add(tag=tag, value=text, data_type=AnalyticsDataType.TEXT, kwargs=kwargs)
+        self._add(tag=tag, value=text, data_type=AnalyticsDataType.TEXT, global_step=global_step, kwargs=kwargs)
 
-    def add_image(self, tag: str, image, **kwargs):
+    def add_image(self, tag: str, image, global_step: Optional[int] = None, **kwargs):
         """Sends an image.
 
         Args:
             tag (str): Data identifier.
             image: Image to send.
+            global_step (optional, int): Global step value.
             **kwargs: Additional arguments to pass to the receiver side.
         """
-        self._add(tag=tag, value=image, data_type=AnalyticsDataType.IMAGE, kwargs=kwargs)
+        self._add(tag=tag, value=image, data_type=AnalyticsDataType.IMAGE, global_step=global_step, kwargs=kwargs)
 
     def _log(self, tag: LogMessageTag, msg: str, event_type: str, *args, **kwargs):
         """Logs a message.
@@ -210,6 +226,18 @@ class AnalyticsSender(Widget):
         """Logs a message with tag LogMessageTag.CRITICAL."""
         self._log(tag=LogMessageTag.CRITICAL, msg=msg, event_type=_LOG_CRITICAL_EVENT_TYPE, args=args, kwargs=kwargs)
 
+    def flush(self):
+        """Flushes out the message.
+
+        This is doing nothing, it is defined for mimic the PyTorch SummaryWriter behavior.
+        """
+        pass
+
+    def close(self):
+        """Close resources."""
+        if self.engine:
+            self.engine = None
+
 
 class AnalyticsReceiver(Widget, ABC):
     def __init__(self, events: Optional[List[str]] = None):
@@ -223,6 +251,7 @@ class AnalyticsReceiver(Widget, ABC):
             events = [_ANALYTIC_EVENT_TYPE, f"fed.{_ANALYTIC_EVENT_TYPE}"]
         self.events = events
         self._save_lock = Lock()
+        self._end = False
 
     @abstractmethod
     def initialize(self, fl_ctx: FLContext):
@@ -250,32 +279,34 @@ class AnalyticsReceiver(Widget, ABC):
 
         Args:
             fl_ctx (FLContext): fl context.
-
         """
         pass
 
     def handle_event(self, event_type: str, fl_ctx: FLContext):
         if event_type == EventType.START_RUN:
             self.initialize(fl_ctx)
-        elif event_type in self.events:
+        elif event_type in self.events and not self._end:
             data = fl_ctx.get_prop(FLContextKey.EVENT_DATA, None)
             if data is None:
-                self.log_error(fl_ctx, "Missing event data.")
+                self.log_error(fl_ctx, "Missing event data.", fire_event=False)
                 return
             if not isinstance(data, Shareable):
-                self.log_error(fl_ctx, f"Expect shareable but get {type(data)}")
+                self.log_error(
+                    fl_ctx, f"Expect data to be an instance of shareable but get {type(data)}", fire_event=False
+                )
                 return
-            record_origin = fl_ctx.get_identity_name()
 
             # if fed event use peer name to save
             if fl_ctx.get_prop(FLContextKey.EVENT_SCOPE) == EventScope.FEDERATION:
-                peer_name = data.get_peer_prop(ReservedKey.IDENTITY_NAME, None)
-                record_origin = peer_name
+                record_origin = data.get_peer_prop(ReservedKey.IDENTITY_NAME, None)
+            else:
+                record_origin = fl_ctx.get_identity_name()
 
             if record_origin is None:
-                self.log_error(fl_ctx, "record_origin can't be None.")
+                self.log_error(fl_ctx, "record_origin can't be None.", fire_event=False)
                 return
             with self._save_lock:
                 self.save(shareable=data, fl_ctx=fl_ctx, record_origin=record_origin)
         elif event_type == EventType.END_RUN:
+            self._end = True
             self.finalize(fl_ctx)


Test patch:
diff --git a/test/test_analytix.py b/test/test_analytix.py
new file mode 100644
index 00000000..1b61d292
--- /dev/null
+++ b/test/test_analytix.py
@@ -0,0 +1,94 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.
+#
+# 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 pytest
+
+from nvflare.apis.analytix import _DATA_TYPE_KEY, _KWARGS_KEY, AnalyticsData, AnalyticsDataType
+from nvflare.apis.dxo import DXO, DataKind
+
+FROM_DXO_TEST_CASES = [
+    ("hello", 3.0, AnalyticsDataType.SCALAR),
+    ("world", "text", AnalyticsDataType.TEXT),
+    ("dict", {"key": 1.0}, AnalyticsDataType.SCALARS),
+]
+
+TO_DXO_TEST_CASES = [
+    AnalyticsData(tag="hello", value=3.0, data_type=AnalyticsDataType.SCALAR),
+    AnalyticsData(tag="world", value="text", data_type=AnalyticsDataType.TEXT),
+    AnalyticsData(tag="dict", value={"key": 1.0}, data_type=AnalyticsDataType.SCALARS),
+]
+
+FROM_DXO_INVALID_TEST_CASES = [
+    (dict(), TypeError, f"expect dxo to be an instance of DXO, but got {type(dict())}."),
+    (
+        DXO(data_kind=DataKind.WEIGHTS, data={"w": 1.0}),
+        TypeError,
+        f"expect data_type to be an instance of AnalyticsDataType, but got {type(None)}.",
+    ),
+]
+
+INVALID_TEST_CASES = [
+    (
+        dict(),
+        1.0,
+        AnalyticsDataType.SCALAR,
+        None,
+        TypeError,
+        f"expect tag to be an instance of str, but got {type(dict())}.",
+    ),
+    (
+        "tag",
+        1.0,
+        "scalar",
+        None,
+        TypeError,
+        f"expect data_type to be an instance of AnalyticsDataType, but got {type('')}.",
+    ),
+    (
+        "tag",
+        1.0,
+        AnalyticsDataType.SCALAR,
+        [1],
+        TypeError,
+        f"expect kwargs to be an instance of dict, but got {type(list())}.",
+    ),
+]
+
+
+class TestAnalytix:
+    @pytest.mark.parametrize("tag,value,data_type,kwargs,expected_error,expected_msg", INVALID_TEST_CASES)
+    def test_invalid(self, tag, value, data_type, kwargs, expected_error, expected_msg):
+        with pytest.raises(expected_error, match=expected_msg):
+            _ = AnalyticsData(tag=tag, value=value, data_type=data_type, kwargs=kwargs)
+
+    @pytest.mark.parametrize("tag,value,data_type", FROM_DXO_TEST_CASES)
+    def test_from_dxo(self, tag, value, data_type):
+        dxo = DXO(data_kind=DataKind.ANALYTIC, data={tag: value})
+        dxo.set_meta_prop(_DATA_TYPE_KEY, data_type)
+        result = AnalyticsData.from_dxo(dxo)
+        assert result.tag == tag
+        assert result.value == value
+
+    @pytest.mark.parametrize("data", TO_DXO_TEST_CASES)
+    def test_to_dxo(self, data: AnalyticsData):
+        result = data.to_dxo()
+        assert result.data_kind == DataKind.ANALYTIC
+        assert result.data == {data.tag: data.value}
+        assert result.get_meta_prop(_DATA_TYPE_KEY) == data.data_type
+        assert result.get_meta_prop(_KWARGS_KEY) == data.kwargs
+
+    @pytest.mark.parametrize("dxo,expected_error,expected_msg", FROM_DXO_INVALID_TEST_CASES)
+    def test_from_dxo_invalid(self, dxo, expected_error, expected_msg):
+        with pytest.raises(expected_error, match=expected_msg):
+            _ = AnalyticsData.from_dxo(dxo)
diff --git a/test/test_streaming.py b/test/test_streaming.py
new file mode 100644
index 00000000..b1a105c1
--- /dev/null
+++ b/test/test_streaming.py
@@ -0,0 +1,54 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.
+#
+# 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 pytest
+
+from nvflare.apis.dxo import DXO, DataKind
+from nvflare.apis.fl_component import FLComponent
+from nvflare.apis.fl_context import FLContext
+from nvflare.app_common.widgets.streaming import send_analytic_dxo, write_scalar, write_scalars, write_text, write_image
+
+INVALID_TEST_CASES = [
+    (list(), dict(), FLContext(), TypeError, f"expect comp to be an instance of FLComponent, but got {type(list())}"),
+    (FLComponent(), dict(), FLContext(), TypeError, f"expect dxo to be an instance of DXO, but got {type(dict())}"),
+    (
+        FLComponent(),
+        DXO(data={"k": "v"}, data_kind=DataKind.ANALYTIC),
+        list(),
+        TypeError,
+        f"expect fl_ctx to be an instance of FLContext, but got {type(list())}",
+    ),
+]
+
+INVALID_WRITE_TEST_CASES = [
+    (write_scalar, list(), 1.0, TypeError, f"expect tag to be an instance of str, but got {type(list())}"),
+    (write_scalar, "tag", list(), TypeError, f"expect value to be an instance of float, but got {type(list())}"),
+    (write_scalars, list(), 1.0, TypeError, f"expect tag to be an instance of str, but got {type(list())}"),
+    (write_scalars, "tag", 1.0, TypeError, f"expect value to be an instance of dict, but got {type(1.0)}"),
+    (write_text, list(), 1.0, TypeError, f"expect tag to be an instance of str, but got {type(list())}"),
+    (write_text, "tag", 1.0, TypeError, f"expect value to be an instance of str, but got {type(1.0)}"),
+    (write_image, list(), 1.0, TypeError, f"expect tag to be an instance of str, but got {type(list())}"),
+]
+
+
+class TestStreaming:
+    @pytest.mark.parametrize("comp,dxo,fl_ctx,expected_error,expected_msg", INVALID_TEST_CASES)
+    def test_invalid_send_analytic_dxo(self, comp, dxo, fl_ctx, expected_error, expected_msg):
+        with pytest.raises(expected_error, match=expected_msg):
+            send_analytic_dxo(comp=comp, dxo=dxo, fl_ctx=fl_ctx)
+
+    @pytest.mark.parametrize("func,tag,value,expected_error,expected_msg", INVALID_WRITE_TEST_CASES)
+    def test_invalid_write_func(self, func, tag, value, expected_error, expected_msg):
+        with pytest.raises(expected_error, match=expected_msg):
+            func(tag, value)