fkat 0.1.2__py3-none-any.whl

fkat/pytorch/callbacks/logging/validation_metrics.py

@@ -0,0 +1,94 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import json
+import tempfile
+from typing_extensions import override
+
+import fsspec
+import lightning as L
+from lightning.pytorch.utilities import rank_zero_only
+
+from fkat.utils.logging import rank0_logger
+from fkat.pytorch.callbacks.loggers import CallbackLogger
+
+logger = rank0_logger(__name__)
+
+
+class ValidationMetrics(L.Callback):
+    """
+    Saves validation metrics after each validation epoch.
+
+    This callback persists validation metrics in JSON format, creating both
+    a versioned file (with epoch and step) and a "latest" file for easy access.
+
+    Args:
+        output_path (str | None): Directory path where metrics will be saved.
+            Supports any fsspec-compatible filesystem (local, s3://, gcs://, etc.).
+            If None, logs to MLflow artifacts. Defaults to None.
+
+    Example:
+        >>> # MLflow artifacts (default)
+        >>> callback = ValidationMetrics()
+        >>> # Local storage
+        >>> callback = ValidationMetrics(output_path="/tmp/metrics")
+        >>> # S3 storage
+        >>> callback = ValidationMetrics(output_path="s3://my-bucket/metrics")
+        >>> trainer = L.Trainer(callbacks=[callback])
+    """
+
+    def __init__(self, output_path: str | None = None) -> None:
+        self.output_path = output_path
+        self._cb_logger: CallbackLogger | None = None
+
+    @override
+    def setup(self, trainer: "L.Trainer", pl_module: "L.LightningModule", stage: str) -> None:
+        if self.output_path is None:
+            self._cb_logger = CallbackLogger(trainer)
+
+    @override
+    @rank_zero_only
+    def on_validation_end(self, trainer: "L.Trainer", pl_module: "L.LightningModule") -> None:
+        if trainer.sanity_checking:
+            logger.info("Skipping validation metrics save during sanity checking")
+            return
+
+        # Extract metrics
+        metrics_dict = {k: v.item() for k, v in trainer.logged_metrics.items()}
+        metrics_json = json.dumps(metrics_dict)
+
+        # Filenames
+        versioned_name = f"validation-metrics-epoch={pl_module.current_epoch}-step={pl_module.global_step}.json"
+        latest_name = "validation-metrics-latest.json"
+
+        if self.output_path is None:
+            # Log to MLflow artifacts
+            logger.info("Saving validation metrics to MLflow artifacts")
+            with tempfile.TemporaryDirectory() as tmpdir:
+                versioned_file = f"{tmpdir}/{versioned_name}"
+                latest_file = f"{tmpdir}/{latest_name}"
+
+                with open(versioned_file, "w") as f:
+                    f.write(metrics_json)
+                with open(latest_file, "w") as f:
+                    f.write(metrics_json)
+
+                if self._cb_logger:
+                    self._cb_logger.log_artifact(versioned_file, "validation_metrics")
+                    self._cb_logger.log_artifact(latest_file, "validation_metrics")
+
+            logger.info("Validation metrics saved to MLflow artifacts")
+        else:
+            # Use fsspec for filesystem
+            logger.info(f"Saving validation metrics to {self.output_path}")
+            fs, _, paths = fsspec.get_fs_token_paths(self.output_path)
+            base_path = paths[0] if paths else self.output_path
+
+            versioned_file = f"{base_path}/{versioned_name}"
+            latest_file = f"{base_path}/{latest_name}"
+
+            with fs.open(versioned_file, "w") as f:
+                f.write(metrics_json)
+            with fs.open(latest_file, "w") as f:
+                f.write(metrics_json)
+
+            logger.info(f"Validation metrics saved to {versioned_file} and {latest_file}")

fkat/pytorch/callbacks/monitoring/__init__.py

