tensoracle 0.3.1__tar.gz → 0.3.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (69) hide show
  1. {tensoracle-0.3.1 → tensoracle-0.3.2}/PKG-INFO +3 -1
  2. {tensoracle-0.3.1 → tensoracle-0.3.2}/README.md +2 -0
  3. {tensoracle-0.3.1 → tensoracle-0.3.2}/docs/index.md +2 -0
  4. {tensoracle-0.3.1 → tensoracle-0.3.2}/src/tensoracle/__init__.py +9 -1
  5. {tensoracle-0.3.1 → tensoracle-0.3.2}/src/tensoracle/_build.py +35 -1
  6. {tensoracle-0.3.1 → tensoracle-0.3.2}/src/tensoracle/_cache.py +30 -8
  7. {tensoracle-0.3.1 → tensoracle-0.3.2}/src/tensoracle/_execution.py +3 -0
  8. {tensoracle-0.3.1 → tensoracle-0.3.2}/src/tensoracle/_version.py +2 -2
  9. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_build.py +1 -0
  10. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_build_memory.py +1 -1
  11. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_cache.py +5 -4
  12. tensoracle-0.3.2/tests/test_shape_tensor_outputs.py +282 -0
  13. {tensoracle-0.3.1 → tensoracle-0.3.2}/.gitignore +0 -0
  14. {tensoracle-0.3.1 → tensoracle-0.3.2}/LICENSE +0 -0
  15. {tensoracle-0.3.1 → tensoracle-0.3.2}/cpp/cmake/TensoracleCubinPlugin.cmake +0 -0
  16. {tensoracle-0.3.1 → tensoracle-0.3.2}/cpp/include/tensoracle/cubin_plugin.hpp +0 -0
  17. {tensoracle-0.3.1 → tensoracle-0.3.2}/cpp/src/cubin_plugin.cpp +0 -0
  18. {tensoracle-0.3.1 → tensoracle-0.3.2}/docs/api/bundle.md +0 -0
  19. {tensoracle-0.3.1 → tensoracle-0.3.2}/docs/api/index.md +0 -0
  20. {tensoracle-0.3.1 → tensoracle-0.3.2}/docs/api/producer.md +0 -0
  21. {tensoracle-0.3.1 → tensoracle-0.3.2}/docs/optimization.md +0 -0
  22. {tensoracle-0.3.1 → tensoracle-0.3.2}/docs/tensorrt-sdk.md +0 -0
  23. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/CMakeLists.txt +0 -0
  24. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/README.md +0 -0
  25. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/bundle.cpp +0 -0
  26. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/bundle.json +0 -0
  27. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
  28. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
  29. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
  30. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
  31. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
  32. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
  33. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
  34. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
  35. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
  36. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
  37. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
  38. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
  39. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
  40. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
  41. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
  42. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/example_cubin_plugin/__init__.py +0 -0
  43. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors +0 -0
  44. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/kernels.py +0 -0
  45. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/plugin.cpp +0 -0
  46. {tensoracle-0.3.1 → tensoracle-0.3.2}/examples/cubin_plugin/pyproject.toml +0 -0
  47. {tensoracle-0.3.1 → tensoracle-0.3.2}/mkdocs.yml +0 -0
  48. {tensoracle-0.3.1 → tensoracle-0.3.2}/pyproject.toml +0 -0
  49. {tensoracle-0.3.1 → tensoracle-0.3.2}/src/tensoracle/_comparison.py +0 -0
  50. {tensoracle-0.3.1 → tensoracle-0.3.2}/src/tensoracle/_memory.py +0 -0
  51. {tensoracle-0.3.1 → tensoracle-0.3.2}/src/tensoracle/_plugins.py +0 -0
  52. {tensoracle-0.3.1 → tensoracle-0.3.2}/src/tensoracle/bundle.py +0 -0
  53. {tensoracle-0.3.1 → tensoracle-0.3.2}/src/tensoracle/convert.py +0 -0
  54. {tensoracle-0.3.1 → tensoracle-0.3.2}/src/tensoracle/producer.py +0 -0
  55. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/build_memory_runtime.py +0 -0
  56. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/example_runtime.py +0 -0
  57. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_bundle.py +0 -0
  58. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_comparison.py +0 -0
  59. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_convert.py +0 -0
  60. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_cpp_core.py +0 -0
  61. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_distribution.py +0 -0
  62. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_engine_io.py +0 -0
  63. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_example_plugin.py +0 -0
  64. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_execution.py +0 -0
  65. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_plugins.py +0 -0
  66. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_producer.py +0 -0
  67. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_public_mirror_spec.py +0 -0
  68. {tensoracle-0.3.1 → tensoracle-0.3.2}/tests/test_weight_hash.py +0 -0
  69. {tensoracle-0.3.1 → tensoracle-0.3.2}/uv.lock +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: tensoracle
