# Copyright 2025 - Oumi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Collects sub-step/step/epoch timings."""

import copy
import pathlib
import sys
from pprint import pformat
from typing import Optional, Union

import transformers

import wandb  # isort: skip
from oumi.core.callbacks.base_trainer_callback import BaseTrainerCallback
from oumi.core.configs import TrainingParams
from oumi.core.distributed import get_device_rank_info, is_world_process_zero
from oumi.performance.telemetry import TelemetryTracker, TimerContext
from oumi.utils.device_utils import (
    log_nvidia_gpu_runtime_info,
)
from oumi.utils.io_utils import save_json
from oumi.utils.logging import logger

_LOGS_KWARG = "logs"



[docs]
class TelemetryCallback(BaseTrainerCallback):
    """Trainer callback to collect sub-step/step/epoch timings.

    Based on `oumi.performance.telemetry.TelemetryTracker`.
    """

    def __init__(
        self,
        skip_first_steps: int = 1,
        world_process_zero_only: bool = True,
        include_timer_metrics: bool = False,
        track_gpu_temperature: bool = False,
        output_dir: Optional[pathlib.Path] = None,
    ):
        """Initializes the TelemetryCallback.

        Args:
            skip_first_steps: The number of initial steps to exclude from stats.
            world_process_zero_only: Whether to collect stats on the main process only.
            include_timer_metrics: Whether to add timer stats to reported metrics.
                The timings stats can be verbose/distracting, so `False` by default.
                The timings will be written to a file at the end of training regardless
                of the value of this flag.
            track_gpu_temperature:  Whether to record GPU temperature.
            output_dir: If specified, then telemetry stats will be written to
                the directory as JSON files.
        """
        self._telemetry = TelemetryTracker()
        self._microstep_timer: Optional[TimerContext] = None
        self._step_timer: Optional[TimerContext] = None
        self._epoch_timer: Optional[TimerContext] = None

        self._skip_first_steps: int = skip_first_steps
        self._include_timer_metrics = include_timer_metrics
        self._track_gpu_temperature = track_gpu_temperature
        self._output_dir: Optional[pathlib.Path] = output_dir
        self._permanently_disabled: bool = (
            world_process_zero_only and not is_world_process_zero()
        )
        self._world_process_zero_only = world_process_zero_only
        self._step: int = 0

        self._last_metrics_dict: Optional[dict[str, float]] = None


[docs]
    def on_step_begin(
        self,
        args: Union[transformers.TrainingArguments, TrainingParams],
        state: Optional[transformers.TrainerState] = None,
        control: Optional[transformers.TrainerControl] = None,
        **kwargs,
    ):
        """Event called at the beginning of a training step.

        If using gradient accumulation, one training step might take several inputs.
        """
        self._step += 1
        if self._callback_disabled():
            return

        self._complete_previous_microstep_if_needed()
        self._start_microstep()
        self._complete_previous_step_if_needed()
        self._start_step()



[docs]
    def on_substep_end(
        self,
        args: Union[transformers.TrainingArguments, TrainingParams],
        state: Optional[transformers.TrainerState] = None,
        control: Optional[transformers.TrainerControl] = None,
        **kwargs,
    ):
        """Event called at the end of a substep during gradient accumulation."""
        if self._callback_disabled():
            return

        self._complete_previous_microstep_if_needed()
        self._start_microstep()



[docs]
    def on_step_end(
        self,
        args: Union[transformers.TrainingArguments, TrainingParams],
        state: Optional[transformers.TrainerState] = None,
        control: Optional[transformers.TrainerControl] = None,
        **kwargs,
    ):
        """Event called at the end of each train step.

        Note that this will be called after all gradient accumulation substeps.
        """
        if self._callback_disabled():
            return

        self._complete_previous_microstep_if_needed()
        self._complete_previous_step_if_needed()
        if self._track_gpu_temperature:
            self._telemetry.record_gpu_temperature()



[docs]
    def on_epoch_begin(
        self,
        args: Union[transformers.TrainingArguments, TrainingParams],
        state: Optional[transformers.TrainerState] = None,
        control: Optional[transformers.TrainerControl] = None,
        **kwargs,
    ):
        """Event called at the beginning of an epoch."""
        if self._permanently_disabled:
            return

        self._complete_previous_epoch_if_needed()
        self._start_epoch()



[docs]
    def on_epoch_end(
        self,
        args: Union[transformers.TrainingArguments, TrainingParams],
        state: Optional[transformers.TrainerState] = None,
        control: Optional[transformers.TrainerControl] = None,
        **kwargs,
    ):
        """Event called at the end of an epoch."""
        if self._permanently_disabled:
            return
        self._complete_previous_epoch_if_needed()

        log_nvidia_gpu_runtime_info(log_prefix="On epoch end:")



