tensoracle 0.5.0__tar.gz → 0.5.1__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-0.5.0 → tensoracle-0.5.1}/PKG-INFO +4 -2
- {tensoracle-0.5.0 → tensoracle-0.5.1}/README.md +3 -1
- {tensoracle-0.5.0 → tensoracle-0.5.1}/docs/index.md +2 -2
- {tensoracle-0.5.0 → tensoracle-0.5.1}/docs/optimization.md +2 -2
- tensoracle-0.5.1/docs/precision.md +154 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/mkdocs.yml +1 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/src/tensoracle/_cache.py +32 -3
- {tensoracle-0.5.0 → tensoracle-0.5.1}/src/tensoracle/_version.py +2 -2
- tensoracle-0.5.1/tests/precision_runtime.py +164 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_cache.py +5 -5
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_distribution.py +31 -7
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_engine_io.py +1 -1
- tensoracle-0.5.1/tests/test_gpu_cache.py +204 -0
- tensoracle-0.5.1/tests/test_precision.py +111 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_weight_hash.py +13 -13
- {tensoracle-0.5.0 → tensoracle-0.5.1}/.gitignore +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/LICENSE +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/build_backend.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/cpp/CMakeLists.txt +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/cpp/cmake/TensoracleCubinPlugin.cmake +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/cpp/include/tensoracle/cubin_plugin.hpp +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/cpp/src/cubin_plugin.cpp +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/cpp/tests/core.cpp +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/docs/api/bundle.md +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/docs/api/index.md +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/docs/api/producer.md +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/docs/tensorrt-sdk.md +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/CMakeLists.txt +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/README.md +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/bundle.cpp +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/bundle.json +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm120/bias.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm120/dot.specialized.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm120/dot.tf32-control.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm120/dot.unspecialized.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm120/scale.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm86/bias.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm86/dot.specialized.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm86/dot.tf32-control.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm86/dot.unspecialized.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm86/scale.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm89/bias.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm89/dot.specialized.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm89/dot.tf32-control.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm89/dot.unspecialized.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/bundle/cubins/sm89/scale.cubin +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/example_cubin_plugin/__init__.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/example_cubin_plugin/weights.safetensors +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/kernels.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/plugin.cpp +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/examples/cubin_plugin/pyproject.toml +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/pyproject.toml +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/src/tensoracle/__init__.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/src/tensoracle/_build.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/src/tensoracle/_comparison.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/src/tensoracle/_execution.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/src/tensoracle/_memory.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/src/tensoracle/_plugins.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/src/tensoracle/bundle.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/src/tensoracle/convert.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/src/tensoracle/producer.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/build_memory_runtime.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/cubin_io/check.cpp +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/cubin_io/kernels.cu +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/cubin_io/plugin.cpp +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/example_runtime.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/plugin_io_runtime.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_build.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_build_memory.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_bundle.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_comparison.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_convert.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_cpp_core.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_example_plugin.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_execution.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_plugin_io.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_plugins.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_producer.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/tests/test_shape_tensor_outputs.py +0 -0
- {tensoracle-0.5.0 → tensoracle-0.5.1}/uv.lock +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: tensoracle
|
|
3
|
-
Version: 0.5.
|
|
3
|
+
Version: 0.5.1
|
|
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
|
|
@@ -59,7 +59,7 @@ Choose the path that matches your work:
|
|
|
59
59
|
|
|
60
60
|
The supported targets are Linux x86_64 (including WSL2), TensorRT 11, and GPU generations sm86, sm89, and sm120. The producer and native build tools are separate from the `tensoracle` runtime dependencies. Engine building, engine reuse, output comparison, and cubin production do not require torch. Converting `.pth` checkpoints requires the `convert` extra described below.
|
|
61
61
|
|
|
62
|
-
Follow the [four-stage optimization workflow](https://github.com/sync-dev-org/tensoracle/blob/main/docs/optimization.md) to establish output quality, evaluate
|
|
62
|
+
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.
|
|
63
63
|
|
|
64
64
|
## Convert a PyTorch checkpoint
|
|
65
65
|
|
|
@@ -114,6 +114,8 @@ Use `InputShape.fixed((1, 64, 2))` for one input size; the graph may declare tha
|
|
|
114
114
|
|
|
115
115
|
The engine's precision is determined by the tensor types that the graph function places in the network. For an fp16 graph, use `trt.float16` tensors and `np.float16` weights where needed; the graph can also contain fp32 tensors. `EngineBuild` declares input shapes and builder settings. TF32 is always disabled; `BuilderSettings(tf32=True)` raises `ValueError`. The function returns serialized plan bytes. Invalid declarations raise `ValueError`, and a failed TensorRT build raises `RuntimeError`.
|
|
116
116
|
|
|
117
|
+
The same API builds bf16, int8, and fp8 graphs. Cast tensors to `trt.bfloat16` for bf16; for int8/fp8, place `add_quantize` / `add_dequantize` layers and load their scales from safetensors. The application determines scales with an external quantization library and chooses output quality limits with `compare_outputs`. fp8 targets sm89 and sm120; bf16 and int8 also target sm86. Follow [Build bf16, int8, and fp8 engines](https://github.com/sync-dev-org/tensoracle/blob/main/docs/precision.md) for a complete example, NumPy dtype constraints, and comparison against an fp32 engine. Use different engine names or graph versions for different precisions.
|
|
118
|
+
|
|
117
119
|
## Save and reopen an engine
|
|
118
120
|
|
|
119
121
|
`open_engine` builds on the first call and reuses the saved engine on later calls, including calls from another process. It invokes `graph` only when building. Declare the expected input and output names, TensorRT types, and build-time shapes:
|
|
@@ -29,7 +29,7 @@ Choose the path that matches your work:
|
|
|
29
29
|
|
|
30
30
|
The supported targets are Linux x86_64 (including WSL2), TensorRT 11, and GPU generations sm86, sm89, and sm120. The producer and native build tools are separate from the `tensoracle` runtime dependencies. Engine building, engine reuse, output comparison, and cubin production do not require torch. Converting `.pth` checkpoints requires the `convert` extra described below.
|
|
31
31
|
|
|
32
|
-
Follow the [four-stage optimization workflow](https://github.com/sync-dev-org/tensoracle/blob/main/docs/optimization.md) to establish output quality, evaluate
|
|
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
34
|
## Convert a PyTorch checkpoint
|
|
35
35
|
|
|
@@ -84,6 +84,8 @@ Use `InputShape.fixed((1, 64, 2))` for one input size; the graph may declare tha
|
|
|
84
84
|
|
|
85
85
|
The engine's precision is determined by the tensor types that the graph function places in the network. For an fp16 graph, use `trt.float16` tensors and `np.float16` weights where needed; the graph can also contain fp32 tensors. `EngineBuild` declares input shapes and builder settings. TF32 is always disabled; `BuilderSettings(tf32=True)` raises `ValueError`. The function returns serialized plan bytes. Invalid declarations raise `ValueError`, and a failed TensorRT build raises `RuntimeError`.
|
|
86
86
|
|
|
87
|
+
The same API builds bf16, int8, and fp8 graphs. Cast tensors to `trt.bfloat16` for bf16; for int8/fp8, place `add_quantize` / `add_dequantize` layers and load their scales from safetensors. The application determines scales with an external quantization library and chooses output quality limits with `compare_outputs`. fp8 targets sm89 and sm120; bf16 and int8 also target sm86. Follow [Build bf16, int8, and fp8 engines](https://github.com/sync-dev-org/tensoracle/blob/main/docs/precision.md) for a complete example, NumPy dtype constraints, and comparison against an fp32 engine. Use different engine names or graph versions for different precisions.
|
|
88
|
+
|
|
87
89
|
## Save and reopen an engine
|
|
88
90
|
|
|
89
91
|
`open_engine` builds on the first call and reuses the saved engine on later calls, including calls from another process. It invokes `graph` only when building. Declare the expected input and output names, TensorRT types, and build-time shapes:
|
|
@@ -2,7 +2,7 @@
|
|
|
2
2
|
|
|
3
3
|
`tensoracle` builds TensorRT engines from model graphs supplied through TensorRT's Python network API. A model package owns its graph, safetensors weights, and output quality limits. `tensoracle` supplies engine building, saved engine reuse, execution, output comparison, and tools for cubin plugins. These workflows do not require torch. Only the optional `.pth` checkpoint converter requires torch, installed with `tensoracle[convert]`.
|
|
4
4
|
|
|
5
|
-
Follow the [four-stage optimization workflow](optimization.md) to establish output quality, evaluate
|
|
5
|
+
Follow the [four-stage optimization workflow](optimization.md) to establish output quality, evaluate lower precision and plugins, and measure speed.
|
|
6
6
|
|
|
7
7
|
## Terms
|
|
8
8
|
|
|
@@ -161,7 +161,7 @@ Networks with outputs marked as shape tensors emit `tensoracle.ShapeTensorOutput
|
|
|
161
161
|
|
|
162
162
|
### Build a standard graph
|
|
163
163
|
|
|
164
|
-
Define `graph(network)` with TensorRT's Python network API and load model weights from safetensors. Describe each graph input with `InputShape.fixed` or `InputShape.ranged`, and pass an `EngineBuild` declaration to `build_serialized_engine`. `EngineBuild` declares input shapes and builder settings; the graph function determines precision through the tensor types it places in the network. The supported precisions are fp32 (with TF32 disabled) and
|
|
164
|
+
Define `graph(network)` with TensorRT's Python network API and load model weights from safetensors. Describe each graph input with `InputShape.fixed` or `InputShape.ranged`, and pass an `EngineBuild` declaration to `build_serialized_engine`. `EngineBuild` declares input shapes and builder settings; the graph function determines precision through the tensor types it places in the network. The supported precisions are fp32 (with TF32 disabled), fp16, bf16, int8, and fp8. The [bf16, int8, and fp8 guide](precision.md) shows tensor casts, explicit Q/DQ layers, and scales supplied through safetensors; fp8 targets sm89 and sm120. `open_engine` uses the same declaration plus a name, graph version, weight files, and `EngineIO` to save and reuse a checked engine. Use a different `name` or `graph_version` for engines with different tensor types, such as fp32 and fp16. The [graph example](https://github.com/sync-dev-org/tensoracle#define-and-build-a-standard-graph) shows the exact calls.
|
|
165
165
|
|
|
166
166
|
The application runs the engine and passes a CPU NumPy candidate output and a corresponding reference output to `compare_outputs`. The application chooses the acceptance limits. See [Compare output arrays](https://github.com/sync-dev-org/tensoracle#compare-output-arrays).
|
|
167
167
|
|
|
@@ -10,9 +10,9 @@ Use the same representative inputs, shapes, and corresponding outputs throughout
|
|
|
10
10
|
|
|
11
11
|
## 2. Accept lower precision by quality
|
|
12
12
|
|
|
13
|
-
- **tensoracle tools:** Build an fp16
|
|
13
|
+
- **tensoracle tools:** Build an fp16, bf16, int8, or fp8 TensorRT engine using the same `EngineBuild` declaration for input shapes and builder settings. The graph function determines precision through its tensor types. For int8 and fp8, the graph places explicit quantize/dequantize layers and reads their scales from safetensors; the application prepares those scales with an external quantization library. Follow [Build bf16, int8, and fp8 engines](precision.md) for a complete graph example and the GPU targets (fp8 requires sm89 or sm120). When saving with `open_engine`, use a different `name` or `graph_version` from the fp32 engine. Use `compare_outputs` with each lower precision output as the candidate and its accepted fp32 TensorRT output as the reference. It reports non-finite counts and, for finite floating point arrays, maximum absolute error, relative L2 error, cosine similarity, and, when both tolerances are supplied, counts and fractions outside tolerance.
|
|
14
14
|
- **You provide:** Choose limits for those model-independent measurements. Compute any model-specific measurements separately, such as PSNR for images or application-specific time-series quality measures, and choose their limits. Run both engines on the same inputs.
|
|
15
|
-
- **Gate:** Accept the
|
|
15
|
+
- **Gate:** Accept the lower precision engine only when its outputs meet all applicable limits, including the model-specific ones. Otherwise keep the accepted fp32 engine or revise the candidate and repeat this stage.
|
|
16
16
|
|
|
17
17
|
## 3. Replace a remaining bottleneck with a plugin
|
|
18
18
|
|
|
@@ -0,0 +1,154 @@
|
|
|
1
|
+
# Build bf16, int8, and fp8 engines
|
|
2
|
+
|
|
3
|
+
The graph function selects tensor types. Use the same `EngineBuild`, `open_engine`, and execution API for fp32, fp16, bf16, int8, and fp8. For bf16, cast tensors to `trt.bfloat16`. For int8 and fp8, place explicit quantize/dequantize (Q/DQ) layers and their scales in the graph. TensorRT 11's [`add_quantize` and `add_dequantize` operators](https://docs.nvidia.com/deeplearning/tensorrt/11.4.0/_static/operators/Quantize.html) accept the quantized type and dequantized type explicitly.
|
|
4
|
+
|
|
5
|
+
| Graph precision | Tensor placement in this example | Supported tensoracle GPU targets |
|
|
6
|
+
| --- | --- | --- |
|
|
7
|
+
| fp32 | fp32 matrix multiplication and bias, with TF32 disabled | sm86, sm89, sm120 |
|
|
8
|
+
| bf16 | Cast features, matrix, and bias to bf16; cast the result to fp32 | sm86, sm89, sm120 |
|
|
9
|
+
| int8 | Q/DQ the features and matrix with `trt.int8` and explicit scales | sm86, sm89, sm120 |
|
|
10
|
+
| fp8 | Q/DQ the features and matrix with `trt.fp8` (E4M3) and explicit scales | sm89, sm120 |
|
|
11
|
+
|
|
12
|
+
Do not build the fp8 graph on sm86. The graph keeps fp32 input and output buffers for every precision; the internal tensors determine its arithmetic. TensorRT chooses implementations for the typed graph, so an int8 or fp8 graph can retain higher precision operations such as the bias addition. This example does not establish a speed improvement.
|
|
13
|
+
|
|
14
|
+
## Prepare weights and scales outside tensoracle
|
|
15
|
+
|
|
16
|
+
For a real model, use an external quantization library, such as [NVIDIA Model Optimizer](https://github.com/NVIDIA/Model-Optimizer), to determine scales from representative data or quantization-aware training. The model package translates those results into its safetensors file and graph. tensoracle does not calculate scales, calibrate a model, or translate a quantizer's export format.
|
|
17
|
+
|
|
18
|
+
Store the fp32 weights and positive fp32 scales together in safetensors. In this example, `matrix` has shape `(32, 16)`, `bias` has shape `(1, 16)`, and each scale has shape `(1,)`. The keys `int8.input_scale`, `int8.weight_scale`, `fp8.input_scale`, and `fp8.weight_scale` distinguish the two quantization recipes. Both Q and DQ use the same scale for each tensor.
|
|
19
|
+
|
|
20
|
+
The following small fixture uses maximum absolute values to choose scales. Its inputs and weights match the GPU test; this scale calculation is only a demonstration for that fixture, not a calibration recipe for other models:
|
|
21
|
+
|
|
22
|
+
```python
|
|
23
|
+
from pathlib import Path
|
|
24
|
+
|
|
25
|
+
import numpy as np
|
|
26
|
+
from safetensors.numpy import save_file
|
|
27
|
+
|
|
28
|
+
path = Path("linear.safetensors")
|
|
29
|
+
features = np.random.default_rng(11).uniform(-1, 1, (16, 32)).astype(np.float32)
|
|
30
|
+
rng = np.random.default_rng(29)
|
|
31
|
+
matrix = rng.uniform(-0.25, 0.25, (32, 16)).astype(np.float32)
|
|
32
|
+
bias = rng.uniform(-0.1, 0.1, (1, 16)).astype(np.float32)
|
|
33
|
+
weights = {"matrix": matrix, "bias": bias}
|
|
34
|
+
for precision, maximum in (("int8", 127), ("fp8", 448)):
|
|
35
|
+
weights[f"{precision}.input_scale"] = np.array(
|
|
36
|
+
[np.max(np.abs(features)) / maximum], dtype=np.float32
|
|
37
|
+
)
|
|
38
|
+
weights[f"{precision}.weight_scale"] = np.array(
|
|
39
|
+
[np.max(np.abs(matrix)) / maximum], dtype=np.float32
|
|
40
|
+
)
|
|
41
|
+
save_file(weights, path)
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
[`safetensors.numpy` 0.8.0](https://github.com/huggingface/safetensors/blob/v0.8.0/bindings/python/py_src/safetensors/numpy.py) with NumPy 2.5 does not directly load bf16 or fp8 tensors into NumPy arrays: bf16 loading raises `TypeError`, and FP8 E4M3 loading raises `AttributeError`. A safetensors file can store those types, but its NumPy reader cannot represent them in this environment. Keep this example's weights and scales in fp32, then cast or quantize in TensorRT. If an existing checkpoint stores bf16 or fp8 weights, convert them with a reader that supports that dtype outside this NumPy graph path. Installing torch is not required for this example.
|
|
45
|
+
|
|
46
|
+
## Define the typed graph
|
|
47
|
+
|
|
48
|
+
This graph computes `features @ matrix + bias`. The bf16 branch casts all three operands before arithmetic and returns fp32. The int8/fp8 branch applies per-tensor Q/DQ to the features and matrix; its bias stays fp32. Keeping the stored matrix in fp32 lets TensorRT quantize it with the supplied scale during building.
|
|
49
|
+
|
|
50
|
+
```python
|
|
51
|
+
import tensorrt as trt
|
|
52
|
+
from safetensors.numpy import load_file
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def graph(path, precision):
|
|
56
|
+
def define(network):
|
|
57
|
+
weights = load_file(path)
|
|
58
|
+
|
|
59
|
+
def constant(array):
|
|
60
|
+
array = np.ascontiguousarray(array, dtype=np.float32)
|
|
61
|
+
return network.add_constant(array.shape, array).get_output(0)
|
|
62
|
+
|
|
63
|
+
features = network.add_input("features", trt.float32, (16, 32))
|
|
64
|
+
matrix = constant(weights["matrix"])
|
|
65
|
+
bias = constant(weights["bias"])
|
|
66
|
+
if precision == "bf16":
|
|
67
|
+
features = network.add_cast(features, trt.bfloat16).get_output(0)
|
|
68
|
+
matrix = network.add_cast(matrix, trt.bfloat16).get_output(0)
|
|
69
|
+
bias = network.add_cast(bias, trt.bfloat16).get_output(0)
|
|
70
|
+
elif precision in ("int8", "fp8"):
|
|
71
|
+
quantized_type = trt.int8 if precision == "int8" else trt.fp8
|
|
72
|
+
|
|
73
|
+
def qdq(tensor, key):
|
|
74
|
+
scale = constant(weights[f"{precision}.{key}_scale"])
|
|
75
|
+
quantized = network.add_quantize(tensor, scale, quantized_type).get_output(0)
|
|
76
|
+
return network.add_dequantize(quantized, scale, trt.float32).get_output(0)
|
|
77
|
+
|
|
78
|
+
features = qdq(features, "input")
|
|
79
|
+
matrix = qdq(matrix, "weight")
|
|
80
|
+
value = network.add_matrix_multiply(
|
|
81
|
+
features, trt.MatrixOperation.NONE, matrix, trt.MatrixOperation.NONE
|
|
82
|
+
).get_output(0)
|
|
83
|
+
value = network.add_elementwise(value, bias, trt.ElementWiseOperation.SUM).get_output(0)
|
|
84
|
+
if precision == "bf16":
|
|
85
|
+
value = network.add_cast(value, trt.float32).get_output(0)
|
|
86
|
+
value.name = "result"
|
|
87
|
+
network.mark_output(value)
|
|
88
|
+
|
|
89
|
+
return define
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
## Build, save, and reuse each precision
|
|
93
|
+
|
|
94
|
+
Use a distinct engine name or graph version whenever tensor types or the graph recipe change. The same weight file serves all four graphs here, so its digest cannot distinguish their precision. Scales are part of that file and therefore part of the saved engine identity.
|
|
95
|
+
|
|
96
|
+
```python
|
|
97
|
+
from tensoracle import (
|
|
98
|
+
BuilderSettings, EngineBuild, EngineIO, InputShape, TensorSpec, open_engine,
|
|
99
|
+
)
|
|
100
|
+
|
|
101
|
+
declaration = EngineBuild(
|
|
102
|
+
{"features": InputShape.fixed((16, 32))},
|
|
103
|
+
BuilderSettings(optimization_level=0),
|
|
104
|
+
)
|
|
105
|
+
io = EngineIO(
|
|
106
|
+
{"features": TensorSpec(trt.float32, (16, 32))},
|
|
107
|
+
{"result": TensorSpec(trt.float32, (16, 16))},
|
|
108
|
+
)
|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
def open_linear(precision):
|
|
112
|
+
return open_engine(
|
|
113
|
+
graph(path, precision),
|
|
114
|
+
declaration,
|
|
115
|
+
name=f"linear-{precision}",
|
|
116
|
+
graph_version="1",
|
|
117
|
+
weights=(path,),
|
|
118
|
+
io=io,
|
|
119
|
+
)
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
reference_handle = open_linear("fp32")
|
|
123
|
+
candidate_handle = open_linear("bf16") # Or "int8"; "fp8" on sm89/sm120.
|
|
124
|
+
```
|
|
125
|
+
|
|
126
|
+
TF32 remains disabled by tensoracle, including for the fp32 reference. Engines are saved under `XDG_CACHE_HOME/tensoracle`, or `~/.cache/tensoracle`. Repeat the same call with unchanged build conditions in another process to reuse the plan without invoking the graph function. Keep the weight file and the engine names/versions stable across the restart. Follow [Execute an engine](index.md#execute-an-engine) to allocate fp32 input/output buffers and enqueue each handle with the same inputs.
|
|
127
|
+
|
|
128
|
+
## Compare outputs and choose limits
|
|
129
|
+
|
|
130
|
+
First establish the fp32 engine against fixed fp32 outputs from the official model implementation, as described in [stage 1](optimization.md#1-establish-the-fp32-reference). Then run the candidate and accepted fp32 engine with identical inputs. Pass their CPU NumPy outputs to `compare_outputs` and choose quality limits in the application:
|
|
131
|
+
|
|
132
|
+
```python
|
|
133
|
+
from tensoracle import compare_outputs
|
|
134
|
+
|
|
135
|
+
# candidate_output and fp32_output are the application's engine outputs.
|
|
136
|
+
comparison = compare_outputs(candidate_output, fp32_output, atol=0.01, rtol=0)
|
|
137
|
+
accepted = (
|
|
138
|
+
comparison.candidate_nonfinite_count == 0
|
|
139
|
+
and comparison.reference_nonfinite_count == 0
|
|
140
|
+
and comparison.outside_tolerance_count == 0
|
|
141
|
+
and comparison.relative_l2_error is not None
|
|
142
|
+
and comparison.relative_l2_error < 0.01
|
|
143
|
+
)
|
|
144
|
+
```
|
|
145
|
+
|
|
146
|
+
These example limits belong to the small bf16 fixture. For the fixture above, a TensorRT 11.4 run on sm120 measured:
|
|
147
|
+
|
|
148
|
+
| Candidate | Maximum absolute error | Relative L2 error | Test `atol` (`rtol=0`) | Test relative L2 limit |
|
|
149
|
+
| --- | --- | --- | --- | --- |
|
|
150
|
+
| bf16 | 0.00357 | 0.00294 | 0.01 | 0.01 |
|
|
151
|
+
| int8 | 0.00710 | 0.00572 | 0.02 | 0.02 |
|
|
152
|
+
| fp8 | 0.0501 | 0.0382 | 0.12 | 0.08 |
|
|
153
|
+
|
|
154
|
+
The test allows at least twice the observed errors for differences in GPU tactics. These are not tensoracle defaults or quality guarantees for other models. Choose representative data, per-output limits, and model-specific quality measurements for the real application. The GPU test also checks that the two engines occupy separate cache entries and reproduce their outputs in another process without calling either graph.
|
|
@@ -236,7 +236,8 @@ def _write_weight_record(path: Path, record: dict) -> None:
|
|
|
236
236
|
pass
|
|
237
237
|
|
|
238
238
|
|
|
239
|
-
def
|
|
239
|
+
def _gpu_identity() -> dict[str, str | int]:
|
|
240
|
+
"""Describe device compatibility without binding the cache to a physical GPU."""
|
|
240
241
|
library_name = ctypes.util.find_library("cuda") or ctypes.util.find_library("nvcuda")
|
|
241
242
|
if library_name is None:
|
|
242
243
|
raise RuntimeError("CUDA driver library was not found")
|
|
@@ -245,6 +246,8 @@ def _gpu_name() -> str:
|
|
|
245
246
|
cuda.cuCtxGetDevice.argtypes = [ctypes.POINTER(ctypes.c_int)]
|
|
246
247
|
cuda.cuDeviceGet.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int]
|
|
247
248
|
cuda.cuDeviceGetName.argtypes = [ctypes.c_char_p, ctypes.c_int, ctypes.c_int]
|
|
249
|
+
cuda.cuDeviceGetAttribute.argtypes = [ctypes.POINTER(ctypes.c_int), ctypes.c_int, ctypes.c_int]
|
|
250
|
+
cuda.cuDeviceTotalMem_v2.argtypes = [ctypes.POINTER(ctypes.c_size_t), ctypes.c_int]
|
|
248
251
|
if cuda.cuInit(0) != 0:
|
|
249
252
|
raise RuntimeError("CUDA driver initialization failed")
|
|
250
253
|
device = ctypes.c_int()
|
|
@@ -256,7 +259,33 @@ def _gpu_name() -> str:
|
|
|
256
259
|
name = ctypes.create_string_buffer(256)
|
|
257
260
|
if cuda.cuDeviceGetName(name, len(name), device.value) != 0:
|
|
258
261
|
raise RuntimeError("CUDA device name lookup failed")
|
|
259
|
-
|
|
262
|
+
properties: dict[str, str | int] = {"name": name.value.decode("utf-8")}
|
|
263
|
+
# TensorRT's documented compatibility checks, plus ECC and device memory:
|
|
264
|
+
# https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/engine-compatibility.html#compatibility-checks
|
|
265
|
+
# ECC can distinguish otherwise identical devices with different usable memory.
|
|
266
|
+
# Attribute numbers are CUdevice_attribute values from the CUDA driver API.
|
|
267
|
+
for field, attribute in (
|
|
268
|
+
("compute_capability_major", 75),
|
|
269
|
+
("compute_capability_minor", 76),
|
|
270
|
+
("global_memory_bus_width", 37),
|
|
271
|
+
("l2_cache_size", 38),
|
|
272
|
+
("max_shared_memory_per_block", 8),
|
|
273
|
+
("max_shared_memory_per_block_optin", 97),
|
|
274
|
+
("max_shared_memory_per_multiprocessor", 81),
|
|
275
|
+
("texture_alignment", 14),
|
|
276
|
+
("multiprocessor_count", 16),
|
|
277
|
+
("integrated", 18),
|
|
278
|
+
("ecc_enabled", 32),
|
|
279
|
+
):
|
|
280
|
+
value = ctypes.c_int()
|
|
281
|
+
if cuda.cuDeviceGetAttribute(ctypes.byref(value), attribute, device.value) != 0:
|
|
282
|
+
raise RuntimeError(f"CUDA device attribute lookup failed: {field}")
|
|
283
|
+
properties[field] = value.value
|
|
284
|
+
total_memory = ctypes.c_size_t()
|
|
285
|
+
if cuda.cuDeviceTotalMem_v2(ctypes.byref(total_memory), device.value) != 0:
|
|
286
|
+
raise RuntimeError("CUDA device total memory lookup failed")
|
|
287
|
+
properties["total_memory"] = total_memory.value
|
|
288
|
+
return properties
|
|
260
289
|
|
|
261
290
|
|
|
262
291
|
def _conditions(
|
|
@@ -286,7 +315,7 @@ def _conditions(
|
|
|
286
315
|
for name, shape in declaration.inputs.items()
|
|
287
316
|
},
|
|
288
317
|
"builder_settings": asdict(declaration.settings),
|
|
289
|
-
"
|
|
318
|
+
"gpu": _gpu_identity(),
|
|
290
319
|
"os": platform.system(),
|
|
291
320
|
"architecture": platform.machine(),
|
|
292
321
|
"tensorrt_version": trt.__version__,
|
|
@@ -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.5.
|
|
22
|
-
__version_tuple__ = version_tuple = (0, 5,
|
|
21
|
+
__version__ = version = '0.5.1'
|
|
22
|
+
__version_tuple__ = version_tuple = (0, 5, 1)
|
|
23
23
|
|
|
24
24
|
__commit_id__ = commit_id = None
|
|
@@ -0,0 +1,164 @@
|
|
|
1
|
+
"""Build and execute the precision guide's linear graph, including after a restart."""
|
|
2
|
+
|
|
3
|
+
import argparse
|
|
4
|
+
import json
|
|
5
|
+
from pathlib import Path
|
|
6
|
+
|
|
7
|
+
import numpy as np
|
|
8
|
+
import tensorrt as trt
|
|
9
|
+
from cuda.bindings import runtime as cudart
|
|
10
|
+
from safetensors.numpy import load_file, save_file
|
|
11
|
+
|
|
12
|
+
from tensoracle import (
|
|
13
|
+
BuilderSettings,
|
|
14
|
+
EngineBuild,
|
|
15
|
+
EngineIO,
|
|
16
|
+
InputShape,
|
|
17
|
+
TensorSpec,
|
|
18
|
+
open_engine,
|
|
19
|
+
)
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def sample_input():
|
|
23
|
+
return np.random.default_rng(11).uniform(-1, 1, (16, 32)).astype(np.float32)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def write_weights(path):
|
|
27
|
+
rng = np.random.default_rng(29)
|
|
28
|
+
matrix = rng.uniform(-0.25, 0.25, (32, 16)).astype(np.float32)
|
|
29
|
+
bias = rng.uniform(-0.1, 0.1, (1, 16)).astype(np.float32)
|
|
30
|
+
weights = {"matrix": matrix, "bias": bias}
|
|
31
|
+
# A small test fixture, not a model calibration algorithm or a package API.
|
|
32
|
+
for precision, maximum in (("int8", 127), ("fp8", 448)):
|
|
33
|
+
weights[f"{precision}.input_scale"] = np.array(
|
|
34
|
+
[np.max(np.abs(sample_input())) / maximum], dtype=np.float32
|
|
35
|
+
)
|
|
36
|
+
weights[f"{precision}.weight_scale"] = np.array(
|
|
37
|
+
[np.max(np.abs(matrix)) / maximum], dtype=np.float32
|
|
38
|
+
)
|
|
39
|
+
save_file(weights, path)
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def graph(path, precision):
|
|
43
|
+
def define(network):
|
|
44
|
+
weights = load_file(path)
|
|
45
|
+
|
|
46
|
+
def constant(array):
|
|
47
|
+
array = np.ascontiguousarray(array, dtype=np.float32)
|
|
48
|
+
return network.add_constant(array.shape, array).get_output(0)
|
|
49
|
+
|
|
50
|
+
features = network.add_input("features", trt.float32, (16, 32))
|
|
51
|
+
matrix = constant(weights["matrix"])
|
|
52
|
+
bias = constant(weights["bias"])
|
|
53
|
+
if precision == "bf16":
|
|
54
|
+
features = network.add_cast(features, trt.bfloat16).get_output(0)
|
|
55
|
+
matrix = network.add_cast(matrix, trt.bfloat16).get_output(0)
|
|
56
|
+
bias = network.add_cast(bias, trt.bfloat16).get_output(0)
|
|
57
|
+
elif precision in ("int8", "fp8"):
|
|
58
|
+
quantized_type = trt.int8 if precision == "int8" else trt.fp8
|
|
59
|
+
|
|
60
|
+
def qdq(tensor, key):
|
|
61
|
+
scale = constant(weights[f"{precision}.{key}_scale"])
|
|
62
|
+
quantized = network.add_quantize(tensor, scale, quantized_type).get_output(0)
|
|
63
|
+
return network.add_dequantize(quantized, scale, trt.float32).get_output(0)
|
|
64
|
+
|
|
65
|
+
features = qdq(features, "input")
|
|
66
|
+
matrix = qdq(matrix, "weight")
|
|
67
|
+
value = network.add_matrix_multiply(
|
|
68
|
+
features, trt.MatrixOperation.NONE, matrix, trt.MatrixOperation.NONE
|
|
69
|
+
).get_output(0)
|
|
70
|
+
value = network.add_elementwise(value, bias, trt.ElementWiseOperation.SUM).get_output(0)
|
|
71
|
+
if precision == "bf16":
|
|
72
|
+
value = network.add_cast(value, trt.float32).get_output(0)
|
|
73
|
+
value.name = "result"
|
|
74
|
+
network.mark_output(value)
|
|
75
|
+
|
|
76
|
+
return define
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
def open_linear(path, precision, *, reuse=False):
|
|
80
|
+
def forbidden(network):
|
|
81
|
+
raise AssertionError("a saved precision engine called its graph after a process restart")
|
|
82
|
+
|
|
83
|
+
return open_engine(
|
|
84
|
+
forbidden if reuse else graph(path, precision),
|
|
85
|
+
EngineBuild(
|
|
86
|
+
{"features": InputShape.fixed((16, 32))},
|
|
87
|
+
BuilderSettings(optimization_level=0),
|
|
88
|
+
),
|
|
89
|
+
name=f"linear-{precision}",
|
|
90
|
+
graph_version="1",
|
|
91
|
+
weights=(path,),
|
|
92
|
+
io=EngineIO(
|
|
93
|
+
{"features": TensorSpec(trt.float32, (16, 32))},
|
|
94
|
+
{"result": TensorSpec(trt.float32, (16, 16))},
|
|
95
|
+
),
|
|
96
|
+
)
|
|
97
|
+
|
|
98
|
+
|
|
99
|
+
def _cuda(result):
|
|
100
|
+
status, *values = result
|
|
101
|
+
if status != cudart.cudaError_t.cudaSuccess:
|
|
102
|
+
raise RuntimeError(f"CUDA runtime returned {status}")
|
|
103
|
+
return values[0] if values else None
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def run(handle):
|
|
107
|
+
features = sample_input()
|
|
108
|
+
output = np.empty((16, 16), dtype=np.float32)
|
|
109
|
+
source = target = stream = None
|
|
110
|
+
with handle.create_execution() as execution:
|
|
111
|
+
assert execution.get_tensor_shape("result") == output.shape
|
|
112
|
+
assert execution.get_tensor_dtype("result") == trt.float32
|
|
113
|
+
try:
|
|
114
|
+
source = _cuda(cudart.cudaMalloc(features.nbytes))
|
|
115
|
+
target = _cuda(cudart.cudaMalloc(output.nbytes))
|
|
116
|
+
stream = _cuda(cudart.cudaStreamCreate())
|
|
117
|
+
_cuda(
|
|
118
|
+
cudart.cudaMemcpyAsync(
|
|
119
|
+
source,
|
|
120
|
+
features.ctypes.data,
|
|
121
|
+
features.nbytes,
|
|
122
|
+
cudart.cudaMemcpyKind.cudaMemcpyHostToDevice,
|
|
123
|
+
stream,
|
|
124
|
+
)
|
|
125
|
+
)
|
|
126
|
+
execution.enqueue({"features": int(source), "result": int(target)}, stream=int(stream))
|
|
127
|
+
_cuda(
|
|
128
|
+
cudart.cudaMemcpyAsync(
|
|
129
|
+
output.ctypes.data,
|
|
130
|
+
target,
|
|
131
|
+
output.nbytes,
|
|
132
|
+
cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
|
|
133
|
+
stream,
|
|
134
|
+
)
|
|
135
|
+
)
|
|
136
|
+
finally:
|
|
137
|
+
if stream is not None:
|
|
138
|
+
_cuda(cudart.cudaStreamSynchronize(stream))
|
|
139
|
+
_cuda(cudart.cudaStreamDestroy(stream))
|
|
140
|
+
if target is not None:
|
|
141
|
+
_cuda(cudart.cudaFree(target))
|
|
142
|
+
if source is not None:
|
|
143
|
+
_cuda(cudart.cudaFree(source))
|
|
144
|
+
return output
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def main():
|
|
148
|
+
parser = argparse.ArgumentParser()
|
|
149
|
+
parser.add_argument("weights", type=Path)
|
|
150
|
+
parser.add_argument("precision", choices=("bf16", "int8", "fp8"))
|
|
151
|
+
parser.add_argument("outputs", type=Path)
|
|
152
|
+
args = parser.parse_args()
|
|
153
|
+
outputs = {}
|
|
154
|
+
paths = {}
|
|
155
|
+
for precision in ("fp32", args.precision):
|
|
156
|
+
handle = open_linear(args.weights, precision, reuse=True)
|
|
157
|
+
paths[precision] = str(handle.path)
|
|
158
|
+
outputs[precision] = run(handle)
|
|
159
|
+
np.savez(args.outputs, **outputs)
|
|
160
|
+
print(json.dumps(paths))
|
|
161
|
+
|
|
162
|
+
|
|
163
|
+
if __name__ == "__main__":
|
|
164
|
+
main()
|
|
@@ -95,7 +95,7 @@ def test_save_large_built_plan_keeps_other_threads_running(tmp_path, monkeypatch
|
|
|
95
95
|
|
|
96
96
|
plan = b"\0" * (512 * 1024 * 1024)
|
|
97
97
|
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
|
|
98
|
-
monkeypatch.setattr(cache, "
|
|
98
|
+
monkeypatch.setattr(cache, "_gpu_identity", lambda: {"name": "GPU A"})
|
|
99
99
|
builder = SimpleNamespace(
|
|
100
100
|
create_network=lambda flags: SimpleNamespace(num_inputs=0, num_outputs=0),
|
|
101
101
|
create_builder_config=lambda: SimpleNamespace(
|
|
@@ -234,7 +234,7 @@ def test_saved_engine_deserialization_failure_guides_caller_to_check_plugins(tmp
|
|
|
234
234
|
|
|
235
235
|
cache_root = tmp_path / "cache"
|
|
236
236
|
monkeypatch.setenv("XDG_CACHE_HOME", str(cache_root))
|
|
237
|
-
monkeypatch.setattr(cache, "
|
|
237
|
+
monkeypatch.setattr(cache, "_gpu_identity", lambda: {"name": "GPU A"})
|
|
238
238
|
weights = tmp_path / "weights.safetensors"
|
|
239
239
|
_weights(weights)
|
|
240
240
|
declaration = tensoracle.EngineBuild({"features": tensoracle.InputShape.fixed((1, 3, 2))})
|
|
@@ -414,7 +414,7 @@ def test_build_failure_removes_placeholder_and_partial_files(tmp_path, monkeypat
|
|
|
414
414
|
import tensoracle._cache as cache
|
|
415
415
|
|
|
416
416
|
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
|
|
417
|
-
monkeypatch.setattr(cache, "
|
|
417
|
+
monkeypatch.setattr(cache, "_gpu_identity", lambda: {"name": "GPU A"})
|
|
418
418
|
monkeypatch.setattr(
|
|
419
419
|
cache, "_build_serialized_engine", lambda graph, declaration, plugins, logger: graph(None)
|
|
420
420
|
)
|
|
@@ -474,7 +474,7 @@ def test_each_declared_condition_selects_a_distinct_cache_entry(tmp_path, monkey
|
|
|
474
474
|
import tensoracle._cache as cache
|
|
475
475
|
|
|
476
476
|
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
|
|
477
|
-
monkeypatch.setattr(cache, "
|
|
477
|
+
monkeypatch.setattr(cache, "_gpu_identity", lambda: {"name": "GPU A"})
|
|
478
478
|
weights = tmp_path / "weights.safetensors"
|
|
479
479
|
_weights(weights)
|
|
480
480
|
same_contents = tmp_path / "same.safetensors"
|
|
@@ -627,7 +627,7 @@ def test_each_declared_condition_selects_a_distinct_cache_entry(tmp_path, monkey
|
|
|
627
627
|
]
|
|
628
628
|
paths = {open_case(**variant).path for variant in variants}
|
|
629
629
|
with monkeypatch.context() as patch:
|
|
630
|
-
patch.setattr(cache, "
|
|
630
|
+
patch.setattr(cache, "_gpu_identity", lambda: {"name": "GPU B"})
|
|
631
631
|
paths.add(open_case().path)
|
|
632
632
|
with monkeypatch.context() as patch:
|
|
633
633
|
patch.setattr(cache.platform, "system", lambda: "Other OS")
|
|
@@ -130,30 +130,53 @@ def test_public_snapshot_includes_only_the_pages_workflow(publication_config, tm
|
|
|
130
130
|
assert paths - {PAGES_WORKFLOW} <= excluded
|
|
131
131
|
|
|
132
132
|
|
|
133
|
+
@pytest.mark.req("SP-TENSORACLE-024")
|
|
133
134
|
def test_pages_workflow_builds_the_public_snapshot(publication_config):
|
|
134
|
-
"""Only the public repository builds
|
|
135
|
+
"""Only the public repository builds and deploys docs with limited permissions and no secrets."""
|
|
135
136
|
workflow_path = ROOT / PAGES_WORKFLOW
|
|
136
137
|
assert workflow_path.is_file()
|
|
137
|
-
|
|
138
|
+
workflow_source = workflow_path.read_text(encoding="utf-8")
|
|
139
|
+
assert "secrets." not in workflow_source
|
|
140
|
+
workflow = yaml.safe_load(workflow_source)
|
|
138
141
|
triggers = workflow.get("on", workflow.get(True))
|
|
139
142
|
assert triggers == {"push": {"branches": ["main"]}, "workflow_dispatch": None}
|
|
140
143
|
|
|
144
|
+
declared_permissions = set()
|
|
145
|
+
for scope in (workflow, *workflow["jobs"].values()):
|
|
146
|
+
permissions = scope.get("permissions", {})
|
|
147
|
+
assert isinstance(permissions, dict)
|
|
148
|
+
declared_permissions.update(permissions.items())
|
|
149
|
+
assert declared_permissions == {
|
|
150
|
+
("contents", "read"),
|
|
151
|
+
("pages", "write"),
|
|
152
|
+
("id-token", "write"),
|
|
153
|
+
}
|
|
154
|
+
|
|
141
155
|
steps = []
|
|
142
156
|
for job in workflow["jobs"].values():
|
|
143
157
|
assert job["if"] == "github.repository == 'sync-dev-org/tensoracle'"
|
|
144
158
|
steps.extend(job["steps"])
|
|
145
159
|
commands = [step["run"] for step in steps if "run" in step]
|
|
146
|
-
assert
|
|
160
|
+
assert commands == [
|
|
161
|
+
"uv sync --locked --only-group docs",
|
|
162
|
+
"uv run --no-sync mkdocs build --strict --config-file ../mkdocs.yml",
|
|
163
|
+
]
|
|
147
164
|
for job in workflow["jobs"].values():
|
|
148
165
|
if any("run" in step for step in job["steps"]):
|
|
149
166
|
assert job["defaults"]["run"]["working-directory"] == "python"
|
|
150
|
-
assert "uv run --no-sync mkdocs build --strict --config-file ../mkdocs.yml" in commands
|
|
151
167
|
actions = {}
|
|
152
168
|
for step in steps:
|
|
153
169
|
if "uses" in step:
|
|
154
170
|
action, version = step["uses"].split("@")
|
|
155
|
-
assert re.fullmatch(r"[0-9a-f]{40}
|
|
171
|
+
assert re.fullmatch(r"[0-9a-f]{40}", version)
|
|
156
172
|
actions[action] = step
|
|
173
|
+
assert actions.keys() <= {
|
|
174
|
+
"actions/checkout",
|
|
175
|
+
"astral-sh/setup-uv",
|
|
176
|
+
"actions/configure-pages",
|
|
177
|
+
"actions/upload-pages-artifact",
|
|
178
|
+
"actions/deploy-pages",
|
|
179
|
+
}
|
|
157
180
|
assert {
|
|
158
181
|
"actions/configure-pages",
|
|
159
182
|
"actions/upload-pages-artifact",
|
|
@@ -163,8 +186,9 @@ def test_pages_workflow_builds_the_public_snapshot(publication_config):
|
|
|
163
186
|
deployment = next(job for job in workflow["jobs"].values() if "environment" in job)
|
|
164
187
|
assert deployment["needs"] == "build"
|
|
165
188
|
assert deployment["environment"]["name"] == "github-pages"
|
|
166
|
-
|
|
167
|
-
assert
|
|
189
|
+
deployment_permissions = deployment.get("permissions", workflow["permissions"])
|
|
190
|
+
assert deployment_permissions["pages"] == "write"
|
|
191
|
+
assert deployment_permissions["id-token"] == "write"
|
|
168
192
|
|
|
169
193
|
|
|
170
194
|
def test_docs_dependencies_are_separate_from_gpu_development_packages(publication_config):
|
|
@@ -50,7 +50,7 @@ def _cached_open(tmp_path, monkeypatch, engine, declaration, io):
|
|
|
50
50
|
import tensoracle._cache as cache
|
|
51
51
|
|
|
52
52
|
monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
|
|
53
|
-
monkeypatch.setattr(cache, "
|
|
53
|
+
monkeypatch.setattr(cache, "_gpu_identity", lambda: {"name": "GPU A"})
|
|
54
54
|
weights = tmp_path / "weights.safetensors"
|
|
55
55
|
_weights(weights)
|
|
56
56
|
conditions = cache._conditions("checked", "1", (weights,), (), declaration)
|