Validator that checks for custom user objects in configs. Custom user objects break or don't work well with serialization/visualization/codegen.

Fiddle-Config Team · copybara-github · commit c882b09ac0a1 · 2023-07-25T13:13:14.000-07:00
PiperOrigin-RevId: 550975396
diff --git a/fiddle/_src/daglish.py b/fiddle/_src/daglish.py
@@ -438,7 +438,7 @@ def initial_state(self) -> State:
     The initial state has a reference to the traversal, the empty path, and no
     parent states.
     """
-    return State(self, (), self.root_obj, _parent=None)  # pytype: disable=attribute-error
+    return State(self, (), self.root_obj, parent=None)
 
   @classmethod
   def begin(cls, fn: Callable[..., Any], root_obj: Any) -> State:
@@ -484,8 +484,9 @@ class State:
     current_path: A path that can be followed to the current object. In the case
       of shared objects, there will be other paths to the current object, and
       often these are determined by a somewhat arbitrary DAG traversal order.
+    parent: The parent state.
   """
-  __slots__ = ("traversal", "current_path", "_value", "_parent")
+  __slots__ = ("traversal", "current_path", "_value", "parent")
 
   traversal: Traversal
   current_path: Path
@@ -494,7 +495,7 @@ class State:
   # accessors.
   _value: Any  # pylint: disable=invalid-name
 
-  _parent: Optional[State]
+  parent: Optional[State]
 
   @property
   def _object_id(self) -> int:
@@ -504,12 +505,23 @@ def _object_id(self) -> int:
   def _is_memoizable(self) -> bool:
     return is_memoizable(self._value)
 
+  @property
+  def original_value(self):
+    """Original value constructed with this state.
+
+    Generally please don't use this value, it's much more clear to use the
+    first argument of your `traverse(value, state)` function, especially since
+    for post-order traversals, you'll often write `value =
+    state.map_children(value)`.
+    """
+    return self._value
+
   @property
   def ancestors_inclusive(self) -> Iterable[State]:
     """Gets ancestors, including the current state."""
     yield self
