[BC] Add an interactive only mode for env.py

compiler_opt/rl/env.py

@@ -290,18 +290,22 @@ def _get_scores() -> dict[str, float]:
 def _get_clang_generator(
     clang_path: str,
     task_type: Type[MLGOTask],
+    interactive_only: bool = False,
 ) -> Generator[Optional[Tuple[ClangProcess, InteractiveClang]],
                Optional[corpus.LoadedModuleSpec], None]:
-  """Returns a generator for creating InteractiveClang objects.
-
-  TODO: fix this docstring
+  """Returns a tuple of generators for creating InteractiveClang objects.
 
   Args:
     clang_path: Path to the clang binary to use within InteractiveClang.
     task_type: Type of the MLGO task to use.
+    interactive_only: If set to true the returned tuple of generators is
+      iclang, iclang instead of iclang, clang
 
   Returns:
-    The generator for InteractiveClang objects.
+    A tuple of generators created with clang_session. First argument of
+    the tuple is always an interactive clang session. The second argumnet
+    is a default clang session if interactive_only is False and otherwise
+    the exact same interactive clang session object as the first argument.
   """
   while True:
     # The following line should be type-hinted as follows:
@@ -311,9 +315,12 @@ def _get_clang_generator(
     module = yield
     with clang_session(
         clang_path, module, task_type, interactive=True) as iclang:
-      with clang_session(
-          clang_path, module, task_type, interactive=False) as clang:
-        yield iclang, clang
+      if interactive_only:
+        yield iclang, iclang
+      else:
+        with clang_session(
+            clang_path, module, task_type, interactive=False) as clang:
+          yield iclang, clang
 
 
 class MLGOEnvironmentBase:
@@ -332,8 +339,10 @@ def __init__(
       task_type: Type[MLGOTask],
       obs_spec,
       action_spec,
+      interactive_only: bool = False,
   ):
-    self._clang_generator = _get_clang_generator(clang_path, task_type)
+    self._clang_generator = _get_clang_generator(
+        clang_path, task_type, interactive_only=interactive_only)
     self._obs_spec = obs_spec
     self._action_spec = action_spec
 

compiler_opt/rl/env_test.py

@@ -214,6 +214,33 @@ def test_env(self, mock_popen):
 
       step = test_env.step(np.array([1], dtype=np.int64))
       self.assertEqual(step.step_type, env.StepType.LAST)
+      self.assertNotEqual(test_env._iclang, test_env._clang)  # pylint: disable=protected-access
+
+  @mock.patch('subprocess.Popen')
+  def test_env_interactive_only(self, mock_popen):
+    mock_popen.side_effect = mock_interactive_clang
+
+    test_env = env.MLGOEnvironmentBase(
+        clang_path=_CLANG_PATH,
+        task_type=MockTask,
+        obs_spec={},
+        action_spec={},
+        interactive_only=True,
+    )
+
+    for env_itr in range(3):
+      del env_itr
+      step = test_env.reset(_MOCK_MODULE)
+      self.assertEqual(step.step_type, env.StepType.FIRST)
+
+      for step_itr in range(_NUM_STEPS - 1):
+        del step_itr
+        step = test_env.step(np.array([1], dtype=np.int64))
+        self.assertEqual(step.step_type, env.StepType.MID)
+
+      step = test_env.step(np.array([1], dtype=np.int64))
+      self.assertEqual(step.step_type, env.StepType.LAST)
+      self.assertEqual(test_env._iclang, test_env._clang)  # pylint: disable=protected-access
 
 
 if __name__ == '__main__':