fkat 0.1.2__py3-none-any.whl

fkat/pytorch/schedule/mlflow.py

@@ -0,0 +1,143 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import logging
+import lightning as L
+from typing_extensions import override
+
+from fkat.pytorch.schedule import Schedule
+from fkat.pytorch.callbacks.loggers import CallbackLogger
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+class HasTag(Schedule):
+    """
+    A schedule that activates only when a specific MLflow tag is present AND the trigger schedule is satisfied.
+
+    This schedule combines another schedule (the trigger schedule) with MLflow tag validation.
+    It allows callbacks to be dynamically enabled or disabled through experiment configuration rather than
+    code changes, which is particularly useful for performance-intensive callbacks like FLOP measurement
+    or detailed logging that should only run conditionally.
+
+    The schedule checks two conditions:
+    1. If the trigger schedule is satisfied
+    2. If the specified MLflow tag exists in the current MLflow run
+
+    Both conditions must be true for the schedule to activate. If the trigger schedule doesn't activate or
+    the trainer is not provided, the schedule will never activate.
+
+    Note:
+        - The trainer can be optionally provided to the ``check`` method for MLflow tag validation. If trainer is None,
+          the schedule will never activate.
+        - Tag checking occurs only when the trigger schedule condition is already satisfied,
+          minimizing MLflow API calls.
+        - If an exception occurs during tag checking, it will be logged and the schedule will not activate.
+
+    Example:
+        Python code example::
+
+            # Create a schedule that checks every 5 batches if the 'enable_flops' tag exists
+            from fkat.pytorch.schedule import Every
+
+            trigger = Every(n_batches=5)
+            flops_schedule = HasTag(tag="enable_flops", schedule=trigger)
+            flops_callback = Flops(schedule=flops_schedule)
+            trainer = L.Trainer(callbacks=[flops_callback])
+
+        Hydra configuration example:
+
+        .. code-block:: yaml
+
+            # In your config.yaml file
+            callbacks:
+              - _target_: fkat.pytorch.callbacks.profiling.Flops
+                schedule:
+                  _target_: fkat.pytorch.schedule.mlflow.HasTag
+                  tag: ENABLE_FLOPS
+                  schedule:
+                    _target_: fkat.pytorch.schedule.Every
+                    n_steps: 20
+
+            # Another example using Fixed schedule
+            callbacks:
+              - _target_: fkat.pytorch.callbacks.heartbeat.Heartbeat
+                schedule:
+                  _target_: fkat.pytorch.schedule.mlflow.HasTag
+                  tag: ENABLE_HEARTBEAT
+                  schedule:
+                    _target_: fkat.pytorch.schedule.Fixed
+                    warmup_steps: 100
+                    active_steps: 1000
+
+            # Example with Elapsed time-based schedule
+            callbacks:
+              - _target_: fkat.pytorch.callbacks.custom_logging.DetailedMetrics
+                schedule:
+                  _target_: fkat.pytorch.schedule.mlflow.HasTag
+                  tag: DETAILED_LOGGING
+                  schedule:
+                    _target_: fkat.pytorch.schedule.Elapsed
+                    interval: ${timedelta:minutes=15}
+    """
+
+    def __init__(self, tag: str, schedule: Schedule) -> None:
+        """
+        Initialize a new MLflow HasTag schedule.
+
+        Args:
+            tag (str): The name of the tag that must be present in the MLflow run.
+            schedule (Schedule): The schedule that determines when to check for the tag.
+                This can be any implementation of the Schedule protocol (e.g., Every, Fixed, Elapsed).
+        """
+        self._tag: str = tag
+        self._schedule: Schedule = schedule
+        self._cb_logger: CallbackLogger | 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:
+        """
+        Check if the schedule should activate based on the trigger schedule and MLflow tag presence.
+
+        This method first checks if the trigger schedule is satisfied.
+        If this condition is met, it then checks if the specified MLflow tag exists in the current run.
+        Both conditions must be true for the method to return True.
+
+        Args:
+            stage (str, optional): Current training stage (e.g., "train", "validate", "test").
+                Passed to the trigger schedule.
+            batch_idx (int, optional): Current batch index within the epoch.
+                Passed to the trigger schedule.
+            step (int, optional): Current global step (cumulative across epochs).
+                Passed to the trigger schedule.
+            trainer (L.Trainer, optional): The Lightning Trainer instance.
+                Required for MLflow tag validation.
+
+        Returns:
+            bool: True if both the trigger schedule is satisfied AND the specified tag is present,
+                  False otherwise.
+
+        Note:
+            - The trainer must be provided for MLflow tag validation.
+            - Tag checking occurs only when the trigger schedule is already satisfied.
+            - If an exception occurs during tag checking, it will be logged as a warning
+              and the method will return False.
+        """
+        triggered = self._schedule.check(stage=stage, batch_idx=batch_idx, step=step, trainer=trainer)
+        if not triggered or trainer is None:
+            return False
+        try:
+            if self._cb_logger is None:
+                self._cb_logger = CallbackLogger(trainer)
+
+            tags = self._cb_logger.tags()
+            return self._tag in tags
+        except Exception as e:
+            logger.warning(f"Error when checking if tag {self._tag} exists: {e}")
+            return False