-    if self._parent is not None:
-      yield from self._parent.ancestors_inclusive
+    if self.parent is not None:
+      yield from self.parent.ancestors_inclusive
 
   def get_all_paths(self, allow_caching: bool = True) -> List[Path]:
     """Gets all paths to the current value.
diff --git a/fiddle/_src/history.py b/fiddle/_src/history.py
@@ -55,10 +55,20 @@ class Location:
   line_number: int
   function_name: Optional[str]
 
-  def __str__(self) -> str:
+  def format(self, max_filename_parts: Optional[int] = None):
+    filename = self.filename
+    if max_filename_parts is not None:
+      filename_parts = filename.split(os.path.sep)
+      if len(filename_parts) > max_filename_parts:
+        filename = os.path.sep.join(
+            ["...", *filename_parts[-max_filename_parts:]]
+        )
     if self.function_name is None:
-      return f"{self.filename}:{self.line_number}"
-    return f"{self.filename}:{self.line_number}:{self.function_name}"
+      return f"{filename}:{self.line_number}"
+    return f"{filename}:{self.line_number}:{self.function_name}"
+
+  def __str__(self) -> str:
+    return self.format()
 
   def __deepcopy__(self, memo):
     del memo  # unused
diff --git a/fiddle/_src/history_test.py b/fiddle/_src/history_test.py
@@ -41,6 +41,22 @@ def test_location_formatting(self):
         function_name="make_config")
     self.assertEqual(str(location), "my_other_file.py:321:make_config")
 
+  def test_location_formatting_concise(self):
+    location = history.Location(
+        filename="foo/bar/baz/my_file.py", line_number=123, function_name=None
+    )
+    self.assertEqual(location.format(3), ".../bar/baz/my_file.py:123")
+    self.assertEqual(location.format(2), ".../baz/my_file.py:123")
+    location = history.Location(
+        filename="foo/bar/baz/my_other_file.py",
+        line_number=321,
+        function_name="make_config",
+    )
+    self.assertEqual(
+        location.format(4), "foo/bar/baz/my_other_file.py:321:make_config"
+    )
+    self.assertEqual(location.format(1), ".../my_other_file.py:321:make_config")
+
   def test_entry_simple(self):
     entry = history.new_value("x", 1)
     self.assertEqual(entry.param_name, "x")
diff --git a/fiddle/_src/validation/no_custom_objects.py b/fiddle/_src/validation/no_custom_objects.py
@@ -0,0 +1,94 @@
+# coding=utf-8
+# Copyright 2022 The Fiddle-Config Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Checks that custom objects/instances are not present in a configuration."""
+
+import dataclasses
+from typing import Any, List, Optional
+
+import fiddle as fdl
+from fiddle import daglish
+from fiddle import history
+
+
+def _concise_history(entries: List[history.HistoryEntry]) -> str:
+  """Returns a concise string for a history entry."""
+  if not entries:
+    return "(no history)"
+  set_or_deleted = (
+      "Set in "
+      if entries[-1].new_value is not history.DELETED
+      else "Deleted in "
+  )
+  return set_or_deleted + entries[-1].location.format(max_filename_parts=3)
+
+
+def _get_history_from_state(
+    state: daglish.State,
+) -> Optional[List[history.HistoryEntry]]:
+  """Returns the history from a Buildable for a state."""
+  while state.current_path and state.parent is not None:
+    attr = state.current_path[-1]
+    state = state.parent
+    if isinstance(state.original_value, fdl.Buildable):
+      assert isinstance(attr, daglish.Attr)
+      entries = state.original_value.__argument_history__[attr.name]
+      return entries if entries else None
+  return None
+
+
+def get_config_errors(config: Any) -> List[str]:
+  """Returns a list of errors found in the given config.
+
+  Args:
+    config: Fiddle config object, or nested structure of configs.
+  """
+  errors = []
+
+  def history_str(state):
+    return ", " + _concise_history(_get_history_from_state(state))
+
+  def traverse(value, state: daglish.State):
+    path_str = daglish.path_str(state.current_path)
+    if isinstance(value, tuple) and hasattr(type(value), "_fields"):
+      errors.append(f"Found namedtuple at {path_str}{history_str(state)}")
+    elif dataclasses.is_dataclass(value):
+      errors.append(f"Found dataclass at {path_str}{history_str(state)}")
+    elif (not state.is_traversable(value)) and not daglish.is_unshareable(
+        value
+    ):
+      errors.append(f"Found {type(value)} at {path_str}{history_str(state)}")
+    return state.map_children(value)
+
+  daglish.MemoizedTraversal.run(traverse, config)
+  return errors
+
+
+def check_no_custom_objects(config: Any) -> None:
+  """Checks that no custom objects are present in the given config.
+
+  Args:
+    config: Fiddle config object, or nested structure of configs.
+
+  Raises:
+    ValueError: If the configuration contains custom objects.
+  """
+  errors = get_config_errors(config)
+  if errors:
+    raise ValueError(
+        "Custom objects were found in the config. In general, you should "
+        "be able to convert these to fdl.Config's. Custom objects:\n  "
+        + "\n  ".join(errors),
+    )
diff --git a/fiddle/_src/validation/no_custom_objects_test.py b/fiddle/_src/validation/no_custom_objects_test.py
@@ -0,0 +1,174 @@
+# coding=utf-8
+# Copyright 2022 The Fiddle-Config Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for no_custom_objects."""
+
+import dataclasses
+import typing
+
+from absl.testing import absltest
+import fiddle as fdl
+from fiddle import daglish
+from fiddle._src.testing.example import fake_encoder_decoder
+from fiddle._src.validation import no_custom_objects
+
+
+def foo(x):
+  return x
+
+
+class MyNamedtuple(typing.NamedTuple):
+  a: int
+  b: str
+
+
+@dataclasses.dataclass(frozen=True)
+class MyDataclass:
+  a: int
+  b: str
+
+
+class NoCustomObjectsTest(absltest.TestCase):
+
+  def test_empty_history(self):
+    self.assertEqual(
+        no_custom_objects._concise_history(entries=[]), "(no history)"
+    )
+
+  def test_concise_history(self):
+    config = fake_encoder_decoder.fixture.as_buildable()
+    config.encoder.attention.dtype = "float64"
+    with self.subTest("overridden"):
+      history = no_custom_objects._concise_history(
+          entries=config.encoder.attention.__argument_history__["dtype"]
+      )
+      self.assertRegex(history, r"Set in .+:\d+:test_concise_history")
+    with self.subTest("not_overridden"):
+      history = no_custom_objects._concise_history(
+          entries=config.encoder.__argument_history__["attention"]
+      )
+      self.assertRegex(history, r"Set in .+fake_encoder.+:\d+:fixture")
+    with self.subTest("deleted"):
+      del config.encoder.attention.dtype
+      history = no_custom_objects._concise_history(
+          entries=config.encoder.attention.__argument_history__["dtype"]
+      )
+      self.assertRegex(history, r"Deleted in .+:\d+:test_concise_history")
+
+  def test_get_history_from_state(self):
+    config = fdl.Config(foo, {"a": {"b": 1}})
+    traversal = daglish.MemoizedTraversal(NotImplemented, config)  # pytype: disable=wrong-arg-types
+    state = traversal.initial_state()
+    state = daglish.State(
+        state.traversal,
+        (*state.current_path, daglish.Attr("x")),
+        config.x,
+        state,
+    )
+    state = daglish.State(
+        state.traversal,
+        (*state.current_path, daglish.Key("a")),
+        config.x["a"],
+        state,
+    )
+    history1 = no_custom_objects._get_history_from_state(state=state)
+    state = daglish.State(
+        state.traversal,
+        (*state.current_path, daglish.Key("b")),
+        config.x["a"]["b"],
+        state,
+    )
+    history2 = no_custom_objects._get_history_from_state(state=state)
+    self.assertNotEmpty(history2)
+    self.assertIs(
+        history1,
+        history2,
+        msg=(
+            "Since dictionaries do not have state, _get_history_from_state"
+            " should traverse to an ancestor to find the history."
+        ),
+    )
+
+  def test_get_config_errors_empty(self):
+    config = fake_encoder_decoder.fixture.as_buildable()
+    self.assertEmpty(no_custom_objects.get_config_errors(config=config))
+
+  def test_get_config_errors_namedtuple(self):
+    config = fake_encoder_decoder.fixture.as_buildable()
+    config.encoder.attention = MyNamedtuple(1, "a")
+    errors = no_custom_objects.get_config_errors(config=config)
+    self.assertLen(errors, 1)
+    self.assertRegex(
+        errors[0],
+        r"Found.*namedtuple.*at \.encoder\.attention.*Set"
+        r" in.*:\d+:test_get_config_errors_namedtuple",
+    )
+
+  def test_get_config_errors_dataclass(self):
+    config = fake_encoder_decoder.fixture.as_buildable()
+    config.encoder.attention = MyDataclass(1, "a")
+    errors = no_custom_objects.get_config_errors(config=config)
+    self.assertLen(errors, 1)
+    self.assertRegex(
+        errors[0],
+        r"Found.*dataclass.*at \.encoder\.attention.*Set"
+        r" in.*:\d+:test_get_config_errors_dataclass",
+    )
+
+  def test_get_config_errors_not_empty(self):
+    config = fake_encoder_decoder.fixture.as_buildable()
+    config.encoder.attention = object()
+    errors = no_custom_objects.get_config_errors(config=config)
+    self.assertLen(errors, 1)
+    self.assertRegex(
+        errors[0],
+        r"Found.*object.*at \.encoder\.attention.*Set"
+        r" in.*:\d+:test_get_config_errors_not_empty",
+    )
+
+  def test_check_no_custom_objects_okay(self):
+    config = fake_encoder_decoder.fixture.as_buildable()
+    no_custom_objects.check_no_custom_objects(config)
+
+  def test_check_no_custom_objects_error(self):
+    config = fake_encoder_decoder.fixture.as_buildable()
+    config.encoder.attention = object()
+    config.decoder.self_attention = object()
+    with self.assertRaisesRegex(
+        ValueError,
+        r"Custom objects were found.*Custom objects:\n  Found.*object.*at"
+        r" \.encoder\.attention, Set in"
+        r" .*:\d+:test_check_no_custom_objects_error\n  Found.*object.*at"
+        r" \.decoder\.self_attention, Set in"
+        r" .*:\d+:test_check_no_custom_objects_error",
+    ):
+      no_custom_objects.check_no_custom_objects(config=config)
+
+  def test_no_history_custom_objects_error(self):
+    config = {
+        "encoder_attention": object(),
+        "decoder_self_attention": object(),
+    }
+    with self.assertRaisesRegex(
+        ValueError,
+        r"Custom objects were found.*Custom objects:"
+        r"\n  Found.*object.*at \['encoder_attention'\],.*no history.*"
+        r"\n  Found.*object.*at \['decoder_self_attention'\],.*no history.*",
+    ):
+      no_custom_objects.check_no_custom_objects(config=config)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/fiddle/validation/no_custom_objects.py b/fiddle/validation/no_custom_objects.py