@@ -0,0 +1,14 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from .crash import CrashDetector
+from .dp import DpSyncMonitor
+from .hardware_stats import HardwareStats
+from .shutdown import GracefulShutdown
+
+
+__all__ = [
+    "CrashDetector",
+    "DpSyncMonitor",
+    "GracefulShutdown",
+    "HardwareStats",
+]

fkat/pytorch/callbacks/monitoring/crash.py

@@ -0,0 +1,162 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import os
+import logging
+import traceback
+import multiprocessing
+import datetime as dt
+import tempfile
+from pathlib import Path
+from typing_extensions import override
+
+import lightning as L
+from lightning.pytorch.utilities import rank_zero_only
+
+from fkat.pytorch.loggers import LightningLogger
+from fkat.pytorch.callbacks.loggers import CallbackLogger
+
+log = logging.getLogger(__name__)
+
+
+def _monitor_process(queue: multiprocessing.Queue, parent_pid: int, rank: int) -> None:
+    """Monitor parent process and report crash info."""
+    try:
+        _, status = os.waitpid(parent_pid, 0)
+        exit_code = os.WEXITSTATUS(status) if os.WIFEXITED(status) else -1
+        signal_num = os.WTERMSIG(status) if os.WIFSIGNALED(status) else None
+
+        crash_info = {
+            "pid": parent_pid,
+            "rank": rank,
+            "exit_code": exit_code,
+            "signal": signal_num,
+            "timestamp": str(dt.datetime.now(dt.timezone.utc)),
+        }
+        queue.put(crash_info)
+    except Exception as e:
+        log.error(f"Error monitoring process {parent_pid}: {e}")
+
+
+class CrashDetector(L.Callback):
+    """
+    Detects process crashes and logs detailed error information.
+
+    Monitors the training process and any spawned subprocesses for crashes.
+    Captures PID, rank, error details, and stack traces, logging them to
+    the configured Lightning logger.
+
+    Args:
+        error_tag: Tag for error messages (default: "error")
+        crash_info_tag: Tag for crash details (default: "crash_info")
+
+    Example:
+        >>> from fkat.pytorch.callbacks.monitoring import CrashDetector
+        >>> callback = CrashDetector()
+        >>> trainer = L.Trainer(callbacks=[callback])
+    """
+
+    def __init__(
+        self,
+        error_tag: str = "error",
+        crash_info_tag: str = "crash_info",
+    ) -> None:
+        self.error_tag = error_tag
+        self.crash_info_tag = crash_info_tag
+        self._cb_logger: LightningLogger | None = None
+        self._processes: list[multiprocessing.Process] = []
+        self._queue: multiprocessing.Queue | None = None
+
+    @override
+    def setup(self, trainer: "L.Trainer", pl_module: "L.LightningModule", stage: str) -> None:
+        """Initialize crash detection."""
+        if trainer.local_rank == 0:
+            self._cb_logger = CallbackLogger(trainer)
+            self._queue = multiprocessing.Queue()
+
+            # Monitor main process
+            process = multiprocessing.Process(
+                target=_monitor_process, args=(self._queue, os.getpid(), trainer.global_rank)
+            )
+            process.daemon = True
+            process.start()
+            self._processes.append(process)
+
+    @override
+    @rank_zero_only
+    def on_exception(self, trainer: "L.Trainer", pl_module: "L.LightningModule", exception: BaseException) -> None:
+        """Log exception details."""
+        if not self._cb_logger:
+            return
+
+        exc_type = type(exception)
+        stacktrace = "".join(traceback.format_exception(exc_type, exception, exception.__traceback__))
+
+        error_msg = f"[{exc_type.__name__}]: {exception}"
+        crash_info = {
+            "pid": os.getpid(),
+            "rank": trainer.global_rank,
+            "error": error_msg,
+            "stacktrace": stacktrace,
+            "timestamp": str(dt.datetime.now(dt.timezone.utc)),
+        }
+
+        log.error(f"Exception: {error_msg}\n{stacktrace}")
+        self._cb_logger.log_tag(self.error_tag, error_msg)
+        self._cb_logger.log_tag(self.crash_info_tag, str(crash_info))
+
+        # Log to MLflow artifacts if available
+        self._log_to_mlflow_artifact(trainer, crash_info)
+
+    @override
+    def teardown(self, trainer: "L.Trainer", pl_module: "L.LightningModule", stage: str) -> None:
+        """Check for crashes and cleanup."""
+        if trainer.global_rank == 0 and self._queue:
+            # Check for any crash reports
+            while not self._queue.empty():
+                crash_info = self._queue.get_nowait()
+                log.warning(f"Detected crash: {crash_info}")
+                if self._cb_logger:
+                    self._cb_logger.log_tag(self.crash_info_tag, str(crash_info))
+                    self._log_to_mlflow_artifact(trainer, crash_info)
+
+        self._terminate_monitors()
+
+    def _log_to_mlflow_artifact(self, trainer: "L.Trainer", crash_info: dict) -> None:
+        """Log crash info to MLflow artifacts."""
+        try:
+            from lightning.pytorch.loggers import MLFlowLogger
+
+            mlflow_logger = next((logger for logger in trainer.loggers if isinstance(logger, MLFlowLogger)), None)
+            if not mlflow_logger:
+                return
+
+            # Create filename: rank0-timestamp.txt
+            rank = crash_info.get("rank", 0)
+            timestamp = crash_info.get("timestamp", "unknown")
+            # Convert timestamp to filename-safe format
+            timestamp_safe = timestamp.replace(" ", "_").replace(":", "-")
+            filename = f"rank{rank}-{timestamp_safe}.txt"
+
+            # Create temp file with crash info
+            temp_dir = Path(tempfile.gettempdir())
+            temp_path = temp_dir / filename
+
+            with open(temp_path, "w") as f:
+                f.write("Crash Information\n")
+                f.write("=" * 80 + "\n\n")
+                for key, value in crash_info.items():
+                    f.write(f"{key}: {value}\n")
+
+            # Log as artifact
+            mlflow_logger.experiment.log_artifact(mlflow_logger.run_id, str(temp_path), "crashes")
+            temp_path.unlink()
+            log.info(f"Logged crash info to MLflow artifacts: {filename}")
+        except Exception as e:
+            log.warning(f"Failed to log crash info to MLflow: {e}")
+
+    def _terminate_monitors(self) -> None:
+        """Terminate all monitoring processes."""
+        for process in self._processes:
+            if process.is_alive():
+                process.kill()
+        self._processes.clear()

