Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 14 additions & 1 deletion py/torch_tensorrt/dynamo/_engine_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,19 @@
]


def _canonicalize_setting_value(value: Any) -> str:
"""Stringify a setting so the result does not depend on iteration order.

``str()`` of a set lays elements out in hash order, and Python randomizes
string hashing per process unless PYTHONHASHSEED is pinned. Hashing that
directly would give the same compilation settings a different cache key on
every run, so unordered collections are sorted first.
"""
if isinstance(value, (set, frozenset)):
return str(sorted(str(element) for element in value))
return str(value)


class BaseEngineCache(ABC):
@abstractmethod
def __init__(
Expand Down Expand Up @@ -88,7 +101,7 @@ def canonicalize_graph(graph: torch.fx.Graph) -> str:
input_specs_hash = sha256_hash(input_specs_data)

invariant_engine_specs = [
str(getattr(settings, field))
_canonicalize_setting_value(getattr(settings, field))
for field in sorted(_SETTINGS_TO_BE_ENGINE_INVARIANT)
]
with io.BytesIO() as stream:
Expand Down
21 changes: 20 additions & 1 deletion tests/py/dynamo/models/test_engine_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,10 @@
import torch_tensorrt as torch_trt
from torch.testing._internal.common_utils import TestCase
from torch_tensorrt.dynamo._defaults import TIMING_CACHE_PATH
from torch_tensorrt.dynamo._engine_cache import BaseEngineCache
from torch_tensorrt.dynamo._engine_cache import (
BaseEngineCache,
_canonicalize_setting_value,
)
from torch_tensorrt.dynamo._settings import CompilationSettings
from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity

Expand Down Expand Up @@ -60,6 +63,22 @@ def load(self, hash: str, prefix: str = "blob") -> Optional[bytes]:


class TestHashFunction(TestCase):
def test_unordered_settings_are_hashed_in_a_stable_order(self):
"""Engine-invariant settings that are sets must not depend on hash order.

str() of a set lays elements out by hash, and string hashing is
randomized per process, so hashing one directly would give the same
settings a different cache key on every run.
"""
self.assertEqual(
_canonicalize_setting_value({"b", "a", "c"}), "['a', 'b', 'c']"
)
self.assertEqual(
_canonicalize_setting_value(frozenset({"b", "a"})),
_canonicalize_setting_value({"a", "b"}),
)
self.assertEqual(_canonicalize_setting_value(True), "True")

@unittest.skipIf(
not importlib.util.find_spec("torchvision"), "torchvision not installed"
)
Expand Down
Loading