[Utils] Replace preserve_attr with patch_attr (#1187)

## Purpose ##
* Provide explicit patching functionality to `preserve_attr`
* This function is very similar to `unittest.mock.patch`, except that
using this functionality does not require using the `unittest` library
in source code

## Changes ##
* Replace `preserve_attr` with `patch_attr`
* Replace usage in `src/llmcompressor/pipelines/sequential/helpers.py`,
which helps with clarity
* Add unit test mark to utils/helpers pytests

## Testing ##
* Added tests

---------

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

Author: kylesayrs

src/llmcompressor/pipelines/sequential/helpers.py

@@ -13,7 +13,7 @@
 from transformers.utils.fx import HFTracer
 
 from llmcompressor.modifiers.utils.hooks import HooksMixin
-from llmcompressor.utils.helpers import calibration_forward_context, preserve_attr
+from llmcompressor.utils.helpers import calibration_forward_context, patch_attr
 
 __all__ = ["trace_subgraphs", "Subgraph"]
 
@@ -132,15 +132,14 @@ def is_leaf_module(self, module: Module, module_qualified_name: str) -> bool:
 
         def trace(self, root: Union[Module, Callable], *args, **kwargs) -> Graph:
             if isinstance(root, Module):
-                with preserve_attr(type(root), "forward"):
-                    # due to a bug in Tracer.create_args_for_root (_patch_function),
-                    # we must unwrap function wrappers prior to tracing, for example
-                    # the `deprecate_kwarg` by transformers which wraps forward
-
-                    # we override the class method because the
-                    # class method is the one being traced
-                    type(root).forward = inspect.unwrap(type(root).forward)
-
+                # due to a bug in Tracer.create_args_for_root (_patch_function),
+                # we must unwrap function wrappers prior to tracing, for example
+                # the `deprecate_kwarg` by transformers which wraps forward
+                unwrapped_forward = inspect.unwrap(type(root).forward)
+
+                # we override the class method because the
+                # class method is the one being traced
+                with patch_attr(type(root), "forward", unwrapped_forward):
                     return super().trace(root, *args, **kwargs)
 
             else:

src/llmcompressor/utils/helpers.py

@@ -65,7 +65,7 @@
     "DisableQuantization",
     "eval_context",
     "calibration_forward_context",
-    "preserve_attr",
+    "patch_attr",
 ]
 
 
@@ -1051,9 +1051,29 @@ def calibration_forward_context(model: PreTrainedModel):
 
 
 @contextlib.contextmanager
-def preserve_attr(base: object, attr: str):
-    value = getattr(base, attr)
+def patch_attr(base: object, attr: str, value: Any):
+    """
+    Patch the value of an object attribute. Original value is restored upon exit
+
+    :param base: object which has the attribute to patch
+    :param attr: name of the the attribute to patch
+    :param value: used to replace original value
+
+    Usage:
+    >>> from types import SimpleNamespace
+    >>> obj = SimpleNamespace()
+    >>> with patch_attr(obj, "attribute", "value"):
+    ...     assert obj.attribute == "value"
+    >>> assert not hasattr(obj, "attribute")
+    """
+    _sentinel = object()
+    original_value = getattr(base, attr, _sentinel)
+
+    setattr(base, attr, value)
     try:
         yield
     finally:
-        setattr(base, attr, value)
+        if original_value is not _sentinel:
+            setattr(base, attr, original_value)
+        else:
+            delattr(base, attr)

tests/llmcompressor/utils/test_helpers.py

@@ -11,10 +11,12 @@
     flatten_iterable,
     getattr_chain,
     interpolate,
+    patch_attr,
     validate_str_iterable,
 )
 
 
+@pytest.mark.unit
 @pytest.mark.parametrize(
     "test_list,output",
     [
@@ -29,6 +31,7 @@ def test_flatten_iterable(test_list, output):
     assert flattened == output
 
 
+@pytest.mark.unit
 @pytest.mark.parametrize(
     "test_bool,output",
     [
@@ -53,6 +56,7 @@ def test_convert_to_bool(test_bool, output):
     assert converted == output
 
 
+@pytest.mark.unit
 @pytest.mark.parametrize(
     "test_list,output",
     [
@@ -68,11 +72,13 @@ def test_validate_str_iterable(test_list, output):
     assert validated == output
 
 
+@pytest.mark.unit
 def test_validate_str_iterable_negative():
     with pytest.raises(ValueError):
         validate_str_iterable("will fail", "")
 
 
+@pytest.mark.unit
 @pytest.mark.parametrize(
     "x_cur,x0,x1,y0,y1,inter_func,out",
     [
@@ -92,6 +98,7 @@ def test_interpolate(x_cur, x0, x1, y0, y1, inter_func, out):
     assert abs(out - interpolated) < 0.01
 
 
+@pytest.mark.unit
 def test_getattr_chain():
     base = SimpleNamespace()
     base.a = None
@@ -123,13 +130,15 @@ def test_getattr_chain():
         getattr_chain(base, "b.d.dne")
 
 
+@pytest.mark.unit
 def test_DisableQuantization():
     model = torch.nn.Linear(1, 1)
     with DisableQuantization(model):
         assert not model.quantization_enabled
     assert model.quantization_enabled
 
 
+@pytest.mark.unit
 def test_calibration_forward_context():
     model = torch.nn.Linear(1, 1)
     model.config = SimpleNamespace()
@@ -143,3 +152,21 @@ def test_calibration_forward_context():
     assert torch.is_grad_enabled()
     assert model.config.use_cache
     assert model.training
+
+
+@pytest.mark.unit
+def test_patch_attr():
+    # patch, original value
+    obj = SimpleNamespace()
+    obj.attribute = "original"
+    with patch_attr(obj, "attribute", "patched"):
+        assert obj.attribute == "patched"
+        obj.attribute = "modified"
+    assert obj.attribute == "original"
+
+    # patch, no original attribute
+    obj = SimpleNamespace()
+    with patch_attr(obj, "attribute", "patched"):
+        assert obj.attribute == "patched"
+        obj.attribute = "modified"
+    assert not hasattr(obj, "attribute")