fkat/pytorch/callbacks/monitoring/dp.py

@@ -0,0 +1,130 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import os
+import time
+from typing import Any, Protocol
+from typing_extensions import override
+
+import lightning as L
+import torch
+import torch.distributed as dist
+
+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 DpGroupStrategy(Protocol):
+    """Protocol for getting DP group info for the current rank."""
+
+    def dp_group_info(self) -> tuple[int, int]:
+        """Return (group_id, rank_in_group) for the current rank."""
+        ...
+
+
+class DistDpGroup(DpGroupStrategy):
+    """Calculates DP group info based on dp_size using distributed rank."""
+
+    def __init__(self, dp_size: int) -> None:
+        self.dp_size = dp_size
+
+    def dp_group_info(self) -> tuple[int, int]:
+        return divmod(dist.get_rank(), self.dp_size)  # type: ignore[possibly-unbound-attribute]
+
+
+class EnvDpGroup(DpGroupStrategy):
+    """Calculates DP group info based on dp_size using environment variables."""
+
+    def __init__(self, dp_size: int) -> None:
+        self.dp_size = dp_size
+
+    def dp_group_info(self) -> tuple[int, int]:
+        rank = int(os.environ.get("RANK", 0))
+        return divmod(rank, self.dp_size)
+
+
+class MegatronDpGroup(DpGroupStrategy):
+    """Gets DP group info from Megatron parallel_state."""
+
+    def dp_group_info(self) -> tuple[int, int]:
+        from megatron.core import parallel_state  # type: ignore[import-not-found]
+
+        group = parallel_state.get_data_parallel_group()
+        rank_in_group = dist.get_rank(group)  # type: ignore[possibly-unbound-attribute]
+        # For Megatron, we need to calculate group_id differently
+        # This assumes we can derive it from global rank and group size
+        global_rank = dist.get_rank()  # Get global rank separately  # type: ignore[possibly-unbound-attribute]
+        group_size = dist.get_world_size(group)  # type: ignore[possibly-unbound-attribute]
+        group_id = global_rank // group_size
+        return group_id, rank_in_group
+
+
+class DpSyncMonitor(L.Callback):
+    """
+    Monitors time for each DP group to reach synchronization point.
+    Measures from batch start to before optimizer step to identify slow/fast groups.
+    """
+
+    def __init__(self, dp_group: DpGroupStrategy, schedule: Schedule | None = None) -> None:
+        """
+        Initialize the DP synchronization monitor.
+
+        Args:
+            dp_group: Strategy for determining DP group info (required).
+            schedule: Controls when logging occurs. Defaults to ``Never``.
+        """
+        self.dp_group = dp_group
+        self.schedule = schedule or Never()
+        self.batch_start_time: float | None = None
+        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)
+
+    @override
+    def on_train_batch_start(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        """Start timing when batch processing begins."""
+        self.batch_start_time = time.perf_counter()
+
+    @override
+    def on_before_optimizer_step(
+        self,
+        trainer: "L.Trainer",
+        pl_module: "L.LightningModule",
+        optimizer: torch.optim.Optimizer,
+    ) -> None:
+        """End timing when ready for sync (before optimizer step) and log if needed."""
+        if self.batch_start_time is not None:
+            sync_time_s = time.perf_counter() - self.batch_start_time
+            # Log immediately since we're at the sync point, before any DP comms
+            self._log_statistics(trainer, "train", 0, sync_time_s)
+        self.batch_start_time = None
+
+    def _log_statistics(self, trainer: "L.Trainer", stage: str, batch_idx: int, sync_time_s: float) -> None:
+        """Log current group timing if schedule permits and this is DP group rank 0."""
+        group_id, rank_in_group = self.dp_group.dp_group_info()
+        if rank_in_group != 0:
+            return
+
+        if not self.schedule.check(stage=stage, batch_idx=batch_idx, step=trainer.global_step, trainer=trainer):
+            return
+
+        if self._cb_logger:
+            metrics = {f"dp_sync/group{group_id}/sync_s": sync_time_s}
+            self._cb_logger.log_batch(metrics=metrics, timestamp=int(time.time()), step=trainer.global_step)

fkat/pytorch/callbacks/monitoring/hardware_stats.py

@@ -0,0 +1,135 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import logging
+from time import time_ns
+from typing import Any
+from typing_extensions import override
+
+import lightning as L
+from lightning.pytorch.utilities import rank_zero_only
+import psutil
+import torch
+
+from fkat.pytorch.callbacks.loggers import CallbackLogger
+from fkat.pytorch.schedule import Schedule, Every
+
+logger = logging.getLogger(__name__)
+
+
+class HardwareStats(L.Callback):
+    """Monitor and log hardware usage (CPU, RAM, and GPU) during training.
+
+    Args:
+        accelerator: Hardware accelerator to monitor ("gpu" or "cpu").
+            If None or unsupported, only CPU/RAM are monitored.
+        schedule: Controls when hardware stats are logged. Defaults to every batch.
+    """
+
+    def __init__(self, accelerator: str | None = None, schedule: Schedule | None = None) -> None:
+        if accelerator not in ("gpu", "cpu", None):
+            logger.warning(f"Unsupported accelerator: {accelerator}. Monitoring CPU/RAM only.")
+            accelerator = "cpu"
+        self.accelerator = accelerator
+        self.schedule = schedule or Every(n_batches=1)
+        self._cb_logger: CallbackLogger | None = None
+        self._total_gpu_memory_gb: float | None = None
+
+    @override
+    def setup(self, trainer: L.Trainer, pl_module: L.LightningModule, stage: str) -> None:
+        self._cb_logger = CallbackLogger(trainer)
+        if self.accelerator == "gpu":
+            if torch.cuda.is_available():
+                self._total_gpu_memory_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+                torch.cuda.reset_peak_memory_stats()
+            else:
+                logger.warning("GPU accelerator requested but CUDA not available. Monitoring CPU/RAM only.")
+
+    @rank_zero_only
+    @override
+    def on_train_epoch_start(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
+        if self.accelerator == "gpu" and torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats()
+
+    def _get_stats(self) -> dict[str, float]:
+        stats = {
+            "cpu_usage_percent": psutil.cpu_percent(),
+            "ram_used_percent": psutil.virtual_memory().percent,
+            "ram_used_gb": psutil.virtual_memory().used / 1e9,
+        }
+        if self.accelerator == "gpu" and torch.cuda.is_available():
+            stats.update(
+                {
+                    "gpu_memory_reserved_gb": torch.cuda.memory_reserved() / 1e9,
+                    "gpu_memory_allocated_gb": torch.cuda.memory_allocated() / 1e9,
+                    "gpu_peak_memory_allocated_gb": torch.cuda.max_memory_allocated() / 1e9,
+                    "gpu_memory_total_per_rank_gb": self._total_gpu_memory_gb or 0.0,
+                }
+            )
+        return stats
+
+    def _log_stats(self, trainer: L.Trainer) -> None:
+        if self._cb_logger:
+            self._cb_logger.log_batch(metrics=self._get_stats(), step=trainer.global_step, timestamp=int(time_ns()))
+        if self.accelerator == "gpu" and torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats()
+
+    @rank_zero_only
+    @override
+    def on_train_batch_start(
+        self, trainer: L.Trainer, pl_module: L.LightningModule, batch: Any, batch_idx: int
+    ) -> None:
+        if self.schedule.check(stage="train", batch_idx=batch_idx, step=trainer.global_step, trainer=trainer):
+            self._log_stats(trainer)
+
+    @rank_zero_only
+    @override
+    def on_train_batch_end(
+        self, trainer: L.Trainer, pl_module: L.LightningModule, outputs: Any, batch: Any, batch_idx: int
+    ) -> None:
+        if self.schedule.check(stage="train", batch_idx=batch_idx, step=trainer.global_step, trainer=trainer):
+            self._log_stats(trainer)
+
+    @rank_zero_only
+    @override
+    def on_validation_batch_start(
+        self, trainer: L.Trainer, pl_module: L.LightningModule, batch: Any, batch_idx: int, dataloader_idx: int = 0
+    ) -> None:
+        if self.schedule.check(stage="validation", batch_idx=batch_idx, step=trainer.global_step, trainer=trainer):
+            self._log_stats(trainer)
+
+    @rank_zero_only
+    @override
+    def on_validation_batch_end(
+        self,
+        trainer: L.Trainer,
+        pl_module: L.LightningModule,
+        outputs: Any,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        if self.schedule.check(stage="validation", batch_idx=batch_idx, step=trainer.global_step, trainer=trainer):
+            self._log_stats(trainer)
+
+    @rank_zero_only
+    @override
+    def on_validation_epoch_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
+        self._log_stats(trainer)
+
+    @rank_zero_only
+    @override
+    def on_before_zero_grad(self, trainer: L.Trainer, pl_module: L.LightningModule, optimizer: Any) -> None:
+        if self.schedule.check(
+            stage="train", batch_idx=trainer.fit_loop.batch_idx, step=trainer.global_step, trainer=trainer
+        ):
+            self._log_stats(trainer)
+
+    @rank_zero_only
+    @override
+    def teardown(self, trainer: L.Trainer, pl_module: L.LightningModule, stage: str) -> None:
+        self._log_stats(trainer)
+
+    @rank_zero_only
+    @override
+    def on_exception(self, trainer: L.Trainer, pl_module: L.LightningModule, exception: BaseException) -> None:
+        self._log_stats(trainer)