[RLlib] - Add timers to env step, forward pass, and complete connector pipelines runs. #51160
Conversation
Signed-off-by: simonsays1980 <[email protected]>
Signed-off-by: simonsays1980 <[email protected]>
Signed-off-by: simonsays1980 <[email protected]>
Signed-off-by: simonsays1980 <[email protected]>
Signed-off-by: simonsays1980 <[email protected]>
| @@ -320,7 +327,8 @@ def _sample( | |||
| actions = to_env.pop(Columns.ACTIONS) | |||
| actions_for_env = to_env.pop(Columns.ACTIONS_FOR_ENV, actions) | |||
| # Try stepping the environment. | |||
| results = self._try_env_step(actions_for_env) | |||
| with self.metrics.log_time(ENV_STEP_TIMER): | |||
There was a problem hiding this comment.
The metrics context can go inside _try_env_step so we don't have to duplicate the call to it for single- and multi-agent env runners.
| shared_data=shared_data, | ||
| metrics=self.metrics, | ||
| ) | ||
| with self.metrics.log_time(MODULE_TO_ENV_TIMER): |
There was a problem hiding this comment.
We time all connector calls of each ConnectorPipeline.
That's why we pass the metrics object.
We should also encapsulate the timing of the overall connector pipeline call in there so that we don't duplicate this code everywhere (deduplication will make this easier to maintain)
Each connectorpipeline can own its name for logging purposes and we can also take this opportunity to prepend a meaningful prefix to each connector logging entry.
For example log the ClipReward time could be logged as ".....env_to_agent_connector_pipeline.clip_reward" and the overall env to agent pipeline could be logged as "....env_to_agent_connector_pipeline" or something like this.
| ENV_STEP_TIMER = "env_step_timer" | ||
| ENV_TO_MODULE_TIMER = "env_to_module_timer" | ||
| RLMODULE_INFERENCE_TIMER = "rlmodule_inference_timer" | ||
| MODULE_TO_ENV_TIMER = "module_to_env_timer" |
There was a problem hiding this comment.
We should take this opportunity to make the names of our metrics more expressive.
Metrics should self-document as much as possible for users who do not have a good overview of RLlib.
"env_step_timer" is good imo. "env_to_module_timer" should be "env_to_module_connector_pipeline_timer" or something similar. Same goes for the others 👍
There was a problem hiding this comment.
Changes in this file are nice.
Why are these changes needed?
Momentarily, we time solely the complete sampling process as well as single connector pieces inside of the connector pipelines. Important measurements including the inference/exploration forward pass of the
RLModule, the times for an environment step, or the complete time for a run of theEnvToModulePipelineandModuleToEnvPipelineare not timed. This PR proposes necessary changes to include these measurements into our metrics.Related issue number
Checks
git commit -s) in this PR.scripts/format.shto lint the changes in this PR.method in Tune, I've added it in
doc/source/tune/api/under thecorresponding
.rstfile.