Fix to issue #648

compiler_gym/wrappers/time_limit.py

@@ -2,7 +2,7 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-from typing import Optional
+from typing import Iterable, Optional
 
 from compiler_gym.envs import CompilerEnv
 from compiler_gym.util.gym_type_hints import ActionType
@@ -32,12 +32,23 @@ def __init__(self, env: CompilerEnv, max_episode_steps: Optional[int] = None):
         self._max_episode_steps = max_episode_steps
         self._elapsed_steps = None
 
-    def step(self, action: ActionType, **kwargs):
+    # def step(self, action: ActionType, **kwargs):
+    #     assert (
+    #         self._elapsed_steps is not None
+    #     ), "Cannot call env.step() before calling reset()"
+    #     observation, reward, done, info = self.env.step(action, **kwargs)
+    #     self._elapsed_steps += 1
+    #     if self._elapsed_steps >= self._max_episode_steps:
+    #         info["TimeLimit.truncated"] = not done
+    #         done = True
+    #     return observation, reward, done, info
+
+    def multistep(self, actions: Iterable[ActionType], **kwargs):
         assert (
             self._elapsed_steps is not None
         ), "Cannot call env.step() before calling reset()"
-        observation, reward, done, info = self.env.step(action, **kwargs)
-        self._elapsed_steps += 1
+        observation, reward, done, info = self.env.multistep(actions, **kwargs)
+        self._elapsed_steps += len(actions)
         if self._elapsed_steps >= self._max_episode_steps:
             info["TimeLimit.truncated"] = not done
             done = True

tests/wrappers/time_limit_wrappers_test.py

@@ -4,7 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 """Unit tests for //compiler_gym/wrappers."""
 from compiler_gym.envs.llvm import LlvmEnv
-from compiler_gym.wrappers import TimeLimit
+from compiler_gym.wrappers import CycleOverBenchmarks, TimeLimit
 from tests.test_main import main
 
 pytest_plugins = ["tests.pytest_plugins.llvm"]
@@ -81,5 +81,21 @@ def test_time_limit_fork(env: LlvmEnv):
         fkd.close()
 
 
+# @pytest.mark.xfail(strict=True, reason="https://github.com/facebookresearch/CompilerGym/issues/648")
+def test_time_limit(env: LlvmEnv):
+    """Check CycleOverBenchmarks does not break TimeLimit"""
+    env = TimeLimit(env, max_episode_steps=1)
+    env = CycleOverBenchmarks(
+        env,
+        benchmarks=[
+            "benchmark://cbench-v1/crc32",
+        ],
+    )
+    env.reset()
+    _, _, done, _ = env.step(0)
+
+    assert done
+
+
 if __name__ == "__main__":
     main()