[llvm] Fix the number of runtime observations in reward reset.

This adds a `resend_on_reset` flag to the `send_params()` method that
enables parameters to be resent immediately after a service is
reset. This is required by the LLVM environment to ensure that the
runtime observation parameters are sent before the reward space is
reset.

Fixes #756.

Author: ChrisCummins

compiler_gym/envs/llvm/llvm_env.py

@@ -28,7 +28,7 @@
     CostFunctionReward,
     NormalizedReward,
 )
-from compiler_gym.errors import BenchmarkInitError
+from compiler_gym.errors import BenchmarkInitError, SessionNotFound
 from compiler_gym.service.client_service_compiler_env import ClientServiceCompilerEnv
 from compiler_gym.spaces import Box, Commandline
 from compiler_gym.spaces import Dict as DictSpace
@@ -363,7 +363,7 @@ def __init__(
 
     def reset(self, *args, **kwargs):
         try:
-            observation = super().reset(*args, **kwargs)
+            return super().reset(*args, **kwargs)
         except ValueError as e:
             # Catch and re-raise some known benchmark initialization errors with
             # a more informative error type.
@@ -379,15 +379,6 @@ def reset(self, *args, **kwargs):
                 raise BenchmarkInitError(str(e)) from e
             raise
 
-        # Resend the runtimes-per-observation session parameter, if it is a
-        # non-default value.
-        if self._runtimes_per_observation_count is not None:
-            self.runtime_observation_count = self._runtimes_per_observation_count
-        if self._runtimes_warmup_per_observation_count is not None:
-            self.runtime_warmup_runs_count = self._runtimes_warmup_per_observation_count
-
-        return observation
-
     def make_benchmark(
         self,
         inputs: Union[
@@ -612,11 +603,15 @@ def runtime_observation_count(self) -> int:
 
     @runtime_observation_count.setter
     def runtime_observation_count(self, n: int) -> None:
-        if self.in_episode:
-            self.send_param("llvm.set_runtimes_per_observation_count", str(n))
-        # NOTE(cummins): Keep this after the send_param() call because
-        # send_param() will raise an error if the valid is invalid.
-        self._runtimes_per_observation_count = n
+        try:
+            self.send_param(
+                "llvm.set_runtimes_per_observation_count", str(n), resend_on_reset=True
+            )
+            # NOTE(cummins): Keep this after the send_param() call because
+            # send_param() will raise an error if the valid is invalid.
+            self._runtimes_per_observation_count = n
+        except SessionNotFound:
+            pass  # Not in session yet, will be sent on reset().
 
     @property
     def runtime_warmup_runs_count(self) -> int:
@@ -648,13 +643,17 @@ def runtime_warmup_runs_count(self) -> int:
 
     @runtime_warmup_runs_count.setter
     def runtime_warmup_runs_count(self, n: int) -> None:
-        if self.in_episode:
+        try:
             self.send_param(
-                "llvm.set_warmup_runs_count_per_runtime_observation", str(n)
+                "llvm.set_warmup_runs_count_per_runtime_observation",
+                str(n),
+                resend_on_reset=True,
             )
-        # NOTE(cummins): Keep this after the send_param() call because
-        # send_param() will raise an error if the valid is invalid.
-        self._runtimes_warmup_per_observation_count = n
+            # NOTE(cummins): Keep this after the send_param() call because
+            # send_param() will raise an error if the valid is invalid.
+            self._runtimes_warmup_per_observation_count = n
+        except SessionNotFound:
+            pass  # Not in session yet, will be sent on reset().
 
     def fork(self):
         fkd = super().fork()

compiler_gym/service/client_service_compiler_env.py

@@ -202,6 +202,7 @@ def __init__(
 
         self._service_endpoint: Union[str, Path] = service
         self._connection_settings = connection_settings or ConnectionOpts()
+        self._params_to_send_on_reset: List[SessionParameter] = []
 
         self.service = service_connection or CompilerGymServiceConnection(
             endpoint=self._service_endpoint,
@@ -788,6 +789,12 @@ def _call_with_error(
                 reply.new_action_space
             )
 
+        # Re-send any session parameters that we marked as needing to be
+        # re-sent on reset(). Do this before any other initialization as they
+        # may affect the behavior of subsequent service calls.
+        if self._params_to_send_on_reset:
+            self.send_params(*[(p.key, p.value) for p in self._params_to_send_on_reset])
+
         self.reward.reset(benchmark=self.benchmark, observation_view=self.observation)
         if self.reward_space:
             self.episode_reward = 0.0
@@ -1236,7 +1243,7 @@ def validate(self, state: Optional[CompilerEnvState] = None) -> ValidationResult
             **validation,
         )
 
-    def send_param(self, key: str, value: str) -> str:
+    def send_param(self, key: str, value: str, resend_on_reset: bool = False) -> str:
         """Send a single <key, value> parameter to the compiler service.
 
         See :meth:`send_params() <compiler_gym.envs.ClientServiceCompilerEnv.send_params>`
@@ -1246,14 +1253,19 @@ def send_param(self, key: str, value: str) -> str:
 
         :param value: The parameter value.
 
+        :param resend_on_reset: Whether to resend this parameter to the compiler
+            service on :code:`reset()`.
+
         :return: The response from the compiler service.
 
         :raises SessionNotFound: If called before :meth:`reset()
             <compiler_gym.envs.ClientServiceCompilerEnv.reset>`.
         """
-        return self.send_params((key, value))[0]
+        return self.send_params((key, value), resend_on_reset=resend_on_reset)[0]
 
-    def send_params(self, *params: Iterable[Tuple[str, str]]) -> List[str]:
+    def send_params(
+        self, *params: Iterable[Tuple[str, str]], resend_on_reset: bool = False
+    ) -> List[str]:
         """Send a list of <key, value> parameters to the compiler service.
 
         This provides a mechanism to send messages to the backend compilation
@@ -1270,17 +1282,25 @@ def send_params(self, *params: Iterable[Tuple[str, str]]) -> List[str]:
         :param params: A list of parameters, where each parameter is a
             :code:`(key, value)` tuple.
 
+        :param resend_on_reset: Whether to resend this parameter to the compiler
+            service on :code:`reset()`.
+
         :return: A list of string responses, one per parameter.
 
         :raises SessionNotFound: If called before :meth:`reset()
             <compiler_gym.envs.ClientServiceCompilerEnv.reset>`.
         """
+        params_to_send = [SessionParameter(key=k, value=v) for (k, v) in params]
+
+        if resend_on_reset:
+            self._params_to_send_on_reset += params_to_send
+
         if not self.in_episode:
             raise SessionNotFound("Must call reset() before send_params()")
 
         request = SendSessionParameterRequest(
             session_id=self._session_id,
-            parameter=[SessionParameter(key=k, value=v) for (k, v) in params],
+            parameter=params_to_send,
         )
         reply: SendSessionParameterReply = self.service(
             self.service.stub.SendSessionParameter, request

tests/llvm/runtime_test.py

@@ -191,10 +191,10 @@ def test_correct_number_of_observations_during_reset(
 
     # Check that the number of observations that you are receive during reset()
     # matches the amount that you asked for.
-    # FIXME(github.com/facebookresearch/CompilerGym/issues/756): This is broken.
-    # Only a single observation is received, irrespective of how many you ask
-    # for.
-    assert len(env.reward.spaces["runtimeseries"].last_runtime_observation) == 1
+    assert (
+        len(env.reward.spaces["runtimeseries"].last_runtime_observation)
+        == runtime_observation_count
+    )
 
     # Check that the number of observations that you are receive during step()
     # matches the amount that you asked for.