[GitHub] [systemds] Baunsgaard commented on a diff in pull request #1668: [SYSTEMDS-2835] Add a federated version of adult_neural and mnist tests

[GitBox <gi...@apache.org>]

Baunsgaard commented on code in PR #1668:
URL: https://github.com/apache/systemds/pull/1668#discussion_r928818389


##########
src/main/python/systemds/context/systemds_context.py:
##########
@@ -306,6 +311,71 @@ def __get_open_port(self):
         s.close()
         return port
 
+    def _execution_completed(self, script: 'DMLScript'):
+        """
+        Should/will be called after execution of a script.
+        Used to update statistics.
+        :param script: The script that got executed
+        """
+        if self._capture_statistics:
+            self._statistics += script.prepared_script.statistics()
+
+    def capture_stats(self, enable: bool = True):
+        """
+        Enable (or disable) capturing of execution statistics.
+        :param enable: if `True` enable capturing, else disable it
+        """
+        self._capture_statistics = enable
+        self.java_gateway.entry_point.getConnection().setStatistics(enable)
+
+    @contextmanager
+    def capture_stats_context(self):
+        """
+        Context for capturing statistics. Should be used in a `with` statement.
+        Afterwards capturing will be reset to the state it was before.
+
+        Example:
+        ```Python
+        with sds.capture_stats_context():

Review Comment:
   i do not like the requirement of with here.
   
   preferably i think something like:
   
   sds.capture_stats()
   
   ...
   x.compute().
   ...
   sds.get_stats()
   
   is perfectly fine.
   this remove the need for this capture_stats_context.
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: dev-unsubscribe@systemds.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org