fkat/pytorch/utilities.py

@@ -0,0 +1,49 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import os
+from functools import wraps
+from typing import TypeVar, ParamSpec
+from collections.abc import Callable
+from typing_extensions import overload
+
+
+def get_rank() -> int | None:
+    return int(os.getenv("RANK", "0"))
+
+
+def get_local_rank() -> int | None:
+    return int(os.getenv("LOCAL_RANK", "0"))
+
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+@overload
+def local_rank_zero_only(fn: Callable[P, T]) -> Callable[P, T | None]: ...
+
+
+@overload
+def local_rank_zero_only(fn: Callable[P, T], default: T) -> Callable[P, T]: ...
+
+
+def local_rank_zero_only(fn: Callable[P, T], default: T | None = None) -> Callable[P, T | None]:
+    """Wrap a function to call internal function only in rank zero.
+
+    Function that can be used as a decorator to enable a function/method being called only on global rank 0.
+
+    """
+
+    @wraps(fn)
+    def wrapped_fn(*args: P.args, **kwargs: P.kwargs) -> T | None:
+        local_rank = getattr(local_rank_zero_only, "local_rank", None)
+        if local_rank is None:
+            raise RuntimeError("The `local_rank_zero_only.local_rank` needs to be set before use")
+        if local_rank == 0:
+            return fn(*args, **kwargs)
+        return default
+
+    return wrapped_fn
+
+
+local_rank_zero_only.local_rank = getattr(local_rank_zero_only, "local_rank", get_local_rank() or 0)  # type: ignore[attr-defined]

fkat/test.py

@@ -0,0 +1,31 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+#!/usr/bin/env python
+
+"""
+The ``fkat.test`` entrypoint processes the provided config,
+instatiates the ``trainer``, ``model`` and ``data`` sections and calls ``trainer.test()``.
+"""
+
+import hydra
+import lightning as L
+from omegaconf import DictConfig
+
+from fkat import initialize, run_main
+
+
+@hydra.main(version_base="1.3")
+def main(cfg: DictConfig) -> None:
+    s = initialize(cfg)
+    kwargs = {
+        "ckpt_path": s.ckpt_path,
+    }
+    if isinstance(s.data, L.LightningDataModule):
+        kwargs["datamodule"] = s.data
+    else:
+        kwargs["test_dataloaders"] = s.data.test_dataloader() if s.data else None
+    s.trainer.test(s.model, **kwargs)
+
+
+if __name__ == "__main__":
+    run_main(main)

fkat/train.py

@@ -0,0 +1,32 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+#!/usr/bin/env python
+
+"""
+The ``fkat.train`` entrypoint processes the provided config,
+instatiates the ``trainer``, ``model`` and ``data`` sections and calls ``trainer.fit()``.
+"""
+
+import hydra
+import lightning as L
+from omegaconf import DictConfig
+
+from fkat import initialize, run_main
+
+
+@hydra.main(version_base="1.3")
+def main(cfg: DictConfig) -> None:
+    s = initialize(cfg)
+    kwargs = {
+        "ckpt_path": s.ckpt_path,
+    }
+    if isinstance(s.data, L.LightningDataModule):
+        kwargs["datamodule"] = s.data
+    else:
+        kwargs["train_dataloaders"] = s.data.train_dataloader() if s.data else None
+        kwargs["val_dataloaders"] = s.data.val_dataloader() if s.data else None
+    s.trainer.fit(s.model, **kwargs)
+
+
+if __name__ == "__main__":
+    run_main(main)

