fkat 0.1.2__py3-none-any.whl

fkat/pytorch/loggers.py

@@ -0,0 +1,284 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any, Protocol, TYPE_CHECKING
+from typing_extensions import override
+
+import lightning as L
+from lightning.pytorch.utilities import rank_zero_only
+from mlflow.entities import Metric, RunTag, Param
+from mlflow.tracking import MlflowClient  # type: ignore[possibly-unbound-import]
+
+if TYPE_CHECKING:
+    from lightning.pytorch.loggers import MLFlowLogger, TensorBoardLogger, WandbLogger
+
+from fkat.utils import assert_not_none
+
+
+def _is_logger_type(logger: Any, logger_name: str) -> bool:
+    """Check if logger matches type from lightning or pytorch_lightning."""
+    module = type(logger).__module__
+    return type(logger).__name__ == logger_name and (
+        module.startswith("lightning.pytorch.loggers") or module.startswith("pytorch_lightning.loggers")
+    )
+
+
+class LightningLogger(Protocol):
+    """
+    Protocol defining the interface for logging that handle metrics, tags, and artifacts.
+    """
+
+    def tags(self) -> dict[str, Any]:
+        """Get current tags"""
+        ...
+
+    def log_tag(self, key: str, value: str) -> None:
+        """
+        Log a single key-value tag.
+
+        Args:
+            key (str): The identifier/name of the tag
+            value (str): The value associated with the tag
+        """
+        ...
+
+    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 multiple metrics and/or tags in a single batch operation.
+
+        Args:
+            metrics (dict[str, float], optional): Dictionary mapping metric names to their float values
+            params (dict[str, Any], optional): Dictionary mapping params names to their values
+            tags (dict[str, str], optional): Dictionary mapping tag names to their string values
+            timestamp (int, optional): Unix timestamp for when the batch was logged
+            step (int, optional): Training step or iteration number
+        """
+        ...
+
+    def log_artifact(self, local_path: str, artifact_path: str | None = None) -> None:
+        """
+        Log a local file as an artifact.
+
+        Args:
+            local_path (str): Path to the file on the local filesystem to be logged
+            artifact_path (str, optional): Remote path where the artifact should be stored
+                If None, a default location should be used
+        """
+        ...
+
+
+class MLFlowLogger:
+    """
+    Mlflow logger class that supports rank_zero logging of tags, metrics and distributed
+    logging of artifacts.
+
+    Args:
+        logger (MLFlowLogger): PTL MLFlow logger object
+    """
+
+    def __init__(
+        self,
+        logger: "MLFlowLogger | None" = None,
+        client: MlflowClient | None = None,
+        synchronous: bool | None = None,
+        run_id: str | None = None,
+    ) -> None:
+        super().__init__()
+        if logger:
+            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
+
+    @rank_zero_only
+    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
+
+    @rank_zero_only
+    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 TensorBoardLogger(LightningLogger):
+    """TensorBoard logger with rank_zero logging."""
+
+    def __init__(self, logger: "TensorBoardLogger") -> None:
+        self._logger = logger
+
+    @rank_zero_only
+    def log_tag(self, key: str, value: str) -> None:
+        self._logger.experiment.add_text(key, value)
+
+    def tags(self) -> dict[str, Any]:
+        return {}
+
+    @rank_zero_only
+    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:
+        from pathlib import Path
+        import shutil
+
+        dest = Path(self._logger.log_dir) / (artifact_path or Path(local_path).name)
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(local_path, dest)
+
+
+class WandbLogger(LightningLogger):
+    """WandB logger with rank_zero logging."""
+
+    def __init__(self, logger: "WandbLogger") -> None:
+        self._logger = logger
+
+    @rank_zero_only
+    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)
+
+    @rank_zero_only
+    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)
+
+        if self._logger.experiment.settings.mode == "offline":
+            import shutil
+            from pathlib import Path
+
+            src = Path(local_path).absolute()
+            dest = Path(self._logger.experiment.settings.files_dir) / src.name
+
+            if dest.is_symlink():
+                dest.unlink()
+                shutil.copy2(src, dest)
+
+
+class CompositeLogger(LightningLogger):
+    """
+    A wrapper on top of the collection of :class:`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(MLFlowLogger(logger=logger))  # type: ignore[arg-type]
+                elif _is_logger_type(logger, "TensorBoardLogger"):
+                    self.loggers.append(TensorBoardLogger(logger=logger))  # type: ignore[arg-type]
+                elif _is_logger_type(logger, "WandbLogger"):
+                    self.loggers.append(WandbLogger(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/schedule/__init__.py

@@ -0,0 +1,27 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from .base import (
+    Schedule,
+    Never,
+    Always,
+    Fixed,
+    Every,
+    Elapsed,
+    GlobalRank,
+    LocalRank,
+    InvertedSchedule,
+    CombinedSchedule,
+)
+
+__all__ = [
+    "Schedule",
+    "Never",
+    "Always",
+    "Fixed",
+    "Every",
+    "Elapsed",
+    "GlobalRank",
+    "LocalRank",
+    "InvertedSchedule",
+    "CombinedSchedule",
+]

fkat/pytorch/schedule/base.py

@@ -0,0 +1,308 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import datetime as dt
+import operator
+from collections.abc import Callable, Sequence
+from functools import reduce
+from typing import Protocol
+from typing_extensions import override
+
+import lightning as L
+
+
+class Schedule(Protocol):
+    """
+    Protocol defining a generic PyTorch schedule
+    """
+
+    def check(
+        self,
+        *,
+        stage: str | None = None,
+        batch_idx: int | None = None,
+        step: int | None = None,
+        trainer: L.Trainer | None = None,
+    ) -> bool:
+        """Checks the schedule for the given moment.
+
+        Args:
+            stage (str): current trainer stage. eg. train/test/validate/predict
+            batch_idx (Optional[int]): current batch_idx
+            step (Optional[int]): current step
+            trainer (Optional[Trainer]): lightning trainer of callback
+        Returns:
+            bool: True if schedule passed the check, False otherwise
+        """
+        ...
+
+    def __and__(self, first: "Schedule", second: "Schedule") -> "CombinedSchedule":
+        return self._combine(operator.and_, first, second)
+
+    def __or__(self, first: "Schedule", second: "Schedule") -> "CombinedSchedule":
+        return self._combine(operator.or_, first, second)
+
+    def __invert__(self, other: "Schedule") -> "InvertedSchedule":
+        return InvertedSchedule(other)
+
+    def _combine(self, fn: Callable[[bool, bool], bool], first: "Schedule", second: "Schedule") -> "CombinedSchedule":
+        return CombinedSchedule(fn, (first, second))
+
+
+class InvertedSchedule(Schedule):
+    def __init__(self, other: Schedule) -> None:
+        self.other = other
+
+    @override
+    def check(
+        self,
+        *,
+        stage: str | None = None,
+        batch_idx: int | None = None,
+        step: int | None = None,
+        trainer: L.Trainer | None = None,
+    ) -> bool:
+        return not self.other.check(stage=stage, batch_idx=batch_idx, step=step, trainer=trainer)
+
+
+class CombinedSchedule(Schedule):
+    def __init__(self, fn: Callable[[bool, bool], bool], schedules: Sequence[Schedule]) -> None:
+        self.fn = fn
+        self.schedules = schedules
+
+    @override
+    def check(
+        self,
+        *,
+        stage: str | None = None,
+        batch_idx: int | None = None,
+        step: int | None = None,
+        trainer: L.Trainer | None = None,
+    ) -> bool:
+        return reduce(
+            self.fn, (s.check(stage=stage, batch_idx=batch_idx, step=step, trainer=trainer) for s in self.schedules)
+        )
+
+
+class GlobalRank(Schedule):
+    """
+    A schedule that only executes on specific global ranks in a distributed training setup.
+
+    This is useful for operations that should only be performed on certain ranks,
+    such as logging, checkpointing, or other operations that would be redundant
+    or conflicting if performed on all ranks.
+
+    Attributes:
+        ranks (tuple[int, ...]): The global ranks on which this schedule should execute.
+    """
+
+    def __init__(self, *ranks: int) -> None:
+        """
+        Initialize a GlobalRank schedule.
+
+        Args:
+            *ranks: Variable number of integer ranks. The schedule will only execute
+                   on these global ranks.
+        """
+        self.ranks = ranks
+
+    @override
+    def check(
+        self,
+        *,
+        stage: str | None = None,
+        batch_idx: int | None = None,
+        step: int | None = None,
+        trainer: L.Trainer | None = None,
+    ) -> bool:
+        """
+        Check if the current global rank is in the specified ranks.
+
+        Args:
+            stage: Current trainer stage (ignored by this schedule).
+            batch_idx: Current batch index (ignored by this schedule).
+            step: Current step (ignored by this schedule).
+            trainer: Lightning trainer instance, used to get the global rank.
+
+        Returns:
+            bool: True if the trainer's global rank is in the specified ranks, False otherwise.
+                 Always returns False if trainer is None.
+        """
+        if trainer is None:
+            return False
+        return trainer.global_rank in self.ranks
+
+
+class LocalRank(Schedule):
+    """
+    A schedule that only executes on specific local ranks in a distributed training setup.
+
+    This is useful for node-specific operations that should only be performed on certain
+    ranks within each node, such as local logging or monitoring.
+
+    Attributes:
+        ranks (tuple[int, ...]): The local ranks on which this schedule should execute.
+    """
+
+    def __init__(self, *ranks: int) -> None:
+        """
+        Initialize a LocalRank schedule.
+
+        Args:
+            *ranks: Variable number of integer ranks. The schedule will only execute
+                   on these local ranks.
+        """
+        self.ranks = ranks
+
+    @override
+    def check(
+        self,
+        *,
+        stage: str | None = None,
+        batch_idx: int | None = None,
+        step: int | None = None,
+        trainer: L.Trainer | None = None,
+    ) -> bool:
+        """
+        Check if the current local rank is in the specified ranks.
+
+        Args:
+            stage: Current trainer stage (ignored by this schedule).
+            batch_idx: Current batch index (ignored by this schedule).
+            step: Current step (ignored by this schedule).
+            trainer: Lightning trainer instance, used to get the local rank.
+
+        Returns:
+            bool: True if the trainer's local rank is in the specified ranks, False otherwise.
+                 Always returns False if trainer is None.
+        """
+        if trainer is None:
+            return False
+        return trainer.local_rank in self.ranks
+
+
+class Never(Schedule):
+    """
+    A schedule for an event that never happens.
+    """
+
+    @override
+    def check(
+        self,
+        *,
+        stage: str | None = None,
+        batch_idx: int | None = None,
+        step: int | None = None,
+        trainer: L.Trainer | None = None,
+    ) -> bool:
+        return False
+
+
+class Always(Schedule):
+    """
+    A schedule for an event that always happens.
+    """
+
+    @override
+    def check(
+        self,
+        *,
+        stage: str | None = None,
+        batch_idx: int | None = None,
+        step: int | None = None,
+        trainer: L.Trainer | None = None,
+    ) -> bool:
+        return True
+
+
+class Fixed(Schedule):
+    """
+    A schedule for an event that happens after a warmup period
+    and lasts for a fixed number of steps.
+
+    Attributes:
+        warmup_steps (int): Number of initial steps to skip before logging starts.
+        active_steps (int): Number of steps to log after warmup period.
+    """
+
+    def __init__(self, warmup_steps: int, active_steps: int) -> None:
+        self._warmup_steps: int = warmup_steps
+        self._active_steps: int = active_steps
+
+    @override
+    def check(
+        self,
+        *,
+        stage: str | None = None,
+        batch_idx: int | None = None,
+        step: int | None = None,
+        trainer: L.Trainer | None = None,
+    ) -> bool:
+        assert step is not None, "step must be provided"
+        if step < self._warmup_steps:
+            return False
+
+        if step - self._warmup_steps >= self._active_steps:
+            return False
+        return True
+
+
+class Every(Schedule):
+    """
+    A schedule for an event that happens every specified number of batches and/or steps.
+
+    Attributes:
+        n_batches (Optional[int]): A positive number of batches between logging events.
+            Defaults to 0 - use only n_steps
+        n_steps (Optional[int]): A positive number of (train) steps between logging events.
+            Defaults to 0 - use only n_batches
+        stage (Optional[str]): The stage this schedule applies to ('train', 'validation', 'test', 'predict').
+            If None, applies to all stages.
+    """
+
+    def __init__(self, *, n_batches: int = 0, n_steps: int = 0, stage: str | None = None) -> None:
+        assert n_batches or n_steps, "either n_batches or n_steps has to be a positive number"
+        self._n_batches = n_batches
+        self._n_steps = n_steps
+        self._stage = stage
+
+    @override
+    def check(
+        self,
+        *,
+        stage: str | None = None,
+        batch_idx: int | None = None,
+        step: int | None = None,
+        trainer: L.Trainer | None = None,
+    ) -> bool:
+        # If stage is specified and doesn't match, return False
+        if self._stage is not None and stage != self._stage:
+            return False
+
+        return (batch_idx is not None and self._n_batches > 0 and batch_idx % self._n_batches == 0) or (
+            step is not None and self._n_steps > 0 and step % self._n_steps == 0
+        )
+
+
+class Elapsed(Schedule):
+    """
+    A schedule for an event that happens after the provided time interval has elapsed.
+    """
+
+    def __init__(self, interval: dt.timedelta) -> None:
+        self.interval = interval
+        self.last_triggered: dt.datetime | None = None
+
+    @override
+    def check(
+        self,
+        *,
+        stage: str | None = None,
+        batch_idx: int | None = None,
+        step: int | None = None,
+        trainer: L.Trainer | None = None,
+    ) -> bool:
+        now = dt.datetime.now(dt.timezone.utc)
+        if self.last_triggered is None or now - self.last_triggered >= self.interval:
+            self.last_triggered = now
+            return True
+        return False