remove session function

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/llmcompressor/core/__init__.py

@@ -16,7 +16,6 @@
     finalize,
     initialize,
     reset_session,
-    update_state,
 )
 from llmcompressor.core.state import Data, Hardware, ModifiedState, State
 
@@ -37,7 +36,6 @@
     "active_session",
     "reset_session",
     "initialize",
-    "update_state",
     "finalize",
     "apply",
     "callbacks",

src/llmcompressor/core/lifecycle.py

@@ -84,14 +84,14 @@ def initialize(
         :return: List of data returned from initialization of modifiers
         :rtype: List[Any]
         """
-        self.state.update(**kwargs)
         if self.initialized_:
             raise ValueError(
                 "Initialize was called twice. To update state values after "
                 "initialization, please use "
                 "`llmcompressor.core.session_functions.update_state`"
             )
 
+        self.state.update(**kwargs)
         logger.debug("Initializing compression lifecycle")
         self.recipe_container.append(recipe, recipe_stage, recipe_args)
         self.modifiers = self.recipe_container.get_modifiers()
@@ -111,13 +111,6 @@ def initialize(
 
         return mod_data
 
-    def update_state(self, **kwargs):
-        """
-        Update the lifecycle state with new values
-        """
-        logger.info(f"Updated state with {kwargs}")
-        self.state.update(**kwargs)
-
     def finalize(self, **kwargs) -> List[Any]:
         """
         Finalize the compression lifecycle.

src/llmcompressor/core/session.py

@@ -143,18 +143,6 @@ def initialize(
             modifier_data=mod_data,
         )
 
-    def update_state(self, **kwargs) -> ModifiedState:
-        """
-        Update the lifecycle state with new values
-        """
-        self._lifecycle.update_state(**kwargs)
-
-        return ModifiedState(
-            model=self.state.model,
-            optimizer=self.state.optimizer,
-            loss=self.state.loss,
-        )
-
     def finalize(self, **kwargs) -> ModifiedState:
         """
         Finalize the session for compression. This will run the finalize method

src/llmcompressor/core/session_functions.py

@@ -12,7 +12,6 @@
     "active_session",
     "reset_session",
     "initialize",
-    "update_state",
     "finalize",
     "callbacks",
     "LifecycleCallbacks",
@@ -122,13 +121,6 @@ def initialize(
     )
 
 
-def update_state(**kwargs) -> ModifiedState:
-    """
-    Update the session state with new values
-    """
-    return active_session().update_state(**kwargs)
-
-
 def finalize(**kwargs) -> ModifiedState:
     """
     Method to finalize the active session for sparsification

src/llmcompressor/transformers/finetune/session_mixin.py

@@ -17,7 +17,6 @@
     create_session,
     finalize,
     initialize,
-    update_state,
 )
 from llmcompressor.metrics import LoggerManager
 from llmcompressor.modifiers.distillation.utils.pytorch.model_wrapper import (
@@ -223,7 +222,7 @@ def create_optimizer(self):
                 len(self.train_dataset) / total_batch_size
             )
 
-        update_state(
+        active_session().state.update(
             optimizer=self.optimizer, steps_per_epoch=self.total_steps_per_epoch
         )