Method maybe_activate_profiler

src/MaxText/profiler.py:51–59

Conditionally activates the profiler based on the current step. This method checks if the current training step matches the step designated for starting an initial profile, or if it meets the criteria for activating a new periodic profile.

(self, step, state)

      ManagedMLDiagnostics(config)  # Initialize the MLRun instance.

  def maybe_activate_profiler(self, step, state):
    """Conditionally activates the profiler based on the current step.
    This method checks if the current training step matches the step designated
    for starting an initial profile, or if it meets the criteria for
    activating a new periodic profile.
    """
    if self.mode != "" and (step == self.start_initial_profile_step or self.should_activate_periodic_profile(step)):
      optional_postfix = f"step_{step}" if self.profile_period > 0 else ""
      self.activate(blocking_object=state, optional_postfix=optional_postfix)

  def activate(self, blocking_object=None, optional_postfix=""):
    """Start the profiler.