llmcompressor.metrics

class BaseLogger(ABC):
    """
    Base class that all modifier loggers must implement.

    :param name: name given to the metrics, used for identification
    :param enabled: True to log, False otherwise
    """

    def __init__(self, name: str, enabled: bool = True):
        self._name = name
        self._enabled = enabled

    @property
    def name(self) -> str:
        """
        :return: name given to the metrics, used for identification
        """
        return self._name

    @property
    def enabled(self) -> bool:
        """
        :return: True to log, False otherwise
        """
        return self._enabled

    @enabled.setter
    def enabled(self, value: bool):
        """
        :param value: True to log, False otherwise
        """
        self._enabled = value

    def __repr__(self):
        return f"{self.__class__.__name__}(name={self._name}, enabled={self._enabled})"

    def log_hyperparams(self, params: Dict[str, float]) -> bool:
        """
        :param params: Each key-value pair in the dictionary is the name of the
            hyper parameter and it's corresponding value.
        :return: True if logged, False otherwise.
        """
        return False

    def log_scalar(
        self,
        tag: str,
        value: float,
        step: Optional[int] = None,
        wall_time: Optional[float] = None,
        **kwargs,
    ) -> bool:
        """
        :param tag: identifying tag to log the value with
        :param value: value to save
        :param step: global step for when the value was taken
        :param wall_time: global wall time for when the value was taken
        :param kwargs: additional logging arguments to support Python and custom loggers
        :return: True if logged, False otherwise.
        """
        return False

    def log_scalars(
        self,
        tag: str,
        values: Dict[str, float],
        step: Optional[int] = None,
        wall_time: Optional[float] = None,
        **kwargs,
    ) -> bool:
        """
        :param tag: identifying tag to log the values with
        :param values: values to save
        :param step: global step for when the values were taken
        :param wall_time: global wall time for when the values were taken
        :param kwargs: additional logging arguments to support Python and custom loggers
        :return: True if logged, False otherwise.
        """
        return False

    def log_string(
        self,
        tag: str,
        string: str,
        step: Optional[int] = None,
        wall_time: Optional[float] = None,
        **kwargs,
    ) -> bool:
        """
        :param tag: identifying tag to log the values with
        :param values: values to save
        :param step: global step for when the values were taken
        :param wall_time: global wall time for when the values were taken
        :param kwargs: additional logging arguments to support Python and custom loggers
        :return: True if logged, False otherwise.
        """
        return False

    def save(
        self,
        file_path: str,
        **kwargs,
    ) -> bool:
        """
        :param file_path: path to a file to be saved
        :param kwargs: additional arguments that a specific metrics might use
        :return: True if saved, False otherwise
        """
        return False

class LambdaLogger(BaseLogger):
    """
    Logger that handles calling back to a lambda function with any logs.

    :param lambda_func: the lambda function to call back into with any logs.
        The expected call sequence is (tag, value, values, step, wall_time) -> bool
        The return type is True if logged and False otherwise.
    :param name: name given to the metrics, used for identification;
        defaults to lambda
    :param enabled: True to log, False otherwise
    """

    def __init__(
        self,
        lambda_func: Callable[
            [
                Optional[str],
                Optional[Union[float, str]],
                Optional[Dict[str, float]],
                Optional[int],
                Optional[float],
                Optional[int],
            ],
            bool,
        ],
        name: str = "lambda",
        enabled: bool = True,
    ):
        super().__init__(name, enabled)
        self._lambda_func = lambda_func
        assert lambda_func, "lambda_func must be set to a callable function"

    @property
    def lambda_func(
        self,
    ) -> Callable[
        [
            Optional[str],
            Optional[Union[float, str]],
            Optional[Dict[str, float]],
            Optional[int],
            Optional[float],
            Optional[int],
        ],
        bool,
    ]:
        """
        :return: the lambda function to call back into with any logs.
            The expected call sequence is (tag, value, values, step, wall_time)
        """
        return self._lambda_func

    def log_hyperparams(
        self,
        params: Dict,
        level: Optional[int] = None,
    ) -> bool:
        """
        :param params: Each key-value pair in the dictionary is the name of the
            hyper parameter and it's corresponding value.
        :return: True if logged, False otherwise.
        """
        if not self.enabled:
            return False

        return self._lambda_func(
            tag=None,
            value=None,
            values=params,
            step=None,
            wall_time=None,
            level=level,
        )

    def log_scalar(
        self,
        tag: str,
        value: float,
        step: Optional[int] = None,
        wall_time: Optional[float] = None,
        level: Optional[int] = None,
    ) -> bool:
        """
        :param tag: identifying tag to log the value with
        :param value: value to save
        :param step: global step for when the value was taken
        :param wall_time: global wall time for when the value was taken,
            defaults to time.time()
        :param kwargs: additional logging arguments to support Python and custom loggers
        :return: True if logged, False otherwise.
        """
        if not wall_time:
            wall_time = time.time()

        return self._lambda_func(
            tag=tag,
            value=value,
            values=None,
            step=step,
            wall_time=wall_time,
            level=level,
        )

    def log_scalars(
        self,
        tag: str,
        values: Dict[str, float],
        step: Optional[int] = None,
        wall_time: Optional[float] = None,
        level: Optional[int] = None,
    ) -> bool:
        """
        :param tag: identifying tag to log the values with
        :param values: values to save
        :param step: global step for when the values were taken
        :param wall_time: global wall time for when the values were taken,
            defaults to time.time()
        :param kwargs: additional logging arguments to support Python and custom loggers
        :return: True if logged, False otherwise.
        """
        if not wall_time:
            wall_time = time.time()

        return self._lambda_func(
            tag=tag,
            value=None,
            values=values,
            step=step,
            wall_time=wall_time,
            level=level,
        )

