kev-inn commented on code in PR #1668:
URL: https://github.com/apache/systemds/pull/1668#discussion_r929456088
##########
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:
This is completely optional, just offers another (arguably neater) way for
the common pattern of
```Python
sds.capture_stats()
# some computation/computations
...
sds.capture_stats(False)
```
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
