fkat 0.1.2__py3-none-any.whl

fkat/pytorch/callbacks/gc.py

@@ -0,0 +1,146 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import gc
+from time import perf_counter
+from typing import Any
+from typing_extensions import override
+
+import lightning as L
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+
+from fkat.pytorch.schedule import (
+    Schedule,
+    Never,
+)
+from fkat.pytorch.loggers import LightningLogger
+from fkat.pytorch.callbacks.loggers import CallbackLogger
+
+
+class ManualGc(L.Callback):
+    def __init__(
+        self,
+        schedule: Schedule | None = None,
+        stats: dict[str, list[str]] | None = None,
+    ) -> None:
+        """
+        PyTorch Lightning callback for manual garbage collection (GC) control.
+
+        This callback allows fine-grained control over Python's garbage collection during training,
+        validation, testing, and prediction. It can disable automatic garbage collection and instead
+        perform manual collection at specified batch intervals.
+
+        Args:
+            schedule (Optional[Schedule]): When to invoke manual GC, defaults to class:`Never`
+            stats (Optional[dict[str, list[str]]]): The list of stats to log per generation.
+                Defaults to all stats for all generations
+
+        Example:
+            >>> trainer = Trainer(callbacks=[ManualGc()])
+        """
+        self.schedule = schedule or Never()
+        self.stats = stats or dict(
+            zip("012", (["collections", "collected", "uncollected"] for _ in range(3)), strict=True)
+        )
+        self._cb_logger: LightningLogger | None = None
+
+    @override
+    def setup(self, trainer: "L.Trainer", pl_module: "L.LightningModule", stage: str) -> None:
+        self._cb_logger = self._cb_logger or CallbackLogger(trainer)
+        if not isinstance(self.schedule, Never):
+            gc.disable()
+
+    @override
+    def on_train_epoch_end(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        """Perform GC after training epoch."""
+        self.maybe_collect()
+
+    @override
+    def on_train_batch_end(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        outputs: "STEP_OUTPUT",
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        """Perform GC after training batch if needed."""
+        self.maybe_gc(trainer, "train", batch_idx)
+
+    @override
+    def on_validation_epoch_end(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        """Perform GC after validation epoch."""
+        self.maybe_collect()
+
+    @override
+    def on_validation_batch_end(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        outputs: "STEP_OUTPUT",
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Perform GC after validation batch if needed."""
+        self.maybe_gc(trainer, "validation", batch_idx)
+
+    @override
+    def on_test_epoch_end(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        """Perform GC after test epoch."""
+        self.maybe_collect()
+
+    @override
+    def on_predict_batch_end(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        outputs: Any,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Perform GC after prediction batch if needed."""
+        self.maybe_gc(trainer, "predict", batch_idx)
+
+    @override
+    def on_predict_epoch_end(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        """Perform GC after predict epoch."""
+        self.maybe_collect()
+
+    @override
+    def on_test_batch_end(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        outputs: "STEP_OUTPUT",
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Perform GC after test batch if needed."""
+        self.maybe_gc(trainer, "test", batch_idx)
+
+    def maybe_collect(self) -> None:
+        if not isinstance(self.schedule, Never):
+            gc.collect()
+
+    def maybe_gc(self, trainer: "L.Trainer", stage: str, batch_idx: int) -> None:
+        """
+        Perform garbage collection if conditions are met.
+
+        Args:
+            batch_idx (int): Current batch index
+        """
+        if self.schedule.check(stage=stage, batch_idx=batch_idx, step=trainer.global_step, trainer=trainer):
+            now = perf_counter()
+            gc.collect()
+            if self._cb_logger is not None:
+                metrics = {f"gc/rank{trainer.global_rank}/time": perf_counter() - now}
+                for i, stats in enumerate(gc.get_stats()):
+                    gen = str(i)
+                    for key, value in stats.items():
+                        if key in self.stats[gen]:
+                            metrics[f"gc/rank{trainer.global_rank}/gen{gen}/{key}"] = float(value)
+                self._cb_logger.log_batch(
+                    metrics=metrics, timestamp=int(now or perf_counter()), step=trainer.global_step
+                )

fkat/pytorch/callbacks/loggers.py

@@ -0,0 +1,211 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any, TYPE_CHECKING
+from typing_extensions import override
+
+import lightning as L
+from lightning.pytorch.loggers import TensorBoardLogger, WandbLogger
+from mlflow.entities import Metric, RunTag, Param
+from mlflow.tracking import MlflowClient  # type: ignore[possibly-unbound-import]
+
+if TYPE_CHECKING:
+    pass
+
+from fkat.pytorch.loggers import LightningLogger, _is_logger_type
+from fkat.utils import assert_not_none
+from fkat.utils.logging import rank0_logger
+from fkat.utils.mlflow import broadcast_mlflow_run_id, mlflow_logger
+
+log = rank0_logger(__name__)
+
+
+class MLFlowCallbackLogger(LightningLogger):
+    """
+    Mlflow logger class that supports distributed logging of tags, metrics and artifacts.
+
+    Args:
+        trainer (L.Trainer): PTL trainer object
+    """
+
+    def __init__(
+        self,
+        trainer: "L.Trainer | None" = None,
+        client: MlflowClient | None = None,
+        synchronous: bool | None = None,
+        run_id: str | None = None,
+    ) -> None:
+        super().__init__()
+        if trainer:
+            # Initialize logger and broadcast run_ids to all ranks
+            logger = assert_not_none(mlflow_logger(trainer))
+            broadcast_mlflow_run_id(logger, trainer)  # type: ignore[arg-type]
+            # Set client and run_id
+            self._client: MlflowClient = assert_not_none(getattr(logger, "_mlflow_client", None))
+            self._synchronous = getattr(logger, "_log_batch_kwargs", {}).get("synchronous")
+            self._run_id: str = assert_not_none(getattr(logger, "_run_id", None))
+        else:
+            assert client
+            self._client = client
+            self._synchronous = synchronous
+            assert run_id
+            self._run_id = run_id
+
+    def log_tag(self, key: str, value: str) -> None:
+        self._client.set_tag(run_id=self._run_id, key=key, value=value, synchronous=self._synchronous)
+
+    def tags(self) -> dict[str, Any]:
+        run = self._client.get_run(self._run_id)
+        return run.data.tags
+
+    def log_batch(
+        self,
+        metrics: dict[str, float] | None = None,
+        params: dict[str, Any] | None = None,
+        tags: dict[str, str] | None = None,
+        timestamp: int | None = None,
+        step: int | None = None,
+    ) -> None:
+        ms = [Metric(k, v, timestamp, step) for k, v in metrics.items()] if metrics else []
+        ps = [Param(k, v) for k, v in params.items()] if params else []
+        ts = [RunTag(k, v) for k, v in tags.items()] if tags else []
+        self._client.log_batch(
+            run_id=self._run_id,
+            metrics=ms,
+            params=ps,
+            tags=ts,
+            synchronous=self._synchronous,
+        )
+
+    def log_artifact(self, local_path: str, artifact_path: str | None = None) -> None:
+        # TODO: log directly to s3 uri
+        # TODO: support async logging
+        self._client.log_artifact(
+            run_id=self._run_id,
+            local_path=local_path,
+            artifact_path=artifact_path,
+        )
+
+
+class TensorBoardCallbackLogger(LightningLogger):
+    """TensorBoard logger for distributed logging."""
+
+    def __init__(self, logger: TensorBoardLogger) -> None:
+        self._logger = logger
+
+    def log_tag(self, key: str, value: str) -> None:
+        self._logger.experiment.add_text(key, value)
+
+    def tags(self) -> dict[str, Any]:
+        return {}
+
+    def log_batch(
+        self,
+        metrics: dict[str, float] | None = None,
+        params: dict[str, Any] | None = None,
+        tags: dict[str, str] | None = None,
+        timestamp: int | None = None,
+        step: int | None = None,
+    ) -> None:
+        if metrics:
+            for k, v in metrics.items():
+                self._logger.experiment.add_scalar(k, v, step)
+        if tags:
+            for k, v in tags.items():
+                self._logger.experiment.add_text(k, v, step)
+
+    def log_artifact(self, local_path: str, artifact_path: str | None = None) -> None:
+        pass
+
+
+class WandbCallbackLogger(LightningLogger):
+    """WandB logger for distributed logging."""
+
+    def __init__(self, logger: WandbLogger) -> None:
+        self._logger = logger
+
+    def log_tag(self, key: str, value: str) -> None:
+        self._logger.experiment.config.update({key: value})
+
+    def tags(self) -> dict[str, Any]:
+        return dict(self._logger.experiment.config)
+
+    def log_batch(
+        self,
+        metrics: dict[str, float] | None = None,
+        params: dict[str, Any] | None = None,
+        tags: dict[str, str] | None = None,
+        timestamp: int | None = None,
+        step: int | None = None,
+    ) -> None:
+        log_dict = {}
+        if metrics:
+            log_dict.update(metrics)
+        if tags:
+            log_dict.update(tags)
+        if log_dict:
+            self._logger.experiment.log(log_dict, step=step)
+
+    def log_artifact(self, local_path: str, artifact_path: str | None = None) -> None:
+        self._logger.experiment.save(local_path)
+
+
+class CallbackLogger(LightningLogger):
+    """
+    A wrapper on top of the collection of Logger instances,
+    providing methods to log metrics, artifacts, and tags across all registered loggers
+    simultaneously.
+
+    Attributes:
+        loggers (list[LightningLogger]): List of loggers
+
+    Args:
+        trainer (L.Trainer): PyTorch Lightning trainer instance used to initialize loggers
+    """
+
+    loggers: list[LightningLogger]
+
+    def __init__(self, trainer: "L.Trainer | None", loggers: list[LightningLogger] | None = None) -> None:
+        if trainer:
+            self.loggers = []
+            for logger in trainer.loggers:
+                if _is_logger_type(logger, "MLFlowLogger"):
+                    self.loggers.append(MLFlowCallbackLogger(trainer=trainer))
+                elif _is_logger_type(logger, "TensorBoardLogger"):
+                    self.loggers.append(TensorBoardCallbackLogger(logger=logger))  # type: ignore[arg-type]
+                elif _is_logger_type(logger, "WandbLogger"):
+                    self.loggers.append(WandbCallbackLogger(logger=logger))  # type: ignore[arg-type]
+        else:
+            assert loggers
+            self.loggers = loggers
+
+    def __str__(self) -> str:
+        return str([type(obj).__name__ for obj in self.loggers])
+
+    @override
+    def log_artifact(self, local_path: str, artifact_path: str | None = None) -> None:
+        for logger in self.loggers:
+            logger.log_artifact(local_path=local_path, artifact_path=artifact_path)
+
+    @override
+    def log_batch(
+        self,
+        metrics: dict[str, float] | None = None,
+        params: dict[str, Any] | None = None,
+        tags: dict[str, str] | None = None,
+        timestamp: int | None = None,
+        step: int | None = None,
+    ) -> None:
+        for logger in self.loggers:
+            logger.log_batch(metrics=metrics, tags=tags, timestamp=timestamp, step=step)
+
+    @override
+    def tags(self) -> dict[str, Any]:
+        tags = {}
+        for logger in self.loggers:
+            tags.update(logger.tags())
+        return tags
+
+    @override
+    def log_tag(self, key: str, value: str) -> None:
+        for logger in self.loggers:
+            logger.log_tag(key=key, value=value)

fkat/pytorch/callbacks/logging/__init__.py

@@ -0,0 +1,12 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from .heartbeat import Heartbeat
+from .throughput import Throughput
+from .validation_metrics import ValidationMetrics
+
+
+__all__ = [
+    "Heartbeat",
+    "Throughput",
+    "ValidationMetrics",
+]

fkat/pytorch/callbacks/logging/heartbeat.py

@@ -0,0 +1,76 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import datetime as dt
+from typing import Any
+from typing_extensions import override
+
+import lightning as L
+from lightning.pytorch.utilities import rank_zero_only
+
+from fkat.pytorch.schedule import (
+    Schedule,
+    Elapsed,
+)
+from fkat.pytorch.loggers import LightningLogger
+from fkat.pytorch.callbacks.loggers import CallbackLogger
+from fkat.utils.logging import rank0_logger
+
+log = rank0_logger(__name__)
+
+
+class Heartbeat(L.Callback):
+    """Publishes tags indicating the time and step of the last heartbeat with the provided schedule."""
+
+    def __init__(
+        self,
+        schedule: Schedule | None = None,
+        last_check_in_time_tag: str = "last_check_in_time",
+        last_check_in_step_tag: str = "last_check_in_step",
+    ) -> None:
+        self.last_check_in_time_tag = last_check_in_time_tag
+        self.last_check_in_step_tag = last_check_in_step_tag
+        self.schedule = schedule or Elapsed(interval=dt.timedelta(minutes=15))
+        self._cb_logger: LightningLogger | None = None
+
+    @override
+    def setup(self, trainer: "L.Trainer", pl_module: "L.LightningModule", stage: str) -> None:
+        self._cb_logger = CallbackLogger(trainer)
+
+    def _publish_tags(self, stage: str, batch_idx: int, trainer: "L.Trainer") -> None:
+        if self.schedule.check(stage=stage, batch_idx=batch_idx, step=trainer.global_step, trainer=trainer):
+            assert self._cb_logger
+            time = dt.datetime.now(dt.timezone.utc)
+            self._cb_logger.log_batch(
+                tags={
+                    self.last_check_in_time_tag: str(time),
+                    self.last_check_in_step_tag: str(trainer.global_step),
+                }
+            )
+
+    @override
+    @rank_zero_only
+    def on_train_batch_start(
+        self, trainer: "L.Trainer", pl_module: "L.LightningModule", batch: Any, batch_idx: int
+    ) -> None:
+        self._publish_tags("train", batch_idx, trainer)
+
+    @override
+    @rank_zero_only
+    def on_test_batch_start(
+        self, trainer: "L.Trainer", pl_module: "L.LightningModule", batch: Any, batch_idx: int, dataloader_idx: int = 0
+    ) -> None:
+        self._publish_tags("test", batch_idx, trainer)
+
+    @override
+    @rank_zero_only
+    def on_validation_batch_start(
+        self, trainer: "L.Trainer", pl_module: "L.LightningModule", batch: Any, batch_idx: int, dataloader_idx: int = 0
+    ) -> None:
+        self._publish_tags("validation", batch_idx, trainer)
+
+    @override
+    @rank_zero_only
+    def on_predict_batch_start(
+        self, trainer: "L.Trainer", pl_module: "L.LightningModule", batch: Any, batch_idx: int, dataloader_idx: int = 0
+    ) -> None:
+        self._publish_tags("predict", batch_idx, trainer)

fkat/pytorch/callbacks/logging/throughput.py

@@ -0,0 +1,253 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from time import time
+from typing import Any, TYPE_CHECKING
+
+import lightning as L
+from lightning.pytorch.callbacks import LearningRateFinder, BatchSizeFinder
+from lightning.pytorch.utilities.data import extract_batch_size
+from lightning.pytorch.utilities import rank_zero_only
+
+if TYPE_CHECKING:
+    from lightning.pytorch.utilities.types import STEP_OUTPUT
+from typing_extensions import override
+
+import torch
+from fkat.pytorch.schedule import (
+    Schedule,
+    Never,
+)
+from fkat.pytorch.loggers import LightningLogger
+from fkat.pytorch.callbacks.loggers import CallbackLogger
+from fkat.utils.logging import rank0_logger
+
+logger = rank0_logger(__name__)
+
+
+class Throughput(L.Callback):
+    def __init__(self, dp_ranks: int | None = None, schedule: Schedule | None = None) -> None:
+        """
+        Throughput logging callback that measures the time spent processing the microbatches.
+        Args:
+            schedule (Optional[Schedule]): Controls when logging occurs. Defaults to ``Never``.
+        """
+        self.schedule = schedule or Never()
+        self.dp_ranks: int | None = dp_ranks
+        self.was_last_step_val = False
+        self.publish = False
+        self.step_start_time: dict[str, float] = {}
+        self.step_time: dict[str, float] = {}
+        self.total_time: dict[str, float] = {}
+        self.step_samples: dict[str, float] = {}
+        self.total_samples: dict[str, float] = {}
+        self.epoch_start_time: dict[str, float] = {}
+        self._cb_logger: LightningLogger | None = None
+
+    @override
+    def setup(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        stage: str,
+    ) -> None:
+        if not self._cb_logger:
+            self._cb_logger = CallbackLogger(trainer)
+            # ignoring special callbacks used for tuning
+            callbacks = [
+                c
+                for c in trainer.callbacks  # type: ignore[attr-defined]
+                if not isinstance(c, LearningRateFinder | BatchSizeFinder)
+            ]
+            tput_callbacks = [i for i, c in enumerate(callbacks) if isinstance(c, Throughput)]
+            assert len(tput_callbacks) == 1, "There can only be one Throughput logging callback in operation"
+        self.dp_ranks: int = self.dp_ranks or trainer.world_size
+        self.step_start_time = {}
+        self.step_time = {}
+        self.total_time = {}
+        self.step_samples = {}
+        self.total_samples = {}
+        self.epoch_start_time = {}
+
+    def _start_epoch(self, stage: str) -> None:
+        self.epoch_start_time[stage] = self.step_start_time[stage] = time()
+
+    @override
+    @rank_zero_only
+    def on_train_epoch_start(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self._start_epoch("train")
+
+    @override
+    @rank_zero_only
+    def on_validation_epoch_start(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self._start_epoch("validation")
+
+    @override
+    @rank_zero_only
+    def on_test_epoch_start(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self._start_epoch("test")
+
+    @override
+    @rank_zero_only
+    def on_predict_epoch_start(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self._start_epoch("predict")
+
+    def _report_epoch(self, trainer: "L.Trainer", stage: str) -> None:
+        self.step_start_time[stage] = (now := time())
+        metrics = {
+            f"{stage}/epochs/epoch_time": now - self.epoch_start_time[stage],
+        }
+        if self._cb_logger:
+            self._cb_logger.log_batch(metrics=metrics, timestamp=int(now), step=trainer.global_step)
+        else:
+            logger.info(metrics)
+
+    @override
+    @rank_zero_only
+    def on_train_epoch_end(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self._report_epoch(trainer, "train")
+
+    @override
+    @rank_zero_only
+    def on_validation_epoch_end(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self._report_epoch(trainer, "validation")
+
+    @override
+    @rank_zero_only
+    def on_test_epoch_end(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self._report_epoch(trainer, "test")
+
+    @override
+    @rank_zero_only
+    def on_predict_epoch_end(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self._report_epoch(trainer, "predict")
+
+    def _update(self, stage: str, batch: Any, batch_idx: int, step: int | None, trainer: "L.Trainer") -> None:
+        # because of other callbacks we want to only measure within batch start/end
+        # and make sure this callback is the first in the list
+        now = time()
+        self.step_time[stage] = self.step_time.get(stage, 0) + (now - self.step_start_time[stage])
+        self.step_start_time[stage] = now
+        num_samples = extract_batch_size(batch) if batch else 0
+        self.step_samples[stage] = self.step_samples.get(stage, 0) + num_samples
+        # train data points have to be at step boundaries or we will have multiple datapoints for the same step
+        self.publish = stage != "train" and self.schedule.check(
+            stage=stage, batch_idx=batch_idx, step=step, trainer=trainer
+        )
+
+    @override
+    @rank_zero_only
+    def on_before_zero_grad(
+        self, trainer: "L.Trainer", pl_module: "L.LightningModule", optimizer: "torch.optim.Optimizer"
+    ) -> None:
+        """
+        Report metrics for individual steps during training.
+        """
+        self.publish = True  # always log on step
+
+    @override
+    @rank_zero_only
+    def on_train_batch_end(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        outputs: "STEP_OUTPUT",
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        # throughput is the first callback so it's safe to capture time here
+        self._update("train", batch, batch_idx, trainer.global_step, trainer)
+        self._report(trainer, "train")
+
+    @override
+    @rank_zero_only
+    def on_validation_batch_end(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        outputs: "STEP_OUTPUT",
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        self._update("validation", batch, batch_idx, trainer.global_step, trainer)
+        self._report(trainer, "validation")
+
+    @override
+    @rank_zero_only
+    def on_test_batch_end(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        outputs: "STEP_OUTPUT",
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        self._update("test", batch, batch_idx, trainer.global_step, trainer)
+        self._report(trainer, "test")
+
+    @override
+    @rank_zero_only
+    def on_predict_batch_end(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        outputs: Any,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        self._update("predict", batch, batch_idx, trainer.global_step, trainer)
+        self._report(trainer, "predict")
+
+    def _report(self, trainer: "L.Trainer", stage: str) -> None:
+        if not self.publish:
+            return
+        if not self.step_time.get(stage):
+            # can end up here outside of training loop e.g. when initializing precision plugin
+            return
+        self.total_time[stage] = self.total_time.get(stage, 0) + self.step_time[stage]
+        self.total_samples[stage] = self.total_samples.get(stage, 0) + self.step_samples[stage]
+        rank0_avg_tput = self.total_samples[stage] / self.total_time[stage]
+        assert self.dp_ranks
+        metrics = {
+            f"{stage}/throughput/running_avg_rank0": rank0_avg_tput,
+            f"{stage}/throughput/running_avg": self.dp_ranks * rank0_avg_tput,
+        }
+        if stage == "train":
+            # we only have steps during fit
+            metrics[f"{stage}/steps/step_time"] = self.step_time[stage]
+        rank0_tput = self.step_samples[stage] / self.step_time[stage]
+        metrics[f"{stage}/throughput/current_rank0"] = rank0_tput
+        metrics[f"{stage}/throughput/current"] = self.dp_ranks * rank0_tput
+        if self._cb_logger:
+            self._cb_logger.log_batch(metrics=metrics, timestamp=int(time()), step=trainer.global_step)
+        else:
+            logger.info(metrics)
+        self.step_time[stage] = 0.0
+        self.step_samples[stage] = 0
+
+    @override
+    @rank_zero_only
+    def on_train_end(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self.publish = True
+        self._report(trainer, "train")
+
+    @override
+    @rank_zero_only
+    def on_validation_end(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self.publish = True
+        self._report(trainer, "validation")
+
+    @override
+    @rank_zero_only
+    def on_test_end(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self.publish = True
+        self._report(trainer, "test")
+
+    @override
+    @rank_zero_only
+    def on_predict_end(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self.publish = True
+        self._report(trainer, "predict")