class LoggerManager(ABC):
    """
    Wrapper around loggers that handles log scheduling and handing off logs to intended
    loggers.

    :param loggers: list of loggers assigned to this manager
    :param log_frequency: number of stes or fraction of steps to wait between logs
    :param mode: The logging mode to use, either "on_change" or "exact",
        "on_change" will log when the model has been updated since the last log,
        "exact" will log at the given frequency regardless of model updates.
        Defaults to "exact"
    :param frequency_type: The frequency type to use, either "epoch", "step", or "batch"
        controls what the frequency manager is tracking, e.g. if the frequency type
        is "epoch", then the frequency manager will track the number of epochs that
        have passed since the last log, if the frequency type is "step", then the
        frequency manager will track the number of optimizer steps
    """

    def __init__(
        self,
        loggers: Optional[List[BaseLogger]] = None,
        log_frequency: Union[float, None] = 0.1,
        log_python: bool = True,
        name: str = "manager",
        mode: LoggingModeType = "exact",
        frequency_type: FrequencyType = "epoch",
    ):
        self._name = name
        self._loggers = (
            loggers
            or SparsificationGroupLogger(
                python=log_python,
                name=name,
                tensorboard=False,
                wandb_=False,
            ).loggers
        )

        self.frequency_manager = FrequencyManager(
            mode=mode,
            frequency_type=frequency_type,
            log_frequency=log_frequency,
        )

        self.system = SystemLoggingWraper(
            loggers=self._loggers, frequency_manager=self.frequency_manager
        )
        self.metric = MetricLoggingWrapper(
            loggers=self._loggers, frequency_manager=self.frequency_manager
        )

    def __len__(self):
        return len(self.loggers)

    def __iter__(self):
        return iter(self.loggers)

    def add_logger(self, logger: BaseLogger):
        """
        add a BaseLogger implementation to the loggers of this manager

        :param logger: metrics object to add
        """
        if not isinstance(logger, BaseLogger):
            raise ValueError(f"metrics {type(logger)} must be of type BaseLogger")
        self._loggers.append(logger)

    def log_ready(
        self, current_log_step, last_log_step=None, check_model_update: bool = False
    ):
        """
        Check if there is a metrics that is ready to accept a log

        :param current_log_step: current step log is requested at
        :param last_log_step: last time a log was recorder for this object. (Deprecated)
        :param check_model_update: if True, will check if the model has been updated,
            if False, will only check the log frequency
        :return: True if a metrics is ready to accept a log.
        """
        log_enabled = any(logger.enabled for logger in self.loggers)
        if last_log_step is not None:
            self.frequency_manager.log_written(step=last_log_step)

        return log_enabled and self.frequency_manager.log_ready(
            current_log_step=current_log_step,
            check_model_update=check_model_update,
        )

    def log_written(self, step: LogStepType):
        """
        Update the frequency manager with the last log step written

        :param step: step that was last logged
        """
        self.frequency_manager.log_written(step=step)

    def model_updated(self, step: LogStepType):
        """
        Update the frequency manager with the last model update step

        :param step: step that was last logged
        """
        self.frequency_manager.model_updated(step=step)

    @staticmethod
    def epoch_to_step(epoch, steps_per_epoch):
        return round(epoch) if steps_per_epoch <= 0 else round(epoch * steps_per_epoch)

    @property
    def loggers(self) -> List[BaseLogger]:
        """
        :return: list of loggers assigned to this manager
        """
        return self._loggers

    @loggers.setter
    def loggers(self, value: List[BaseLogger]):
        """
        :param value: list of loggers assigned to this manager
        """
        self._loggers = value

    @property
    def log_frequency(self) -> Union[str, float, None]:
        """
        :return: number of epochs or fraction of epochs to wait between logs
        """
        return self.frequency_manager._log_frequency

    @log_frequency.setter
    def log_frequency(self, value: Union[str, float, None]):
        """
        :param value: number of epochs or fraction of epochs to wait between logs
        """
        self.frequency_manager._log_frequency = value

    @property
    def name(self) -> str:
        """
        :return: name given to the metrics, used for identification
        """
        return self._name

    @property
    def wandb(self) -> Optional[ModuleType]:
        """
        :return: wandb module if initialized
        """
        for log in self.loggers:
            if isinstance(log, WANDBLogger) and log.enabled:
                return log.wandb
        return None

    def log_scalar(
        self,
        tag: str,
        value: float,
        step: Optional[int] = None,
        wall_time: Optional[float] = None,
        log_types: Union[str, List[str]] = ALL_TOKEN,
        level: Optional[int] = None,
    ):
        """
        (Note: this method is deprecated and will be removed in a future version,
        use LoggerManager().metric.log_scalar instead)

        :param tag: identifying tag to log the value with
        :param value: value to save
        :param step: global step for when the value was taken
        :param wall_time: global wall time for when the value was taken
        :param kwargs: additional logging arguments to support Python and custom loggers
        :return: True if logged, False otherwise.
        """

        self.metric.log_scalar(
            tag=tag,
            value=value,
            step=step,
            wall_time=wall_time,
            log_types=log_types,
            level=level,
        )

    def log_scalars(
        self,
        tag: str,
        values: Dict[str, float],
        step: Optional[int] = None,
        wall_time: Optional[float] = None,
        log_types: Union[str, List[str]] = ALL_TOKEN,
        level: Optional[int] = None,
    ):
        """
        (Note: this method is deprecated and will be removed in a future version,
        use LoggerManager().metric.log_scalars instead)

        :param tag: identifying tag to log the values with
        :param values: values to save
        :param step: global step for when the values were taken
        :param wall_time: global wall time for when the values were taken
        :param kwargs: additional logging arguments to support Python and custom loggers
        :return: True if logged, False otherwise.
        """

        self.metric.log_scalars(
            tag=tag,
            values=values,
            step=step,
            wall_time=wall_time,
            log_types=log_types,
            level=level,
        )

    def log_hyperparams(
        self,
        params: Dict,
        log_types: Union[str, List[str]] = ALL_TOKEN,
        level: Optional[int] = None,
    ):
        """
        (Note: this method is deprecated and will be removed in a future version,
        use LoggerManager().metric.log_hyperparams instead)

        :param params: Each key-value pair in the dictionary is the name of the
            hyper parameter and it's corresponding value.
        """

        self.metric.log_hyperparams(
            params=params,
            log_types=log_types,
            level=level,
        )

    def log_string(
        self,
        tag: str,
        string: str,
        step: Optional[int] = None,
        wall_time: Optional[float] = None,
        log_types: Union[str, List[str]] = ALL_TOKEN,
        level: Optional[int] = None,
    ):
        """
        (Note: this method is deprecated and will be removed in a future version,
        use LoggerManager().system.log_string instead)

        :param tag: identifying tag to log the values with
        :param values: values to save
        :param step: global step for when the values were taken
        :param wall_time: global wall time for when the values were taken
        :param kwargs: additional logging arguments to support Python and custom loggers
        :return: True if logged, False otherwise.
        """
        self.system.log_string(
            tag=tag,
            string=string,
            step=step,
            wall_time=wall_time,
            log_types=log_types,
            level=level,
        )

    def save(
        self,
        file_path: str,
        **kwargs,
    ):
        """
        :param file_path: path to a file to be saved
        :param kwargs: additional arguments that a specific metrics might use
        """
        for log in self._loggers:
            if log.enabled:
                log.save(file_path, **kwargs)

    @contextmanager
    def time(self, tag: Optional[str] = None, *args, **kwargs):
        """
        Context manager to log the time it takes to run the block of code

        Usage:
        >>> with LoggerManager().time("my_block"):
        >>>    time.sleep(1)

        :param tag: identifying tag to log the values with
        """

        start = time.time()
        yield
        elapsed = time.time() - start
        if not tag:
            tag = f"{DEFAULT_TAG}_time_secs"
        self.log_scalar(tag=tag, value=float(f"{elapsed:.3f}"), *args, **kwargs)