capture solver output

Author: dangunter

idaes/core/util/structfs/fsrunner.py

@@ -81,7 +81,7 @@ class BaseFlowsheetRunner(Runner):
         "check_model_numerics",
     )
 
-    def __init__(self, solver=None, tee=False):
+    def __init__(self, solver=None, tee=True):
         self.build_step = self.STEPS[0]
         self._solver, self._tee = solver, tee
         self._ann = {}
@@ -230,7 +230,7 @@ def __init__(self, runner):
 
             # check DoF after build, initial solve, and optimization solve
             self._a = runner.add_action(
-                "dof",
+                "degrees_of_freedom",
                 UnitDofChecker,
                 "fs",
                 ["build", "solve_initial", "solve_optimization"],
@@ -258,7 +258,7 @@ class Timings:
         def __init__(self, runner):
             from .runner_actions import Timer
 
-            self._a: Timer = runner.add_action("t", Timer)
+            self._a: Timer = runner.add_action("timings", Timer)
 
         @property
         def values(self) -> list[dict]:
@@ -280,9 +280,12 @@ def _ipython_display_(self):
             self._a._ipython_display_()
 
     def __init__(self, **kwargs):
+        from .runner_actions import CaptureSolverOutput
+
         super().__init__(**kwargs)
         self.dof = self.DegreesOfFreedom(self)
         self.timings = self.Timings(self)
+        self.add_action("capture_solver_output", CaptureSolverOutput)
 
     def build(self):
         """Run just the build step"""

idaes/core/util/structfs/runner_actions.py

@@ -325,3 +325,27 @@ def _get_dof(block, fix_inlets: bool = True):
                 inlet.free()
 
         return dof
+
+
+class CaptureSolverOutput(Action):
+    def __init__(self, runner, **kwargs):
+        super().__init__(runner, **kwargs)
+        self._logs = {}
+        self._solver_out = None
+
+    def before_step(self, step_name: str):
+        if self._is_solve_step(step_name):
+            self._solver_out = StringIO()
+            self._save_stdout, sys.stdout = sys.stdout, self._solver_out
+
+    def after_step(self, step_name: str):
+        if self._is_solve_step(step_name):
+            self._logs[step_name] = self._solver_out.getvalue()
+            self._solver_out = None
+            sys.stdout = self._save_stdout
+
+    def _is_solve_step(self, name: str):
+        return name.startswith("solve")
+
+    def report(self):
+        return {"solver_logs": self._logs}