llmcompressor.core.session_functions

class LifecycleCallbacks:
    """
    A class for invoking lifecycle events for the active session
    """

    @classmethod
    def event(cls, event_type: EventType, **kwargs) -> ModifiedState:
        """
        Invoke an event for the active session

        :param event_type: the event type to invoke
        :param kwargs: additional kwargs to pass to the current session's event method
        :return: the modified state of the active session after invoking the event
        """
        if event_type in [EventType.INITIALIZE, EventType.FINALIZE]:
            raise ValueError(
                f"Cannot invoke {event_type} event. "
                f"Use the corresponding method instead."
            )

        # skip event callbacks if no recipe was provided
        if not active_session().lifecycle.recipe_container.check_any_recipe_exists():
            return

        return active_session().event(event_type, **kwargs)

    @classmethod
    def batch_start(cls, batch_data: Optional[Any] = None, **kwargs) -> ModifiedState:
        """
        Invoke a batch start event for the active session

        :param batch_data: the batch data to use for the event
        :param kwargs: additional kwargs to pass to the current session's event method
        :return: the modified state of the active session after invoking the event
        """
        return cls.event(EventType.BATCH_START, batch_data=batch_data, **kwargs)

    @classmethod
    def loss_calculated(cls, loss: Optional[Any] = None, **kwargs) -> ModifiedState:
        """
        Invoke a loss calculated event for the active session

        :param loss: the loss to use for the event
        :param kwargs: additional kwargs to pass to the current session's event method
        :return: the modified state of the active session after invoking the event
        """
        # log loss if loss calculated
        active_session()._log_loss(event_type=EventType.LOSS_CALCULATED, loss=loss)
        return cls.event(EventType.LOSS_CALCULATED, loss=loss, **kwargs)

    @classmethod
    def optim_pre_step(cls, **kwargs) -> ModifiedState:
        """
        Invoke an optimizer pre-step event for the active session

        :param kwargs: additional kwargs to pass to the current session's event method
        :return: the modified state of the active session after invoking the event
        """
        return cls.event(EventType.OPTIM_PRE_STEP, **kwargs)

    @classmethod
    def optim_post_step(cls, **kwargs) -> ModifiedState:
        """
        Invoke an optimizer post-step event for the active session

        :param kwargs: additional kwargs to pass to the current session's event method
        :return: the modified state of the active session after invoking the event
        """
        return cls.event(EventType.OPTIM_POST_STEP, **kwargs)

    @classmethod
    def batch_end(cls, **kwargs) -> ModifiedState:
        """
        Invoke a batch end event for the active session

        :param kwargs: additional kwargs to pass to the current session's event method
        :return: the modified state of the active session after invoking the event
        """
        active_session()._log_model_info()
        return cls.event(EventType.BATCH_END, **kwargs)

    @classmethod
    def calibration_epoch_start(cls, **kwargs) -> ModifiedState:
        """
        Invoke a epoch start event for the active session during calibration. This event
        should be called before calibration starts for one epoch

        see `src/llmcompressor/pipelines/basic/pipeline.py` for usage example
        """
        return cls.event(EventType.CALIBRATION_EPOCH_START, **kwargs)

    @classmethod
    def sequential_epoch_end(cls, **kwargs) -> ModifiedState:
        """
        Invoke a sequential epoch end event for the active session. This event should be
        called after one sequential layer has been calibrated/trained for one epoch

        This is called after a sequential layer has been calibrated with one batch, see
        `src/llmcompressor/pipelines/sequential/pipeline.py` for usage example
        """
        return cls.event(EventType.SEQUENTIAL_EPOCH_END, **kwargs)

    @classmethod
    def calibration_epoch_end(cls, **kwargs) -> ModifiedState:
        """
        Invoke a epoch end event for the active session during calibration. This event
        should be called after the model has been calibrated for one epoch

        see `src/llmcompressor/pipelines/basic/pipeline.py` for usage example
        """
        return cls.event(EventType.CALIBRATION_EPOCH_END, **kwargs)