3
- Version: 0.3.1
3
+ Version: 0.3.2
4
4
  Summary: Build, save, reuse, and execute TensorRT engines from safetensors weights
5
5
  Project-URL: documentation, https://sync-dev-org.github.io/tensoracle/
6
6
  Project-URL: repository, https://github.com/sync-dev-org/tensoracle
@@ -253,6 +253,8 @@ The caller selects the CUDA device and context, orders inputs written on other s
253
253
 
254
254
  Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
255
255
 
256
+ Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutputWarning` during building and when creating an execution unit (`create_execution()` or `EngineExecution(handle)`). The warning names the affected tensors: `create_execution()` cannot handle these engines, and calling it crashes the process. Building continues by default. To stop either operation safely, turn this warning into an exception with `warnings.filterwarnings("error", category=tensoracle.ShapeTensorOutputWarning)` after importing `warnings` and `tensoracle`. Saved engines retain the shape output names outside the cache conditions; older files without these names are reused without an execution warning.
257
+
256
258
  ## Build a cubin plugin
257
259
 
258
260
  Follow [Prepare a TensorRT SDK](https://github.com/sync-dev-org/tensoracle/blob/main/docs/tensorrt-sdk.md) to acquire and verify the headers and native library for a build SDK, including the fixed inputs for SDK 11.2.1.2 and the `TENSORRT_ROOT` setup.
@@ -223,6 +223,8 @@ The caller selects the CUDA device and context, orders inputs written on other s
223
223
 
224
224
  Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
225
225
 
226
+ Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutputWarning` during building and when creating an execution unit (`create_execution()` or `EngineExecution(handle)`). The warning names the affected tensors: `create_execution()` cannot handle these engines, and calling it crashes the process. Building continues by default. To stop either operation safely, turn this warning into an exception with `warnings.filterwarnings("error", category=tensoracle.ShapeTensorOutputWarning)` after importing `warnings` and `tensoracle`. Saved engines retain the shape output names outside the cache conditions; older files without these names are reused without an execution warning.
227
+
226
228
  ## Build a cubin plugin
227
229
 
