tensoracle 1.6.0__tar.gz → 1.6.1__tar.gz

{tensoracle-1.6.0 → tensoracle-1.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: tensoracle
-Version: 1.6.0
+Version: 1.6.1
 Summary: Build, save, reuse, and execute TensorRT engines from safetensors weights
 Project-URL: documentation, https://sync-dev-org.github.io/tensoracle/
 Project-URL: repository, https://github.com/sync-dev-org/tensoracle

{tensoracle-1.6.0 → tensoracle-1.6.1}/src/tensoracle/_build.py

@@ -114,7 +114,7 @@ def _build_serialized_engine(
 
     register_libraries(plugins)
 
-    builder = network = config = profile = serialized = tensor = None
+    builder = network = config = profile = serialized = tensor = timing = None
     inputs = {}
     try:
         builder = trt.Builder(logger or trt.Logger(trt.Logger.ERROR))
@@ -161,12 +161,22 @@ def _build_serialized_engine(
         if profile is not None:
             config.add_optimization_profile(profile)
 
-        serialized = builder.build_serialized_network(network, config)
-        if serialized is None:
-            raise RuntimeError("TensorRT failed to build the serialized engine")
+        timing = config.create_timing_cache(b"")
+        if timing is None:
+            raise RuntimeError("TensorRT failed to create a timing cache")
+        if not config.set_timing_cache(timing, ignore_mismatch=False):
+            raise RuntimeError("TensorRT rejected the timing cache")
+        # A cold build can serialize large zero-filled buffers. Rebuild with the
+        # populated in-memory cache, releasing the first plan before the second.
+        for _ in range(2):
+            serialized = None
+            serialized = builder.build_serialized_network(network, config)
+            if serialized is None:
+                raise RuntimeError("TensorRT failed to build the serialized engine")
         return bytes(serialized), output_names
     finally:
         # Tensor references also keep the network alive, including in exception frames.
         inputs.clear()
-        del tensor, profile, serialized, config, network, builder
+        # TensorRT requires the timing cache to outlive the builder using it.
+        del tensor, profile, serialized, config, network, builder, timing
         release_host_memory()

{tensoracle-1.6.0 → tensoracle-1.6.1}/src/tensoracle/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '1.6.0'
-__version_tuple__ = version_tuple = (1, 6, 0)
+__version__ = version = '1.6.1'
+__version_tuple__ = version_tuple = (1, 6, 1)
 
 __commit_id__ = commit_id = None

{tensoracle-1.6.0 → tensoracle-1.6.1}/tests/test_build.py

@@ -219,6 +219,12 @@ def test_tf32_is_disabled_and_other_builder_settings_are_applied(
         def set_memory_pool_limit(self, pool, limit):
             self.workspace_calls.append((pool, limit))
 
+        def create_timing_cache(self, data):
+            return object()
+
+        def set_timing_cache(self, timing, ignore_mismatch):
+            return True
+
     class Network:
         num_inputs = 1
         num_outputs = 0

{tensoracle-1.6.0 → tensoracle-1.6.1}/tests/test_build_memory.py

@@ -68,6 +68,8 @@ def test_unavailable_host_memory_release_preserves_the_build_result(monkeypatch,
         create_builder_config=lambda: SimpleNamespace(
             clear_flag=lambda flag: None,
             set_memory_pool_limit=lambda pool, limit: None,
+            create_timing_cache=lambda data: object(),
+            set_timing_cache=lambda timing, ignore_mismatch: True,
         ),
         build_serialized_network=serialize,
     )

{tensoracle-1.6.0 → tensoracle-1.6.1}/tests/test_cache.py

@@ -101,6 +101,8 @@ def test_save_large_built_plan_keeps_other_threads_running(tmp_path, monkeypatch
         create_builder_config=lambda: SimpleNamespace(
             clear_flag=lambda flag: None,
             set_memory_pool_limit=lambda pool, limit: None,
+            create_timing_cache=lambda data: object(),
+            set_timing_cache=lambda timing, ignore_mismatch: True,
         ),
         build_serialized_network=lambda network, config: plan,
     )
@@ -498,7 +500,7 @@ def test_each_declared_condition_selects_a_distinct_cache_entry(tmp_path, monkey
 
     monkeypatch.setattr(trt, "get_plugin_registry", PluginRegistry)
 
-    state = {"builds": 0, "graph_calls": 0, "shape": (1, 3, 2), "range": None}
+    state = {"graph_calls": 0, "shape": (1, 3, 2), "range": None}
 
     class FakeEngine:
         num_io_tensors = 2
@@ -546,6 +548,12 @@ def test_each_declared_condition_selects_a_distinct_cache_entry(tmp_path, monkey
         def add_optimization_profile(self, profile):
             pass
 
+        def create_timing_cache(self, data):
+            return object()
+
+        def set_timing_cache(self, timing, ignore_mismatch):
+            return True
+
     class FakeBuilder:
         def __init__(self, logger):
             pass
@@ -560,7 +568,6 @@ def test_each_declared_condition_selects_a_distinct_cache_entry(tmp_path, monkey
             return FakeProfile()
 
         def build_serialized_network(self, network, config):
-            state["builds"] += 1
             return b"small plan"
 
     monkeypatch.setattr(trt, "Builder", FakeBuilder)
@@ -607,9 +614,9 @@ def test_each_declared_condition_selects_a_distinct_cache_entry(tmp_path, monkey
         )
 
     first = open_case()
-    assert state["builds"] == 1
+    assert state["graph_calls"] == 1
     assert open_case(weight_files=(same_contents,)).path == first.path
-    assert state["builds"] == 1
+    assert state["graph_calls"] == 1
     variants = [
         {"name": "decoder"},
         {"graph_version": "2"},
@@ -636,12 +643,12 @@ def test_each_declared_condition_selects_a_distinct_cache_entry(tmp_path, monkey
         paths.add(open_case().path)
     assert first.path not in paths
     assert len(paths) == len(variants) + 4
-    assert state["builds"] == len(paths) + 1
+    assert state["graph_calls"] == len(paths) + 1
     assert (
         open_case(graph=lambda network: pytest.fail("graph changed without version bump")).path
         == first.path
     )
-    assert state["builds"] == len(paths) + 1
+    assert state["graph_calls"] == len(paths) + 1
 
 
 @pytest.mark.gpu

{tensoracle-1.6.0 → tensoracle-1.6.1}/tests/test_cache_management.py

@@ -20,7 +20,10 @@ def open_cached(tmp_path, monkeypatch):
     builder = SimpleNamespace(
         create_network=lambda flags: SimpleNamespace(num_inputs=0, num_outputs=0),
         create_builder_config=lambda: SimpleNamespace(
-            clear_flag=lambda flag: None, set_memory_pool_limit=lambda pool, limit: None
+            clear_flag=lambda flag: None,
+            set_memory_pool_limit=lambda pool, limit: None,
+            create_timing_cache=lambda data: object(),
+            set_timing_cache=lambda timing, ignore_mismatch: True,
         ),
         build_serialized_network=lambda network, config: b"TensorRT plan",
     )

{tensoracle-1.6.0 → tensoracle-1.6.1}/tests/test_gpu_cache.py

@@ -1,15 +1,80 @@
-"""Separate saved engines by the CUDA device properties used for compatibility."""
+"""Save compact engines and separate them by CUDA device compatibility."""
 
 import ctypes
+import time
 from types import SimpleNamespace
 
+import numpy as np
 import pytest
 import tensorrt as trt
+from safetensors.numpy import save_file
 
 import tensoracle
 import tensoracle._cache as cache
 
 
+@pytest.mark.gpu
+@pytest.mark.req("SP-TENSORACLE-009")
+def test_compact_saved_engine_executes_and_reopens_without_calling_graph(tmp_path, monkeypatch):
+    """A LeakyReLU engine stays below 1 MiB and computes correctly before and after reopening."""
+    import cupy as cp
+
+    monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
+    weights = tmp_path / "weights.safetensors"
+    save_file({"slope": np.array([0.125], dtype=np.float16)}, weights)
+    shape = (1, 64, 135, 240)
+    graph_calls = []
+
+    def graph(network):
+        graph_calls.append(1)
+        features = network.add_input("features", trt.float16, shape)
+        layer = network.add_activation(features, trt.ActivationType.LEAKY_RELU)
+        layer.alpha = 0.125
+        result = layer.get_output(0)
+        result.name = "result"
+        network.mark_output(result)
+
+    def open_engine(callback):
+        return tensoracle.open_engine(
+            callback,
+            tensoracle.EngineBuild({"features": tensoracle.InputShape.fixed(shape)}),
+            name="compact-leaky-relu",
+            graph_version="1",
+            weights=(weights,),
+            io=tensoracle.EngineIO(
+                {"features": tensoracle.TensorSpec(trt.float16, shape)},
+                {"result": tensoracle.TensorSpec(trt.float16, shape)},
+            ),
+        )
+
+    started = time.perf_counter()
+    first = open_engine(graph)
+    elapsed = time.perf_counter() - started
+    size = first.path.stat().st_size
+    print(f"LeakyReLU saved engine: {size} bytes, build/open: {elapsed:.6f} s")
+    assert graph_calls == [1]
+    assert size < 1024 * 1024
+
+    def forbidden(network):
+        pytest.fail("saved engine invoked the graph callback")
+
+    reopened = open_engine(forbidden)
+    assert reopened.path == first.path
+    assert reopened.path.stat().st_size == size
+    assert graph_calls == [1]
+    host = np.linspace(-1, 1, np.prod(shape), dtype=np.float32).astype(np.float16).reshape(shape)
+    expected = np.where(host >= 0, host, host * np.float16(0.125))
+    value = cp.asarray(host)
+    for handle in (first, reopened):
+        output = cp.asarray(np.full_like(host, np.nan))
+        with handle.create_execution() as execution:
+            execution.enqueue(
+                {"features": value.data.ptr, "result": output.data.ptr},
+                stream=cp.cuda.get_current_stream().ptr,
+            )
+            np.testing.assert_array_equal(cp.asnumpy(output), expected)
+
+
 @pytest.fixture
 def cuda_driver(monkeypatch):
     properties = {
@@ -128,17 +193,18 @@ def test_changed_gpu_property_builds_a_separate_saved_engine(
     monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
     builds = []
 
-    def serialize(network, config):
+    def graph(network):
         builds.append(cuda_driver.device)
-        return b"serialized engine"
 
     builder = SimpleNamespace(
         create_network=lambda flags: SimpleNamespace(num_inputs=0, num_outputs=0),
         create_builder_config=lambda: SimpleNamespace(
             clear_flag=lambda flag: None,
             set_memory_pool_limit=lambda pool, limit: None,
+            create_timing_cache=lambda data: object(),
+            set_timing_cache=lambda timing, ignore_mismatch: True,
         ),
-        build_serialized_network=serialize,
+        build_serialized_network=lambda network, config: b"serialized engine",
     )
     monkeypatch.setattr(trt, "Builder", lambda logger: builder)
     monkeypatch.setattr(
@@ -153,7 +219,7 @@ def test_changed_gpu_property_builds_a_separate_saved_engine(
 
     def open_engine():
         return tensoracle.open_engine(
-            lambda network: None,
+            graph,
             tensoracle.EngineBuild({}),
             name="encoder",
             graph_version="1",