2.7 KiB
Profiling in Python
As part of the ML-Agents tookit, we provide a lightweight profiling system, in order to identity hotspots in the training process and help spot regressions from changes.
Timers are hierarchical, meaning that the time tracked in a block of code can be further split into other blocks if desired. This also means that a function that is called from multiple places in the code will appear in multiple places in the timing output.
All timers operate using a "global" instance by default, but this can be overridden if necessary (mainly for testing).
Adding Profiling
There are two ways to indicate code should be included in profiling. The simplest way is to add the @timed
decorator to a function or method of interested.
class TrainerController:
# ....
@timed
def advance(self, env: EnvManager) -> int:
# do stuff
You can also used the hierarchical_timer
context manager.
with hierarchical_timer("communicator.exchange"):
outputs = self.communicator.exchange(step_input)
The context manager may be easier than the @timed
decorator for profiling different parts of a large function, or
profiling calls to abstract methods that might not use decorator.
Output
By default, at the end of training, timers are collected and written in json format to
{summaries_dir}/{run_id}_timers.json
. The output consists of node objects with the following keys:
- total (float): The total time in seconds spent in the block, including child calls.
- count (int): The number of times the block was called.
- self (float): The total time in seconds spent in the block, excluding child calls.
- children (dictionary): A dictionary of child nodes, keyed by the node name.
- is_parallel (bool): Indicates that the block of code was executed in multiple threads or processes (see below). This is optional and defaults to false.
Parallel execution
Subprocesses
For code that executes in multiple processes (for example, SubprocessEnvManager), we periodically send the timer
information back to the "main" process, aggregate the timers there, and flush them in the subprocess. Note that
(depending on the number of processes) this can result in timers where the total time may exceed the parent's total
time. This is analogous to the difference between "real" and "user" values reported from the unix time
command. In the
timer output, blocks that were run in parallel are indicated by the is_parallel
flag.
Threads
Timers currently use time.perf_counter()
to track time spent, which may not give accurate results for multiple
threads. If this is problematic, set threaded: false
in your trainer configuration.