[docs]
    def on_log(
        self,
        args: Union[transformers.TrainingArguments, TrainingParams],
        state: Optional[transformers.TrainerState] = None,
        control: Optional[transformers.TrainerControl] = None,
        **kwargs,
    ):
        """Event called after logging the last logs."""
        if self._callback_disabled():
            return

        device_rank_info = get_device_rank_info()
        basename = f"telemetry_rank{device_rank_info.rank:03}"

        summary = self._telemetry.get_summary()
        if (
            self._include_timer_metrics
            and "timers" in summary
            and _LOGS_KWARG in kwargs
        ):
            for name, stats in summary["timers"].items():
                for stats_key in ("mean", "median", "std_dev", "min", "max", "count"):
                    if stats_key in stats:
                        metric_name = f"{basename}_{name}_{stats_key}"
                        kwargs[_LOGS_KWARG][metric_name] = float(stats[stats_key])

        if (
            self._track_gpu_temperature
            and "gpu_temperature" in summary
            and summary["gpu_temperature"]
            and _LOGS_KWARG in kwargs
        ):
            stats = summary["gpu_temperature"]
            for stats_key in ("mean", "median", "std_dev", "min", "max", "count"):
                metric_name = f"{basename}_gpu_temperature_{stats_key}"
                kwargs[_LOGS_KWARG][metric_name] = float(stats[stats_key])

        if _LOGS_KWARG in kwargs and is_world_process_zero():
            self._last_metrics_dict = copy.deepcopy(kwargs[_LOGS_KWARG])



[docs]
    def on_train_end(
        self,
        args: Union[transformers.TrainingArguments, TrainingParams],
        state: Optional[transformers.TrainerState] = None,
        control: Optional[transformers.TrainerControl] = None,
        **kwargs,
    ):
        """Event called at the end of training."""
        if self._callback_disabled() or not self._output_dir:
            return

        device_rank_info = get_device_rank_info()

        if is_world_process_zero():
            metrics_dict = self._last_metrics_dict or {}
            save_json(
                metrics_dict,
                self._output_dir
                / f"telemetry_callback_metrics_rank{device_rank_info.rank:04}.json",
            )
            if wandb.run:
                save_json(
                    {
                        "id": wandb.run.id,
                        "name": wandb.run.name,
                        "url": wandb.run.get_url(),
                    },
                    self._output_dir
                    / f"telemetry_callback_wandb_rank{device_rank_info.rank:04}.json",
                )

        if self._world_process_zero_only:
            if is_world_process_zero():
                summary = self._telemetry.get_summary()
                telemetry_file = (
                    self._output_dir
                    / f"telemetry_callback_rank{device_rank_info.rank:04}.json"
                )
                logger.info(f"Saving telemetry callback summary to {telemetry_file}...")
                save_json(summary, telemetry_file)
        else:
            # The function has to be called by all ranks.
            summaries = self._telemetry.get_summaries_from_all_ranks()
            if is_world_process_zero():
                summaries_dict = {
                    f"rank{rank:04}": summary for rank, summary in enumerate(summaries)
                }
                telemetry_file = self._output_dir / "telemetry_callback_all_ranks.json"
                logger.info(
                    "Saving telemetry callback summaries "
                    f"for all ranks to {telemetry_file}..."
                )
                save_json(summaries_dict, telemetry_file)

                gpu_temperature_info_dict = (
                    self._telemetry.compute_cross_rank_summaries(
                        summaries,
                        measurement_names={
                            "gpu_temperature": {"max", "mean", "median"},
                        },
                    )
                )
                logger.info(
                    f"GPU temperature summary:\n{pformat(gpu_temperature_info_dict)}"
                )
                save_json(
                    gpu_temperature_info_dict,
                    self._output_dir
                    / "telemetry_callback_gpu_temperature_summary.json",
                )


    def _callback_disabled(self) -> bool:
        """Check if the callback should be disabled."""
        if self._permanently_disabled:
            return True
        if self._skip_first_steps > 0 and self._step <= self._skip_first_steps:
            return True
        return False

    @staticmethod
    def _exit_timer_if_needed(timer: Optional[TimerContext]) -> Optional[TimerContext]:
        if timer is not None:
            timer.__exit__(*sys.exc_info())
        return None

    def _start_timer(self, timer_name: str) -> TimerContext:
        timer: TimerContext = self._telemetry.timer(timer_name)
        timer.__enter__()
        return timer

    def _complete_previous_microstep_if_needed(self):
        self._microstep_timer = TelemetryCallback._exit_timer_if_needed(
            self._microstep_timer
        )

    def _start_microstep(self):
        self._microstep_timer = self._start_timer("microsteps")

    def _complete_previous_step_if_needed(self):
        self._step_timer = TelemetryCallback._exit_timer_if_needed(self._step_timer)

    def _start_step(self):
        self._step_timer = self._start_timer("steps")

    def _complete_previous_epoch_if_needed(self):
        self._epoch_timer = TelemetryCallback._exit_timer_if_needed(self._epoch_timer)

    def _start_epoch(self):
        self._epoch_timer = self._start_timer("epochs")