[torch/utils][Code Clean] Clean asserts in torch/utils/*.py (pytorch#165410)

Including:
- `torch/utils/*.py`

Fixes part of pytorch#164878

Pull Request resolved: pytorch#165410
Approved by: https://github.com/albanD, https://github.com/cyyever

Author: KarhouTam

torch/utils/_appending_byte_serializer.py

@@ -38,19 +38,22 @@ def to_bytes(self) -> bytes:
         digest = zlib.crc32(self._data[CHECKSUM_DIGEST_SIZE:]).to_bytes(
             4, byteorder="big", signed=False
         )
-        assert len(digest) == CHECKSUM_DIGEST_SIZE
+        if len(digest) != CHECKSUM_DIGEST_SIZE:
+            raise AssertionError("Computed checksum digest has unexpected size")
         self._data[0:CHECKSUM_DIGEST_SIZE] = digest
         return bytes(self._data)
 
 
 class BytesReader:
     def __init__(self, data: bytes) -> None:
         # Check for data corruption
-        assert len(data) >= CHECKSUM_DIGEST_SIZE
+        if len(data) < CHECKSUM_DIGEST_SIZE:
+            raise AssertionError("Input data is too short to contain checksum")
         digest = zlib.crc32(data[CHECKSUM_DIGEST_SIZE:]).to_bytes(
             4, byteorder="big", signed=False
         )
-        assert len(digest) == CHECKSUM_DIGEST_SIZE
+        if len(digest) != CHECKSUM_DIGEST_SIZE:
+            raise AssertionError("Computed checksum digest has unexpected size")
         if data[0:CHECKSUM_DIGEST_SIZE] != digest:
             raise RuntimeError(
                 "Bytes object is corrupted, checksum does not match. "
@@ -120,7 +123,11 @@ def to_bytes(self) -> bytes:
     @staticmethod
     def to_list(data: bytes, *, deserialize_fn: Callable[[BytesReader], T]) -> list[T]:
         reader = BytesReader(data)
-        assert reader.read_uint64() == _ENCODING_VERSION
+        if reader.read_uint64() != _ENCODING_VERSION:
+            raise AssertionError(
+                f"Encoding version mismatch in AppendingByteSerializer.to_list, \
+                    got {reader.read_uint64()}"
+            )
 
         result: list[T] = []
         while not reader.is_finished():

torch/utils/_config_module.py

@@ -85,12 +85,16 @@ def __post_init__(self) -> None:
         )
 
         if self.alias is not None:
-            assert (
-                self.default is _UNSET_SENTINEL
-                and self.justknob is None
-                and self.env_name_default is None
-                and self.env_name_force is None
-            ), "if alias is set, none of {default, justknob and env var} can be set"
+            if (
+                self.default is not _UNSET_SENTINEL
+                or self.justknob is not None
+                or self.env_name_default is not None
+                or self.env_name_force is not None
+            ):
+                raise AssertionError(
+                    "if alias is set, none of {default, justknob, \
+                        env_name_default and env_name_force} can be set"
+                )
 
     @staticmethod
     def string_or_list_of_string_to_list(
@@ -100,7 +104,8 @@ def string_or_list_of_string_to_list(
             return None
         if isinstance(val, str):
             return [val]
-        assert isinstance(val, list)
+        if not isinstance(val, list):
+            raise AssertionError(f"val is not a list, got {type(val)}")
         return val
 
 
@@ -193,7 +198,10 @@ def visit(
                 if dest is module:
                     delattr(module, key)
             elif isinstance(value, type):
-                assert value.__module__ == module.__name__
+                if value.__module__ != module.__name__:
+                    raise AssertionError(
+                        f"subconfig class {value} must be defined in module {module.__name__}"
+                    )
                 # a subconfig with `class Blah:` syntax
                 proxy = SubConfigProxy(module, f"{name}.")
                 visit(value, proxy, f"{name}.")
@@ -234,10 +242,8 @@ def get_assignments_with_compile_ignored_comments(module: ModuleType) -> set[str
             prev_name = ""
             maybe_current = token.string.strip()
             if COMPILE_IGNORED_MARKER in maybe_current:
-                assert current_comment == (
-                    "",
-                    -1,
-                ), f"unconsumed {COMPILE_IGNORED_MARKER}"
+                if current_comment != ("", -1):
+                    raise AssertionError(f"unconsumed {COMPILE_IGNORED_MARKER}")
                 current_comment = maybe_current, token.start[0]
         elif token.type == tokenize.NAME:
             # Only accept the first name token, to handle if you have
@@ -254,7 +260,8 @@ def get_assignments_with_compile_ignored_comments(module: ModuleType) -> set[str
                 assignments.add(prev_name)
                 current_comment = "", -1  # reset
             prev_name = ""
-    assert current_comment == ("", -1), f"unconsumed {COMPILE_IGNORED_MARKER}"
+    if current_comment != ("", -1):
+        raise AssertionError(f"unconsumed {COMPILE_IGNORED_MARKER}")
     return assignments
 
 
@@ -306,20 +313,22 @@ def __init__(self, config: _Config):
 
         # Ensure justknobs and envvars are allowlisted types
         if self.justknob is not None and self.default is not None:
-            assert isinstance(self.default, bool), (
-                f"justknobs only support booleans, {self.default} is not a boolean"
-            )
+            if not isinstance(self.default, bool):
+                raise AssertionError(
+                    f"justknobs only support booleans, {self.default} is not a boolean"
+                )
         if self.value_type is not None and (
             config.env_name_default is not None or config.env_name_force is not None
         ):
-            assert self.value_type in (
+            if self.value_type not in (
                 bool,
                 str,
                 Optional[bool],
                 Optional[str],
-            ), (
-                f"envvar configs only support (optional) booleans or strings, {self.value_type} is neither"
-            )
+            ):
+                raise AssertionError(
+                    f"envvar configs only support (optional) booleans or strings, {self.value_type} is neither"
+                )
 
 
 class ConfigModule(ModuleType):
@@ -417,7 +426,10 @@ def _get_alias_val(self, entry: _ConfigEntry) -> Any:
 
     def _set_alias_val(self, entry: _ConfigEntry, val: Any) -> None:
         data = self._get_alias_module_and_name(entry)
-        assert data is not None
+        if data is None:
+            raise AssertionError(
+                "alias data should not be None when setting alias value"
+            )
         module, constant_name = data
         setattr(module, constant_name, val)
 
@@ -642,19 +654,32 @@ def foo(...):
         changes: dict[str, Any]
         if arg1 is not None:
             if arg2 is not None:
-                assert isinstance(arg1, str)
+                if not isinstance(arg1, str):
+                    raise AssertionError(
+                        "first argument must be a string when passing 2 positional args to patch"
+                    )
                 # patch("key", True) syntax
                 changes = {arg1: arg2}
             else:
-                assert isinstance(arg1, dict)
+                if not isinstance(arg1, dict):
+                    raise AssertionError(
+                        "first argument must be a dict when passing a single positional arg to patch"
+                    )
                 # patch({"key": True}) syntax
                 changes = arg1
-            assert not kwargs
+            if kwargs:
+                raise AssertionError(
+                    "cannot pass both positional and keyword arguments to patch"
+                )
         else:
             # patch(key=True) syntax
             changes = kwargs
-            assert arg2 is None
-        assert isinstance(changes, dict), f"expected `dict` got {type(changes)}"
+            if arg2 is not None:
+                raise AssertionError(
+                    "second positional argument is only valid when first argument is a key string"
+                )
+        if not isinstance(changes, dict):
+            raise AssertionError(f"expected `dict` got {type(changes)}")
         prior: dict[str, Any] = {}
         config = self
 
@@ -663,7 +688,10 @@ def __init__(self) -> None:
                 self.changes = changes
 
             def __enter__(self) -> None:
-                assert not prior
+                if prior:
+                    raise AssertionError(
+                        "prior should be empty when entering ConfigPatch"
+                    )
                 for key in self.changes.keys():
                     # KeyError on invalid entry
                     prior[key] = config.__getattr__(key)

torch/utils/_config_typing.pyi

@@ -21,7 +21,8 @@ This file should be imported into any file that uses install_config_module like
 Note that the import should happen before the call to install_config_module(), otherwise runtime errors may occur.
 """
 
-assert TYPE_CHECKING, "Do not use at runtime"
+if not TYPE_CHECKING:  # noqa: PYI002
+    raise AssertionError("Do not use at runtime")  # noqa: W291
 
 def save_config() -> bytes: ...
 def save_config_portable(*, ignore_private_configs: bool = True) -> dict[str, Any]: ...

torch/utils/_content_store.py

@@ -217,7 +217,10 @@ def read_storage(self, h: str, *, device=None) -> torch.UntypedStorage:
             weights_only=True,
             map_location=device,
         )._untyped_storage
-        assert s is not None
+        if s is None:
+            raise AssertionError(
+                f"expected storage for hash {h} in {os.path.join(self.loc, 'storages')}, got None"
+            )
         if self.storage_cache is not None:
             self.storage_cache[device][h] = StorageWeakRef(s)
         return s

torch/utils/_contextlib.py

@@ -86,13 +86,14 @@ def context_decorator(ctx, func):
     be a multi-shot context manager that can be directly invoked multiple times)
     or a callable that produces a context manager.
     """
-    assert not (callable(ctx) and hasattr(ctx, "__enter__")), (
-        f"Passed in {ctx} is both callable and also a valid context manager "
-        "(has __enter__), making it ambiguous which interface to use.  If you "
-        "intended to pass a context manager factory, rewrite your call as "
-        "context_decorator(lambda: ctx()); if you intended to pass a context "
-        "manager directly, rewrite your call as context_decorator(lambda: ctx)"
-    )
+    if callable(ctx) and hasattr(ctx, "__enter__"):
+        raise AssertionError(
+            f"Passed in {ctx} is both callable and also a valid context manager "
+            "(has __enter__), making it ambiguous which interface to use.  If you "
+            "intended to pass a context manager factory, rewrite your call as "
+            "context_decorator(lambda: ctx()); if you intended to pass a context "
+            "manager directly, rewrite your call as context_decorator(lambda: ctx)"
+        )
 
     if not callable(ctx):
 

torch/utils/_cxx_pytree.py

@@ -931,7 +931,10 @@ def _broadcast_to_and_flatten(
     treespec: TreeSpec,
     is_leaf: Optional[Callable[[PyTree], bool]] = None,
 ) -> Optional[list[Any]]:
-    assert _is_pytreespec_instance(treespec)
+    if not _is_pytreespec_instance(treespec):
+        raise AssertionError(
+            f"_broadcast_to_and_flatten: Expected `treespec` to be instance of PyTreeSpec but got {type(treespec)}"
+        )
     full_tree = tree_unflatten([0] * treespec.num_leaves, treespec)
     try:
         return broadcast_prefix(tree, full_tree, is_leaf=is_leaf)

torch/utils/_device.py

@@ -87,12 +87,18 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         # or else someone else has popped it!
         for _ in range(_len_torch_function_stack() - 1):
             mode = _pop_mode()
-            assert not isinstance(mode, DeviceContext)
+            if isinstance(mode, DeviceContext):
+                raise AssertionError(
+                    "Found nested DeviceContext on the mode stack where none expected"
+                )
             cur_stack.append(mode)
 
         if _len_torch_function_stack() > 0:
             mode = _pop_mode()
-            assert isinstance(mode, DeviceContext)
+            if not isinstance(mode, DeviceContext):
+                raise AssertionError(
+                    "Expected a DeviceContext at the bottom of the mode stack"
+                )
 
         for mode in reversed(cur_stack):
             _push_mode(mode)

torch/utils/_functools.py

@@ -31,7 +31,8 @@ def cache_method(
 
     @functools.wraps(f)
     def wrap(self: _C, *args: _P.args, **kwargs: _P.kwargs) -> _T:
-        assert not kwargs
+        if kwargs:
+            raise AssertionError("cache_method does not accept keyword arguments")
         if not (cache := getattr(self, cache_name, None)):
             cache = {}
             setattr(self, cache_name, cache)