228
230
  Follow [Prepare a TensorRT SDK](https://github.com/sync-dev-org/tensoracle/blob/main/docs/tensorrt-sdk.md) to acquire and verify the headers and native library for a build SDK, including the fixed inputs for SDK 11.2.1.2 and the `TENSORRT_ROOT` setup.
@@ -155,6 +155,8 @@ The caller selects the CUDA device and context, orders inputs written on other s
155
155
 
156
156
  Tensoracle neither allocates or frees device buffers, copies data, nor calls stream or device synchronization. TensorRT itself may allocate internal resources or act synchronously for some networks and plugins. Execution units reject HOST-location tensors, externally supplied shape tensor values, and data-dependent output sizes at creation. Ordinary networks that compute shapes internally from input dimensions are supported. CUDA graphs, profile switching, and caller-owned TensorRT workspace memory are outside this API.
157
157
 
158
+ Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutputWarning` during building and when creating an execution unit (`create_execution()` or `EngineExecution(handle)`). The warning names the affected tensors: `create_execution()` cannot handle these engines, and calling it crashes the process. Building continues by default. To stop either operation safely, turn this warning into an exception with `warnings.filterwarnings("error", category=tensoracle.ShapeTensorOutputWarning)` after importing `warnings` and `tensoracle`. Saved engines retain the shape output names outside the cache conditions; older files without these names are reused without an execution warning.
159
+
158
160
  `TensorSpec` accepts optional `location=trt.TensorLocation.DEVICE` and `format=trt.TensorFormat.LINEAR`. `open_engine` checks these only when declared, together with tensor names, modes, dtypes, build-time shapes, and profiles. An `EngineIOError` names the tensor, field, observed value, expected value, and saved file path. The execution unit does not accept a second declaration.
159
161
 
160
162
  ### Build a standard graph
@@ -15,6 +15,7 @@ __all__ = [
15
15
  "InputShape",
16
16
  "OutputComparison",
17
17
  "PluginLibraryError",
18
+ "ShapeTensorOutputWarning",
18
19
  "TensorSpec",
19
20
  "build_serialized_engine",
20
21
  "compare_outputs",
@@ -23,7 +24,14 @@ __all__ = [
23
24
 
24
25
  _MODULES = {
25
26
  **dict.fromkeys(
26
- ("BuilderSettings", "EngineBuild", "InputShape", "build_serialized_engine"), "._build"
27
+ (
28
+ "BuilderSettings",
29
+ "EngineBuild",
30
+ "InputShape",
31
+ "ShapeTensorOutputWarning",
32
+ "build_serialized_engine",
33
+ ),
34
+ "._build",
27
35
  ),
28
36
  **dict.fromkeys(
29
37
  (
@@ -1,5 +1,6 @@
1
1
  """Build serialized TensorRT engines from caller-defined network graphs."""
2
2
 
3
+ import warnings
3
4
  from collections.abc import Callable, Mapping, Sequence
4
5
  from dataclasses import dataclass, field
5
6
  from pathlib import Path
@@ -11,6 +12,21 @@ from ._memory import release_host_memory
11
12
  from ._plugins import register_libraries
12
13
 
13
14
 
15
+ class ShapeTensorOutputWarning(UserWarning):
16
+ """Shape tensor outputs cannot be handled by tensoracle execution units."""
17
+
18
+
19
+ def _warn_shape_tensor_outputs(names: tuple[str, ...], *, stacklevel: int) -> None:
20
+ if names:
21
+ warnings.warn(
22
+ f"TensorRT outputs {', '.join(repr(name) for name in names)} are shape tensors; "
23
+ "tensoracle.create_execution() cannot handle this engine and calling it will "
24
+ "crash the process.",
25
+ ShapeTensorOutputWarning,
26
+ stacklevel=stacklevel + 1,
27
+ )
28
+
29
+
14
30
  @dataclass(frozen=True)
15
31
  class InputShape:
16
32
  """The minimum, preferred, and maximum dimensions of one input."""
@@ -76,6 +92,17 @@ def build_serialized_engine(
76
92
  logger: trt.ILogger | None = None,
77
93
  ) -> bytes:
78
94
  """Register ``plugins`` before calling ``graph`` and serializing its engine."""
95
+ plan, _ = _build_serialized_engine(graph, declaration, plugins=plugins, logger=logger)
96
+ return plan
97
+
98
+
99
+ def _build_serialized_engine(
100
+ graph: Callable[[trt.INetworkDefinition], None],
101
+ declaration: EngineBuild,
102
+ *,
103
+ plugins: Sequence[Path] = (),
104
+ logger: trt.ILogger | None = None,
105
+ ) -> tuple[bytes, tuple[str, ...]]:
79
106
  if declaration.precision not in ("fp32", "fp16"):
80
107
  raise ValueError(f"unsupported precision: {declaration.precision}")
81
108
  if declaration.settings.tf32:
@@ -93,6 +120,13 @@ def build_serialized_engine(
93
120
  builder = trt.Builder(logger or trt.Logger(trt.Logger.ERROR))
94
121
  network = builder.create_network(0)
95
122
  graph(network)
123
+ shape_tensor_outputs = []
124
+ for index in range(network.num_outputs):
125
+ tensor = network.get_output(index)
126
+ if tensor.is_shape_tensor:
127
+ shape_tensor_outputs.append(tensor.name)
128
+ output_names = tuple(shape_tensor_outputs)
129
+ _warn_shape_tensor_outputs(output_names, stacklevel=3)
96
130
  config = builder.create_builder_config()
97
131
  config.clear_flag(trt.BuilderFlag.TF32)
98
132
  config.set_memory_pool_limit(
@@ -129,7 +163,7 @@ def build_serialized_engine(
129
163
  serialized = builder.build_serialized_network(network, config)
130
164
  if serialized is None:
131
165
  raise RuntimeError("TensorRT failed to build the serialized engine")
132
- return bytes(serialized)
166
+ return bytes(serialized), output_names
133
167
  finally:
134
168
  # Tensor references also keep the network alive, including in exception frames.
135
169
  inputs.clear()
@@ -9,12 +9,13 @@ import platform
9
9
  import tempfile
10
10
  from collections.abc import Callable, Mapping, Sequence
11
11
  from dataclasses import asdict, dataclass
12
+ from dataclasses import field as dataclass_field
12
13
  from pathlib import Path
13
14
  from typing import TYPE_CHECKING
14
15
 
15
16
  import tensorrt as trt
16
17
 
17
- from ._build import EngineBuild, build_serialized_engine
18
+ from ._build import EngineBuild, _build_serialized_engine
18
19
  from ._plugins import register_libraries
19
20
 
20
21
  if TYPE_CHECKING:
@@ -49,6 +50,7 @@ class EngineHandle:
49
50
  engine: trt.ICudaEngine
50
51
  runtime: trt.Runtime
51
52
  path: Path
53
+ _shape_tensor_outputs: tuple[str, ...] = dataclass_field(default=(), repr=False)
52
54
 
53
55
  def create_execution(self) -> "EngineExecution":
54
56
  """Create an independent execution context with explicit close semantics."""
@@ -103,16 +105,29 @@ def open_engine(
103
105
  try:
104
106
  fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
105
107
  except FileExistsError:
106
- plan = _read_plan(path, conditions)
107
- return _open_plan(plan, path, declaration, io, actual_logger)
108
+ plan, shape_tensor_outputs = _read_plan(path, conditions)
109
+ return _open_plan(
110
+ plan,
111
+ path,
112
+ declaration,
113
+ io,
114
+ actual_logger,
115
+ shape_tensor_outputs=shape_tensor_outputs,
116
+ )
108
117
  else:
109
118
  os.close(fd)
110
119
 
111
120
  temporary: Path | None = None
112
121
  try:
113
- plan = build_serialized_engine(graph, declaration, plugins=plugins, logger=actual_logger)
114
- handle = _open_plan(plan, path, declaration, io, actual_logger)
122
+ plan, shape_tensor_outputs = _build_serialized_engine(
123
+ graph, declaration, plugins=plugins, logger=actual_logger
124
+ )
125
+ handle = _open_plan(
126
+ plan, path, declaration, io, actual_logger, shape_tensor_outputs=shape_tensor_outputs
127
+ )
115
128
  header = {"conditions": conditions, "plan_sha256": hashlib.sha256(plan).hexdigest()}
129
+ if shape_tensor_outputs:
130
+ header["shape_tensor_outputs"] = list(shape_tensor_outputs)
116
131
  header_bytes = _json_bytes(header)
117
132
  fd, filename = tempfile.mkstemp(prefix=f".{digest}.", dir=path.parent)
118
133
  temporary = Path(filename)
@@ -279,7 +294,7 @@ def _conditions(
279
294
  }
280
295
 
281
296
 
282
- def _read_plan(path: Path, conditions: dict) -> memoryview:
297
+ def _read_plan(path: Path, conditions: dict) -> tuple[memoryview, tuple[str, ...]]:
283
298
  data = path.read_bytes()
284
299
  offset = len(_MAGIC)
285
300
  if not data.startswith(_MAGIC) or len(data) < offset + 8:
@@ -294,12 +309,17 @@ def _read_plan(path: Path, conditions: dict) -> memoryview:
294
309
  raise CorruptEngineError(f"invalid TensorRT engine header: {path}") from error
295
310
  if not isinstance(header, dict) or not isinstance(header.get("conditions"), dict):
296
311
  raise CorruptEngineError(f"invalid TensorRT engine header: {path}")
312
+ shape_tensor_outputs = header.get("shape_tensor_outputs", [])
313
+ if not isinstance(shape_tensor_outputs, list) or any(
314
+ not isinstance(name, str) for name in shape_tensor_outputs
315
+ ):
316
+ raise CorruptEngineError(f"invalid TensorRT engine header: {path}")
297
317
  if header["conditions"] != conditions:
298
318
  raise EngineConditionError(f"TensorRT engine conditions differ: {path}")
299
319
  plan = memoryview(data)[header_end:]
300
320
  if header.get("plan_sha256") != hashlib.sha256(plan).hexdigest():
301
321
  raise CorruptEngineError(f"TensorRT engine checksum differs: {path}")
302
- return plan
322
+ return plan, tuple(shape_tensor_outputs)
303
323
 
304
324
 
305
325
  def _open_plan(
@@ -308,6 +328,8 @@ def _open_plan(
308
328
  declaration: EngineBuild,
309
329
  io: EngineIO,
310
330
  logger: trt.ILogger,
331
+ *,
332
+ shape_tensor_outputs: tuple[str, ...] = (),
311
333
  ) -> EngineHandle:
312
334
  runtime = trt.Runtime(logger)
313
335
  engine = runtime.deserialize_cuda_engine(plan)
@@ -349,7 +371,7 @@ def _open_plan(
349
371
  actual_range = (io.inputs[name].shape,) * 3
350
372
  if actual_range != expected_range:
351
373
  _io_mismatch(path, name, "profile range", actual_range, expected_range)
352
- return EngineHandle(engine, runtime, path)
374
+ return EngineHandle(engine, runtime, path, _shape_tensor_outputs=shape_tensor_outputs)
353
375
 
354
376
 
355
377
  def _io_mismatch(path: Path, name: str, field: str, observed: object, expected: object) -> None:
@@ -7,6 +7,8 @@ from typing import TYPE_CHECKING, Self
7
7
 
8
8
  import tensorrt as trt
9
9
 
10
+ from ._build import _warn_shape_tensor_outputs
11
+
10
12
  if TYPE_CHECKING:
11
13
  from ._cache import EngineHandle
12
14
 
@@ -51,6 +53,7 @@ class EngineExecution:
51
53
  self._context = engine.create_execution_context()
52
54
  if self._context is None:
53
55
  raise RuntimeError("TensorRT failed to create an execution context")
56
+ _warn_shape_tensor_outputs(handle._shape_tensor_outputs, stacklevel=2)
54
57
  # Probe output dimensions before accepting the engine. A ranged engine
55
58
  # still requires an explicit caller setting after this capability check.
56
59
  self._set_shapes(initial_shapes, creating=True)
@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
18
18
  commit_id: str | None
19
19
  __commit_id__: str | None
20
20
 
21
- __version__ = version = '0.3.1'
22
- __version_tuple__ = version_tuple = (0, 3, 1)
21
+ __version__ = version = '0.3.2'
22
+ __version_tuple__ = version_tuple = (0, 3, 2)
23
23
 
24
24
  __commit_id__ = commit_id = None
@@ -202,6 +202,7 @@ def test_tf32_is_disabled_and_other_builder_settings_are_applied(monkeypatch):
202
202
 
203
203
  class Network:
204
204
  num_inputs = 1
205
+ num_outputs = 0
205
206
 
206
207
  def get_input(self, index):
207
208
  return type("Input", (), {"name": "features", "shape": (1, 3, 2)})()
@@ -64,7 +64,7 @@ def test_unavailable_host_memory_release_preserves_the_build_result(monkeypatch,
64
64
  return b"serialized plan"
65
65
 
66
66
  builder = SimpleNamespace(
67
- create_network=lambda flags: SimpleNamespace(num_inputs=0),
67
+ create_network=lambda flags: SimpleNamespace(num_inputs=0, num_outputs=0),
68
68
  create_builder_config=lambda: SimpleNamespace(
69
69
  clear_flag=lambda flag: None,
70
70
  set_memory_pool_limit=lambda pool, limit: None,
@@ -72,7 +72,7 @@ def test_read_large_saved_plan_keeps_other_threads_running(tmp_path):
72
72
 
73
73
  measurements = []
74
74
  for _ in range(3):
75
- loaded, measured = _heartbeat_during(lambda: cache._read_plan(path, conditions))
75
+ (loaded, _), measured = _heartbeat_during(lambda: cache._read_plan(path, conditions))
76
76
  assert len(loaded) == 512 * 1024 * 1024
77
77
  assert hashlib.sha256(loaded).hexdigest() == digest
78
78
  measurements.append(measured)
@@ -97,7 +97,7 @@ def test_save_large_built_plan_keeps_other_threads_running(tmp_path, monkeypatch
97
97
  monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
98
98
  monkeypatch.setattr(cache, "_gpu_name", lambda: "GPU A")
99
99
  builder = SimpleNamespace(
100
- create_network=lambda flags: SimpleNamespace(num_inputs=0),
100
+ create_network=lambda flags: SimpleNamespace(num_inputs=0, num_outputs=0),
101
101
  create_builder_config=lambda: SimpleNamespace(
102
102
  clear_flag=lambda flag: None,
103
103
  set_memory_pool_limit=lambda pool, limit: None,
@@ -131,7 +131,7 @@ def test_save_large_built_plan_keeps_other_threads_running(tmp_path, monkeypatch
131
131
  for _ in range(3):
132
132
  handle, measured = _heartbeat_during(save)
133
133
  conditions = cache._conditions("heartbeat", "1", (weights,), (), declaration)
134
- loaded = cache._read_plan(handle.path, conditions)
134
+ loaded, _ = cache._read_plan(handle.path, conditions)
135
135
  assert len(loaded) == len(plan)
136
136
  assert hashlib.sha256(loaded).digest() == hashlib.sha256(plan).digest()
137
137
  measurements.append(measured)
@@ -422,7 +422,7 @@ def test_build_failure_removes_placeholder_and_partial_files(tmp_path, monkeypat
422
422
  monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
423
423
  monkeypatch.setattr(cache, "_gpu_name", lambda: "GPU A")
424
424
  monkeypatch.setattr(
425
- cache, "build_serialized_engine", lambda graph, declaration, plugins, logger: graph(None)
425
+ cache, "_build_serialized_engine", lambda graph, declaration, plugins, logger: graph(None)
426
426
  )
427
427
  weights = tmp_path / "weights.safetensors"
428
428
  _weights(weights)
@@ -537,6 +537,7 @@ def test_each_declared_condition_selects_a_distinct_cache_entry(tmp_path, monkey
537
537
 
538
538
  class FakeNetwork:
539
539
  num_inputs = 1
540
+ num_outputs = 0
540
541
 
541
542
  def get_input(self, index):
542
543
  return type("Input", (), {"name": "features", "shape": state["shape"]})()
@@ -0,0 +1,282 @@
1
+ """Warn about shape tensor outputs before unsafe TensorRT shape inference."""
2
+
3
+ import hashlib
4
+ import json
5
+ import subprocess
6
+ import sys
7
+ import warnings
8
+
9
+ import numpy as np
10
+ import pytest
11
+ import tensorrt as trt
12
+ from safetensors.numpy import save_file
13
+
14
+ import tensoracle
15
+ from test_execution import Context, Engine, _handle
16
+
17
+
18
+ def _graph(network, *, output_kind="shape"):
19
+ features = network.add_input("features", trt.float32, (3, 2))
20
+ result = network.add_identity(features).get_output(0)
21
+ result.name = "result"
22
+ network.mark_output(result)
23
+ if output_kind != "ordinary":
24
+ dimensions = network.add_shape(features).get_output(0)
25
+ dimensions.name = "dimensions"
26
+ if output_kind == "shape":
27
+ assert network.mark_output_for_shapes(dimensions)
28
+ else:
29
+ network.mark_output(dimensions)
30
+
31
+
32
+ def _declaration():
33
+ return tensoracle.EngineBuild(
34
+ "fp32",
35
+ {"features": tensoracle.InputShape.fixed((3, 2))},
36
+ tensoracle.BuilderSettings(optimization_level=0),
37
+ )
38
+
39
+
40
+ def _io(output_kind="shape"):
41
+ outputs = {"result": tensoracle.TensorSpec(trt.float32, (3, 2))}
42
+ if output_kind != "ordinary":
43
+ outputs["dimensions"] = tensoracle.TensorSpec(trt.int64, (2,))
44
+ return tensoracle.EngineIO({"features": tensoracle.TensorSpec(trt.float32, (3, 2))}, outputs)
45
+
46
+
47
+ @pytest.fixture
48
+ def weights(tmp_path, monkeypatch):
49
+ monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
50
+ path = tmp_path / "weights.safetensors"
51
+ save_file({"unused": np.zeros(1, dtype=np.float32)}, path)
52
+ return path
53
+
54
+
55
+ def _open(weights, *, output_kind="shape", graph=None, force_rebuild=False):
56
+ return tensoracle.open_engine(
57
+ graph if graph is not None else lambda network: _graph(network, output_kind=output_kind),
58
+ _declaration(),
59
+ name="shape-output",
60
+ graph_version="1",
61
+ weights=(weights,),
62
+ io=_io(output_kind),
63
+ force_rebuild=force_rebuild,
64
+ )
65
+
66
+
67
+ def _read_file(path):
68
+ data = path.read_bytes()
69
+ offset = len(b"TENSORACLE\x01")
70
+ header_end = offset + 8 + int.from_bytes(data[offset : offset + 8], "big")
71
+ return json.loads(data[offset + 8 : header_end]), data[header_end:]
72
+
73
+
74
+ def _write_file(path, header, plan):
75
+ encoded = json.dumps(header, sort_keys=True, separators=(",", ":")).encode()
76
+ path.write_bytes(b"TENSORACLE\x01" + len(encoded).to_bytes(8, "big") + encoded + plan)
77
+
78
+
79
+ def _assert_warning(message):
80
+ assert "dimensions" in message
81
+ assert "shape tensor" in message
82
+ assert "create_execution()" in message
83
+ assert "cannot" in message
84
+ assert "process" in message and "crash" in message
85
+
86
+
87
+ @pytest.mark.gpu
88
+ @pytest.mark.req("REQ-TENSORACLE-018")
89
+ @pytest.mark.req("REQ-TENSORACLE-002")
90
+ def test_shape_output_build_warns_once_and_returns_plan():
91
+ """A shape tensor output warns about the execution crash while returning plan bytes."""
92
+ with warnings.catch_warnings(record=True) as caught:
93
+ warnings.simplefilter("always")
94
+ plan = tensoracle.build_serialized_engine(_graph, _declaration())
95
+ assert isinstance(plan, bytes) and plan
96
+ assert len(caught) == 1
97
+ assert caught[0].category is tensoracle.ShapeTensorOutputWarning
98
+ assert issubclass(caught[0].category, UserWarning)
99
+ assert "ShapeTensorOutputWarning" in tensoracle.__all__
100
+ _assert_warning(str(caught[0].message))
101
+
102
+
103
+ @pytest.mark.gpu
104
+ @pytest.mark.req("REQ-TENSORACLE-007")
105
+ @pytest.mark.req("REQ-TENSORACLE-018")
106
+ @pytest.mark.req("REQ-TENSORACLE-002")
107
+ def test_shape_output_cache_records_names_without_changing_conditions(weights):
108
+ """Shape output names are saved outside conditions and do not change the cache identity."""
109
+ with warnings.catch_warnings(record=True) as caught:
110
+ warnings.simplefilter("always")
111
+ handle = _open(weights)
112
+ assert len(caught) == 1
113
+ assert caught[0].category is tensoracle.ShapeTensorOutputWarning
114
+ _assert_warning(str(caught[0].message))
115
+ header, plan = _read_file(handle.path)
116
+ assert header["shape_tensor_outputs"] == ["dimensions"]
117
+ assert "shape_tensor_outputs" not in header["conditions"]
118
+ assert header["plan_sha256"] == hashlib.sha256(plan).hexdigest()
119
+ digest = hashlib.sha256(
120
+ json.dumps(header["conditions"], sort_keys=True, separators=(",", ":")).encode()
121
+ ).hexdigest()
122
+ assert handle.path.name == f"{digest}.engine"
123
+ ordinary = _open(weights, output_kind="ordinary", force_rebuild=True)
124
+ ordinary_header, _ = _read_file(ordinary.path)
125
+ assert ordinary.path == handle.path
126
+ assert ordinary_header["conditions"] == header["conditions"]
127
+
128
+
129
+ @pytest.mark.gpu
130
+ @pytest.mark.req("REQ-TENSORACLE-018")
131
+ @pytest.mark.req("REQ-TENSORACLE-002")
132
+ @pytest.mark.parametrize("entry", ["create_execution", "EngineExecution"])
133
+ def test_saved_shape_output_warns_before_inference_in_another_process(weights, entry):
134
+ """A later process can turn the execution warning into an exception before unsafe inference."""
135
+ category = tensoracle.ShapeTensorOutputWarning
136
+ with warnings.catch_warnings():
137
+ warnings.simplefilter("ignore", category)
138
+ handle = _open(weights)
139
+ child = """
140
+ import json
141
+ import sys
142
+ import warnings
143
+ from pathlib import Path
144
+ import tensorrt as trt
145
+ import tensoracle
146
+
147
+ def forbidden(network):
148
+ raise AssertionError('saved engine was rebuilt')
149
+
150
+ handle = tensoracle.open_engine(
151
+ forbidden,
152
+ tensoracle.EngineBuild('fp32', {'features': tensoracle.InputShape.fixed((3, 2))},
153
+ tensoracle.BuilderSettings(optimization_level=0)),
154
+ name='shape-output', graph_version='1', weights=(Path(sys.argv[1]),),
155
+ io=tensoracle.EngineIO({'features': tensoracle.TensorSpec(trt.float32, (3, 2))},
156
+ {'result': tensoracle.TensorSpec(trt.float32, (3, 2)),
157
+ 'dimensions': tensoracle.TensorSpec(trt.int64, (2,))}),
158
+ )
159
+ warnings.filterwarnings('error', category=tensoracle.ShapeTensorOutputWarning)
160
+ try:
161
+ if sys.argv[2] == 'create_execution':
162
+ handle.create_execution()
163
+ else:
164
+ tensoracle.EngineExecution(handle)
165
+ except tensoracle.ShapeTensorOutputWarning as warning:
166
+ print(json.dumps({'path': str(handle.path), 'message': str(warning)}))
167
+ else:
168
+ raise AssertionError('execution did not raise the warning')
169
+ """
170
+ result = subprocess.run(
171
+ [sys.executable, "-c", child, str(weights), entry],
172
+ capture_output=True,
173
+ text=True,
174
+ check=False,
175
+ timeout=60,
176
+ )
177
+ assert result.returncode == 0, result.stdout + result.stderr
178
+ observed = json.loads(result.stdout.splitlines()[-1])
179
+ assert observed["path"] == str(handle.path)
180
+ _assert_warning(observed["message"])
181
+
182
+
183
+ @pytest.mark.gpu
184
+ @pytest.mark.req("REQ-TENSORACLE-018")
185
+ @pytest.mark.req("REQ-TENSORACLE-002")
186
+ @pytest.mark.req("REQ-TENSORACLE-007")
187
+ @pytest.mark.parametrize("entry", ["build_serialized_engine", "open_engine"])
188
+ def test_shape_output_warning_as_error_stops_build_without_cache_files(weights, entry):
189
+ """Turning the build warning into an error leaves no engine or temporary engine file."""
190
+ category = tensoracle.ShapeTensorOutputWarning
191
+ with warnings.catch_warnings():
192
+ warnings.filterwarnings("error", category=category)
193
+ with pytest.raises(category) as caught:
194
+ if entry == "build_serialized_engine":
195
+ tensoracle.build_serialized_engine(_graph, _declaration())
196
+ else:
197
+ _open(weights)
198
+ _assert_warning(str(caught.value))
199
+ cache_directory = weights.parent / "cache" / "tensoracle"
200
+ assert all(path.name.startswith("weight-") for path in cache_directory.glob("*"))
201
+
202
+
203
+ @pytest.mark.gpu
204
+ @pytest.mark.req("REQ-TENSORACLE-007")
205
+ @pytest.mark.req("REQ-TENSORACLE-018")
206
+ def test_saved_engine_without_shape_output_metadata_is_reused_without_warning(weights):
207
+ """An older header loads without rebuilding or warning when creating an execution unit."""
208
+ first = _open(weights, output_kind="ordinary")
209
+ header, plan = _read_file(first.path)
210
+ header.pop("shape_tensor_outputs", None)
211
+ _write_file(first.path, header, plan)
212
+ legacy_bytes = first.path.read_bytes()
213
+
214
+ def forbidden(network):
215
+ pytest.fail("saved engine was rebuilt")
216
+
217
+ with warnings.catch_warnings(record=True) as caught:
218
+ warnings.simplefilter("always")
219
+ reopened = _open(weights, output_kind="ordinary", graph=forbidden)
220
+ with reopened.create_execution() as execution:
221
+ assert execution.get_tensor_shape("result") == (3, 2)
222
+ assert caught == []
223
+ assert reopened.path == first.path
224
+ assert reopened.path.read_bytes() == legacy_bytes
225
+
226
+
227
+ @pytest.mark.gpu
228
+ @pytest.mark.req("REQ-TENSORACLE-018")
229
+ @pytest.mark.parametrize("output_kind", ["ordinary", "shape-layer"])
230
+ def test_ordinary_outputs_including_shape_layer_outputs_never_warn(weights, output_kind):
231
+ """Ordinary marked outputs stay supported even when their values come from a shape layer."""
232
+ with warnings.catch_warnings(record=True) as caught:
233
+ warnings.simplefilter("always")
234
+ assert tensoracle.build_serialized_engine(
235
+ lambda network: _graph(network, output_kind=output_kind), _declaration()
236
+ )
237
+ handle = _open(weights, output_kind=output_kind)
238
+ reopened = _open(weights, output_kind=output_kind)
239
+ for checked in (handle, reopened):
240
+ with checked.create_execution() as execution:
241
+ assert execution.get_tensor_shape("result") == (3, 2)
242
+ if output_kind == "shape-layer":
243
+ assert execution.get_tensor_shape("dimensions") == (2,)
244
+ assert caught == []
245
+
246
+
247
+ @pytest.mark.gpu
248
+ @pytest.mark.req("REQ-TENSORACLE-008")
249
+ @pytest.mark.req("REQ-TENSORACLE-002")
250
+ @pytest.mark.parametrize("metadata", [None, "dimensions", {}, [1], ["dimensions", False]])
251
+ def test_invalid_shape_output_metadata_is_a_corrupt_engine(weights, metadata):
252
+ """A present shape output record must be a list of tensor names to load a saved engine."""
253
+ handle = _open(weights, output_kind="ordinary")
254
+ header, plan = _read_file(handle.path)
255
+ header["shape_tensor_outputs"] = metadata
256
+ _write_file(handle.path, header, plan)
257
+ with pytest.raises(tensoracle.CorruptEngineError, match="header"):
258
+ _open(weights, output_kind="ordinary")
259
+
260
+
261
+ @pytest.mark.req("REQ-TENSORACLE-018")
262
+ @pytest.mark.req("REQ-TENSORACLE-002")
263
+ def test_execution_warning_as_error_releases_context_before_inference(monkeypatch):
264
+ """A warning raised as an error releases the new context without probing output shapes."""
265
+ category = tensoracle.ShapeTensorOutputWarning
266
+ engine = Engine(fixed=True)
267
+ monkeypatch.setattr(
268
+ Context, "infer_shapes", lambda context: pytest.fail("inference was called")
269
+ )
270
+ handle = _handle(engine)
271
+ object.__setattr__(handle, "_shape_tensor_outputs", ("dimensions",))
272
+ execution = tensoracle.EngineExecution.__new__(tensoracle.EngineExecution)
273
+ with warnings.catch_warnings():
274
+ warnings.filterwarnings("error", category=category)
275
+ with pytest.raises(category) as caught:
276
+ execution.__init__(handle)
277
+ _assert_warning(str(caught.value))
278
+ assert len(engine.contexts) == 1
279
+ assert engine.latest is None
280
+ with pytest.raises(RuntimeError, match="closed"):
281
+ execution.get_tensor_shape("result")
282
+ assert execution._handle is None
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes