From bd864961336d8fa797bb97be75cacab0f8262d04 Mon Sep 17 00:00:00 2001
From: tvmarino <145081464+tvmarino@users.noreply.github.com>
Date: Wed, 2 Oct 2024 15:33:30 -0400
Subject: [PATCH] [BC] Add an interactive only mode for env.py (#377)

Added interactive only mode for env.py which
compiles only using iclang instead of both clang
and iclang.
---
 compiler_opt/rl/env.py      | 26 ++++++++++++++++++--------
 compiler_opt/rl/env_test.py | 27 +++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/compiler_opt/rl/env.py b/compiler_opt/rl/env.py
index 0b40f1b9..6f9ee41a 100644
--- a/compiler_opt/rl/env.py
+++ b/compiler_opt/rl/env.py
@@ -290,18 +290,23 @@ 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 generator of tuples. Each element of the tuple is 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 +316,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 +340,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
 
diff --git a/compiler_opt/rl/env_test.py b/compiler_opt/rl/env_test.py
index f6d3c63b..67a951df 100644
--- a/compiler_opt/rl/env_test.py
+++ b/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__':