fkat/utils/__init__.py

@@ -0,0 +1,28 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from datetime import datetime, timezone
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+def assert_not_none(obj: T | None, name: str = "obj") -> T:
+    assert obj is not None, f"{name} cannot be None"
+    return obj
+
+
+def safe_timestamp(dt: datetime | None = None) -> str:
+    """
+    Generate a filesystem-safe timestamp string.
+
+    Format: YYYY-MM-DD_HH-MM-SS-mmm (e.g., 2026-01-14_16-09-15-123)
+
+    Args:
+        dt: datetime object to format. If None, uses current UTC time.
+
+    Returns:
+        Formatted timestamp string safe for use in filenames
+    """
+    if dt is None:
+        dt = datetime.now(timezone.utc)
+    return dt.strftime("%Y-%m-%d_%H-%M-%S-") + f"{dt.microsecond // 1000:03d}"

fkat/utils/aws/__init__.py

@@ -0,0 +1,3 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+

fkat/utils/aws/imds.py

@@ -0,0 +1,137 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+EC2 Instance Metadata Service related methods
+"""
+
+import logging
+import socket
+from dataclasses import dataclass
+from functools import lru_cache
+
+import requests
+from requests import HTTPError
+
+from fkat.utils.logging import rank0_logger
+
+IMDS_URL = "http://169.254.169.254/latest"
+IMDS_METADATA_URL = f"{IMDS_URL}/meta-data"
+IMDS_V2_TOKEN_URL = f"{IMDS_URL}/api/token"
+NULL = "_NULL_"  # sentinel value used to mark null (not-available) values
+
+
+log: logging.Logger = rank0_logger(__name__)
+
+Token = str
+
+
+@dataclass
+class InstanceMetadata:
+    """
+    Struct representing the instance metadata as fetched from IMDS on the current host.
+    Use :py:func:`fkat.utils.aws.imds.instance_metadata` to get a filled-out
+    instance of this object.
+    """
+
+    instance_id: str
+    instance_type: str
+    hostname: str
+    public_hostname: str
+    local_hostname: str
+    local_ipv4: str
+    availability_zone: str
+    region: str
+    ami_id: str
+
+
+@lru_cache
+def fetch(metadata: str = "", token: Token | None = None) -> str | None:
+    """
+    Fetches the specified ``metadata`` from EC2's Instance MetaData Service (IMDS) running
+    on the current host, by sending an HTTP GET request to ``http://169.254.169.254/latest/meta-data/<metadata>``.
+
+    To get a list of all valid values of ``metadata`` run this method with no arguments then split
+    the return value by new-line.
+
+
+    Arguments:
+        metadata: Name of the instance metadata to query (e.g. ``instance-type``)
+        token: IMDS token
+
+    Returns:
+        the specified ``metadata`` or ``None`` if IMDS cannot be reached
+    """
+
+    try:
+        response = requests.get(f"{IMDS_METADATA_URL}/{metadata}", headers={"X-aws-ec2-metadata-token": token or ""})
+    except Exception as e:
+        log.warning("Error querying IMDSV2 instance metadata won't be available", exc_info=e)
+        return None
+
+    if response.ok:
+        return response.text
+    else:  # response NOT ok
+        try:
+            response.raise_for_status()
+        except HTTPError:
+            return None
+
+    raise AssertionError("Unreachable code!")
+
+
+def token(timeout: int = 60) -> Token | None:
+    """
+    Fetches IMDS ``token`` from EC2's Instance MetaData Service (IMDSV2) running
+    on the current host, by sending an HTTP GET request to ``http://169.254.169.254/latest/meta-data/<metadata>``.
+
+    Arguments:
+        timeout: request timeout
+
+    Returns:
+        the specified ``token`` or ``""`` if IMDSV2 cannot be reached
+    """
+
+    try:
+        response = requests.put(
+            IMDS_V2_TOKEN_URL, headers={"X-aws-ec2-metadata-token-ttl-seconds": "21600"}, timeout=timeout
+        )
+    except Exception as e:
+        log.warning("Error querying IMDSV2 instance token won't be available", exc_info=e)
+        return ""
+
+    if response.ok:
+        return response.text
+    else:  # response NOT ok
+        error_code_msg = f"IMDS Token response is not ok with status code {response.status_code}"
+        log.warning("Error querying IMDSV2 instance token won't be available", exc_info=Exception(error_code_msg))
+        return ""
+
+    raise AssertionError("Unreachable code!")
+
+
+@lru_cache
+def instance_metadata() -> InstanceMetadata:
+    """
+    Fetches IMDS instance metadata for the current host from EC2's Instance Metadata Service (IMDS),
+    which typically runs on localhost at ``http://169.254.169.254``.
+    If IMDS cannot be reached for any reason returns an instance of :py:class:`InstanceMetadata`
+    where all the fields are empty strings.
+
+    .. note::
+        This method is memoized (value is cached) hence, only the first call
+        will actually hit IMDS, and subsequent calls will return the memoized
+        value. Therefore, it is ok to call this function multiple times.
+
+    """
+    tkn = token()
+    return InstanceMetadata(
+        instance_id=fetch("instance-id", tkn) or "localhost",
+        instance_type=fetch("instance-type", tkn) or NULL,
+        availability_zone=fetch("placement/availability-zone", tkn) or NULL,
+        region=fetch("placement/region", tkn) or NULL,
+        hostname=fetch("hostname", tkn) or socket.gethostname(),
+        local_ipv4=fetch("local-ipv4", tkn) or NULL,
+        public_hostname=fetch("public-hostname", tkn) or NULL,
+        local_hostname=fetch("local-hostname", tkn) or NULL,
+        ami_id=fetch("ami-id", tkn) or NULL,
+    )

fkat/utils/boto3.py

@@ -0,0 +1,24 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Literal
+
+import boto3
+from botocore.config import Config
+
+
+def session(
+    max_attempts: int = 6,
+    mode: Literal["legacy", "standard", "adaptive"] = "standard",
+    clients: list[str] | None = None,
+) -> boto3.Session:
+    config = Config(
+        retries={
+            "max_attempts": max_attempts,
+            "mode": mode,
+        }
+    )
+    session = boto3.Session()
+    if clients:
+        for client in clients:
+            session.client(client, config=config)  # type: ignore
+    return session

fkat/utils/config.py

@@ -0,0 +1,194 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import os
+from typing import Any, Protocol
+from tempfile import TemporaryDirectory
+
+from lightning import Trainer as LightningTrainer
+from lightning.pytorch.utilities import rank_zero_only
+from omegaconf import OmegaConf
+
+from fkat.utils.mlflow import mlflow_logger
+
+
+class Trainer(Protocol):
+    """
+    Protocol defining the interface for training models.
+
+    This protocol establishes the required methods that any trainer implementation
+    must provide. It serves as a structural subtyping interface for objects that
+    can train, evaluate, and make predictions with machine learning models.
+
+    Implementations of this protocol should handle the complete training lifecycle,
+    including model fitting, prediction, testing, and validation.
+
+    Note:
+        As a Protocol class, this is not meant to be instantiated directly but
+        rather used for type checking and interface definition.
+    """
+
+    def fit(self, *args: Any, **kwargs: Any) -> Any:
+        """
+        Train a model on the provided data.
+
+        This method handles the model training process, including data loading,
+        optimization, and potentially checkpointing.
+
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            Any: Training results, which may include training history,
+                 trained model, or other relevant information.
+        """
+        ...
+
+    def predict(self, *args: Any, **kwargs: Any) -> Any:
+        """
+        Generate predictions using a trained model.
+
+        This method applies the trained model to new data to produce predictions.
+
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            Any: Model predictions, which could be probabilities, class labels,
+                 regression values, or other outputs depending on the model type.
+        """
+        ...
+
+    def test(self, *args: Any, **kwargs: Any) -> Any:
+        """
+        Evaluate a trained model on test data.
+
+        This method assesses model performance on a test dataset, typically
+        calculating metrics such as accuracy, loss, or other relevant measures.
+
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            Any: Test results, which may include metrics, predictions,
+                 or other evaluation information.
+        """
+        ...
+
+    def validate(self, *args: Any, **kwargs: Any) -> Any:
+        """
+        Evaluate a trained model on validation data.
+
+        This method assesses model performance on a validation dataset, which is
+        typically used during the training process for hyperparameter tuning
+        or early stopping.
+
+        Args:
+            *args: Variable length argument list.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            Any: Validation results, which may include metrics, predictions,
+                 or other evaluation information.
+        """
+        ...
+
+
+class SingletonResolver:
+    """
+    A singleton class that resolves and manages training components.
+
+    This class serves as a central registry for training-related objects such as
+    trainers, data, models, and checkpoint paths. It ensures that these components
+    are accessible throughout the application while maintaining a single instance.
+
+    Example:
+        >>> resolver = SingletonResolver()
+        >>> resolver.trainer = MyTrainer()
+        >>> resolver.model = MyModel()
+        >>> # Access the same instance elsewhere
+        >>> same_resolver = SingletonResolver()
+        >>> assert same_resolver.trainer is resolver.trainer
+    """
+
+    trainer: Trainer
+    """The trainer instance responsible for executing the training process."""
+
+    data: Any | None = None
+    """The dataset or data loader used for training and evaluation.
+    Defaults to None."""
+
+    model: Any | None = None
+    """The model architecture to be trained or evaluated.
+    Defaults to None."""
+
+    ckpt_path: Any | None = None
+    """Path to checkpoint files for model loading/saving.
+    Defaults to None."""
+
+    return_predictions: Any | None = None
+    """Flag or configuration for returning predictions.
+    Defaults to None."""
+
+    tuners: Any | None = None
+    """Hyperparameter tuners or optimization components.
+    Defaults to None."""
+
+
+def register_singleton_resolver() -> Any:
+    resolver = SingletonResolver()
+
+    def resolve(key: str) -> Any:
+        res = resolver
+        for attr in key.split("."):
+            res = getattr(res, attr)
+        return res
+
+    OmegaConf.register_new_resolver("fkat", resolve)
+    return resolver
+
+
+def to_str(cfg: Any) -> str:
+    """
+    Convert a configuration object to a formatted string representation.
+
+    This function takes a configuration object and converts it to a human-readable
+    YAML string. It's useful for logging, debugging, or displaying configuration settings.
+
+    Args:
+        cfg (Any): The configuration object to convert to string.
+
+    Returns:
+        str: A formatted string representation of the configuration.
+
+    Example:
+        >>> config = {"model": {"type": "resnet", "layers": 50}, "batch_size": 32}
+        >>> print(to_str(config))
+        Config:
+        model:
+          type: resnet
+          layers: 50
+        batch_size: 32
+    """
+    return "Config:\n" + OmegaConf.to_yaml(cfg)
+
+
+def to_primitive_container(cfg: Any) -> Any:
+    if OmegaConf.is_config(cfg):
+        return OmegaConf.to_container(cfg)
+    return cfg
+
+
+@rank_zero_only
+def save(cfg: Any, trainer: LightningTrainer) -> None:
+    yaml_str = OmegaConf.to_yaml(cfg)
+    with TemporaryDirectory() as temp_dir:
+        yaml_path = os.path.join(temp_dir, "config.yaml")
+        os.makedirs(os.path.dirname(yaml_path), exist_ok=True)
+        with open(yaml_path, "w") as f:
+            f.write(yaml_str)
+        if mlflow := mlflow_logger(trainer):
+            if mlflow.run_id:
+                mlflow.experiment.log_artifact(mlflow.run_id, yaml_path)

fkat/utils/cuda/__init__.py

@@ -0,0 +1,3 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+

fkat/utils/cuda/preflight/__init__.py

@@ -0,0 +1,3 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+