Merge pull request #277 from optimas-org/capture_logs

Allow `FunctionEvaluator` to redirect `stdout` and `stderr` to a log file

Author: RemiLehe

optimas/evaluators/function_evaluator.py

@@ -20,15 +20,26 @@ class FunctionEvaluator(Evaluator):
         using this option, the current working directory inside the ``function``
         will be changed to the corresponding evaluation directory.
         By default, ``False``.
+    redirect_logs_to_file : bool
+        Whether to redirect the logs (stdout and stderr) of the evaluation
+        function to a file (log.out and log.err). This can be useful to keep the
+        logs of the exploration clean, preventing many processes from writing to the
+        terminal at once. If enabled, `create_evaluation_dirs` will be set to `True`.
 
     """
 
     def __init__(
-        self, function: Callable, create_evaluation_dirs: bool = False
+        self,
+        function: Callable,
+        create_evaluation_dirs: bool = False,
+        redirect_logs_to_file: bool = False,
     ) -> None:
         super().__init__(sim_function=run_function)
         self.function = function
         self._create_evaluation_dirs = create_evaluation_dirs
+        self._redirect_logs_to_file = redirect_logs_to_file
+        if self._redirect_logs_to_file:
+            self._create_evaluation_dirs = True
 
     def get_sim_specs(
         self,
@@ -43,6 +54,7 @@ def get_sim_specs(
         )
         # Add evaluation function to sim_specs.
         sim_specs["user"]["evaluation_func"] = self.function
+        sim_specs["user"]["redirect_logs_to_file"] = self._redirect_logs_to_file
         return sim_specs
 
     def get_libe_specs(self) -> Dict:

optimas/sim_functions.py

@@ -1,5 +1,7 @@
 """Contains the definition of the simulation functions given to libEnsemble."""
 
+from contextlib import redirect_stderr, redirect_stdout
+
 import jinja2
 import numpy as np
 
@@ -122,6 +124,7 @@ def run_function(H, persis_info, sim_specs, libE_info):
     else:
         user_specs = sim_specs["user"]
     evaluation_func = user_specs["evaluation_func"]
+    redirect_logs_to_file = user_specs["redirect_logs_to_file"]
 
     # Prepare the array that is returned to libE
     libE_output = np.zeros(1, dtype=sim_specs["out"])
@@ -130,7 +133,16 @@ def run_function(H, persis_info, sim_specs, libE_info):
             libE_output[name].fill(np.nan)
 
     # Run evaluation.
-    evaluation_func(input_values, libE_output[0])
+    if redirect_logs_to_file:
+        with (
+            open("log.out", "w") as stdout_file,
+            open("log.err", "w") as stderr_file,
+            redirect_stdout(stdout_file),
+            redirect_stderr(stderr_file),
+        ):
+            evaluation_func(input_values, libE_output[0])
+    else:
+        evaluation_func(input_values, libE_output[0])
     calc_status = WORKER_DONE
 
     # If required, fail when the objectives are NaN.

tests/test_function_evaluator.py

@@ -1,4 +1,5 @@
 import os
+import sys
 
 import numpy as np
 import matplotlib.pyplot as plt
@@ -25,6 +26,17 @@ def eval_func(input_params, output_params):
     plt.savefig("fig.png")
 
 
+def eval_func_logs(input_params, output_params):
+    """Evaluation function used for testing"""
+    x0 = input_params["x0"]
+    x1 = input_params["x1"]
+    result = -(x0 + 10 * np.cos(x0)) * (x1 + 5 * np.cos(x1))
+    output_params["f"] = result
+    # write something to stdout and stderr
+    print("This is a message to stdout.")
+    print("This is a message to stderr.", file=sys.stderr)
+
+
 def test_function_evaluator():
     """Test that an exploration runs successfully with a function evaluator."""
 
@@ -93,5 +105,55 @@ def test_function_evaluator():
                     diags.get_evaluation_dir_path(trial_index)
 
 
+def test_function_evaluator_with_logs():
+    """Test a function evaluator with redirected stdout and stderr."""
+
+    # Define variables and objectives.
+    var1 = VaryingParameter("x0", -50.0, 5.0)
+    var2 = VaryingParameter("x1", -5.0, 15.0)
+    obj = Objective("f", minimize=False)
+
+    # Create generator.
+    gen = RandomSamplingGenerator(
+        varying_parameters=[var1, var2],
+        objectives=[obj],
+    )
+
+    # Create function evaluator.
+    ev = FunctionEvaluator(
+        function=eval_func_logs,
+        redirect_logs_to_file=True,
+    )
+
+    # Create exploration.
+    exploration = Exploration(
+        generator=gen,
+        evaluator=ev,
+        max_evals=10,
+        sim_workers=2,
+        exploration_dir_path="./tests_output/test_function_evaluator_logs",
+    )
+
+    # Run exploration.
+    exploration.run()
+
+    # Get diagnostics.
+    diags = ExplorationDiagnostics(exploration)
+
+    # Check that the logs were redirected if specified.
+    for trial_index in diags.history.trial_index:
+        trial_dir = diags.get_evaluation_dir_path(trial_index)
+        assert os.path.exists(os.path.join(trial_dir, "log.out"))
+        assert os.path.exists(os.path.join(trial_dir, "log.err"))
+        # Check contents of log files are as expected
+        with open(os.path.join(trial_dir, "log.out"), "r") as f:
+            log_out_content = f.read()
+            assert "This is a message to stdout." in log_out_content
+        with open(os.path.join(trial_dir, "log.err"), "r") as f:
+            log_err_content = f.read()
+            assert "This is a message to stderr." in log_err_content
+
+
 if __name__ == "__main__":
     test_function_evaluator()
+    test_function_evaluator_with_logs()