fkat 0.1.2__py3-none-any.whl

fkat/pytorch/callbacks/profiling/memray.py

@@ -0,0 +1,212 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import os
+import tempfile
+import atexit
+import signal
+import gzip
+import shutil
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from collections.abc import Sequence
+from typing import Any, TYPE_CHECKING
+from typing_extensions import override
+
+import lightning as L
+
+if TYPE_CHECKING:
+    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
+
+memray = None
+
+
+class Memray(L.Callback):
+    def __init__(
+        self,
+        ranks: Sequence[int] | None = None,
+        flamegraph: bool = False,
+        output_path_prefix: str | None = None,
+        schedule: Schedule | None = None,
+        compress: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """
+        [Memray](https://bloomberg.github.io/memray/api.html) PyTorch Lightning callback.
+        This callbacks traces host RAM (DRAM) allocations and publishes a report to help identify
+        potential memory leaks and investigate OOM errors.
+
+        Args:
+            ranks (Optional[Sequence[int]]): only trace the provided ranks, defaults to all ranks
+            flamegraph (bool): whether to generate [Flamegraph](https://www.brendangregg.com/flamegraphs.html)
+                for the traced allocations, generates HTML report that van be viewed without installing `memray`
+            output_path_prefix (Optional[str]): output path prefix for generated reports,
+                use to persist these files locally, defaults to temporary location that is cleaned as soon as possible
+            schedule (Optional[Schedule]): Controls when logging occurs during training.
+                Defaults to Never - no logging
+            compress (bool): publish reports as compressed files defaults to publishing raw files
+        """
+        self.ranks = ranks
+        self.flamegraph = flamegraph
+        self.compress = compress
+        self.rank: int | None = None
+        self.stage: str | None = None
+        self.kwargs = kwargs
+
+        self.output_path_prefix = output_path_prefix
+        self.schedule = schedule or Never()
+        self._cb_logger: LightningLogger | None = None
+
+        self.executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="Memray")
+
+        global memray
+        import memray  # type: ignore[unresolved-import]
+
+        self.tracker: memray.Tracker | None = None  # type: ignore
+        self.dir = self.tmp_dir = "/tmp"
+
+        signal.signal(signal.SIGTERM, self._terminate)  # terminate signal
+        signal.signal(signal.SIGINT, self._terminate)  # keyboard interrupt
+        atexit.register(self._terminate)
+
+    @override
+    def setup(self, trainer: "L.Trainer", pl_module: "L.LightningModule", stage: str) -> None:
+        self._cb_logger = CallbackLogger(trainer)
+        self.rank = trainer.global_rank
+        self.stage = stage
+
+    @override
+    def on_train_start(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self._start(trainer)
+
+    @override
+    def on_train_batch_end(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        outputs: "STEP_OUTPUT",
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        self._on_batch_end(trainer, "train", batch_idx + 1)
+
+    @override
+    def on_validation_start(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self._start(trainer)
+
+    @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:
+        self._on_batch_end(trainer, "validation", batch_idx + 1)
+
+    @override
+    def on_predict_start(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self._start(trainer)
+
+    @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:
+        self._on_batch_end(trainer, "predict", batch_idx + 1)
+
+    @override
+    def on_test_start(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        self._start(trainer)
+
+    @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:
+        self._on_batch_end(trainer, "test", batch_idx + 1)
+
+    def _on_batch_end(self, trainer: "L.Trainer", stage: str, batch_idx: int) -> None:
+        if self.schedule.check(stage=stage, batch_idx=batch_idx, step=trainer.global_step, trainer=trainer):
+            self._stop(str(batch_idx))
+            self._start(trainer)
+
+    @override
+    def teardown(self, trainer: "L.Trainer", pl_module: "L.LightningModule", stage: str) -> None:
+        self._terminate()
+
+    @override
+    def on_exception(self, trainer: "L.Trainer", pl_module: "L.LightningModule", exception: BaseException) -> None:
+        self._terminate()
+
+    def _terminate(self, *args: Any, **kwargs: Any) -> None:
+        # calling synchronously since this can be called during interpreter shutdown
+        self._stop("last", sync=True)
+        self.executor.shutdown()
+
+    def _start(self, trainer: "L.Trainer") -> None:
+        if self.ranks is not None and trainer.global_rank not in self.ranks:
+            return
+        if not self.tracker:
+            self.tmp_dir = tempfile.mkdtemp()
+            self.dir = self.output_path_prefix or self.tmp_dir
+            path = os.path.join(self.dir, f"rank{trainer.global_rank}.bin")
+            assert memray
+            self.tracker = memray.Tracker(path, **self.kwargs)
+            self.tracker.__enter__()
+        assert self.tracker
+
+    def _stop(self, suffix: str, sync: bool = False) -> None:
+        if not self.tracker:
+            return
+        # create reports synchronously
+        self.tracker.__exit__(None, None, None)
+        self.tracker = None
+        if self.flamegraph:
+            for f in os.listdir(self.dir):
+                results = os.path.join(self.dir, f)
+                from memray.commands.flamegraph import FlamegraphCommand  # type: ignore[unresolved-import]
+
+                # creating this report synchronously because it uses a global memray lock
+                FlamegraphCommand().write_report(Path(results), Path(results + ".html"), True, -1, False)
+        # process reports asynchronously
+        artifacts_path = f"memray/{self.stage}/{suffix}"
+        if sync:
+            self._process(artifacts_path, self.dir, self.tmp_dir)
+        else:
+            self.executor.submit(self._process, artifacts_path, self.dir, self.tmp_dir)
+
+    def _process(
+        self,
+        artifacts_path: str,
+        report_dir: str,
+        tmp_dir: str,
+    ) -> None:
+        assert self._cb_logger
+        for f in os.listdir(report_dir):
+            output_file = os.path.join(report_dir, f)
+            if self.compress:
+                with open(output_file, "rb") as f_in:
+                    output_file = output_file + ".gz"
+                    with gzip.open(output_file, "wb") as f_out:
+                        shutil.copyfileobj(f_in, f_out)
+            self._cb_logger.log_artifact(output_file, artifacts_path)
+        shutil.rmtree(tmp_dir, ignore_errors=True)

fkat/pytorch/callbacks/profiling/torch.py

@@ -0,0 +1,197 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import os
+import gzip
+import shutil
+import tempfile
+import atexit
+import signal
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, TYPE_CHECKING
+from collections.abc import Sequence
+from typing_extensions import override
+
+import lightning as L
+import torch
+
+if TYPE_CHECKING:
+    from lightning.pytorch.utilities.types import STEP_OUTPUT
+
+from fkat.pytorch.schedule import (
+    Schedule,
+    Never,
+)
+from fkat.pytorch.utilities import get_rank
+from fkat.pytorch.loggers import LightningLogger
+from fkat.pytorch.callbacks.loggers import CallbackLogger
+
+
+class PyTorch(L.Callback):
+    def __init__(
+        self,
+        ranks: Sequence[int] | None = None,
+        output_path_prefix: str | None = None,
+        schedule: Schedule | None = None,
+        compress: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        """
+        [PyTorch Profiler](https://docs.pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) Lightning callback.
+        This :class:`L.Callback` continiously traces the training process and publishes a report
+        that helps examining the duration of individual calls through time.
+
+        Args:
+            ranks (Optional[Sequence[int]]): only trace the provided ranks, defaults to all ranks
+            output_path_prefix (Optional[str]): output path prefix for generated reports,
+                use to persist these files locally, defaults to temporary location that is cleaned as soon as possible
+            schedule (Optional[Schedule]): Controls when logging occurs during training.
+                Defaults to :class:`Never` - no intermediate logging
+            compress (bool): compress the report
+                Defaults to ``True``
+            **kwargs (Any): Arbitrary keyword arguments passed as is to PyTorch Profiler
+                except for ``execution_trace_observer`` and ``on_trace_ready``.
+        """
+        self.rank = get_rank()
+        self.compress = compress
+        self.schedule = schedule or Never()
+        self.output_path_prefix = output_path_prefix
+
+        self.trace_observer: torch.profiler.ExecutionTraceObserver | None = None
+        self.trace_file: str | None
+        self.profiler: torch.profiler.profile | None = None
+        if ranks is None or self.rank in ranks:
+            self.trace_file = os.path.join(self.output_path_prefix or tempfile.mkdtemp(), f"rank{self.rank}.json")
+            self.trace_observer = torch.profiler.ExecutionTraceObserver()
+            kwargs.pop("execution_trace_observer", None)
+            kwargs.pop("on_trace_ready", None)
+            self.profiler = torch.profiler.profile(
+                schedule=lambda step: torch.profiler.ProfilerAction.RECORD_AND_SAVE,
+                on_trace_ready=self._publish,
+                execution_trace_observer=self.trace_observer,
+                **kwargs,
+            )
+            self._start_profiler()
+        self._cb_logger: LightningLogger | None = None
+        self.stage: str | None = None
+        self.batch_idx = "?"
+        self.executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="PyTorchProfiler")
+
+        signal.signal(signal.SIGTERM, self._terminate)  # terminate signal
+        signal.signal(signal.SIGINT, self._terminate)  # keyboard interrupt
+        atexit.register(self._terminate)
+
+    @override
+    def setup(self, trainer: "L.Trainer", pl_module: "L.LightningModule", stage: str) -> None:
+        self._cb_logger = CallbackLogger(trainer)
+        self.stage = stage
+
+    def _on_batch_end(self, trainer: "L.Trainer", stage: str, batch_idx: int) -> None:
+        if self.profiler and self.schedule.check(
+            stage=stage, batch_idx=batch_idx + 1, step=trainer.global_step if trainer else None, trainer=trainer
+        ):
+            self.batch_idx = str(batch_idx + 1)
+            self.profiler.step()
+
+    @override
+    def on_train_batch_end(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        outputs: Any,
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        self._on_batch_end(trainer, "train", batch_idx)
+
+    @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:
+        self._on_batch_end(trainer, "validation", batch_idx)
+
+    @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:
+        self._on_batch_end(trainer, "predict", batch_idx)
+
+    @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:
+        self._on_batch_end(trainer, "test", batch_idx)
+
+    def _publish(self, prof: torch.profiler.profile) -> None:
+        # create report synchronously
+        assert self.trace_file
+        prof.export_chrome_trace(self.trace_file)
+        base_path = self.output_path_prefix or os.path.dirname(self.trace_file)
+        output_file = os.path.join(base_path, self.batch_idx, os.path.basename(self.trace_file))
+        os.makedirs(os.path.dirname(output_file), exist_ok=True)
+        shutil.move(self.trace_file, output_file)
+        self._start_profiler()
+        # process report asynchronously
+        artifact_path = f"pt_profiler/{self.stage}/{self.batch_idx}"
+        sync = self.profiler is None
+        if sync:
+            # calling synchronously since this can be called during interpreter shutdown
+            self._process(output_file, artifact_path)
+        else:
+            self.executor.submit(self._process, output_file, artifact_path)
+
+    def _process(self, output_file: str, artifacts_path: str) -> None:
+        assert self._cb_logger
+        if self.compress:
+            with open(output_file, "rb") as f_in:
+                output_file = output_file + ".gz"
+                with gzip.open(output_file, "wb") as f_out:
+                    shutil.copyfileobj(f_in, f_out)
+        self._cb_logger.log_artifact(output_file, artifacts_path)
+        shutil.rmtree(output_file, ignore_errors=True)
+
+    @override
+    def teardown(self, trainer: "L.Trainer", pl_module: "L.LightningModule", stage: str) -> None:
+        self._terminate()
+
+    @override
+    def on_exception(self, trainer: "L.Trainer", pl_module: "L.LightningModule", exception: BaseException) -> None:
+        self._terminate()
+
+    def _terminate(self, *_: Any) -> None:
+        if self.profiler and self.stage:
+            self.batch_idx = "last"
+            self._stop_profiler()
+            self.profiler = None
+            self.executor.shutdown()
+            if self.trace_file:
+                shutil.rmtree(self.trace_file, ignore_errors=True)
+
+    def _start_profiler(self) -> None:
+        assert self.trace_file and self.trace_observer and self.profiler
+        shutil.rmtree(self.trace_file, ignore_errors=True)
+        self.trace_observer.register_callback(self.trace_file)
+        self.profiler.start()
+
+    def _stop_profiler(self) -> None:
+        if self.profiler:
+            self.profiler.stop()
+        if self.trace_observer:
+            self.trace_observer.unregister_callback()

fkat/pytorch/callbacks/profiling/viztracer.py

@@ -0,0 +1,197 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import os
+import json
+import gzip
+import shutil
+import tempfile
+import atexit
+import signal
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, TYPE_CHECKING
+from collections.abc import Sequence
+from typing_extensions import override
+
+import lightning as L
+
+if TYPE_CHECKING:
+    from lightning.pytorch.utilities.types import STEP_OUTPUT
+
+from fkat.pytorch.schedule import (
+    Schedule,
+    Never,
+)
+from fkat.pytorch.utilities import get_rank
+from fkat.pytorch.loggers import LightningLogger
+from fkat.pytorch.callbacks.loggers import CallbackLogger
+
+if TYPE_CHECKING:
+    import viztracer
+else:
+    viztracer = None
+
+
+class VizTracer(L.Callback):
+    def __init__(
+        self,
+        ranks: Sequence[int] | None = None,
+        output_path_prefix: str | None = None,
+        schedule: Schedule | None = None,
+        compress: bool = False,
+        patch: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """
+        [VizTracer](https://viztracer.readthedocs.io/en/latest/) PyTorch Lightning callback.
+        This :class:`L.Callback` continiously traces the training process and publishes a report
+        that helps examining the duration of individual calls through time.
+
+        Args:
+            ranks (Optional[Sequence[int]]): only trace the provided ranks, defaults to all ranks
+            output_path_prefix (Optional[str]): output path prefix for generated reports,
+                use to persist these files locally, defaults to temporary location that is cleaned as soon as possible
+            schedule (Optional[Schedule]): Controls when logging occurs during training.
+                Defaults to :class:`Never` - no logging
+            compress (bool): publish reports as compressed binaries
+                (need to be decompressed via `viztracer --decompress <REPORT>`),
+                if ``True``` saves reports using viztracer's own compression that requires `viztracer` installation,
+                defaults to ``False`` and publishes gzipped HTML reports which require no `viztracer` installation
+            patch (bool): whether to let VizTracer patch internal Python hooks: subprocess, multiprocessing, etc.
+                Defaults to ``False``
+            **kwargs (Any): Arbitrary keyword arguments passed as is to VizTracer.
+        """
+        self.rank = get_rank()
+        self.schedule = schedule or Never()
+        self.output_path_prefix = output_path_prefix
+
+        global viztracer
+        import viztracer
+        from viztracer.vcompressor import VCompressor
+
+        self.compressor = VCompressor() if compress else None
+
+        self.tracer: viztracer.VizTracer | None = None  # type: ignore[no-any-unimported]
+        if ranks is None or self.rank in ranks:
+            kwargs["output_file"] = f"rank{self.rank}.json"
+            kwargs["verbose"] = 0
+            self.tracer = viztracer.VizTracer(**kwargs)
+            assert self.tracer
+            if patch:
+                args = [v for k, v in kwargs.items() for v in ("--" * min(2, len(k)) + k, v)]
+                from viztracer.patch import install_all_hooks
+
+                install_all_hooks(self.tracer, args)
+            self.tracer.start()
+        self._cb_logger: LightningLogger | None = None
+        self.stage: str | None = None
+        self.executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="VizTracer")
+
+        signal.signal(signal.SIGTERM, self._terminate)  # terminate signal
+        signal.signal(signal.SIGINT, self._terminate)  # keyboard interrupt
+        atexit.register(self._terminate)
+
+    @override
+    def setup(self, trainer: "L.Trainer", pl_module: "L.LightningModule", stage: str) -> None:
+        self._cb_logger = CallbackLogger(trainer)
+        self.stage = stage
+
+    def _on_batch_end(self, trainer: "L.Trainer", stage: str, batch_idx: int) -> None:
+        if self.tracer and self.schedule.check(
+            stage=stage, batch_idx=batch_idx, step=trainer.global_step, trainer=trainer
+        ):
+            self._publish(str(batch_idx))
+            self.tracer.start()
+
+    @override
+    def on_train_batch_end(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        outputs: Any,
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        self._on_batch_end(trainer, "train", batch_idx + 1)
+
+    @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:
+        self._on_batch_end(trainer, "validation", batch_idx + 1)
+
+    @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:
+        self._on_batch_end(trainer, "predict", batch_idx + 1)
+
+    @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:
+        self._on_batch_end(trainer, "test", batch_idx + 1)
+
+    def _publish(self, suffix: str, sync: bool = False) -> None:
+        assert self.tracer
+        self.tracer.stop()
+        # create report synchronously
+        tmp_dir = tempfile.mkdtemp()
+        output_stem = os.path.join(self.output_path_prefix or tmp_dir, suffix, f"rank{self.rank}")
+        output_file = output_stem + (".json" if self.compressor else ".html")
+        self.tracer.save(output_file=output_file, verbose=0)
+        self.tracer.clear()
+        # process report asynchronously
+        artifact_path = f"viztracer/{self.stage}/{suffix}"
+        if sync:
+            self._process(tmp_dir, output_file, artifact_path)
+        else:
+            self.executor.submit(self._process, tmp_dir, output_file, artifact_path)
+
+    def _process(self, tmp_dir: str, output_file: str, artifacts_path: str) -> None:
+        if not self.compressor:
+            with open(output_file, "rb") as f_in:
+                output_file = output_file + ".gz"
+                with gzip.open(output_file, "wb") as f_out:
+                    shutil.copyfileobj(f_in, f_out)
+        else:
+            with open(output_file) as f:
+                data = json.load(f)
+            output_file = os.path.splitext(output_file)[0] + ".cvf"
+            self.compressor.compress(data, output_file)
+        assert self._cb_logger
+        self._cb_logger.log_artifact(output_file, artifacts_path)
+        shutil.rmtree(tmp_dir, ignore_errors=True)
+
+    @override
+    def teardown(self, trainer: "L.Trainer", pl_module: "L.LightningModule", stage: str) -> None:
+        self._terminate()
+
+    @override
+    def on_exception(self, trainer: "L.Trainer", pl_module: "L.LightningModule", exception: BaseException) -> None:
+        self._terminate()
+
+    def _terminate(self, *_: Any) -> None:
+        if self.tracer and self.stage:
+            # calling synchronously since this can be called during interpreter shutdown
+            self._publish("last", sync=True)
+            self.tracer.terminate()
+            self.tracer = None
+            self.executor.shutdown()