Bump flame-sdk version and update project dependencies

flame/star/star_model_tester.py

@@ -2,8 +2,10 @@
 import threading
 import uuid
 from typing import Any, Type, Literal, Optional, Union
+import traceback
 
 from flame.star import StarModel, StarLocalDPModel, StarAnalyzer, StarAggregator
+from flame.utils.mock_flame_core import MockFlameCoreSDK
 
 
 class StarModelTester:
@@ -28,6 +30,9 @@ def __init__(self,
         participant_ids = [str(uuid.uuid4()) for _ in range(len(node_roles) + 1)]
 
         threads = []
+        thread_errors = {}
+        results_queue = []
+        MockFlameCoreSDK.stop_event = []  # shared stop event for all threads in case of failure in any thread
         for i, participant_id in enumerate(participant_ids):
             test_kwargs = {
                 'analyzer': analyzer,
@@ -54,13 +59,28 @@ def __init__(self,
                 test_kwargs['epsilon'] = epsilon
                 test_kwargs['sensitivity'] = sensitivity
 
-            results_queue = []
             def run_node(kwargs=test_kwargs, use_dp=use_local_dp):
-                if not use_dp:
-                    flame = StarModel(**kwargs).flame
-                else:
-                    flame = StarLocalDPModel(**kwargs).flame
-                results_queue.append(flame.final_results_storage)
+                try:
+                    if not use_dp:
+                        flame = StarModel(**kwargs).flame
+                    else:
+                        flame = StarLocalDPModel(**kwargs).flame
+                    results_queue.append(flame.final_results_storage)
+                except Exception:
+                    stop_event = MockFlameCoreSDK.stop_event
+                    if not stop_event:
+                        stack_trace = traceback.format_exc()#.replace('\n', '\\n').replace('\t', '\\t')
+                        thread_errors[(kwargs['test_kwargs']['role'],
+                                       kwargs['test_kwargs']['node_id'])] = f"\033[31m{stack_trace}\033[0m"
+                        stop_event.append(kwargs['test_kwargs']['node_id'])
+                        mock = MockFlameCoreSDK(test_kwargs=kwargs['test_kwargs'])
+                        mock.__pop_logs__(failure_message=True)
+                    else:
+                        thread_errors[(kwargs['test_kwargs']['role'],
+                                       kwargs['test_kwargs']['node_id'])] = (Exception("Another thread already failed, "
+                                                                                       "stopping this thread as well."))
+                    return
+
             thread = threading.Thread(target=run_node)
             threads.append(thread)
 
@@ -70,8 +90,14 @@ def run_node(kwargs=test_kwargs, use_dp=use_local_dp):
         for thread in threads:
             thread.join()
 
+
         # write final results
-        self.write_result(results_queue[0], output_type, result_filepath, multiple_results)
+        if results_queue:
+            self.write_result(results_queue[0], output_type, result_filepath, multiple_results)
+        else:
+            print("No results to write. All threads failed with errors:")
+            for (role, node_id), error in thread_errors.items():
+                print(f"\t{(role if role != 'default' else 'analyzer').capitalize()} {node_id}: {error}")
 
 
     @staticmethod

flame/utils/mock_flame_core.py

@@ -44,11 +44,23 @@ def __init__(self, test_kwargs) -> None:
         self.finished: bool = False
 
 
+class IterationTracker:
+    def __init__(self):
+        self.iter = 0
+
+    def increment(self):
+        self.iter += 1
+
+    def get_iterations(self):
+        return self.iter
+
+
 class MockFlameCoreSDK:
-    num_iterations: int = 0
+    num_iterations: IterationTracker = IterationTracker()
     logger: dict[str, list[str]] = {}
     message_broker: dict[str, list[dict[str, Any]]] = {}
     final_results_storage: Optional[Any] = None
+    stop_event: list[tuple[str]] = []
 
     def __init__(self, test_kwargs):
         self.sanity_check(test_kwargs)
@@ -202,6 +214,8 @@ def await_messages(self,
                         break
                 raise KeyError
             except KeyError:
+                if self.stop_event:
+                    raise Exception
                 time.sleep(.01)
                 pass
 
@@ -323,12 +337,17 @@ def _node_finished(self) -> bool:
         self.config.finished = True
         return self.config.finished
 
-    def __pop_logs__(self) -> None:
-        print(f"--- Starting Iteration {self.num_iterations} ---")
+    def __pop_logs__(self, failure_message: bool = False) -> None:
+        print(f"--- Starting Iteration {self.__get_iteration__()} ---")
+        if failure_message:
+            self.flame_log("Exception was raised (see Stacktrace)!", log_type='error')
         for k, v in self.logger.items():
             role, log = self.logger[k]
             print(f"Logs for {'Analyzer' if role == 'default' else role.capitalize()} {k}:")
             self.logger[k] = [role, '']
             print(log, end='')
-        print(f"--- Ending Iteration {self.num_iterations} ---\n")
-        self.num_iterations += 1
+        print(f"--- Ending Iteration {self.__get_iteration__()} ---\n")
+        self.num_iterations.increment()
+
+    def __get_iteration__(self):
+        return self.num_iterations.get_iterations()

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "flame"
-version = "0.6.0"
+version = "0.6.1"
 description = ""
 authors = ["Alexander Röhl <alexander.roehl@uni-tuebingen.de>", "David Hieber <david.hieber@uni-tuebingen.de>"]
 readme = "README.md"
@@ -9,7 +9,7 @@ packages = [{ include = "flame" }]
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"
-flamesdk = {git = "https://github.com/PrivateAIM/python-sdk.git", tag = "0.4.1"}
+flamesdk = {git = "https://github.com/PrivateAIM/python-sdk.git", tag = "0.4.2"}
 opendp = ">=0.12.1,<0.13.0"
 
 [tool.poetry.group.dev.dependencies]