tensoracle 1.4.0__tar.gz → 1.5.0__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.
- {tensoracle-1.4.0 → tensoracle-1.5.0}/PKG-INFO +28 -14
- {tensoracle-1.4.0 → tensoracle-1.5.0}/README.md +27 -13
- tensoracle-1.5.0/docs/cpp-python-repositories.md +65 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/docs/index.md +26 -2
- {tensoracle-1.4.0 → tensoracle-1.5.0}/mkdocs.yml +1 -1
- {tensoracle-1.4.0 → tensoracle-1.5.0}/pyproject.toml +4 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/src/tensoracle/_cache.py +29 -70
- tensoracle-1.5.0/src/tensoracle/_cache_format.py +78 -0
- tensoracle-1.5.0/src/tensoracle/_cli.py +249 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/src/tensoracle/_version.py +2 -2
- tensoracle-1.5.0/tests/test_cache_cli.py +368 -0
- tensoracle-1.5.0/tests/test_cache_management.py +157 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_distribution.py +46 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_precision.py +2 -6
- tensoracle-1.4.0/docs/cpp-python-repositories.md +0 -199
- {tensoracle-1.4.0 → tensoracle-1.5.0}/.gitignore +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/LICENSE +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/build_backend.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/cpp/CMakeLists.txt +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/cpp/cmake/TensoracleCubinPlugin.cmake +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/cpp/cmake/TensoracleCuda.cmake +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/cpp/include/tensoracle/cubin_plugin.hpp +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/cpp/src/cubin_plugin.cpp +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/cpp/tests/core.cpp +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/docs/api/bundle.md +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/docs/api/index.md +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/docs/api/producer.md +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/docs/optimization.md +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/docs/precision.md +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/docs/tensorrt-sdk.md +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/CMakeLists.txt +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/README.md +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/bundle.cpp +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/bundle.json +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/example_cubin_plugin/__init__.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/kernels.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/plugin.cpp +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/examples/cubin_plugin/pyproject.toml +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/src/tensoracle/__init__.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/src/tensoracle/_build.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/src/tensoracle/_comparison.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/src/tensoracle/_execution.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/src/tensoracle/_memory.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/src/tensoracle/_plugins.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/src/tensoracle/bundle.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/src/tensoracle/convert.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/src/tensoracle/producer.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/build_memory_runtime.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/cubin_io/check.cpp +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/cubin_io/kernels.cu +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/cubin_io/plugin.cpp +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/cubins/ptxas-12.9.86/sm120.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/cubins/ptxas-12.9.86/sm86.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/cubins/ptxas-12.9.86/sm89.cubin +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/example_runtime.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/native/CMakeLists.txt +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/native/export.h +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/native_build.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/plugin_io_runtime.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/precision_runtime.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_build.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_build_memory.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_bundle.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_cache.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_cmake_cuda.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_cmake_plugin.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_comparison.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_convert.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_cpp_core.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_cuda_graphs.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_engine_io.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_example_plugin.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_execution.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_gpu_cache.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_plugin_io.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_plugins.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_producer.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_shape_tensor_outputs.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/tests/test_weight_hash.py +0 -0
- {tensoracle-1.4.0 → tensoracle-1.5.0}/uv.lock +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: tensoracle
|
|
3
|
-
Version: 1.
|
|
3
|
+
Version: 1.5.0
|
|
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
|
|
@@ -64,7 +64,7 @@ The supported targets are Linux x86_64 (including WSL2) and Windows x86_64, Tens
|
|
|
64
64
|
|
|
65
65
|
Follow the [four-stage optimization workflow](https://github.com/sync-dev-org/tensoracle/blob/main/docs/optimization.md) to establish output quality, evaluate lower precision and plugins, and measure speed.
|
|
66
66
|
|
|
67
|
-
Follow [Build
|
|
67
|
+
Follow [Build plugins with tensoracle's C++ sources](https://sync-dev-org.github.io/tensoracle/cpp-python-repositories/) to locate the sources in the wheel and use the CMake helper in your plugin build.
|
|
68
68
|
|
|
69
69
|
## Convert a PyTorch checkpoint
|
|
70
70
|
|
|
@@ -156,6 +156,30 @@ For each weight file, `open_engine` stores its SHA-256 alongside the saved engin
|
|
|
156
156
|
|
|
157
157
|
`CorruptEngineError` reports incomplete or damaged saved files and plans that TensorRT cannot deserialize. When deserialization fails, its message asks the caller to check the required plugin libraries in `plugins=`. `EngineConditionError` reports a saved header with different build conditions. These types and `EngineIOError` derive from `EngineCacheError`. The caller chooses whether to rebuild after an error.
|
|
158
158
|
|
|
159
|
+
## Manage saved engines
|
|
160
|
+
|
|
161
|
+
Inspect saved engines and preview or perform cleanup with the installed command:
|
|
162
|
+
|
|
163
|
+
```bash
|
|
164
|
+
tensoracle cache list
|
|
165
|
+
tensoracle cache prune --unused-for 30 --dry-run
|
|
166
|
+
tensoracle cache prune --unused-for 30
|
|
167
|
+
tensoracle cache prune --max-size 10GiB --dry-run
|
|
168
|
+
tensoracle cache prune --max-size 10GiB
|
|
169
|
+
```
|
|
170
|
+
|
|
171
|
+
The commands use the same `XDG_CACHE_HOME` / home `.cache` directory as `open_engine`. They read only filesystem metadata, engine headers, and saved weight digest records. They do not read or checksum the plan body, load the original weights, deserialize engines, import TensorRT, or query a GPU. Listing recognizes the header format; it does not establish that an engine can be loaded.
|
|
172
|
+
|
|
173
|
+
`cache list` prints one JSON object with `directory`, `entries`, and `total_size_bytes`. Each entry contains the file name, size in bytes, recorded time in UTC (`recorded_at`) and exact nanoseconds (`recorded_at_ns`), and the complete header conditions. This distinguishes engines with the same name and graph version but different weights or plugins. `weight_paths` associates each weight digest with paths from valid, checksum-verified `weight-*.json` records whose SHA-256 matches that digest; an empty `paths` array means unknown. These paths are references and do not establish which application uses an engine or whether the weight file still exists. Empty files have status `building or leftover`; unreadable or invalid headers have status `unrecognized` and a reason.
|
|
174
|
+
|
|
175
|
+
An engine's `.engine` modification time records its initial save or the last successful `open_engine` cache read. Failed loads and these management commands leave that time unchanged. A timestamp write failure, including a read-only cache, does not prevent a successful load, and a disappeared file is never recreated to record its time. The record describes the time of `open_engine`, not whether inference continues afterward. Files saved before timestamp recording was introduced, and files whose timestamp could not be written, may not reflect the actual last successful read.
|
|
176
|
+
|
|
177
|
+
`cache prune` requires exactly one of `--unused-for DAYS` and `--max-size SIZE`. Days are finite nonnegative numbers, may be fractional, and mean 24 hours each; only timestamps strictly earlier than the cutoff are selected. Sizes accept nonnegative bytes or the case-insensitive units `B`, `KiB`, `MiB`, `GiB`, and `TiB` (powers of 1024). Fractions such as `1.5GiB` are allowed and fractional bytes are rounded down. Size pruning selects the oldest recorded times first, breaking ties by file name, until the total `.engine` bytes would fit. Empty and unrecognized files count toward that total but are always excluded from deletion. Weight digest records are kept.
|
|
178
|
+
|
|
179
|
+
Before deletion, pruning prints a JSON `plan` containing the selected engines, excluded files, and a notice that removed engines will be rebuilt the next time they are used. It then prints a JSON `result` separating removed, excluded, failed, and already missing files, the bytes actually removed, and the remaining total. `--dry-run` shows the same targets without changing file contents or recorded times; its result also gives the projected total. Exit status is 0 for success, 1 for filesystem failures, failed deletions, or an unattainable size limit, and 2 for invalid usage. Files that disappear before removal are reported in `missing` and do not count as bytes removed. Observed file changes are reported as failed removals with a reason.
|
|
180
|
+
|
|
181
|
+
Run pruning during maintenance when no engine is being built or opened in the same cache directory. Cleanup provides no lock or guarantee of protection against concurrent builds and reads. Inspect the selection with `--dry-run` before performing cleanup. `tensoracle cache prune --help` describes the options.
|
|
182
|
+
|
|
159
183
|
## Use a saved engine for inference
|
|
160
184
|
|
|
161
185
|
An inference application can call `open_engine` with the same build declaration and files used to save the engine. For example, assume the model package exports `graph`, `declaration`, and `io` matching the example above:
|
|
@@ -336,23 +360,13 @@ Install the producer extra when you need to compile cubins from your own kernels
|
|
|
336
360
|
pip install 'tensoracle[producer]'
|
|
337
361
|
```
|
|
338
362
|
|
|
339
|
-
|
|
340
|
-
|
|
341
|
-
```text
|
|
342
|
-
cpp/
|
|
343
|
-
CMakeLists.txt
|
|
344
|
-
python/
|
|
345
|
-
pyproject.toml
|
|
346
|
-
src/
|
|
347
|
-
```
|
|
348
|
-
|
|
349
|
-
Keep C++ sources in `cpp/`, with a CMake project that builds and runs its tests through CTest without Python. Keep the Python project, its src-layout package, and its uv environment in `python/`; run the project's uv commands from that directory. Tensoracle itself follows this layout, and the installed C++ helper paths described below remain inside its Python package.
|
|
363
|
+
Use the [C++ sources and CMake helper supplied in the tensoracle wheel](https://sync-dev-org.github.io/tensoracle/cpp-python-repositories/) to build your native plugin.
|
|
350
364
|
|
|
351
365
|
`tensoracle.producer.compile_bundle` uses the extra's **Triton 3.8.0** with **ptxas >=13.2,<14** and **cuobjdump** from the PyPI `nvidia-cuda-nvcc` and `nvidia-cuda-cuobjdump` wheels. The extra pins those tools, `nvidia-cuda-crt`, and `nvidia-nvvm` to the 13.2 series. It does not read `CUDA_HOME`; explicit `ptxas` and `cuobjdump` paths override the wheel defaults independently. A missing tool identifies its package and searched path and explains how to install the extra or pass a path. The normal import, engine build, and engine reuse paths do not import the producer or Triton. Cubin production does not need a GPU or import torch; run it without importing torch first. The caller supplies Triton functions and `KernelSpec` records with their public argument types, compile time constants, and any divisibility assumptions. The producer writes sm86, sm89, and sm120 cubins, `bundle.json`, and a C++ translation unit that embeds exactly those cubins. The generated C++ entry records carry each public argument's order, type, count, and divisibility rule into the common launcher. It rejects an unexpected parameter ABI or nonzero scratch requirement. The included [example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md#rebuild-the-cubins) has the command to rebuild its cubins from a repository clone.
|
|
352
366
|
|
|
353
367
|
Neither cubin production nor native plugin building requires an OS-installed CUDA Toolkit. For the native build, install PyPI `nvidia-cuda-nvcc==13.2.*` and `nvidia-cuda-runtime==13.2.*`, then pass their `site-packages/nvidia/cu13` directory as `CUDAToolkit_ROOT` (a CMake variable or environment variable). The helper uses its `cuda.h` and links the NVIDIA driver's `libcuda.so` on Linux, including `/usr/lib/wsl/lib/libcuda.so` on WSL, or the wheel's `lib/x64/cuda.lib` on Windows. TensorRT headers still come from the separately prepared SDK. The [SDK guide](https://github.com/sync-dev-org/tensoracle/blob/main/docs/tensorrt-sdk.md#5-pass-the-verified-root-to-the-build-and-tests) gives the setup commands. An explicit OS Toolkit root and the existing default locations remain supported for native builds.
|
|
354
368
|
|
|
355
|
-
The
|
|
369
|
+
The caller implements the graph's meaning, field meanings, weight placement, launch sequence, and variant selection. The common C++ implementation owns typed fields in their original order, handles `IPluginV3` lifecycle methods, checks cubin digests and resources, loads one module per CUDA context, checks the generated argument count and specialization rules, appends the two supported zero scratch pointers, and launches kernels. The caller selects a variant using `Kernel::accepts` and derives the count passed to `Kernel::launch` from its argument array. The common implementation releases modules and device allocations in their owning contexts. The library must export TensorRT's `getCreators`; a static registration constructor is not supported.
|
|
356
370
|
|
|
357
371
|
To record the CUDA inputs used by the plugin build, call `tensoracle_find_cuda` from `TensoracleCuda.cmake` (also included by `TensoracleCubinPlugin.cmake`):
|
|
358
372
|
|
|
@@ -31,7 +31,7 @@ The supported targets are Linux x86_64 (including WSL2) and Windows x86_64, Tens
|
|
|
31
31
|
|
|
32
32
|
Follow the [four-stage optimization workflow](https://github.com/sync-dev-org/tensoracle/blob/main/docs/optimization.md) to establish output quality, evaluate lower precision and plugins, and measure speed.
|
|
33
33
|
|
|
34
|
-
Follow [Build
|
|
34
|
+
Follow [Build plugins with tensoracle's C++ sources](https://sync-dev-org.github.io/tensoracle/cpp-python-repositories/) to locate the sources in the wheel and use the CMake helper in your plugin build.
|
|
35
35
|
|
|
36
36
|
## Convert a PyTorch checkpoint
|
|
37
37
|
|
|
@@ -123,6 +123,30 @@ For each weight file, `open_engine` stores its SHA-256 alongside the saved engin
|
|
|
123
123
|
|
|
124
124
|
`CorruptEngineError` reports incomplete or damaged saved files and plans that TensorRT cannot deserialize. When deserialization fails, its message asks the caller to check the required plugin libraries in `plugins=`. `EngineConditionError` reports a saved header with different build conditions. These types and `EngineIOError` derive from `EngineCacheError`. The caller chooses whether to rebuild after an error.
|
|
125
125
|
|
|
126
|
+
## Manage saved engines
|
|
127
|
+
|
|
128
|
+
Inspect saved engines and preview or perform cleanup with the installed command:
|
|
129
|
+
|
|
130
|
+
```bash
|
|
131
|
+
tensoracle cache list
|
|
132
|
+
tensoracle cache prune --unused-for 30 --dry-run
|
|
133
|
+
tensoracle cache prune --unused-for 30
|
|
134
|
+
tensoracle cache prune --max-size 10GiB --dry-run
|
|
135
|
+
tensoracle cache prune --max-size 10GiB
|
|
136
|
+
```
|
|
137
|
+
|
|
138
|
+
The commands use the same `XDG_CACHE_HOME` / home `.cache` directory as `open_engine`. They read only filesystem metadata, engine headers, and saved weight digest records. They do not read or checksum the plan body, load the original weights, deserialize engines, import TensorRT, or query a GPU. Listing recognizes the header format; it does not establish that an engine can be loaded.
|
|
139
|
+
|
|
140
|
+
`cache list` prints one JSON object with `directory`, `entries`, and `total_size_bytes`. Each entry contains the file name, size in bytes, recorded time in UTC (`recorded_at`) and exact nanoseconds (`recorded_at_ns`), and the complete header conditions. This distinguishes engines with the same name and graph version but different weights or plugins. `weight_paths` associates each weight digest with paths from valid, checksum-verified `weight-*.json` records whose SHA-256 matches that digest; an empty `paths` array means unknown. These paths are references and do not establish which application uses an engine or whether the weight file still exists. Empty files have status `building or leftover`; unreadable or invalid headers have status `unrecognized` and a reason.
|
|
141
|
+
|
|
142
|
+
An engine's `.engine` modification time records its initial save or the last successful `open_engine` cache read. Failed loads and these management commands leave that time unchanged. A timestamp write failure, including a read-only cache, does not prevent a successful load, and a disappeared file is never recreated to record its time. The record describes the time of `open_engine`, not whether inference continues afterward. Files saved before timestamp recording was introduced, and files whose timestamp could not be written, may not reflect the actual last successful read.
|
|
143
|
+
|
|
144
|
+
`cache prune` requires exactly one of `--unused-for DAYS` and `--max-size SIZE`. Days are finite nonnegative numbers, may be fractional, and mean 24 hours each; only timestamps strictly earlier than the cutoff are selected. Sizes accept nonnegative bytes or the case-insensitive units `B`, `KiB`, `MiB`, `GiB`, and `TiB` (powers of 1024). Fractions such as `1.5GiB` are allowed and fractional bytes are rounded down. Size pruning selects the oldest recorded times first, breaking ties by file name, until the total `.engine` bytes would fit. Empty and unrecognized files count toward that total but are always excluded from deletion. Weight digest records are kept.
|
|
145
|
+
|
|
146
|
+
Before deletion, pruning prints a JSON `plan` containing the selected engines, excluded files, and a notice that removed engines will be rebuilt the next time they are used. It then prints a JSON `result` separating removed, excluded, failed, and already missing files, the bytes actually removed, and the remaining total. `--dry-run` shows the same targets without changing file contents or recorded times; its result also gives the projected total. Exit status is 0 for success, 1 for filesystem failures, failed deletions, or an unattainable size limit, and 2 for invalid usage. Files that disappear before removal are reported in `missing` and do not count as bytes removed. Observed file changes are reported as failed removals with a reason.
|
|
147
|
+
|
|
148
|
+
Run pruning during maintenance when no engine is being built or opened in the same cache directory. Cleanup provides no lock or guarantee of protection against concurrent builds and reads. Inspect the selection with `--dry-run` before performing cleanup. `tensoracle cache prune --help` describes the options.
|
|
149
|
+
|
|
126
150
|
## Use a saved engine for inference
|
|
127
151
|
|
|
128
152
|
An inference application can call `open_engine` with the same build declaration and files used to save the engine. For example, assume the model package exports `graph`, `declaration`, and `io` matching the example above:
|
|
@@ -303,23 +327,13 @@ Install the producer extra when you need to compile cubins from your own kernels
|
|
|
303
327
|
pip install 'tensoracle[producer]'
|
|
304
328
|
```
|
|
305
329
|
|
|
306
|
-
|
|
307
|
-
|
|
308
|
-
```text
|
|
309
|
-
cpp/
|
|
310
|
-
CMakeLists.txt
|
|
311
|
-
python/
|
|
312
|
-
pyproject.toml
|
|
313
|
-
src/
|
|
314
|
-
```
|
|
315
|
-
|
|
316
|
-
Keep C++ sources in `cpp/`, with a CMake project that builds and runs its tests through CTest without Python. Keep the Python project, its src-layout package, and its uv environment in `python/`; run the project's uv commands from that directory. Tensoracle itself follows this layout, and the installed C++ helper paths described below remain inside its Python package.
|
|
330
|
+
Use the [C++ sources and CMake helper supplied in the tensoracle wheel](https://sync-dev-org.github.io/tensoracle/cpp-python-repositories/) to build your native plugin.
|
|
317
331
|
|
|
318
332
|
`tensoracle.producer.compile_bundle` uses the extra's **Triton 3.8.0** with **ptxas >=13.2,<14** and **cuobjdump** from the PyPI `nvidia-cuda-nvcc` and `nvidia-cuda-cuobjdump` wheels. The extra pins those tools, `nvidia-cuda-crt`, and `nvidia-nvvm` to the 13.2 series. It does not read `CUDA_HOME`; explicit `ptxas` and `cuobjdump` paths override the wheel defaults independently. A missing tool identifies its package and searched path and explains how to install the extra or pass a path. The normal import, engine build, and engine reuse paths do not import the producer or Triton. Cubin production does not need a GPU or import torch; run it without importing torch first. The caller supplies Triton functions and `KernelSpec` records with their public argument types, compile time constants, and any divisibility assumptions. The producer writes sm86, sm89, and sm120 cubins, `bundle.json`, and a C++ translation unit that embeds exactly those cubins. The generated C++ entry records carry each public argument's order, type, count, and divisibility rule into the common launcher. It rejects an unexpected parameter ABI or nonzero scratch requirement. The included [example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md#rebuild-the-cubins) has the command to rebuild its cubins from a repository clone.
|
|
319
333
|
|
|
320
334
|
Neither cubin production nor native plugin building requires an OS-installed CUDA Toolkit. For the native build, install PyPI `nvidia-cuda-nvcc==13.2.*` and `nvidia-cuda-runtime==13.2.*`, then pass their `site-packages/nvidia/cu13` directory as `CUDAToolkit_ROOT` (a CMake variable or environment variable). The helper uses its `cuda.h` and links the NVIDIA driver's `libcuda.so` on Linux, including `/usr/lib/wsl/lib/libcuda.so` on WSL, or the wheel's `lib/x64/cuda.lib` on Windows. TensorRT headers still come from the separately prepared SDK. The [SDK guide](https://github.com/sync-dev-org/tensoracle/blob/main/docs/tensorrt-sdk.md#5-pass-the-verified-root-to-the-build-and-tests) gives the setup commands. An explicit OS Toolkit root and the existing default locations remain supported for native builds.
|
|
321
335
|
|
|
322
|
-
The
|
|
336
|
+
The caller implements the graph's meaning, field meanings, weight placement, launch sequence, and variant selection. The common C++ implementation owns typed fields in their original order, handles `IPluginV3` lifecycle methods, checks cubin digests and resources, loads one module per CUDA context, checks the generated argument count and specialization rules, appends the two supported zero scratch pointers, and launches kernels. The caller selects a variant using `Kernel::accepts` and derives the count passed to `Kernel::launch` from its argument array. The common implementation releases modules and device allocations in their owning contexts. The library must export TensorRT's `getCreators`; a static registration constructor is not supported.
|
|
323
337
|
|
|
324
338
|
To record the CUDA inputs used by the plugin build, call `tensoracle_find_cuda` from `TensoracleCuda.cmake` (also included by `TensoracleCubinPlugin.cmake`):
|
|
325
339
|
|
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
# Build plugins with tensoracle's C++ sources
|
|
2
|
+
|
|
3
|
+
Tensoracle supplies the shared C++ implementation and CMake helpers for cubin plugins. Your model package supplies the kernels and model-specific plugin behavior, compiles the native plugin, and packages it for the application.
|
|
4
|
+
|
|
5
|
+
## C++ sources in the wheel
|
|
6
|
+
|
|
7
|
+
The tensoracle wheel includes:
|
|
8
|
+
|
|
9
|
+
```text
|
|
10
|
+
tensoracle/cpp/
|
|
11
|
+
include/tensoracle/cubin_plugin.hpp
|
|
12
|
+
src/cubin_plugin.cpp
|
|
13
|
+
cmake/TensoracleCubinPlugin.cmake
|
|
14
|
+
cmake/TensoracleCuda.cmake
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
The wheel contains no native binaries and uses the `py3-none-any` tag. Your plugin's native library is compiled separately and belongs to your model package.
|
|
18
|
+
|
|
19
|
+
## Prepare the build inputs
|
|
20
|
+
|
|
21
|
+
Use CMake 3.20 or later and a C++17 compiler. Follow [Prepare a TensorRT SDK](tensorrt-sdk.md) to obtain and verify the TensorRT 11 headers and library and prepare the PyPI CUDA headers and driver link library. That guide covers Linux, WSL2, and [Windows](tensorrt-sdk.md#windows-x86_64); no OS-installed CUDA Toolkit is required.
|
|
22
|
+
|
|
23
|
+
Generate a cubin bundle with the [cubin producer](api/producer.md). The bundle directory must contain `bundle.cpp`, `bundle.json`, and the cubin files. Cubin production runs on Linux; the same bundle can be used to build a Windows plugin DLL. Follow [Build a cubin plugin](index.md#build-a-cubin-plugin) for the model-specific plugin interface and tensor declarations.
|
|
24
|
+
|
|
25
|
+
## Use the packaged CMake helper
|
|
26
|
+
|
|
27
|
+
In the Python environment where tensoracle is installed, locate its C++ source directory:
|
|
28
|
+
|
|
29
|
+
```bash
|
|
30
|
+
python -c 'from importlib.metadata import distribution; print(distribution("tensoracle").locate_file("tensoracle/cpp"))'
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
Pass that directory as `TENSORACLE_CPP_ROOT` to your plugin's CMake project. Include the helper and give it your plugin source, verified SDK root, and generated bundle directory:
|
|
34
|
+
|
|
35
|
+
```cmake
|
|
36
|
+
cmake_minimum_required(VERSION 3.20)
|
|
37
|
+
project(my_model_plugin LANGUAGES CXX)
|
|
38
|
+
|
|
39
|
+
include("${TENSORACLE_CPP_ROOT}/cmake/TensoracleCubinPlugin.cmake")
|
|
40
|
+
tensoracle_add_cubin_plugin(
|
|
41
|
+
TARGET my_model_plugin
|
|
42
|
+
TENSORRT_ROOT "${TENSORRT_ROOT}"
|
|
43
|
+
BUNDLE_ROOT "${TENSORACLE_BUNDLE_ROOT}"
|
|
44
|
+
SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/plugin.cpp")
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
`tensoracle_add_cubin_plugin` builds a shared library from your `plugin.cpp`, tensoracle's common `cubin_plugin.cpp`, and the bundle's generated `bundle.cpp`. It adds the required include directories and links TensorRT and the CUDA driver.
|
|
48
|
+
|
|
49
|
+
For a Bash build, replace the absolute paths below with your plugin project, installed tensoracle sources, verified SDK, PyPI CUDA root, and generated bundle:
|
|
50
|
+
|
|
51
|
+
```bash
|
|
52
|
+
cmake -S /absolute/path/to/plugin-project -B build/plugin \
|
|
53
|
+
-DCMAKE_BUILD_TYPE=Release \
|
|
54
|
+
-DTENSORACLE_CPP_ROOT=/absolute/path/to/tensoracle/cpp \
|
|
55
|
+
-DTENSORRT_ROOT=/absolute/path/to/TensorRT-SDK \
|
|
56
|
+
-DCUDAToolkit_ROOT=/absolute/path/to/site-packages/nvidia/cu13 \
|
|
57
|
+
-DTENSORACLE_BUNDLE_ROOT=/absolute/path/to/bundle
|
|
58
|
+
cmake --build build/plugin --config Release
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
The [plugin guide](https://github.com/sync-dev-org/tensoracle#build-a-cubin-plugin) documents the required `getCreators` export and the optional `tensoracle_find_cuda` function and `tensoracle_cuda_driver` target for consumers that also need the resolved CUDA inputs.
|
|
62
|
+
|
|
63
|
+
## Package and inspect the plugin
|
|
64
|
+
|
|
65
|
+
Follow the [cubin plugin example](https://github.com/sync-dev-org/tensoracle/blob/main/examples/cubin_plugin/README.md#build-and-inspect-the-wheel) for packaging the native library and bundle in your model wheel and testing it in a runtime environment. Use the [bundle inspection API](api/bundle.md) to check the bundle and its match with the compiled plugin before distribution.
|
|
@@ -4,7 +4,7 @@
|
|
|
4
4
|
|
|
5
5
|
Follow the [four-stage optimization workflow](optimization.md) to establish output quality, evaluate lower precision and plugins, and measure speed.
|
|
6
6
|
|
|
7
|
-
Follow [Build
|
|
7
|
+
Follow [Build plugins with tensoracle's C++ sources](cpp-python-repositories.md) to locate the sources in the wheel and use the CMake helper in your plugin build.
|
|
8
8
|
|
|
9
9
|
## Terms
|
|
10
10
|
|
|
@@ -74,6 +74,30 @@ handle = open_engine(
|
|
|
74
74
|
|
|
75
75
|
`my_model` represents the application's model package. On a valid cache hit, `open_engine` does not call `graph`; on a miss, it builds and saves a new engine. Raise `graph_version` when the graph code changes or when any other caller value that changes the engine contents but is not part of the cache identity changes, such as a plugin manifest or a build recipe. `open_engine` itself does not execute the engine or manage GPU buffers. Create an execution unit from its returned handle to enqueue inference. Saved engines live under `XDG_CACHE_HOME/tensoracle`, or `~/.cache/tensoracle` when `XDG_CACHE_HOME` is unset. On Windows, `~` is `%USERPROFILE%`. See the [complete cache example](https://github.com/sync-dev-org/tensoracle#save-and-reopen-an-engine) for the declaration and expected input and output types.
|
|
76
76
|
|
|
77
|
+
### Manage saved engines
|
|
78
|
+
|
|
79
|
+
Inspect saved engines and preview or perform cleanup with the installed command:
|
|
80
|
+
|
|
81
|
+
```bash
|
|
82
|
+
tensoracle cache list
|
|
83
|
+
tensoracle cache prune --unused-for 30 --dry-run
|
|
84
|
+
tensoracle cache prune --unused-for 30
|
|
85
|
+
tensoracle cache prune --max-size 10GiB --dry-run
|
|
86
|
+
tensoracle cache prune --max-size 10GiB
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
The commands use the same `XDG_CACHE_HOME` / home `.cache` directory as `open_engine`. They read only filesystem metadata, engine headers, and saved weight digest records. They do not read or checksum the plan body, load the original weights, deserialize engines, import TensorRT, or query a GPU. Listing recognizes the header format; it does not establish that an engine can be loaded.
|
|
90
|
+
|
|
91
|
+
`cache list` prints one JSON object with `directory`, `entries`, and `total_size_bytes`. Each entry contains the file name, size in bytes, recorded time in UTC (`recorded_at`) and exact nanoseconds (`recorded_at_ns`), and the complete header conditions. This distinguishes engines with the same name and graph version but different weights or plugins. `weight_paths` associates each weight digest with paths from valid, checksum-verified `weight-*.json` records whose SHA-256 matches that digest; an empty `paths` array means unknown. These paths are references and do not establish which application uses an engine or whether the weight file still exists. Empty files have status `building or leftover`; unreadable or invalid headers have status `unrecognized` and a reason.
|
|
92
|
+
|
|
93
|
+
An engine's `.engine` modification time records its initial save or the last successful `open_engine` cache read. Failed loads and these management commands leave that time unchanged. A timestamp write failure, including a read-only cache, does not prevent a successful load, and a disappeared file is never recreated to record its time. The record describes the time of `open_engine`, not whether inference continues afterward. Files saved before timestamp recording was introduced, and files whose timestamp could not be written, may not reflect the actual last successful read.
|
|
94
|
+
|
|
95
|
+
`cache prune` requires exactly one of `--unused-for DAYS` and `--max-size SIZE`. Days are finite nonnegative numbers, may be fractional, and mean 24 hours each; only timestamps strictly earlier than the cutoff are selected. Sizes accept nonnegative bytes or the case-insensitive units `B`, `KiB`, `MiB`, `GiB`, and `TiB` (powers of 1024). Fractions such as `1.5GiB` are allowed and fractional bytes are rounded down. Size pruning selects the oldest recorded times first, breaking ties by file name, until the total `.engine` bytes would fit. Empty and unrecognized files count toward that total but are always excluded from deletion. Weight digest records are kept.
|
|
96
|
+
|
|
97
|
+
Before deletion, pruning prints a JSON `plan` containing the selected engines, excluded files, and a notice that removed engines will be rebuilt the next time they are used. It then prints a JSON `result` separating removed, excluded, failed, and already missing files, the bytes actually removed, and the remaining total. `--dry-run` shows the same targets without changing file contents or recorded times; its result also gives the projected total. Exit status is 0 for success, 1 for filesystem failures, failed deletions, or an unattainable size limit, and 2 for invalid usage. Files that disappear before removal are reported in `missing` and do not count as bytes removed. Observed file changes are reported as failed removals with a reason.
|
|
98
|
+
|
|
99
|
+
Run pruning during maintenance when no engine is being built or opened in the same cache directory. Cleanup provides no lock or guarantee of protection against concurrent builds and reads. Inspect the selection with `--dry-run` before performing cleanup. `tensoracle cache prune --help` describes the options.
|
|
100
|
+
|
|
77
101
|
### Execute an engine
|
|
78
102
|
|
|
79
103
|
The execution flow is: open the engine, create an execution unit, set all ranged input shapes together, read concrete output shapes and dtypes, allocate output buffers, then enqueue the device addresses on a caller-owned stream. Each `EngineExecution` keeps the handle alive and owns an independent TensorRT context. Fixed inputs and single-point profiles are initialized when the unit is created. Ranged inputs retain their shapes until the next successful `set_input_shapes` call. A partial TensorRT setting failure blocks queries for concrete shapes and enqueue until a complete setting succeeds.
|
|
@@ -234,7 +258,7 @@ The application runs the engine and passes a CPU NumPy candidate output and a co
|
|
|
234
258
|
|
|
235
259
|
### Build a cubin plugin
|
|
236
260
|
|
|
237
|
-
|
|
261
|
+
Use the [C++ sources and CMake helper supplied in the tensoracle wheel](cpp-python-repositories.md) to build your native plugin.
|
|
238
262
|
|
|
239
263
|
Install `tensoracle[producer]` when you need to compile cubins from your own kernels:
|
|
240
264
|
|
|
@@ -35,7 +35,7 @@ nav:
|
|
|
35
35
|
- Optimization workflow: optimization.md
|
|
36
36
|
- bf16, int8, and fp8 engines: precision.md
|
|
37
37
|
- TensorRT SDK setup: tensorrt-sdk.md
|
|
38
|
-
- C++
|
|
38
|
+
- C++ plugin builds: cpp-python-repositories.md
|
|
39
39
|
- API reference:
|
|
40
40
|
- Engine and comparison: api/index.md
|
|
41
41
|
- Cubin bundle inspection: api/bundle.md
|
|
@@ -24,6 +24,9 @@ dependencies = [
|
|
|
24
24
|
"tensorrt-cu13>=11.0,<12",
|
|
25
25
|
]
|
|
26
26
|
|
|
27
|
+
[project.scripts]
|
|
28
|
+
tensoracle = "tensoracle._cli:main"
|
|
29
|
+
|
|
27
30
|
[project.urls]
|
|
28
31
|
documentation = "https://sync-dev-org.github.io/tensoracle/"
|
|
29
32
|
repository = "https://github.com/sync-dev-org/tensoracle"
|
|
@@ -63,6 +66,7 @@ exclude = [
|
|
|
63
66
|
"/.github/",
|
|
64
67
|
"/AGENTS.md",
|
|
65
68
|
"/CLAUDE.md",
|
|
69
|
+
"/intro/",
|
|
66
70
|
"/tools/",
|
|
67
71
|
"/noren.yaml",
|
|
68
72
|
"/noren-words.txt",
|
|
@@ -3,7 +3,6 @@
|
|
|
3
3
|
import ctypes
|
|
4
4
|
import ctypes.util
|
|
5
5
|
import hashlib
|
|
6
|
-
import json
|
|
7
6
|
import os
|
|
8
7
|
import platform
|
|
9
8
|
import tempfile
|
|
@@ -16,14 +15,22 @@ from typing import TYPE_CHECKING
|
|
|
16
15
|
import tensorrt as trt
|
|
17
16
|
|
|
18
17
|
from ._build import EngineBuild, _build_serialized_engine
|
|
18
|
+
from ._cache_format import (
|
|
19
|
+
_MAGIC,
|
|
20
|
+
_WEIGHT_RECORD_VERSION,
|
|
21
|
+
CorruptEngineError,
|
|
22
|
+
EngineCacheError,
|
|
23
|
+
EngineConditionError,
|
|
24
|
+
_cache_directory,
|
|
25
|
+
_json_bytes,
|
|
26
|
+
_read_header,
|
|
27
|
+
_read_weight_record,
|
|
28
|
+
)
|
|
19
29
|
from ._plugins import register_libraries
|
|
20
30
|
|
|
21
31
|
if TYPE_CHECKING:
|
|
22
32
|
from ._execution import EngineExecution
|
|
23
33
|
|
|
24
|
-
_MAGIC = b"TENSORACLE\x01"
|
|
25
|
-
_WEIGHT_RECORD_VERSION = 1
|
|
26
|
-
|
|
27
34
|
|
|
28
35
|
@dataclass(frozen=True)
|
|
29
36
|
class TensorSpec:
|
|
@@ -59,18 +66,6 @@ class EngineHandle:
|
|
|
59
66
|
return EngineExecution(self)
|
|
60
67
|
|
|
61
68
|
|
|
62
|
-
class EngineCacheError(ValueError):
|
|
63
|
-
"""A stored engine cannot be used for this declaration."""
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
class CorruptEngineError(EngineCacheError):
|
|
67
|
-
"""A stored engine is incomplete or corrupt."""
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
class EngineConditionError(EngineCacheError):
|
|
71
|
-
"""A stored engine belongs to different build conditions."""
|
|
72
|
-
|
|
73
|
-
|
|
74
69
|
class EngineIOError(EngineCacheError):
|
|
75
70
|
"""The engine's input or output signature differs from the declaration."""
|
|
76
71
|
|
|
@@ -106,7 +101,7 @@ def open_engine(
|
|
|
106
101
|
fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
|
107
102
|
except FileExistsError:
|
|
108
103
|
plan, shape_tensor_outputs = _read_plan(path, conditions)
|
|
109
|
-
|
|
104
|
+
handle = _open_plan(
|
|
110
105
|
plan,
|
|
111
106
|
path,
|
|
112
107
|
declaration,
|
|
@@ -114,6 +109,11 @@ def open_engine(
|
|
|
114
109
|
actual_logger,
|
|
115
110
|
shape_tensor_outputs=shape_tensor_outputs,
|
|
116
111
|
)
|
|
112
|
+
try:
|
|
113
|
+
os.utime(path, None)
|
|
114
|
+
except OSError:
|
|
115
|
+
pass
|
|
116
|
+
return handle
|
|
117
117
|
else:
|
|
118
118
|
os.close(fd)
|
|
119
119
|
|
|
@@ -145,10 +145,6 @@ def open_engine(
|
|
|
145
145
|
raise
|
|
146
146
|
|
|
147
147
|
|
|
148
|
-
def _json_bytes(value: object) -> bytes:
|
|
149
|
-
return json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8")
|
|
150
|
-
|
|
151
|
-
|
|
152
148
|
def _file_sha256(path: Path) -> str:
|
|
153
149
|
digest = hashlib.sha256()
|
|
154
150
|
with path.open("rb") as stream:
|
|
@@ -157,11 +153,6 @@ def _file_sha256(path: Path) -> str:
|
|
|
157
153
|
return digest.hexdigest()
|
|
158
154
|
|
|
159
155
|
|
|
160
|
-
def _cache_directory() -> Path:
|
|
161
|
-
cache_root = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache")
|
|
162
|
-
return cache_root / "tensoracle"
|
|
163
|
-
|
|
164
|
-
|
|
165
156
|
def _weight_sha256(path: Path, *, force_rebuild: bool) -> str:
|
|
166
157
|
absolute_path = Path(os.path.abspath(path))
|
|
167
158
|
before = absolute_path.stat()
|
|
@@ -175,27 +166,8 @@ def _weight_sha256(path: Path, *, force_rebuild: bool) -> str:
|
|
|
175
166
|
record_path = _cache_directory() / f"weight-{path_digest}.json"
|
|
176
167
|
|
|
177
168
|
if not force_rebuild:
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
except (OSError, UnicodeDecodeError, json.JSONDecodeError):
|
|
181
|
-
record = None
|
|
182
|
-
if (
|
|
183
|
-
isinstance(record, dict)
|
|
184
|
-
and record.keys() == identity.keys() | {"version", "sha256", "checksum"}
|
|
185
|
-
and record["version"] == _WEIGHT_RECORD_VERSION
|
|
186
|
-
and type(record["version"]) is int
|
|
187
|
-
and all(
|
|
188
|
-
type(record[key]) is type(value) and record[key] == value
|
|
189
|
-
for key, value in identity.items()
|
|
190
|
-
)
|
|
191
|
-
and isinstance(record["sha256"], str)
|
|
192
|
-
and len(record["sha256"]) == 64
|
|
193
|
-
and all(character in "0123456789abcdef" for character in record["sha256"])
|
|
194
|
-
and record["checksum"]
|
|
195
|
-
== hashlib.sha256(
|
|
196
|
-
_json_bytes({key: value for key, value in record.items() if key != "checksum"})
|
|
197
|
-
).hexdigest()
|
|
198
|
-
):
|
|
169
|
+
record = _read_weight_record(record_path)
|
|
170
|
+
if record is not None and all(record[key] == value for key, value in identity.items()):
|
|
199
171
|
return record["sha256"]
|
|
200
172
|
|
|
201
173
|
digest = _file_sha256(absolute_path)
|
|
@@ -323,31 +295,18 @@ def _conditions(
|
|
|
323
295
|
|
|
324
296
|
|
|
325
297
|
def _read_plan(path: Path, conditions: dict) -> tuple[memoryview, tuple[str, ...]]:
|
|
326
|
-
|
|
327
|
-
|
|
328
|
-
|
|
329
|
-
|
|
330
|
-
|
|
331
|
-
|
|
332
|
-
|
|
333
|
-
|
|
334
|
-
|
|
335
|
-
header = json.loads(data[offset + 8 : header_end])
|
|
336
|
-
except (UnicodeDecodeError, json.JSONDecodeError) as error:
|
|
337
|
-
raise CorruptEngineError(f"invalid TensorRT engine header: {path}") from error
|
|
338
|
-
if not isinstance(header, dict) or not isinstance(header.get("conditions"), dict):
|
|
339
|
-
raise CorruptEngineError(f"invalid TensorRT engine header: {path}")
|
|
340
|
-
shape_tensor_outputs = header.get("shape_tensor_outputs", [])
|
|
341
|
-
if not isinstance(shape_tensor_outputs, list) or any(
|
|
342
|
-
not isinstance(name, str) for name in shape_tensor_outputs
|
|
343
|
-
):
|
|
344
|
-
raise CorruptEngineError(f"invalid TensorRT engine header: {path}")
|
|
345
|
-
if header["conditions"] != conditions:
|
|
346
|
-
raise EngineConditionError(f"TensorRT engine conditions differ: {path}")
|
|
347
|
-
plan = memoryview(data)[header_end:]
|
|
298
|
+
# A buffered read after the header can concatenate a buffered prefix with a
|
|
299
|
+
# large plan while holding the GIL. FileIO.readall avoids that concatenation
|
|
300
|
+
# and still handles plans larger than a single OS read can return.
|
|
301
|
+
with path.open("rb", buffering=0) as stream:
|
|
302
|
+
file_size = os.fstat(stream.fileno()).st_size
|
|
303
|
+
header = _read_header(stream, path, file_size)
|
|
304
|
+
if header["conditions"] != conditions:
|
|
305
|
+
raise EngineConditionError(f"TensorRT engine conditions differ: {path}")
|
|
306
|
+
plan = memoryview(stream.read())
|
|
348
307
|
if header.get("plan_sha256") != hashlib.sha256(plan).hexdigest():
|
|
349
308
|
raise CorruptEngineError(f"TensorRT engine checksum differs: {path}")
|
|
350
|
-
return plan, tuple(shape_tensor_outputs)
|
|
309
|
+
return plan, tuple(header.get("shape_tensor_outputs", []))
|
|
351
310
|
|
|
352
311
|
|
|
353
312
|
def _open_plan(
|
|
@@ -0,0 +1,78 @@
|
|
|
1
|
+
"""Shared on-disk cache format and weight records, without TensorRT imports."""
|
|
2
|
+
|
|
3
|
+
import hashlib
|
|
4
|
+
import json
|
|
5
|
+
import os
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
from typing import BinaryIO
|
|
8
|
+
|
|
9
|
+
_MAGIC = b"TENSORACLE\x01"
|
|
10
|
+
_WEIGHT_RECORD_VERSION = 1
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class EngineCacheError(ValueError):
|
|
14
|
+
"""A stored engine cannot be used for this declaration."""
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class CorruptEngineError(EngineCacheError):
|
|
18
|
+
"""A stored engine is incomplete or corrupt."""
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class EngineConditionError(EngineCacheError):
|
|
22
|
+
"""A stored engine belongs to different build conditions."""
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def _json_bytes(value: object) -> bytes:
|
|
26
|
+
return json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8")
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def _cache_directory() -> Path:
|
|
30
|
+
cache_root = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache")
|
|
31
|
+
return cache_root / "tensoracle"
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _read_header(stream: BinaryIO, path: Path, file_size: int) -> dict:
|
|
35
|
+
"""Read and validate only the header, leaving the stream at the plan body."""
|
|
36
|
+
prefix = stream.read(len(_MAGIC) + 8)
|
|
37
|
+
if not prefix.startswith(_MAGIC) or len(prefix) != len(_MAGIC) + 8:
|
|
38
|
+
raise CorruptEngineError(f"incomplete TensorRT engine file: {path}")
|
|
39
|
+
header_size = int.from_bytes(prefix[len(_MAGIC) :], "big")
|
|
40
|
+
if header_size == 0 or len(prefix) + header_size >= file_size:
|
|
41
|
+
raise CorruptEngineError(f"invalid TensorRT engine header: {path}")
|
|
42
|
+
try:
|
|
43
|
+
header = json.loads(stream.read(header_size))
|
|
44
|
+
except (UnicodeDecodeError, ValueError, RecursionError) as error:
|
|
45
|
+
raise CorruptEngineError(f"invalid TensorRT engine header: {path}") from error
|
|
46
|
+
if not isinstance(header, dict) or not isinstance(header.get("conditions"), dict):
|
|
47
|
+
raise CorruptEngineError(f"invalid TensorRT engine header: {path}")
|
|
48
|
+
shape_tensor_outputs = header.get("shape_tensor_outputs", [])
|
|
49
|
+
if not isinstance(shape_tensor_outputs, list) or any(
|
|
50
|
+
not isinstance(name, str) for name in shape_tensor_outputs
|
|
51
|
+
):
|
|
52
|
+
raise CorruptEngineError(f"invalid TensorRT engine header: {path}")
|
|
53
|
+
return header
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def _read_weight_record(path: Path) -> dict | None:
|
|
57
|
+
"""Validate a digest record without accessing the original weight file."""
|
|
58
|
+
try:
|
|
59
|
+
record = json.loads(path.read_bytes())
|
|
60
|
+
except (OSError, UnicodeDecodeError, ValueError, RecursionError):
|
|
61
|
+
return None
|
|
62
|
+
if (
|
|
63
|
+
not isinstance(record, dict)
|
|
64
|
+
or record.keys() != {"version", "path", "size", "mtime_ns", "inode", "sha256", "checksum"}
|
|
65
|
+
or type(record["version"]) is not int
|
|
66
|
+
or record["version"] != _WEIGHT_RECORD_VERSION
|
|
67
|
+
or not isinstance(record["path"], str)
|
|
68
|
+
or any(type(record[key]) is not int for key in ("size", "mtime_ns", "inode"))
|
|
69
|
+
or not isinstance(record["sha256"], str)
|
|
70
|
+
or len(record["sha256"]) != 64
|
|
71
|
+
or any(character not in "0123456789abcdef" for character in record["sha256"])
|
|
72
|
+
or record["checksum"]
|
|
73
|
+
!= hashlib.sha256(
|
|
74
|
+
_json_bytes({key: value for key, value in record.items() if key != "checksum"})
|
|
75
|
+
).hexdigest()
|
|
76
|
+
):
|
|
77
|
+
return None
|
|
78
|
+
return record
|