From e9e37c45ef4fd97c663debcc35b2b9cbdc02cbcd Mon Sep 17 00:00:00 2001 From: Frank Lai Date: Tue, 13 Oct 2026 17:18:59 -0700 Subject: [PATCH 01/21] docs: publish downloadable LEAPP integration skill Signed-off-by: Frank Lai --- AGENTS.md | 5 +- SKILL.md | 208 ++++++++++++++++++++++++++++++ docs/source/api/index.rst | 2 +- docs/source/conf.py | 21 +++ docs/source/index.rst | 20 +++ references/function-patching.md | 72 +++++++++++ references/state-and-semantics.md | 106 +++++++++++++++ references/warp.md | 121 +++++++++++++++++ 8 files changed, 552 insertions(+), 3 deletions(-) create mode 100644 SKILL.md create mode 100644 references/function-patching.md create mode 100644 references/state-and-semantics.md create mode 100644 references/warp.md diff --git a/AGENTS.md b/AGENTS.md index 73697815..540456fd 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -7,8 +7,9 @@ SPDX-License-Identifier: Apache-2.0 Instructions for coding agents working **in** nvidia-isaac/leapp: where code lives, which tests to run, and which skill to load. This is not a -guide for annotating a downstream policy; that lives in -`docs/source/` and at https://nvidia-isaac.github.io/leapp/ +guide for annotating a downstream policy; read [SKILL.md](SKILL.md) for +that workflow, with API documentation in `docs/source/` and at +https://nvidia-isaac.github.io/leapp/ Python 3.10+ and PyTorch 2.6.0+. Visualization extras install only on 3.11+. Environment is a local venv (typically `uv`), never conda. diff --git a/SKILL.md b/SKILL.md new file mode 100644 index 00000000..07d1a32f --- /dev/null +++ b/SKILL.md @@ -0,0 +1,208 @@ +--- +name: leapp-integrate +description: Annotate and export downstream LEAPP pipelines with a single-model ONNX default, function patches, state and deployment semantics, runtime validation, and optional Warp setup/capture. Use for application integration rather than tracer internals. +--- + +# Integrate LEAPP into a project + +**Start by exporting the full compute pipeline as one model (one LEAPP node), unless the user requests multiple models or a concrete technical/deployment requirement justifies splitting it.** Existing preprocessing, policy, and postprocessing helpers can all execute within that node; their function boundaries do not require separate exports. + +Preserve the existing computation. **Avoid rewriting pipelines, model implementations, or library call sites to make them traceable. Use `leapp.FunctionPatch` for unsupported module-level functions instead.** Keep adapters in the integration layer, scoped to the behavior the application actually uses. Supported Python helpers can execute unchanged inside a node. + +**Recommend ONNX unless the user specifies another format.** Explicitly set `export_with="onnx"` (the `onnx-dynamo` alias); omitting `export_with` can select no export. Try compatible alternative exporters when the first fails, before considering application changes. + +## Discover docs and establish the environment + +Start with the [online LEAPP documentation](https://nvidia-isaac.github.io/leapp/). **Do not assume LEAPP is installed or that a source checkout exists.** Read only the guides relevant to the application: + +| Need | Online reference | +|------|------------------| +| Public signatures and annotation patterns | [API](https://nvidia-isaac.github.io/leapp/api/index.html), [nodes](https://nvidia-isaac.github.io/leapp/guides/nodes.html) | +| Export formats, parameters, prebuilt models | [Export](https://nvidia-isaac.github.io/leapp/guides/export.html) | +| Unsupported functions and trace loss | [Function patches](https://nvidia-isaac.github.io/leapp/guides/patching.html), [Torch limitations](https://nvidia-isaac.github.io/leapp/tensor_libraries/torch_limitations.html), [NumPy limitations](https://nvidia-isaac.github.io/leapp/tensor_libraries/limitations.html) | +| State and metadata | [State](https://nvidia-isaac.github.io/leapp/guides/graph.html), [semantics](https://nvidia-isaac.github.io/leapp/semantics/usage.html), [temporal axes](https://nvidia-isaac.github.io/leapp/semantics/temporal.html) | +| Validation and deployment | [Debugging](https://nvidia-isaac.github.io/leapp/guides/debugging.html), [runtime](https://nvidia-isaac.github.io/leapp/leapp_runtime.html), [bundle/YAML](https://nvidia-isaac.github.io/leapp/generated_configs.html) | +| Warp | [Setup](https://nvidia-isaac.github.io/leapp/tensor_libraries/installation.html), [capture](https://nvidia-isaac.github.io/leapp/tensor_libraries/warp.html), [limitations](https://nvidia-isaac.github.io/leapp/tensor_libraries/warp_limitations.html) | + +Use the project's established Python environment and dependency manager. If LEAPP is absent, add `leapp` using that workflow (for example, `python -m pip install leapp` inside a virtual environment). Respect project version constraints; do not broadly upgrade dependencies. Warp needs the extras and native build described below. + +Once available, record `leapp.__version__`, `leapp.__file__`, and relevant public signatures. Online docs can be newer than an installed package; verify compatibility before using an API. Resolve conflicting examples against the version's signature or matching upstream source/tests, without requiring a local checkout. Keep the user's single-model/ONNX preferences when docs illustrate other choices. + +## Start with one model; split only for a reason + +1. Run or inspect the original entrypoint. Map external inputs, meaningful computation stages, final outputs, mutable state, and external services. Record shapes, dtypes, devices, layout, and representative inputs. +2. Place one input boundary before the pipeline's compute and one output boundary after its final result. Include preprocessing, feature fusion, policy inference, and action conversion in that same node where applicable. Multiple helpers, branches, model calls, or stateful components alone are not reasons to split. +3. Keep acquisition, simulator stepping, file/network I/O, visualization, and orchestration outside exported compute nodes. Explicitly identify the runtime responsible for any external computation between nodes; a disconnected graph is not a complete exported pipeline. +4. Keep **every input-dependent transformation** inside the exported compute, including normalization, slicing, reshaping, dtype conversion, clipping, and concatenation. These operations must run on new inputs at inference time, not be omitted or frozen to values from tracing. If multiple nodes are needed, include each transformation in a producer or consumer; unannotated math between nodes can become an extra graph input instead of an edge. +5. Split only when requested or justified by a concrete requirement: separately deployed components, different devices/rates required by the target runtime, incompatible export backends, or an unsupported external operation that a faithful function patch cannot capture. For export failures, investigate patches and compatible exporters before partitioning. State the reason for each split, then choose cohesive stage boundaries with clear tensor contracts. +6. When using multiple nodes, use stable names, finalize producers before consumers, and explicitly connect published outputs at branches and joins. Do not nest independently annotated nodes inside an active node. Node-name reuse means re-entry of the same logical stage and annotation origin; independently exported instances with different inputs, weights, or state need distinct names. + +## Annotation API chooser + +| API | Use | +|-----|-----| +| `input_tensors` + `output_tensors` | Primary interface for explicit names, semantic metadata, and nodes spanning existing helpers or files. | +| `@annotate.method(...)` | Shorthand when one existing function naturally defines a node. `node_name` defaults to the function name. | +| `leapp.FunctionPatch(module, name, replacement)` | Traceable adapter for an unsupported function, registered through `start(patching=[...])`; not a node declaration. | +| `annotate.module(name, model)` | Track registered `nn.Module` buffers used in a node; reassigned buffers become feedback. | +| `state_tensors` + `update_state` | Explicit recurrent state and its next value. | +| `annotate.register_buffer` | Frozen constants in the exported model, without feedback. | +| `output_tensors(..., static_outputs=...)` | Deliberately constant output tensors. Never use to conceal lost tracing. | +| `annotate.warp_op` | Explicit Warp segment within an existing node; see the Warp subsection below. | + +Avoid legacy private `annotate._method` for new integrations; investigate public interfaces first. + +## Complete annotation and validation recipe + +This runnable stateless example keeps an original function unchanged, caches multiple trace samples, exports one model, and checks fresh inputs. Replace the illustrative function and sample generator with the application's existing compute and representative data. For mutable pipelines, isolate baseline inputs/state and reset both executions to the same initial state; see [state and deployment semantics](references/state-and-semantics.md). + +```python +import torch +import leapp +from leapp import annotate, InferenceManager + + +def original_pipeline(obs): + features = (obs - obs.mean()) / (obs.std() + 1e-6) + return torch.tanh(features[..., :12]) + + +samples = [torch.randn(1, 32) * scale + 0.2 for scale in (0.3, 0.5, 1, 2, 4)] +expected = [original_pipeline(obs) for obs in samples] + +leapp.start(name="demo_graph", save_path=".", max_cached_io=5) +try: + for obs, reference in zip(samples, expected): + x = annotate.input_tensors("pipeline", {"obs": obs}) + action = original_pipeline(x) + annotate.output_tensors("pipeline", {"action": action}, export_with="onnx") + torch.testing.assert_close(action, reference) +finally: + leapp.stop() +leapp.compile_graph(visualize=True, validate=True, rtol=1e-3, atol=1e-5, strict=True) + +manager = InferenceManager("demo_graph/demo_graph.yaml") +assert set(manager.inputs) == {"pipeline/obs"} +assert set(manager.outputs) == {"pipeline/action"} +for scale in (0.7, 1.5, 3.0): + inputs = manager.get_mock_input() + fresh = torch.randn_like(inputs["pipeline/obs"]) * scale - 0.4 + reference = original_pipeline(fresh) + inputs["pipeline/obs"] = fresh + actual = manager.run_policy(inputs)["pipeline/action"] + torch.testing.assert_close(actual, reference.to(actual.device), rtol=1e-3, atol=1e-5) +``` + +`max_cached_io=5` only sets cache capacity; **execute the annotated pipeline repeatedly on distinct inputs** to collect samples. A mock-input run or parity on the first captured input is insufficient. Preserve application weights, execution mode, state, and random seeds as appropriate. Never call `compile_graph()` before `stop()`. Restore session patches on errors; resolve active-wrapper errors before retrying in that process. + +- Raw I/O uses a top-level **dict of named tensors**. A single input returns its value, not a dict; multiple inputs return a tuple in dict insertion order. Always use these returned values for computation. +- Nested dictionaries, lists, and tuples can remain nested within named values. Inspect generated flattened port names rather than guessing them or rewriting application containers. +- `input_tensors("node", ...)` may be called at multiple sites while assembling that node's initial trace. One `output_tensors()` finalizes it. Later iterations re-enter the same stage to validate/cache I/O and detect feedback; they do not create additional graph nodes. +- `method()` derives return names from source. Use explicit I/O annotations for ambiguous returns or unavailable source rather than rewriting a function solely for name extraction. + +## Export methods and fallback strategy + +Select per node through `export_with` on `output_tensors()` or `method()`. LEAPP emits per-node artifacts plus `.yaml`, optional PNG, and feedback initial-value data when applicable. + +| `export_with` | Method / artifact | When to try | +|---------------|-------------------|-------------| +| `"onnx"`, `"onnx-dynamo"` | Dynamo ONNX exporter → `.onnx` | Default recommendation; modern tensor graphs and supported Warp nodes. These two names select the same backend. | +| `"onnx-torchscript"` | TorchScript-based ONNX exporter → `.onnx` | Alternative ONNX lowering when Dynamo fails or validates incorrectly; often appropriate first for GRU/LSTM. | +| `"jit"`, `"jit-script"` | TorchScript scripting → `.pt` | A TorchScript target or an alternative when ONNX lowering is unsupported. Both names select the same backend. | +| `"jit-trace"` | TorchScript tracing → `.pt` | Scripting cannot handle the captured graph; requires representative input and execution-path coverage. | +| `"pt2"`, `"exported-program"` | `torch.export` → `.pt2` | A requested ExportedProgram target or a compatible modern PyTorch runtime; supports Warp with its runtime adapter. Both names select the same backend. | +| `None` | No compilation; optionally reference an existing model | Deliberate external/prebuilt model integration or an explicitly incomplete placeholder. Not an export-success fallback. | + +When the first export fails: + +1. Identify the failing node and whether the problem is tracing, exporter lowering, missing dependencies, or runtime/numerical validation. An exporter switch cannot recover computation already lost during tracing; use a function patch for an unsupported source operation. +2. For ONNX lowering failures, try the **other ONNX exporter** with unchanged computation. Switching `"onnx"` to `"onnx-dynamo"` is not a new attempt. For recurrent graphs, try `onnx-torchscript` early. Warp supports Dynamo ONNX and PT2, not the legacy ONNX or JIT routes. +3. If both ONNX routes are unsuitable and no user requirement mandates ONNX, try compatible `pt2`, `jit-script`, or `jit-trace` based on the error and target runtime. For JIT scripting failures, try `jit-trace`. Mixed per-node backends are possible when the consuming runtime supports them; report any departure from the ONNX recommendation. +4. Honor a user-specified format. Other formats can help diagnose a failure but do not satisfy a required ONNX/PT2/JIT deliverable. Do not silently substitute formats or drop failing nodes. +5. Change the annotation's backend settings, then retrace in a fresh session/output directory with original state reset. Validate each candidate; keep a concise record of backend, parameters, and failure. Do not retry the same failing configuration without new evidence. If suitable alternatives fail, report the narrow unsupported operation and remaining deployment gap before proposing a scoped application change. + +Useful `backend_params` (check the installed implementation): + +- Both ONNX paths: `opset_version`, `report`; Dynamo also `verify`, `optimize`; legacy ONNX also `prescript`. +- PT2: `strict` controls `torch.export`, distinct from `compile_graph(strict=True)` validation. +- `None`: `{"model_path": "/path/model.onnx", "copy_original_model": True}` references/copies an existing `.onnx`, `.pt`, or `.pt2`. Verify its I/O contract and runtime execution. + +Choose opset and optimizations for the actual runtime. Disabling validation, `skip_validation=True`, dry-run, or `non_traced` is diagnostic only, not evidence of a working export. + +## State, constants, and loops + +For recurrent policies, history buffers, or episode resets, read [state and deployment semantics](references/state-and-semantics.md). + +- `state_tensors` requires an existing node and a **flat dict of named Torch tensors**; nested state is unsupported. `update_state` must use declared names and preserve each state's shape and dtype. One entry returns a value; multiple entries return a tuple in insertion order. +- Only states passed to `update_state()` become explicit feedback; an unupdated state remains a regular input. Keep evolving state distinct from frozen parameters and constants. +- After opening the node, call `annotate.module("node", model, buffer_names=[...])` before model execution to select registered buffers; omit `buffer_names` to track all. It tracks reassignment (`self.h = new_h`), not in-place updates (`self.h.copy_(...)`) or plain attributes. Use explicit state annotations for existing mutation patterns where possible; a pure function patch cannot represent mutation. +- Inferred cross-node feedback needs **at least two iterations** between `start()` and `stop()`. Explicit state/module feedback can be declared in one trace; still validate multiple time steps and initial/reset state. Confirm `pipeline.feedback_flow` and the saved initial values. +- Re-entry does not make a Python data-dependent branch, a loop count, or shape dynamic. Verify the supported input/shape domain; do not promise dynamic shapes from a single successful example. + +## Preserve graph connections and computation + +Current tracing connects a published `(producing node, output port)` to a consumer. Verify the installed version's behavior if using an older release. + +- Equivalent copies can preserve the port: Torch `clone`/`detach`/`contiguous`/device-only transfers/full buffer overwrites, NumPy `np.copy`/`.copy()`/`np.asanyarray`, and full-range `wp.copy`. +- Arithmetic, slicing, reshaping, dtype casts, and partial writes between nodes can clear the port. Move the annotation boundary to include that existing computation; do not relabel the result as its unmodified source. +- `annotate.mirror_leapp_tags(source, target)` is only for **equivalent copies between finished nodes**, with matching values, shape, and dtype. Never use it to repair a transformed value or a missing computation. Assign its return value, especially for raw NumPy destinations that cannot be upgraded in place. +- Use `static_outputs` only for intentional constants (named raw Torch tensors or supported semantic wrappers). Capture unsupported numerical functions with [FunctionPatch](references/function-patching.md). If that cannot preserve behavior, identify an explicit external runtime boundary and report that the export is partial. + +## Diagnose trace loss before changing exporters + +Calls in NumPy, SciPy, scikit-learn, scikit-image, OpenCV, or other libraries **may return ordinary untraced arrays/tensors**, depending on the operation. A library name alone does not determine support: helpers built entirely from supported operations can trace unchanged. If an input-dependent result loses tracing, patch the unsupported module-level function with a faithful [FunctionPatch adapter](references/function-patching.md). For a scikit estimator method, locate the unsupported module-level call or a suitable module-level application helper; `FunctionPatch` does not target arbitrary bound methods. + +| Hazard | What the agent should do | +|--------|--------------------------| +| Unmapped NumPy/native operation | Treat the trace warning as a correctness failure. A plain result mixed back into traced data can become a constant. Patch the operation and compare new inputs against the original. | +| `.item()`, `.tolist()`, `float()`, `int()` | Audit explicitly: Torch extraction can be silent. Keep dynamic data in traced operations; logs alone cannot prove correctness. | +| NumPy indexing down to a scalar or NumPy scalar operands | Check for loss of array tracing or unsupported operand types. Use a faithful adapter preserving the needed array computation. | +| NumPy negative-stride slicing, unsupported `astype` dtype | Eager and exported behavior can disagree. Check the online limitations and implement supported equivalent operations within the adapter. | +| Tensor-dependent Python `if`/`while` or iteration counts | Only the executed path is captured. Use a faithful traceable expression in an adapter where possible; otherwise document the supported domain or external control-flow boundary. Another exporter cannot recover a branch already lost. | + +Supported NumPy/Torch conversions such as `torch.from_numpy()` and `.numpy()` can preserve tracing with global patching enabled. Do not rewrite all NumPy code or treat every conversion as a trace break. + +## Semantic metadata + +For deployment-facing tensors, add known meaning with `TensorSemantics`: observed state versus output targets, real joint/feature order, and temporal axes for action chunks. Read the runnable example in [state and deployment semantics](references/state-and-semantics.md). + +- Semantic annotations use `input_tensors` / `output_tensors`, **not `method`**. Pass a single `TensorSemantics` or a list directly; no dict wrapping or mixing with raw tensors in that list. Names must be unique within the inputs or outputs. +- `element_names` describes axis ordering; use `None` for unnamed dimensions. Put at most one `TemporalAxis(period_ms=...)` directly in the outer axis list, not inside a nested list. +- `extra` fields such as units/frames are flattened into YAML and are project-specific. Avoid collisions with built-in fields and confirm the downstream consumer understands them. +- Use `GraphConfigs(frequency=...)` for graph execution frequency in Hz. Temporal sample spacing is in milliseconds and is independent of graph frequency; LEAPP does not validate their consistency or schedule the application from this metadata. + +## Warp setup and export (experimental) + +Warp is part of this skill. **Before integrating Warp, read [the complete Warp guide](references/warp.md)** for prerequisites, native runtime setup, automatic/explicit capture, limitations, replay, and troubleshooting. No separate skill or LEAPP source checkout is required. + +- Use the matching `leapp[warp-cu12]` or `leapp[warp-cu13]` extra, then run `leapp-build-warp-runtime` and `leapp-build-warp-runtime --status` in that environment. +- Export capture currently requires Linux, a CUDA GPU, the matching CUDA toolkit, and Warp >=1.16. Verify current requirements in the online setup guide. +- Initialize Warp before `start(global_patching=True)`. Run the same Warp path at least twice: discovery then capture. A segment groups Warp calls inside a node; it is not a reason to split the full pipeline into multiple models. +- Use ONNX Dynamo by default; PT2 is the other supported Warp exporter. Deploy the corresponding native runtime dependency along with the bundle. + +## Validate the whole exported pipeline + +1. Compare the original application and annotated execution using the same inputs, weights, state, execution mode, and random seeds where relevant. Check function adapters against original functions separately. +2. Run `compile_graph(validate=True, strict=True)` with justified tolerances. Inspect warnings and generated YAML: expected external inputs, final outputs, forward/feedback edges, initial values, backend types, and existing artifact paths. Explain any unexpected dangling input. +3. Load the bundle with `InferenceManager` and compare **fresh representative inputs not used during tracing** against the original pipeline. A mock-input smoke test alone is insufficient. Include realistic value ranges and relevant branches/shapes; for stateful pipelines compare a sequence from matching initial/reset state. +4. Verify the runtime/device and any custom operators required by deployment. ONNX Runtime defaults to CPU-safe execution; GPU/Warp requirements must be checked in the actual environment. + +Use actual `manager.inputs` / `manager.outputs` keys (`node/port`). Inspect intermediate node values via `manager.value_dict`; seed feedback when needed with `set_input_value(node, port, value)` or a fresh manager. + +## Debugging controls + +- Read `//log.txt`; it includes DEBUG output even when the console is quiet. Search for `Compiled graph module for ` to inspect the captured FX graph: live inputs should be placeholders, expected math should be present, and constants should be intentional. Unused inputs may be pruned; investigate a missing input expected to affect the output. +- On a mismatch, use the reported node, output, sample index, shapes/dtypes, and error statistics. A first-sample pass followed by later failures suggests frozen values or state problems. `strict=False` exposes diagnostic results; it does not turn failures into successful validation. +- `start(verbose=True)` enables detailed tracing logs. Keep `max_cached_io` around `3–5` initially and expand when representative replay coverage needs it. +- `start(dry_run=True)` checks tracing/structure without model export or model validation. Current `compile_graph()` does not accept `dry_run`; inspect signatures on older installations. +- `start(non_traced=["node"])` skips selected exports while retaining I/O/edges. Use to isolate a failure, and restore real exports before declaring completion. +- `global_patching=False` can diagnose patch interference, but disables user `FunctionPatch` and automatic patch-dependent tracing; it is not a transparent compatibility fix. +- `visualize=False` avoids optional PNG rendering. Model export and validation remain required. + +## Deliver the bundle + +Deliver the complete export directory: model artifact(s), graph YAML, `*_initial_values.safetensors` when feedback exists, and all referenced model/external-data files. Preserve paths relative to the YAML. PNG and logs are useful diagnostics; they do not replace runnable artifacts. For Warp, document/build the matching native adapter and its discovery/override configuration on the destination. + +Confirm runtime input/output keys, shapes, dtypes, device expectations, semantics, state initialization/reset behavior, and supported input domain. Preserve YAML positional I/O order for consumers loading the model directly. Where practical, load and run a copied bundle from its delivery location to detect hidden source-path dependencies. + +Report the model count and any reasons for splitting, node boundaries and graph I/O, annotation/adapter changes, chosen backend per node and fallback attempts, generated bundle paths, fresh-input and stateful parity results, and any remaining runtime or input-domain limitations. Distinguish a complete runnable bundle from a traced graph, partial export, or diagnostic-only result. diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst index 631254ff..e8756f0b 100644 --- a/docs/source/api/index.rst +++ b/docs/source/api/index.rst @@ -233,7 +233,7 @@ Graph-level metadata passed to :func:`~leapp.compile_graph`. leapp.compile_graph( graph_configs=GraphConfigs( - frequency_hz=50, + frequency=50, extra={"runtime": "isaac_lab"}, ), ) diff --git a/docs/source/conf.py b/docs/source/conf.py index 7699554f..abee0867 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -4,6 +4,8 @@ import os import sys from datetime import datetime +from pathlib import Path +from zipfile import ZIP_DEFLATED, ZipFile sys.path.insert(0, os.path.abspath('../..')) @@ -93,3 +95,22 @@ html_css_files = ['css/leapp.css'] html_last_updated_fmt = '%b %d, %Y' + + +def package_agent_skill(app, exception): + """Bundle the repository skill and its references with the published HTML.""" + if exception is not None or app.builder.format != 'html': + return + + repo_root = Path(__file__).resolve().parents[2] + download_dir = Path(app.outdir) / '_downloads' + download_dir.mkdir(parents=True, exist_ok=True) + files = [repo_root / 'SKILL.md', *sorted((repo_root / 'references').rglob('*.md'))] + with ZipFile(download_dir / 'leapp-integrate.zip', 'w', ZIP_DEFLATED) as bundle: + for source in files: + name = (Path('leapp-integrate') / source.relative_to(repo_root)).as_posix() + bundle.writestr(name, source.read_bytes()) + + +def setup(app): + app.connect('build-finished', package_agent_skill) diff --git a/docs/source/index.rst b/docs/source/index.rst index 1dd9c1e2..4f2d3ffa 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -44,6 +44,26 @@ traced graph into portable artifacts. Read usage guide +Use LEAPP with a Coding Agent +============================= + +Download the LEAPP integration skill to help your coding agent annotate, +export, and validate a pipeline. The ZIP includes ``SKILL.md`` and supporting +guides for function patches, state, deployment semantics, and Warp. + +.. raw:: html + +

+ + Download agent skill (ZIP) + +

+ +Extract the ZIP and give your agent the ``leapp-integrate/SKILL.md`` file, +or place the ``leapp-integrate`` folder in your agent's skills directory. +Keep the ``references`` folder beside ``SKILL.md`` so the linked guides remain +available. You do not need a local LEAPP checkout to use the skill. + Is LEAPP a Fit? =============== diff --git a/references/function-patching.md b/references/function-patching.md new file mode 100644 index 00000000..cff155c1 --- /dev/null +++ b/references/function-patching.md @@ -0,0 +1,72 @@ +# Function patches for unsupported operations + +Use this when an existing module-level function cannot be traced, including native-library calls nested inside helpers. Preserve the original application implementation and register a small adapter in the integration layer. Do not rewrite the whole pipeline, replace every call site, or convert all NumPy code to Torch merely because one operation is unsupported. + +Consult the online [function-patching guide](https://nvidia-isaac.github.io/leapp/guides/patching.html) and [public API](https://nvidia-isaac.github.io/leapp/api/index.html); no source checkout is needed. NumPy, SciPy, scikit-learn, scikit-image, and OpenCV calls can return untraced values. Patch calls that actually break the trace, not every call from those libraries. Scikit estimator methods are not module-level targets; find the unsupported module function or a pure module-level helper enclosing the computation. + +## Contract + +```python +leapp.start( + name="pipeline", + patching=[leapp.FunctionPatch(module, "function_name", traceable_replacement)], +) +``` + +- `FunctionPatch` is a public `leapp` import. Check that it exists in the active installation before relying on it. Missing support is a version constraint, not a reason to invent a private patch API. +- The target is a callable attribute on a **Python module**, not an arbitrary object or bound instance method. Target the smallest unsupported function that explains the trace break. +- The adapter is called when arguments contain actively traced data (including nested supported containers); ordinary-array calls still use the original function. The adapter's traced operations become the exported computation. +- `global_patching=True` is required (the default). `stop()` restores session patches. Scope them with `try/finally`; do not leave a session active after an exception. +- Patches affect the whole Python process, including other threads. Avoid concurrent sessions with different patch lists. Stop and restart before changing patches. +- Calls must resolve the patched module attribute. References captured earlier by `from package import function`, default arguments, closures, or caches are not automatically rebound. Inspect the actual call path. Prefer patching the module attribute actually used, or a minimal integration-level binding adjustment, over changing every application call. + +## Adapter requirements + +Implement the same computation with supported traced Torch/NumPy operations. Preserve the signature used by the application, accepted argument combinations, output structure, shapes, dtype, device, layout, and numerical semantics (including relevant clipping, interpolation endpoints, or rounding). + +Adapters must be **pure**: mutation, `dst=`/out buffers, input/output aliasing, and calling the patched function from its own replacement are unsupported. Reject unsupported modes explicitly rather than silently approximating them. A patch is not an escape hatch for file I/O, simulator side effects, or stateful native code. + +Keep genuinely dynamic values as traced data; do not extract them into Python scalars or copy raw buffers. Do not call the original unsupported operation on traced inputs and then attach metadata to its result: that still loses the computation. + +## Example: preserve nested OpenCV callers + +Existing helpers may call `cv2.cvtColor(image, cv2.COLOR_BGR2RGB)` in several places. This adapter supports that precise three-channel HWC operation without changing those helpers: + +```python +import cv2 +import torch +import leapp +from leapp import annotate + + +def traced_cvt_color(src, code, dst=None, dstCn=0): + if code != cv2.COLOR_BGR2RGB or dst is not None or dstCn not in (0, 3): + raise ValueError("Only BGR2RGB without a destination buffer is supported") + if src.ndim != 3 or src.shape[-1] != 3: + raise ValueError("Expected a three-channel HWC image") + return torch.from_numpy(src).flip((-1,)).numpy() + + +leapp.start( + name="vision", + patching=[leapp.FunctionPatch(cv2, "cvtColor", traced_cvt_color)], +) +try: + image = annotate.input_tensors("preprocess", {"bgr": sample_bgr}) + features = existing_image_pipeline(image) # Existing helpers stay unchanged. + annotate.output_tensors("preprocess", {"features": features}, export_with="onnx") +finally: + leapp.stop() +leapp.compile_graph(validate=True, strict=True) +``` + +`sample_bgr` and `existing_image_pipeline` are application-provided. The online guide links a complete example; the [upstream example source](https://github.com/nvidia-isaac/leapp/blob/main/examples/function_patch.py) is also available without cloning LEAPP. Match examples to the package version used by the application. + +## Verify before accepting + +1. Compare the adapter with the original function on ordinary representative inputs, including relevant edge cases. Reject unsupported arguments deliberately. +2. Exercise the real nested application call path during tracing. Confirm it reaches the adapter and retains dependence on inputs. +3. Confirm ordinary calls still use the original behavior during the session and the module attribute is restored afterward, including failure cleanup where relevant. +4. Export and replay on **new values** against the unmodified original pipeline. Same-input validation can pass when an unsupported result was captured as a constant. + +If a faithful pure adapter is impractical, inspect whether a compatible exporter can handle an already correctly traced graph. Otherwise expose the operation as an explicit external runtime boundary or propose a narrowly scoped application change, with the unsupported behavior and parity requirements stated. Never label a partial graph as a complete export. diff --git a/references/state-and-semantics.md b/references/state-and-semantics.md new file mode 100644 index 00000000..13e7b6fa --- /dev/null +++ b/references/state-and-semantics.md @@ -0,0 +1,106 @@ +# State and deployment semantics + +Read this for recurrent policies, history windows, episode resets, robot-facing ports, or action chunks. Online references: [state](https://nvidia-isaac.github.io/leapp/guides/graph.html), [API](https://nvidia-isaac.github.io/leapp/api/index.html), [semantics](https://nvidia-isaac.github.io/leapp/semantics/usage.html), [axis names](https://nvidia-isaac.github.io/leapp/semantics/kind_element_names.html), [temporal axes](https://nvidia-isaac.github.io/leapp/semantics/temporal.html), [runtime](https://nvidia-isaac.github.io/leapp/leapp_runtime.html). + +## Explicit state across an episode + +`state_tensors` takes a flat dict of named Torch tensors inside an existing node. Nested state is unsupported; expose tensor leaves under distinct stable names. `update_state` uses those same names and preserves shape and dtype. An unupdated state remains an ordinary input. Constants belong in `annotate.register_buffer`, not feedback. + +This runnable example demonstrates state ownership and sequence parity without modifying the original step function: + +```python +import torch +import leapp +from leapp import annotate, InferenceManager + + +def original_step(obs, hidden): + next_hidden = torch.tanh(obs + 0.5 * hidden) + return 2.0 * next_hidden, next_hidden + + +samples = [torch.randn(1, 3) for _ in range(5)] +hidden = torch.zeros(1, 3) +expected = [] +for obs in samples: + action, hidden = original_step(obs, hidden) + expected.append(action.clone()) + +hidden = torch.zeros(1, 3) +leapp.start("stateful_graph", max_cached_io=5) +try: + for obs, reference in zip(samples, expected): + x = annotate.input_tensors("pipeline", {"obs": obs}) + h = annotate.state_tensors("pipeline", {"hidden": hidden}) + action, next_hidden = original_step(x, h) + annotate.update_state("pipeline", {"hidden": next_hidden}) + annotate.output_tensors("pipeline", {"action": action}, export_with="onnx") + torch.testing.assert_close(action, reference) + hidden = next_hidden +finally: + leapp.stop() +leapp.compile_graph(visualize=False, validate=True, strict=True) + +# Each new manager starts a new episode with the saved feedback initial values. +for episode in range(2): + manager = InferenceManager("stateful_graph/stateful_graph.yaml") + template = manager.get_mock_input()["pipeline/obs"] + hidden = torch.zeros_like(template) + for step in range(4): + fresh = torch.randn_like(template) + episode * 0.2 + expected_action, hidden = original_step(fresh, hidden) + actual = manager.run_policy({"pipeline/obs": fresh})["pipeline/action"] + torch.testing.assert_close(actual, expected_action.to(actual.device), rtol=1e-3, atol=1e-5) +``` + +For an actual application, preserve its initial state, reset rules, and mutation semantics. Clone inputs or use separate model instances when baseline and traced execution could mutate shared data. Do not reset every trace iteration if deployment carries state between ticks. Compare first-tick and later-tick behavior on fresh sequences. + +## Existing module state + +After opening a node, register the module before executing it: + +```python +annotate.module("pipeline", model, buffer_names=["h_state", "encoder.running_mean"]) +``` + +Replace the example names with real registered buffers; omit `buffer_names` to track all. Reassigned buffers become feedback; read-only buffers are constants. Plain attributes and in-place updates such as `self.h.copy_(next_h)` are not captured as automatic state updates. Use explicit state annotations where appropriate rather than rewriting the model or assuming a pure function patch captures mutation. + +Explicit state/module feedback can be declared in one trace. **Inferred cross-node feedback** needs at least two graph iterations in the same session. Validate `pipeline.feedback_flow` and `pipeline.initial_values`; carry the referenced safetensors file with the bundle. Use `manager.feedback_inputs` and `set_input_value(node, port, value)` to inspect/seed feedback. A fresh manager is the portable episode-reset recipe; if using an installed `reset()` method, verify its contract—buffer reset does not necessarily reset private state inside model objects. + +## Semantic deployment contract + +`TensorSemantics` works with `input_tensors` / `output_tensors`, not `method`. Pass one object or a homogeneous list directly. Do not put semantic wrappers inside a dict. Keep names unique per input/output set. + +This complete example names joint order and a four-step action chunk within a single model. In a real project, use the existing action computation, actual joint names, action quantity, and timing: + +```python +import torch +import leapp +from leapp import annotate, TensorSemantics, TemporalAxis, GraphConfigs +from leapp import InputKindEnum, OutputKindEnum + +joint_names = ["hip", "knee", "ankle"] +leapp.start("semantic_graph") +try: + q = annotate.input_tensors("pipeline", TensorSemantics( + "joint_pos", torch.randn(3), kind=InputKindEnum.JOINT_POSITION, + element_names=[joint_names], extra={"units": "rad"}, + )) + actions = q.unsqueeze(0).repeat(4, 1) + annotate.output_tensors("pipeline", TensorSemantics( + "joint_targets", actions, kind=OutputKindEnum.JOINT_POSITION, + element_names=[TemporalAxis(period_ms=20), joint_names], + extra={"units": "rad"}, + ), export_with="onnx") +finally: + leapp.stop() +leapp.compile_graph(visualize=False, validate=True, strict=True, + graph_configs=GraphConfigs(frequency=50)) +``` + +- Match `kind` to observed state, commanded references, or output targets. A custom string is allowed when no standard kind fits; confirm consumer support. +- `element_names` is an outer axis list with per-axis name lists; `None` leaves an axis unnamed. For `[batch, joints]`, use `[None, joint_names]`. Names must reflect actual ordering rather than a guessed robot convention. +- Put at most one `TemporalAxis` directly in that outer list. LEAPP emits `__temporal_axis__` plus `temporal_period_ms`; do not write the reserved sentinel yourself. The consumer must interpret how the chunk is applied. +- `GraphConfigs(frequency=...)` is in Hz; `TemporalAxis(period_ms=...)` is in ms. They are not automatically validated against one another; neither makes `InferenceManager` schedule calls in real time. +- `extra` entries are flattened into the tensor YAML mapping. Units, frames, and external IDs are project-specific fields, not automatic conversion or routing instructions. Agree on them with the consumer and avoid collisions with built-ins such as `name`, `shape`, `dtype`, or `kind`. +- Inspect generated YAML for expected names, order, shapes, temporal metadata, and graph configs, then verify that the deployment runtime consumes that contract. diff --git a/references/warp.md b/references/warp.md new file mode 100644 index 00000000..f1390de9 --- /dev/null +++ b/references/warp.md @@ -0,0 +1,121 @@ +# Warp inside a LEAPP pipeline + +This is the complete Warp integration reference bundled with `leapp-integrate`. Read it before tracing Warp; a separate skill or source checkout is unnecessary. Warp support is experimental. Consult the online [setup](https://nvidia-isaac.github.io/leapp/tensor_libraries/installation.html), [usage](https://nvidia-isaac.github.io/leapp/tensor_libraries/warp.html), [capture model](https://nvidia-isaac.github.io/leapp/tensor_libraries/warp_how_it_works.html), and [limitations](https://nvidia-isaac.github.io/leapp/tensor_libraries/warp_limitations.html) for the release used by the application. + +LEAPP records Warp regions as `leapp::warp_runner` custom operations, embedding APIC archives and compiled kernel modules in an ONNX or PT2 model. It does not translate kernel bodies to Torch operators. Keep the full compute pipeline in one node by default; several Warp regions and Torch operations can share that node. + +## Set up the export environment + +Current export/capture requirements are Linux, a CUDA GPU, Warp >=1.16, and a matching CUDA toolkit with `nvcc`. Native builds also need a C++17 compiler, CMake >=3.18, and network access to download ONNX Runtime headers. PTX is currently targeted at compute capability 7.5: older GPUs and kernels relying on newer hardware features are outside this export contract. Verify version-specific requirements before installing. + +Use the application's environment and package manager. A default `leapp` installation does not install Warp support. Select **one** extra matching the CUDA toolkit (not merely the driver's maximum supported CUDA version): + +```bash +python -m pip install "leapp[warp-cu12]" # CUDA 12; respect the project's LEAPP version pin +# For CUDA 13, use "leapp[warp-cu13]" instead. +leapp-build-warp-runtime +leapp-build-warp-runtime --status +``` + +For an actual LEAPP source checkout managed by uv, `uv sync --extra test --extra warp-cu12` (or `warp-cu13`) is the alternative. Do not assume a downstream project defines these extras. + +The CUDA extras select compatible `cupti-python` and `onnxruntime-gpu` packages with CUDA/cuDNN runtime dependencies. Do not independently install a mismatched CUPTI major or add a competing CPU ONNX Runtime package. CUDA 13's GPU runtime package requires Python >=3.11. Check `nvcc --version`, `python -m pip check`, and `onnxruntime.get_available_providers()` in the export environment; Warp ONNX replay needs `CUDAExecutionProvider`. + +The build command builds **both ONNX and PT2 native adapters**, checking rather than installing dependencies. Default discovery uses the user cache: + +- Linux: `${XDG_CACHE_HOME:-~/.cache}/leapp/warp-runtime/build` +- Windows native runtime builds: `%LOCALAPPDATA%\leapp\warp-runtime\build` + +No environment variable is required for a valid cached build. For isolation between environments, overrides are available: + +```bash +export LEAPP_WARP_ONNX_CUSTOM_OP_LIBRARY="/path/to/libleapp_wrp_onnx_custom_op.so" +export LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY="/path/to/libleapp_wrp_torch_custom_op.so" +``` + +Use actual built libraries in place of these example paths. An invalid override is an error, not a fallback to the cache. The shared cache is not versioned per environment. After changing Torch, Warp, ONNX Runtime, or CUDA, rebuild and restart processes that loaded the libraries. Account for other environments using that cache before replacing a shared build. + +Windows library builds use corresponding `.dll` names and need `LEAPP_WARP_IMPORT_LIBRARY` pointing to `warp.lib`, which the Warp wheel does not ship. This build capability does **not** enable Linux-only LEAPP tracing on Windows. Capture uses CUPTI; replay of an exported bundle is a separate dependency check and is not inherently CUPTI-dependent. + +## Capture and export + +1. Initialize Warp with `wp.init()` before `leapp.start(..., global_patching=True)` and before compiling application kernels under incompatible settings. Keep global patching enabled. +2. Pass application inputs through `annotate.input_tensors` and use returned values. Public Warp calls involving actively traced Warp arrays participate in capture. Work on unrelated raw arrays may be absent from the export. +3. Execute the annotated path **at least twice** before stopping: discovery, then APIC capture. Calls and region ordering must agree. Repeated identical inputs alone are not a correctness test; also validate fresh values on the supported path. +4. Export with `"onnx"` / `"onnx-dynamo"` by default. `"pt2"` / `"exported-program"` is the other supported route. Legacy ONNX and JIT do not support these Warp bundles. + +Save this runnable example to a Python file so Warp can inspect its kernel source. It assumes setup is complete and CUDA device 0 is available: + +```python +import torch +import warp as wp +import leapp +from leapp import annotate, InferenceManager + + +@wp.kernel +def scale_kernel(src: wp.array(dtype=wp.float32), dst: wp.array(dtype=wp.float32)): + i = wp.tid() + dst[i] = 2.0 * src[i] + + +wp.init() +leapp.start("warp_graph", global_patching=True) +try: + for _ in range(2): + source = wp.ones(8, dtype=wp.float32, device="cuda:0") + x = annotate.input_tensors("pipeline", {"x": source}) + with annotate.warp_op("pipeline"): + y = wp.empty_like(x) + wp.launch(scale_kernel, dim=x.size, inputs=[x, y], device="cuda:0") + result = torch.tanh(wp.to_torch(y)) + annotate.output_tensors("pipeline", {"result": result}, export_with="onnx") +finally: + leapp.stop() +leapp.compile_graph(visualize=False, validate=True, strict=True) + +manager = InferenceManager("warp_graph/warp_graph.yaml") +inputs = manager.get_mock_input() +fresh = torch.linspace(-1, 1, 8, device=inputs["pipeline/x"].device) +inputs["pipeline/x"] = fresh +actual = manager.run_policy(inputs)["pipeline/result"] +torch.testing.assert_close(actual, torch.tanh(2 * fresh).to(actual.device), rtol=1e-3, atol=1e-5) +``` + +In the application, use its existing kernels and compare replay with the original Warp pipeline on fresh inputs; the Torch expression above is just a simple oracle for this example. + +## Automatic versus explicit regions + +With global patching, public Warp calls touching active traced Warp values automatically open/continue a region. `wp.to_torch()`, Warp `.numpy()`, synchronization, and consumption by traced Torch operations close automatic regions. `wp.from_torch()` / `wp.from_numpy()` do not themselves close them. Use automatic regions when they match existing code; an explicit context is optional. + +`annotate.warp_op("node")` protects a region until context exit. Only its node name currently controls capture; `inputs`, `outputs`, and extra keyword arguments do not select its boundary. Capture currently uses `cuda:0`, so this context is not a device-selection API. Keep synchronization, host readback, and unrelated CUDA work outside explicit regions. + +A Warp operation must not consume active arrays owned by different nodes. Finish producers and rebind inputs in the consumer first. Equivalent full-range `wp.copy` can preserve a finished node's output identity; transformations/partial copies belong inside annotated compute. + +## Constraints and process effects + +- Boundary arrays need scalar or densely packed homogeneous dtypes. Logical `(B, N)` `wp.vec3` arrays are represented as `(B, N, 3)` float32. Heterogeneous `@wp.struct` arrays are unsupported at these boundaries. +- CPU Warp arrays can execute eagerly but cannot replay in the exported CUDA runner. Successful discovery does not prove export compatibility. +- Private `warp._*` entrypoints are not intercepted. Kernel-language symbols inside `@wp.kernel` belong to the captured kernel, not independent Python patches. +- Data-dependent changes to the sequence of host Warp calls can diverge between passes. This concerns captured call structure, not every conditional inside a kernel body. +- Capture changes process-wide PTX configuration and Torch CUDA caching allocator behavior that can persist after `stop()`. Prefer an isolated export process for a long-lived application; stopping does not restore every CUDA/Warp setting. +- Driver compilation of embedded PTX can add first-load latency; separate it from steady-state replay when measuring performance. + +## Troubleshooting and delivery + +| Symptom | Check | +|---------|-------| +| No Warp backend / runtime missing | Matching extra, `leapp-build-warp-runtime --status`, initialization before start, global patching, and valid overrides. | +| CUPTI initialization fails | CUPTI and toolkit major versions match; do not randomly upgrade packages. | +| Missing second APIC capture pass | Execute the same annotated path twice in one session. Dry-run still needs a second discovery-path replay even though APIC capture is skipped. | +| More regions / segment diverged | Compare host call sequence and boundaries between passes; reset relevant application state without changing intended semantics. | +| Unsupported export backend | Use ONNX Dynamo or PT2; legacy ONNX is not a Warp fallback. | +| Missing CUDA provider / CPU validation fails | Check GPU runtime/provider and device placement; CPU fallback cannot run Warp. | +| Cross-context / owner-token error | Keep one node per active operation; leave explicit regions before another node starts. | +| Replay mismatch | Inspect inputs, mutations, region boundaries, and logs; compare fresh inputs with the original before changing tolerances. | + +Deliver the full LEAPP bundle plus native-runtime setup information. Embedded kernels do not make the model self-contained: the destination needs the adapter for its ONNX/PT2 backend and compatible CUDA/runtime dependencies. Test loading and replay there when available. Do not claim portability from a CPU-only smoke test. + +## Only when maintaining LEAPP itself + +For internal changes, use the tracing-maintenance skill if available. Traced Torch operations (including views/indexing) must close a pending automatic Warp region **before eager execution or extracting FX argument proxies**. Closing only at FX insertion can preserve stale proxies; view operations may have no CUDA callback. Internal paused work and explicit owner-protected regions retain their lifecycle. In source checkouts, relevant regression coverage includes `tests/unit_tests/test_warp_torch_boundaries.py` and `test_torch_view_consumes_pending_warp_output` in the Warp functional suite. This is maintainer guidance, not an application rewrite requirement. From 59d2d352f7ec8efe20849b7d0f791cd90eeefd58 Mon Sep 17 00:00:00 2001 From: Frank Lai Date: Tue, 13 Oct 2026 17:33:45 -0700 Subject: [PATCH 02/21] docs: simplify skill banner and separate state from semantics Signed-off-by: Frank Lai --- SKILL.md | 18 ++-- docs/source/_static/css/leapp.css | 82 +++++++++++------ docs/source/index.rst | 50 +++++------ references/semantics.md | 90 +++++++++++++++++++ .../{state-and-semantics.md => state.md} | 44 +-------- 5 files changed, 183 insertions(+), 101 deletions(-) create mode 100644 references/semantics.md rename references/{state-and-semantics.md => state.md} (53%) diff --git a/SKILL.md b/SKILL.md index 07d1a32f..698681e6 100644 --- a/SKILL.md +++ b/SKILL.md @@ -20,7 +20,8 @@ Start with the [online LEAPP documentation](https://nvidia-isaac.github.io/leapp | Public signatures and annotation patterns | [API](https://nvidia-isaac.github.io/leapp/api/index.html), [nodes](https://nvidia-isaac.github.io/leapp/guides/nodes.html) | | Export formats, parameters, prebuilt models | [Export](https://nvidia-isaac.github.io/leapp/guides/export.html) | | Unsupported functions and trace loss | [Function patches](https://nvidia-isaac.github.io/leapp/guides/patching.html), [Torch limitations](https://nvidia-isaac.github.io/leapp/tensor_libraries/torch_limitations.html), [NumPy limitations](https://nvidia-isaac.github.io/leapp/tensor_libraries/limitations.html) | -| State and metadata | [State](https://nvidia-isaac.github.io/leapp/guides/graph.html), [semantics](https://nvidia-isaac.github.io/leapp/semantics/usage.html), [temporal axes](https://nvidia-isaac.github.io/leapp/semantics/temporal.html) | +| State and resets | [State](https://nvidia-isaac.github.io/leapp/guides/graph.html) | +| Optional deployment metadata | [Semantics](https://nvidia-isaac.github.io/leapp/semantics/usage.html), [standard kinds](https://nvidia-isaac.github.io/leapp/semantics/kind_element_names.html), [temporal axes](https://nvidia-isaac.github.io/leapp/semantics/temporal.html) | | Validation and deployment | [Debugging](https://nvidia-isaac.github.io/leapp/guides/debugging.html), [runtime](https://nvidia-isaac.github.io/leapp/leapp_runtime.html), [bundle/YAML](https://nvidia-isaac.github.io/leapp/generated_configs.html) | | Warp | [Setup](https://nvidia-isaac.github.io/leapp/tensor_libraries/installation.html), [capture](https://nvidia-isaac.github.io/leapp/tensor_libraries/warp.html), [limitations](https://nvidia-isaac.github.io/leapp/tensor_libraries/warp_limitations.html) | @@ -54,7 +55,7 @@ Avoid legacy private `annotate._method` for new integrations; investigate public ## Complete annotation and validation recipe -This runnable stateless example keeps an original function unchanged, caches multiple trace samples, exports one model, and checks fresh inputs. Replace the illustrative function and sample generator with the application's existing compute and representative data. For mutable pipelines, isolate baseline inputs/state and reset both executions to the same initial state; see [state and deployment semantics](references/state-and-semantics.md). +This runnable stateless example keeps an original function unchanged, caches multiple trace samples, exports one model, and checks fresh inputs. Replace the illustrative function and sample generator with the application's existing compute and representative data. For mutable pipelines, isolate baseline inputs/state and reset both executions to the same initial state; see [state and episode lifecycle](references/state.md). ```python import torch @@ -131,7 +132,7 @@ Choose opset and optimizations for the actual runtime. Disabling validation, `sk ## State, constants, and loops -For recurrent policies, history buffers, or episode resets, read [state and deployment semantics](references/state-and-semantics.md). +For recurrent policies, history buffers, or episode resets, read [state and episode lifecycle](references/state.md). - `state_tensors` requires an existing node and a **flat dict of named Torch tensors**; nested state is unsupported. `update_state` must use declared names and preserve each state's shape and dtype. One entry returns a value; multiple entries return a tuple in insertion order. - Only states passed to `update_state()` become explicit feedback; an unupdated state remains a regular input. Keep evolving state distinct from frozen parameters and constants. @@ -162,14 +163,11 @@ Calls in NumPy, SciPy, scikit-learn, scikit-image, OpenCV, or other libraries ** Supported NumPy/Torch conversions such as `torch.from_numpy()` and `.numpy()` can preserve tracing with global patching enabled. Do not rewrite all NumPy code or treat every conversion as a trace break. -## Semantic metadata +## Optional deployment semantics -For deployment-facing tensors, add known meaning with `TensorSemantics`: observed state versus output targets, real joint/feature order, and temporal axes for action chunks. Read the runnable example in [state and deployment semantics](references/state-and-semantics.md). +Add semantics when a user requests them or an identified downstream consumer needs tensor meaning to bind observations, map joint order, interpret commands, or apply action chunks. They are not required for tracing, export, numerical validation, or state feedback. Default to named raw tensors for ordinary internal node connections; do not annotate every intermediate value or add nodes solely for metadata. -- Semantic annotations use `input_tensors` / `output_tensors`, **not `method`**. Pass a single `TensorSemantics` or a list directly; no dict wrapping or mixing with raw tensors in that list. Names must be unique within the inputs or outputs. -- `element_names` describes axis ordering; use `None` for unnamed dimensions. Put at most one `TemporalAxis(period_ms=...)` directly in the outer axis list, not inside a nested list. -- `extra` fields such as units/frames are flattened into YAML and are project-specific. Avoid collisions with built-in fields and confirm the downstream consumer understands them. -- Use `GraphConfigs(frequency=...)` for graph execution frequency in Hz. Temporal sample spacing is in milliseconds and is independent of graph frequency; LEAPP does not validate their consistency or schedule the application from this metadata. +When needed, read [deployment semantics](references/semantics.md). **Prefer standard `InputKindEnum` / `OutputKindEnum` kinds** over custom strings or `extra`. Keep annotations compact: reuse existing ordered names, specify only consumed fields, and avoid repeated wrappers or metadata factories. Do not guess physical meaning for normalized features, latent vectors, or mixed observations. State/reset contracts are covered separately in [state](references/state.md). ## Warp setup and export (experimental) @@ -203,6 +201,6 @@ Use actual `manager.inputs` / `manager.outputs` keys (`node/port`). Inspect inte Deliver the complete export directory: model artifact(s), graph YAML, `*_initial_values.safetensors` when feedback exists, and all referenced model/external-data files. Preserve paths relative to the YAML. PNG and logs are useful diagnostics; they do not replace runnable artifacts. For Warp, document/build the matching native adapter and its discovery/override configuration on the destination. -Confirm runtime input/output keys, shapes, dtypes, device expectations, semantics, state initialization/reset behavior, and supported input domain. Preserve YAML positional I/O order for consumers loading the model directly. Where practical, load and run a copied bundle from its delivery location to detect hidden source-path dependencies. +Confirm runtime input/output keys, shapes, dtypes, device expectations, semantics when used, state initialization/reset behavior, and supported input domain. Preserve YAML positional I/O order for consumers loading the model directly. Where practical, load and run a copied bundle from its delivery location to detect hidden source-path dependencies. Report the model count and any reasons for splitting, node boundaries and graph I/O, annotation/adapter changes, chosen backend per node and fallback attempts, generated bundle paths, fresh-input and stateful parity results, and any remaining runtime or input-domain limitations. Distinguish a complete runnable bundle from a traced graph, partial export, or diagnostic-only result. diff --git a/docs/source/_static/css/leapp.css b/docs/source/_static/css/leapp.css index 3dbbc01f..178e2dad 100644 --- a/docs/source/_static/css/leapp.css +++ b/docs/source/_static/css/leapp.css @@ -140,45 +140,77 @@ a.reference.external::after { } -h1:has(+ .leapp-github-link) { - display: inline-block; - margin-right: 0.8rem; +/* Keep the Sphinx document title for navigation; the banner renders its heading. */ +h1:has(+ .leapp-hero) { + display: none; } -.leapp-github-link { - display: inline-flex; - margin: 0 0 1rem; - vertical-align: 0.25rem; +.leapp-hero { + display: grid; + grid-template-columns: minmax(0, 1fr) auto; + align-items: center; + gap: 2rem; + margin: 0 0 2rem; + padding: 0 0 1.75rem; + border-bottom: 1px solid var(--pst-color-border, #d1d5da); } -.leapp-github-link a { +.leapp-hero h1 { + margin: 0 0 0.65rem; + font-size: clamp(2.5rem, 5vw, 3rem); + font-weight: 700; + line-height: 1.1; +} + +.leapp-hero-copy > p { + margin: 0; + font-size: 1.05rem; + line-height: 1.55; +} + +.leapp-hero-action { + display: grid; + gap: 0.75rem; +} + +.leapp-hero-button { + display: flex; align-items: center; - color: #0969da; - display: inline-flex; - font-family: inherit; + justify-content: center; + gap: 0.6rem; + margin: 0; + padding: 0.8rem 1.1rem; + border-radius: 4px; font-size: 0.95rem; - font-weight: 700; - gap: 0.32rem; + font-weight: 600; + line-height: 1.4; text-decoration: none; + white-space: nowrap; } -.leapp-github-link a:hover { - color: #054da7; - text-decoration: underline; +.leapp-hero-button img, +.leapp-hero-button svg { + flex-shrink: 0; + width: 20px; + height: 20px; +} + +.leapp-hero-button img { + filter: brightness(0) invert(1); } -.leapp-github-link img { - height: 0.95rem; - width: 0.95rem; +.leapp-hero a:focus-visible { + outline: 3px solid var(--pst-color-primary, #006eaf); + outline-offset: 4px; } -@media (max-width: 520px) { - h1:has(+ .leapp-github-link), - .leapp-github-link { - display: block; +@media (max-width: 700px) { + .leapp-hero { + grid-template-columns: minmax(0, 1fr); + gap: 1.5rem; } - .leapp-github-link { - margin: -0.4rem 0 1rem; + .leapp-hero-button { + white-space: normal; } } diff --git a/docs/source/index.rst b/docs/source/index.rst index 4f2d3ffa..5dd6099a 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -4,14 +4,24 @@ LEAPP .. raw:: html -

- - - nvidia-isaac/leapp - -

- -**Lightweight Export Annotations for Policy Pipelines** +
+
+

LEAPP

+

Lightweight Export Annotations
for Policy Pipelines

+
+
+ + + GitHub repository + + + + Download agent skill + +
+
Export the whole policy pipeline, not just the model weights. @@ -44,25 +54,13 @@ traced graph into portable artifacts. Read usage guide -Use LEAPP with a Coding Agent -============================= - -Download the LEAPP integration skill to help your coding agent annotate, -export, and validate a pipeline. The ZIP includes ``SKILL.md`` and supporting -guides for function patches, state, deployment semantics, and Warp. - -.. raw:: html - -

- - Download agent skill (ZIP) - -

+.. dropdown:: Using the agent skill -Extract the ZIP and give your agent the ``leapp-integrate/SKILL.md`` file, -or place the ``leapp-integrate`` folder in your agent's skills directory. -Keep the ``references`` folder beside ``SKILL.md`` so the linked guides remain -available. You do not need a local LEAPP checkout to use the skill. + The download includes ``SKILL.md`` and supporting guides for function patches, + state, deployment semantics, and Warp. Extract the ZIP and give your agent + ``leapp-integrate/SKILL.md``, or place the ``leapp-integrate`` folder in your + agent's skills directory. Keep the ``references`` folder beside ``SKILL.md`` + so the linked guides remain available. No local LEAPP checkout is needed. Is LEAPP a Fit? =============== diff --git a/references/semantics.md b/references/semantics.md new file mode 100644 index 00000000..d5e99378 --- /dev/null +++ b/references/semantics.md @@ -0,0 +1,90 @@ +# Optional deployment semantics + +Read this when a user requests semantic metadata or a deployment consumer needs it. For recurrent buffers and episode resets, use [state](state.md) independently. + +Online references: [usage and passing conventions](https://nvidia-isaac.github.io/leapp/semantics/usage.html), [standard kinds and element names](https://nvidia-isaac.github.io/leapp/semantics/kind_element_names.html), [temporal axes](https://nvidia-isaac.github.io/leapp/semantics/temporal.html), [API](https://nvidia-isaac.github.io/leapp/api/index.html). Check these for the current API; a local source checkout is not required. + +## Decide whether metadata has a consumer + +Semantics describe what a tensor means beyond its name, shape, and dtype. They let a deployment integration distinguish measured joint positions from target commands, match elements to actuators, or interpret a sequence of actions. They do not change numerical computation, repair tracing, or implement state feedback. + +Before adding a field, identify the consumer and the decision it enables. Apply metadata primarily to external graph inputs and outputs where this contract is useful. An internal port needs semantics only when a specific consumer actually inspects it, such as an independently deployed component. A node boundary alone is not a reason to add them. + +| Situation | Guidance | +| --- | --- | +| Robot adapter needs measured positions and actuator target order | Add the corresponding standard kinds and actual joint names at the external ports. | +| Controller consumes an action chunk | Describe its actual time axis and spacing when the controller supports that contract. | +| Preprocessor passes normalized features to an internal policy node | Use a named raw tensor; no semantic wrapper is needed. | +| Hidden state or history stays inside the exported pipeline | Capture state as needed; do not add semantics just because it persists. | +| Consumer already binds named tensors and needs no semantic metadata | Keep named raw tensors, including at graph inputs and outputs. | + +Do not split a model, convert a working `method` annotation, or label every internal edge merely to populate metadata. Preserve requested deployment contracts, but do not invent one when the task is simply export and replay. + +## Use standard kinds first + +Choose the documented `InputKindEnum` or `OutputKindEnum` member matching the real quantity and role. Prefer the enum over a hand-written equivalent string; standard kinds give consumers a shared vocabulary. For example, `InputKindEnum.JOINT_POSITION` describes measured positions, `InputKindEnum.COMMAND_JOINT_POSITION` an incoming reference, and `OutputKindEnum.JOINT_POSITION` a produced target. A tensor named `action` is not necessarily a torque command. + +Use a custom kind only when no standard kind fits **and** the consumer explicitly supports the agreed label. Do not create synonyms for existing kinds or repeat `kind` in `extra`. Unknown meaning should remain unspecified. In particular, normalized joint features, learned latent vectors, and concatenated mixed observations must not be mislabeled as physical joint positions. + +## Keep the annotation small + +- Start with named raw tensors. Add `TensorSemantics` only to ports needing metadata and include only fields the consumer uses; omit empty/default fields. +- Reuse ordered joint/feature names from the application's existing configuration. Define a shared list once if necessary; do not paste large lists into each annotation or reconstruct guessed ordering. +- Keep short wrappers inline, with one readable entry per port. A small import alias is fine if repeated names dominate the code; avoid generic builders, classes, or registries solely to shorten a few calls. +- Keep capture and export lifecycle calls in their existing locations. Metadata should add a small amount of code around the relevant boundary, not duplicate the pipeline. +- Use standard fields before `extra`. Do not add units, frames, IDs, or timing just to make the export look complete. Retain every field genuinely required by the consumer even if it takes more lines. + +For an existing `[batch, joints]` output whose consumer needs target positions and joint order, the annotation can stay compact (`actions` and `joint_names` come from the application): + +```python +from leapp import annotate, TensorSemantics, OutputKindEnum + +annotate.output_tensors("pipeline", TensorSemantics( + "joint_targets", actions, kind=OutputKindEnum.JOINT_POSITION, + element_names=[None, joint_names], +), export_with="onnx") +``` + +## Passing and axis contracts + +`TensorSemantics` works with `input_tensors` / `output_tensors`, not `method`. Pass one object or a homogeneous list directly, never wrappers inside a dict or mixed with raw tensors in a list. Use separate calls for semantic and raw ports when needed; keep names unique within each node's inputs or outputs. Semantic fields are optional, so an empty wrapper adds no value over a named raw tensor. + +`element_names` uses an outer axis list with per-axis name lists; `None` leaves an axis unnamed. For `[batch, joints]`, use `[None, joint_names]`; for `[batch, time, joints]`, use `[None, TemporalAxis(period_ms=...), joint_names]`. Match the actual tensor layout and element order. Do not name batch indices or invent names for latent dimensions to fill the schema. + +Use at most one `TemporalAxis` directly in the outer list. LEAPP emits `__temporal_axis__` and `temporal_period_ms`; do not write the reserved sentinel yourself. Include it when the consumer needs sample timing, not on every tensor with multiple dimensions. + +`GraphConfigs(frequency=...)` records graph execution frequency in Hz; `TemporalAxis(period_ms=...)` records sample spacing in milliseconds. They describe different things and need not be reciprocals (for example, overlapping action chunks). LEAPP does not validate their consistency or make `InferenceManager` schedule calls in real time. Supply actual deployment timing only when needed, never a guessed default. + +## Example when a controller needs action chunks + +This runnable example illustrates an external joint-position contract with four future targets. It is not a template requiring all pipelines to add names, timing, or graph configs. In a real project, keep the existing computation and use its actual joint order, quantity, and timing. No custom fields are needed here. + +```python +import torch +import leapp +from leapp import annotate, TensorSemantics, TemporalAxis, GraphConfigs +from leapp import InputKindEnum, OutputKindEnum + +joint_names = ["hip", "knee", "ankle"] +leapp.start("semantic_graph") +try: + q = annotate.input_tensors("pipeline", TensorSemantics( + "joint_pos", torch.randn(3), kind=InputKindEnum.JOINT_POSITION, + element_names=[joint_names], + )) + actions = q.unsqueeze(0).repeat(4, 1) + annotate.output_tensors("pipeline", TensorSemantics( + "joint_targets", actions, kind=OutputKindEnum.JOINT_POSITION, + element_names=[TemporalAxis(period_ms=20), joint_names], + ), export_with="onnx") +finally: + leapp.stop() +leapp.compile_graph(visualize=False, validate=True, strict=True, + graph_configs=GraphConfigs(frequency=50)) +``` + +## Custom fields and validation + +Use `extra` only for agreed consumer-specific metadata that standard fields cannot express. Entries are flattened into the tensor YAML mapping; graph-level extras are flattened into `pipeline.configs`. Avoid collisions with built-ins such as `name`, `shape`, `dtype`, `type`, `kind`, `element_names`, `temporal_period_ms`, or `__temporal_axis__`. Units, frames, and external IDs are not automatic conversion or routing instructions; arbitrary fields do not become portable merely by appearing in YAML. + +Inspect the exported YAML for the intended standard kinds, actual element order and axes, and any required timing/config fields. Confirm that the deployment consumer recognizes the contract. Numerical parity tests alone cannot detect mislabeled units, joint order, or physical meaning. Report missing consumer information instead of filling it with plausible metadata. diff --git a/references/state-and-semantics.md b/references/state.md similarity index 53% rename from references/state-and-semantics.md rename to references/state.md index 13e7b6fa..3f8d2cef 100644 --- a/references/state-and-semantics.md +++ b/references/state.md @@ -1,6 +1,8 @@ -# State and deployment semantics +# State and episode lifecycle -Read this for recurrent policies, history windows, episode resets, robot-facing ports, or action chunks. Online references: [state](https://nvidia-isaac.github.io/leapp/guides/graph.html), [API](https://nvidia-isaac.github.io/leapp/api/index.html), [semantics](https://nvidia-isaac.github.io/leapp/semantics/usage.html), [axis names](https://nvidia-isaac.github.io/leapp/semantics/kind_element_names.html), [temporal axes](https://nvidia-isaac.github.io/leapp/semantics/temporal.html), [runtime](https://nvidia-isaac.github.io/leapp/leapp_runtime.html). +Read this for recurrent policies, history windows, or episode resets. State controls values carried between calls; semantic metadata describes tensor meaning and is optional independently of state. See [deployment semantics](semantics.md) only when a consumer needs that metadata. + +Online references: [state](https://nvidia-isaac.github.io/leapp/guides/graph.html), [API](https://nvidia-isaac.github.io/leapp/api/index.html), [runtime](https://nvidia-isaac.github.io/leapp/leapp_runtime.html). ## Explicit state across an episode @@ -66,41 +68,3 @@ annotate.module("pipeline", model, buffer_names=["h_state", "encoder.running_mea Replace the example names with real registered buffers; omit `buffer_names` to track all. Reassigned buffers become feedback; read-only buffers are constants. Plain attributes and in-place updates such as `self.h.copy_(next_h)` are not captured as automatic state updates. Use explicit state annotations where appropriate rather than rewriting the model or assuming a pure function patch captures mutation. Explicit state/module feedback can be declared in one trace. **Inferred cross-node feedback** needs at least two graph iterations in the same session. Validate `pipeline.feedback_flow` and `pipeline.initial_values`; carry the referenced safetensors file with the bundle. Use `manager.feedback_inputs` and `set_input_value(node, port, value)` to inspect/seed feedback. A fresh manager is the portable episode-reset recipe; if using an installed `reset()` method, verify its contract—buffer reset does not necessarily reset private state inside model objects. - -## Semantic deployment contract - -`TensorSemantics` works with `input_tensors` / `output_tensors`, not `method`. Pass one object or a homogeneous list directly. Do not put semantic wrappers inside a dict. Keep names unique per input/output set. - -This complete example names joint order and a four-step action chunk within a single model. In a real project, use the existing action computation, actual joint names, action quantity, and timing: - -```python -import torch -import leapp -from leapp import annotate, TensorSemantics, TemporalAxis, GraphConfigs -from leapp import InputKindEnum, OutputKindEnum - -joint_names = ["hip", "knee", "ankle"] -leapp.start("semantic_graph") -try: - q = annotate.input_tensors("pipeline", TensorSemantics( - "joint_pos", torch.randn(3), kind=InputKindEnum.JOINT_POSITION, - element_names=[joint_names], extra={"units": "rad"}, - )) - actions = q.unsqueeze(0).repeat(4, 1) - annotate.output_tensors("pipeline", TensorSemantics( - "joint_targets", actions, kind=OutputKindEnum.JOINT_POSITION, - element_names=[TemporalAxis(period_ms=20), joint_names], - extra={"units": "rad"}, - ), export_with="onnx") -finally: - leapp.stop() -leapp.compile_graph(visualize=False, validate=True, strict=True, - graph_configs=GraphConfigs(frequency=50)) -``` - -- Match `kind` to observed state, commanded references, or output targets. A custom string is allowed when no standard kind fits; confirm consumer support. -- `element_names` is an outer axis list with per-axis name lists; `None` leaves an axis unnamed. For `[batch, joints]`, use `[None, joint_names]`. Names must reflect actual ordering rather than a guessed robot convention. -- Put at most one `TemporalAxis` directly in that outer list. LEAPP emits `__temporal_axis__` plus `temporal_period_ms`; do not write the reserved sentinel yourself. The consumer must interpret how the chunk is applied. -- `GraphConfigs(frequency=...)` is in Hz; `TemporalAxis(period_ms=...)` is in ms. They are not automatically validated against one another; neither makes `InferenceManager` schedule calls in real time. -- `extra` entries are flattened into the tensor YAML mapping. Units, frames, and external IDs are project-specific fields, not automatic conversion or routing instructions. Agree on them with the consumer and avoid collisions with built-ins such as `name`, `shape`, `dtype`, or `kind`. -- Inspect generated YAML for expected names, order, shapes, temporal metadata, and graph configs, then verify that the deployment runtime consumes that contract. From 465df4f2b1a4b35b6c51bcef410d1250bf19e44b Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Wed, 14 Oct 2026 10:21:26 -0700 Subject: [PATCH 03/21] feat: add Rust bundle inspection API Signed-off-by: Marc Alwan --- packages/leapp-runtime-rust/.gitignore | 6 + packages/leapp-runtime-rust/Cargo.lock | 365 ++++++++++ packages/leapp-runtime-rust/Cargo.toml | 28 + packages/leapp-runtime-rust/README.md | 67 ++ packages/leapp-runtime-rust/src/error.rs | 90 +++ packages/leapp-runtime-rust/src/lib.rs | 148 ++++ packages/leapp-runtime-rust/src/load.rs | 305 ++++++++ packages/leapp-runtime-rust/src/schema.rs | 183 +++++ packages/leapp-runtime-rust/src/validate.rs | 128 ++++ packages/leapp-runtime-rust/tests/bundle.rs | 682 ++++++++++++++++++ .../tests/fixtures/feedback.yaml | 19 + .../tests/fixtures/pipeline.yaml | 53 ++ .../tests/fixtures/policy.yaml | 49 ++ 13 files changed, 2123 insertions(+) create mode 100644 packages/leapp-runtime-rust/.gitignore create mode 100644 packages/leapp-runtime-rust/Cargo.lock create mode 100644 packages/leapp-runtime-rust/Cargo.toml create mode 100644 packages/leapp-runtime-rust/README.md create mode 100644 packages/leapp-runtime-rust/src/error.rs create mode 100644 packages/leapp-runtime-rust/src/lib.rs create mode 100644 packages/leapp-runtime-rust/src/load.rs create mode 100644 packages/leapp-runtime-rust/src/schema.rs create mode 100644 packages/leapp-runtime-rust/src/validate.rs create mode 100644 packages/leapp-runtime-rust/tests/bundle.rs create mode 100644 packages/leapp-runtime-rust/tests/fixtures/feedback.yaml create mode 100644 packages/leapp-runtime-rust/tests/fixtures/pipeline.yaml create mode 100644 packages/leapp-runtime-rust/tests/fixtures/policy.yaml diff --git a/packages/leapp-runtime-rust/.gitignore b/packages/leapp-runtime-rust/.gitignore new file mode 100644 index 00000000..128abe75 --- /dev/null +++ b/packages/leapp-runtime-rust/.gitignore @@ -0,0 +1,6 @@ +/target/ +/.venv/ +__pycache__/ +.pytest_cache/ +*.so +*.dylib diff --git a/packages/leapp-runtime-rust/Cargo.lock b/packages/leapp-runtime-rust/Cargo.lock new file mode 100644 index 00000000..aa8bc8a5 --- /dev/null +++ b/packages/leapp-runtime-rust/Cargo.lock @@ -0,0 +1,365 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "bitflags" +version = "2.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cfg-if" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown", + "serde", + "serde_core", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "leapp-runtime" +version = "0.1.0" +dependencies = [ + "anyhow", + "indexmap", + "safetensors", + "serde", + "serde_json", + "serde_path_to_error", + "serde_yaml_ng", + "sha2", + "tempfile", +] + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rustix" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "safetensors" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "172dd94c5a87b5c79f945c863da53b2ebc7ccef4eca24ac63cca66a41aab2178" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + +[[package]] +name = "serde_yaml_ng" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b4db627b98b36d4203a7b458cf3573730f2bb591b28871d916dfa9efabfd41f" +dependencies = [ + "indexmap", + "itoa", + "ryu", + "serde", + "unsafe-libyaml", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "syn" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8593e8e72159ed2257d083c7a454a85cbf854f37a0966d8d483aff8c8a3ebcee" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom", + "once_cell", + "rustix", + "windows-sys", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d245f478577f809a851594d02313b640fb437e0bb33866753cff937863096954" + +[[package]] +name = "unsafe-libyaml" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/packages/leapp-runtime-rust/Cargo.toml b/packages/leapp-runtime-rust/Cargo.toml new file mode 100644 index 00000000..10596e06 --- /dev/null +++ b/packages/leapp-runtime-rust/Cargo.toml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[workspace] +members = [] +default-members = ["."] +resolver = "2" + +[package] +name = "leapp-runtime" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +publish = false +description = "Backend-independent LEAPP bundle loading and inspection" + +[dependencies] +anyhow = "1" +indexmap = { version = "2", features = ["serde"] } +serde = { version = "1", features = ["derive"] } +serde_path_to_error = "0.1" +serde_yaml_ng = "=0.10.0" +serde_json = "1" +sha2 = "0.10" +safetensors = "=0.6.2" + +[dev-dependencies] +tempfile = "3" diff --git a/packages/leapp-runtime-rust/README.md b/packages/leapp-runtime-rust/README.md new file mode 100644 index 00000000..34c72050 --- /dev/null +++ b/packages/leapp-runtime-rust/README.md @@ -0,0 +1,67 @@ + + +# LEAPP Rust bundle API + +This package loads and validates exported LEAPP bundles and exposes their +public inspection API in Rust. Bundle loading does not initialize an inference +backend. It provides `Bundle`, `PortRef`, `PortSpec`, `DataType`, `Metadata`, +`MetadataValue`, `BundleError`, and `ErrorCode`. + +The loader validates graph connections, artifact paths, optional SHA-256 hashes, +and safetensors initial state. It preserves ordered metadata, nested YAML tags, +duplicate mapping keys, and YAML merges. + +## Rust + +```rust +use leapp_runtime::Bundle; + +let bundle = Bundle::open("export/policy.yaml")?; +for output in bundle.pipeline_outputs() { + let spec = bundle.output_info(&output.node, &output.port).unwrap(); + println!("{}: {:?} {:?}", output.name(), spec.dtype, spec.shape); +} +# Ok::<(), leapp_runtime::BundleError>(()) +``` + +`Bundle` is an immutable, reference-counted snapshot. Clones share the snapshot; +inspection is safe across threads. Getters preserve declaration order and +unknown metadata. Exporter parameters retain their original paths and backend +names. Internal paths are resolved relative to the YAML, falling back to the +supplied artifact path. Optional SHA-256 hashes are checked before publication. +Files must be available while opening; their bytes are not retained for later +execution. + +Bundle loading validates graph ports and connections, artifacts, and optional +initial-value files. It does not check whether an inference engine can execute +the graph. Declared external inputs can remain unresolved, matching the source +implementations' inspection contract. + +The loader normalizes known backend aliases, accepts custom nonempty backend +names, and leaves adapter availability to execution. + +## Build and test + +Run from this directory with a Rust toolchain: + +```sh +cargo test --locked +cargo check --locked --workspace +cargo clippy --locked --workspace --all-targets -- -D warnings +cargo fmt --all --check +``` + +The default workspace member is the core crate. Its inspection tests use opaque +artifact bytes and require no inference SDK. + +## Extraction provenance + +This package is the common foundation extracted from the existing Triton-based +and native runtime implementations. Execution engines, sessions, tensor/DLPack +code, graph planners, the Triton converter, and backend adapters remain in those +implementations. The package path and library names match the runtime branches +to support later integration. The extraction has no Triton, CUDA, ONNX Runtime, +PyTorch, or NumPy dependency. diff --git a/packages/leapp-runtime-rust/src/error.rs b/packages/leapp-runtime-rust/src/error.rs new file mode 100644 index 00000000..fb332775 --- /dev/null +++ b/packages/leapp-runtime-rust/src/error.rs @@ -0,0 +1,90 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::{ + fmt, io, + path::{Path, PathBuf}, +}; + +/// Failure categories returned by fallible runtime operations. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(i32)] +pub enum ErrorCode { + Ok = 0, + InvalidArgument = 1, + NotFound = 2, + Unsupported = 3, + Io = 4, + Internal = 5, +} + +/// An owned loading error, including its phase, YAML path and underlying cause. +#[derive(Debug)] +pub struct BundleError { + code: ErrorCode, + phase: &'static str, + path: PathBuf, + source: anyhow::Error, +} + +impl BundleError { + /// Category suitable for programmatic error handling. + pub fn code(&self) -> ErrorCode { + self.code + } + + /// Loading phase that failed, such as `parse` or `artifact`. + pub fn phase(&self) -> &'static str { + self.phase + } + + /// YAML path supplied to the failed open attempt. + pub fn path(&self) -> &Path { + &self.path + } + + /// Full contextual diagnostic, including the original dependency error. + pub fn message(&self) -> String { + self.to_string() + } +} + +impl fmt::Display for BundleError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}: {:#}", self.path.display(), self.source) + } +} + +impl std::error::Error for BundleError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(self.source.as_ref()) + } +} + +pub(crate) fn in_phase( + path: &Path, + phase: &'static str, + operation: impl FnOnce() -> anyhow::Result, +) -> Result { + operation().map_err(|source| { + let code = classify_error(&source); + BundleError { + code, + phase, + path: path.into(), + source, + } + }) +} + +fn classify_error(source: &anyhow::Error) -> ErrorCode { + if let Some(error) = source.downcast_ref::() { + if error.kind() == io::ErrorKind::NotFound { + ErrorCode::NotFound + } else { + ErrorCode::Io + } + } else { + ErrorCode::InvalidArgument + } +} diff --git a/packages/leapp-runtime-rust/src/lib.rs b/packages/leapp-runtime-rust/src/lib.rs new file mode 100644 index 00000000..b12214f5 --- /dev/null +++ b/packages/leapp-runtime-rust/src/lib.rs @@ -0,0 +1,148 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Load and inspect an immutable LEAPP bundle without initializing a backend. +//! +//! ```no_run +//! use leapp_runtime::Bundle; +//! let bundle = Bundle::open("export/policy.yaml")?; +//! for input in bundle.pipeline_inputs() { +//! if let Some(spec) = bundle.input_info(&input.node, &input.port) { +//! println!("{}: {:?}", input.name(), spec.shape); +//! } +//! } +//! # Ok::<(), leapp_runtime::BundleError>(()) +//! ``` +//! Implementation schema types are intentionally private: +//! +//! ```compile_fail +//! use leapp_runtime::{BundleDescription, NodeSpec, Pipeline}; +//! ``` +#![forbid(unsafe_code)] + +mod error; +mod load; +mod schema; +mod validate; + +pub use error::{BundleError, ErrorCode}; +use indexmap::IndexMap; +use schema::{BundleDescription, InitialValueSpec}; +pub use schema::{DataType, Metadata, MetadataValue, PortRef, PortSpec}; +use std::{ + path::{Path, PathBuf}, + sync::Arc, +}; + +#[derive(Debug)] +struct BundleData { + yaml_path: PathBuf, + description: BundleDescription, + inputs: Vec, + outputs: Vec, + // Validated initial descriptors are retained for later bundle conversion. + #[allow(dead_code)] + initial_values: IndexMap, +} + +/// Shared ownership of an immutable, validated bundle snapshot. +/// +/// Cloning shares the snapshot. Inspection borrows are tied to the bundle; +/// concurrent reads are safe. Artifact files must remain available and unchanged +/// for any subsequent execution that uses the resolved paths. +#[derive(Debug, Clone)] +pub struct Bundle(Arc); + +impl Bundle { + /// Parse and validate a YAML bundle, resolving artifacts beside its path. + /// + /// The path is made absolute and lexically normalized before reading YAML + /// and resolving relative artifact paths. Symlinks are not canonicalized. + /// Artifacts are sought beside the YAML first, then at the supplied path. + /// + /// Returns a new bundle on success and an error with a phase and diagnostic + /// on failure. Existing bundles are never modified by an open attempt. + pub fn open(path: impl AsRef) -> Result { + load::read_bundle(path.as_ref()).map(|data| Self(Arc::new(data))) + } + + /// Absolute YAML path, without resolving symlinks. + pub fn yaml_path(&self) -> &Path { + &self.0.yaml_path + } + + /// Model names in declaration order, borrowed from this snapshot. + pub fn model_names(&self) -> Vec<&str> { + self.0 + .description + .models + .keys() + .map(String::as_str) + .collect() + } + + /// Declared input references in order; these need not identify existing ports. + pub fn pipeline_inputs(&self) -> &[PortRef] { + &self.0.inputs + } + + /// Public output references in declaration order. + pub fn pipeline_outputs(&self) -> &[PortRef] { + &self.0.outputs + } + + /// Pipeline configs; an empty mapping when absent. + pub fn pipeline_configs(&self) -> &Metadata { + &self.0.description.pipeline.configs + } + + /// Additional descriptive fields at pipeline level. + pub fn pipeline_metadata(&self) -> &Metadata { + &self.0.description.pipeline.metadata + } + + /// Export provenance from the `system information` mapping. + pub fn system_information(&self) -> &Metadata { + &self.0.description.system_information + } + + /// Additional descriptive fields at the bundle root. + pub fn metadata(&self) -> &Metadata { + &self.0.description.metadata + } + + /// Original exporter parameters, including unknown fields and original paths. + /// Returns `None` when the model is absent. + pub fn model_parameters(&self, node: &str) -> Option<&Metadata> { + Some(&self.0.description.models.get(node)?.original_parameters) + } + + /// Additional model metadata, or `None` when the model is absent. + pub fn model_metadata(&self, node: &str) -> Option<&Metadata> { + Some(&self.0.description.models.get(node)?.metadata) + } + + /// The first input descriptor with this name, including internal ports. + /// Returns `None` if absent; duplicate declarations are permitted. + pub fn input_info(&self, node: &str, port: &str) -> Option<&PortSpec> { + self.0 + .description + .models + .get(node)? + .inputs + .iter() + .find(|p| p.name == port) + } + + /// The first output descriptor with this name, including internal ports. + /// Returns `None` if absent; duplicate declarations are permitted. + pub fn output_info(&self, node: &str, port: &str) -> Option<&PortSpec> { + self.0 + .description + .models + .get(node)? + .outputs + .iter() + .find(|p| p.name == port) + } +} diff --git a/packages/leapp-runtime-rust/src/load.rs b/packages/leapp-runtime-rust/src/load.rs new file mode 100644 index 00000000..2ad31691 --- /dev/null +++ b/packages/leapp-runtime-rust/src/load.rs @@ -0,0 +1,305 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use crate::{ + error::in_phase, + schema::{BundleDescription, NodeParameters}, + validate, BundleData, BundleError, Metadata, MetadataValue, PortRef, +}; +use anyhow::{ensure, Context, Result}; +use indexmap::IndexMap; +use serde::{ + de::{EnumAccess, MapAccess, SeqAccess, VariantAccess, Visitor}, + Deserialize, Deserializer, +}; +use serde_yaml_ng::value::{Tag, TaggedValue}; +use sha2::{Digest, Sha256}; +use std::{ + fmt, fs, io, + path::{Component, Path, PathBuf}, +}; + +pub(crate) fn read_bundle(path: &Path) -> Result { + let (yaml_path, yaml_text) = in_phase(path, "open", || { + ensure!(!path.as_os_str().is_empty(), "YAML path must not be empty"); + let yaml_path = normalize(&std::path::absolute(path)?); + let yaml_text = fs::read_to_string(&yaml_path)?; + Ok((yaml_path, yaml_text)) + })?; + let base = yaml_path.parent().expect("absolute path has a parent"); + let mut description = in_phase(path, "parse", || { + // PyYAML keeps the last duplicate mapping entry and expands YAML merges. + let value = serde_yaml_ng::from_str::(&yaml_text)?.into_value()?; + // Serde's flatten buffer cannot carry YAML tags. Deserialize known + // schema fields separately, then attach the original metadata values. + let mut schema = value.clone(); + let metadata = take_metadata(Some(&mut schema), ROOT_FIELDS); + let pipeline_metadata = take_metadata(schema.get_mut("pipeline"), PIPELINE_FIELDS); + if let Some(models) = schema + .get_mut("models") + .and_then(MetadataValue::as_mapping_mut) + { + for node in models.values_mut() { + take_metadata(Some(node), MODEL_FIELDS); + for field in ["inputs", "outputs"] { + if let Some(ports) = + node.get_mut(field).and_then(MetadataValue::as_sequence_mut) + { + for port in ports { + take_metadata(Some(port), PORT_FIELDS); + } + } + } + } + } + let mut description: BundleDescription = serde_path_to_error::deserialize(schema)?; + description.metadata = metadata; + description.pipeline.metadata = pipeline_metadata; + for (name, node) in &mut description.models { + let mut original = value["models"][name].clone(); + node.metadata = take_metadata(Some(&mut original), MODEL_FIELDS); + for (field, ports) in [("inputs", &mut node.inputs), ("outputs", &mut node.outputs)] { + for (index, port) in ports.iter_mut().enumerate() { + port.metadata = take_metadata(original[field].get_mut(index), PORT_FIELDS); + } + } + node.original_parameters = + serde_yaml_ng::from_value(value["models"][name]["parameters"].clone())?; + } + Ok(description) + })?; + in_phase(path, "validate", || { + validate::validate_structure(&description) + })?; + for (name, node) in &mut description.models { + in_phase(path, "artifact", || { + let parameters = &mut node.parameters; + ensure!( + !parameters.backend.trim().is_empty(), + "backend name must not be empty" + ); + parameters.backend = backend_name(¶meters.backend).into(); + ensure!( + !parameters.model_path.as_os_str().is_empty(), + "model path must not be empty" + ); + let beside_yaml = normalize(&base.join(¶meters.model_path)); + // Python falls back to the supplied path when no adjacent artifact exists. + parameters.model_path = if beside_yaml.exists() { + beside_yaml + } else { + normalize(&std::path::absolute(¶meters.model_path)?) + }; + verify_artifact(parameters).with_context(|| { + format!("node {name} artifact {}", parameters.model_path.display()) + }) + })?; + } + let initial_values = in_phase(path, "initial_state", || { + if let Some(initial) = description + .pipeline + .initial_values + .as_mut() + .filter(|p| !p.as_os_str().is_empty()) + { + *initial = normalize(&base.join(&initial)); + ensure!( + fs::metadata(&initial)?.is_file(), + "initial values must be a regular file" + ); + validate::read_initial_values(initial, &description.models) + .with_context(|| format!("pipeline.initial_values ({})", initial.display())) + } else { + Ok(Default::default()) + } + })?; + let boundary = |map: &crate::schema::BoundaryPorts| { + map.iter() + .flat_map(|(node, ports)| { + ports.iter().map(|port| PortRef { + node: node.clone(), + port: port.clone(), + }) + }) + .collect() + }; + Ok(BundleData { + yaml_path, + inputs: boundary(&description.pipeline.inputs), + outputs: boundary(&description.pipeline.outputs), + description, + initial_values, + }) +} + +const ROOT_FIELDS: &[&str] = &["models", "pipeline", "system information"]; +const PIPELINE_FIELDS: &[&str] = &[ + "data_flow", + "feedback_flow", + "inputs", + "outputs", + "initial_values", + "configs", +]; +const MODEL_FIELDS: &[&str] = &["inputs", "outputs", "parameters"]; +const PORT_FIELDS: &[&str] = &["name", "dtype", "shape", "type"]; + +fn take_metadata(value: Option<&mut MetadataValue>, fields: &[&str]) -> Metadata { + let Some(mapping) = value.and_then(MetadataValue::as_mapping_mut) else { + // Leave malformed structure to the typed deserializer's path diagnostics. + return Metadata::new(); + }; + let mut metadata = Metadata::new(); + mapping.retain(|key, value| { + if let Some(key) = key.as_str() { + if !fields.contains(&key) { + metadata.insert(key.into(), value.clone()); + return false; + } + } + true + }); + metadata +} + +// Inspection normalizes known exporter aliases without requiring an adapter. +pub(crate) fn backend_name(name: &str) -> &str { + match name { + "onnx" | "onnx-dynamo" | "onnx-torchscript" => "onnx", + "jit" | "jit-script" | "jit-trace" => "jit", + "pt2" | "exported-program" => "pt2", + "trt" | "tensorrt" => "trt", + other => other, + } +} + +fn verify_artifact(parameters: &NodeParameters) -> Result<()> { + ensure!( + fs::metadata(¶meters.model_path)?.is_file(), + "artifact is not a regular file" + ); + if let Some(expected) = ¶meters.sha256sum { + let mut artifact = fs::File::open(¶meters.model_path)?; + let mut hash = Sha256::new(); + io::copy(&mut artifact, &mut hash)?; + let actual = format!("{:x}", hash.finalize()); + ensure!( + actual.eq_ignore_ascii_case(expected), + "SHA256 mismatch: expected {expected}, got {actual}" + ); + } + Ok(()) +} + +fn normalize(path: &Path) -> PathBuf { + let mut result = PathBuf::new(); + for component in path.components() { + match component { + Component::CurDir => {} + Component::ParentDir => { + result.pop(); + } + other => result.push(other.as_os_str()), + } + } + result +} + +// Decode mappings before Value so duplicate keys follow Python's last-value rule. +enum YamlValue { + Mapping(IndexMap), + Sequence(Vec), + Scalar(MetadataValue), + Tagged(String, Box), +} + +// An untagged serde enum cannot accept YAML tags. Read each value directly so +// metadata retains tags while mappings still use Python's last-value rule. +impl<'de> Deserialize<'de> for YamlValue { + fn deserialize>(deserializer: D) -> Result { + struct YamlVisitor; + + impl<'de> Visitor<'de> for YamlVisitor { + type Value = YamlValue; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a YAML value") + } + + fn visit_bool(self, value: bool) -> Result { + Ok(YamlValue::Scalar(MetadataValue::Bool(value))) + } + + fn visit_i64(self, value: i64) -> Result { + Ok(YamlValue::Scalar(MetadataValue::Number(value.into()))) + } + + fn visit_u64(self, value: u64) -> Result { + Ok(YamlValue::Scalar(MetadataValue::Number(value.into()))) + } + + fn visit_f64(self, value: f64) -> Result { + Ok(YamlValue::Scalar(MetadataValue::Number(value.into()))) + } + + fn visit_str(self, value: &str) -> Result { + Ok(YamlValue::Scalar(MetadataValue::String(value.into()))) + } + + fn visit_unit(self) -> Result { + Ok(YamlValue::Scalar(MetadataValue::Null)) + } + + fn visit_seq>(self, mut sequence: A) -> Result { + let mut values = Vec::new(); + while let Some(value) = sequence.next_element()? { + values.push(value); + } + Ok(YamlValue::Sequence(values)) + } + + fn visit_map>(self, mut mapping: A) -> Result { + let mut values = IndexMap::new(); + while let Some((key, value)) = mapping.next_entry()? { + values.insert(key, value); + } + Ok(YamlValue::Mapping(values)) + } + + fn visit_enum>(self, tagged: A) -> Result { + let (tag, value) = tagged.variant::()?; + Ok(YamlValue::Tagged(tag, Box::new(value.newtype_variant()?))) + } + } + + deserializer.deserialize_any(YamlVisitor) + } +} + +impl YamlValue { + fn into_value(self) -> Result { + match self { + Self::Mapping(items) => { + // Expand inherited maps first so chained merges retain their fields. + let mut value = MetadataValue::Mapping( + items + .into_iter() + .map(|(key, value)| value.into_value().map(|value| (key, value))) + .collect::>()?, + ); + value.apply_merge()?; + Ok(value) + } + Self::Sequence(items) => items + .into_iter() + .map(Self::into_value) + .collect::>() + .map(MetadataValue::Sequence), + Self::Scalar(value) => Ok(value), + Self::Tagged(tag, value) => Ok(MetadataValue::Tagged(Box::new(TaggedValue { + tag: Tag::new(tag), + value: value.into_value()?, + }))), + } + } +} diff --git a/packages/leapp-runtime-rust/src/schema.rs b/packages/leapp-runtime-rust/src/schema.rs new file mode 100644 index 00000000..7fe669cb --- /dev/null +++ b/packages/leapp-runtime-rust/src/schema.rs @@ -0,0 +1,183 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use anyhow::{ensure, Context, Result}; +use indexmap::IndexMap; +use serde::{de::DeserializeOwned, Deserialize, Deserializer}; +use std::{fmt, path::PathBuf}; + +use safetensors::Dtype as InitialValueDtype; +pub use serde_yaml_ng::Value as MetadataValue; + +pub type Metadata = IndexMap; +pub(crate) type Connections = IndexMap>; +pub(crate) type BoundaryPorts = IndexMap>; + +/// Ordered YAML fields, without artifact resolution or structural validation. +/// Kept private; loading publishes only a validated inspection snapshot. +#[derive(Debug, Deserialize)] +pub(crate) struct BundleDescription { + #[serde(deserialize_with = "deserialize_non_null")] + pub models: IndexMap, + pub pipeline: Pipeline, + #[serde( + rename = "system information", + deserialize_with = "deserialize_non_null" + )] + pub system_information: Metadata, + #[serde(flatten)] + pub metadata: Metadata, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct Pipeline { + #[serde(default, deserialize_with = "deserialize_list_map")] + pub data_flow: Connections, + #[serde(default, deserialize_with = "deserialize_list_map")] + pub feedback_flow: Connections, + #[serde(deserialize_with = "deserialize_list_map")] + pub inputs: BoundaryPorts, + #[serde(deserialize_with = "deserialize_list_map")] + pub outputs: BoundaryPorts, + pub initial_values: Option, + #[serde(default, deserialize_with = "deserialize_non_null")] + pub configs: Metadata, + #[serde(flatten)] + pub metadata: Metadata, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct NodeSpec { + #[serde(deserialize_with = "deserialize_non_null")] + pub inputs: Vec, + #[serde(deserialize_with = "deserialize_non_null")] + pub outputs: Vec, + pub parameters: NodeParameters, + #[serde(skip)] + pub original_parameters: Metadata, + #[serde(flatten)] + pub metadata: Metadata, +} + +#[derive(Debug, Deserialize)] +pub(crate) struct NodeParameters { + pub backend: String, + pub model_path: PathBuf, + pub sha256sum: Option, +} + +/// Tensor interface and additional descriptive fields for one model port. +#[derive(Debug, Deserialize)] +pub struct PortSpec { + pub name: String, + pub dtype: DataType, + #[serde(deserialize_with = "deserialize_shape")] + /// Empty means scalar; 0 is an empty dimension and -1 is dynamic. + pub shape: Vec, + #[serde(default, rename = "type")] + pub tensor_type: String, + #[serde(flatten)] + pub metadata: Metadata, +} + +/// Scalar types supported in LEAPP port declarations. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum DataType { + Float16, + Float32, + Float64, + Bfloat16, + Int8, + Int16, + Int32, + Int64, + Uint8, + Uint16, + Uint32, + Uint64, + Bool, +} + +/// Actual metadata retained for later bundle conversion. +#[allow(dead_code)] +#[derive(Debug)] +pub(crate) struct InitialValueSpec { + pub dtype: InitialValueDtype, + pub shape: Vec, +} + +/// A qualified model and port identifier. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize)] +#[serde(try_from = "String")] +pub struct PortRef { + pub node: String, + pub port: String, +} + +impl PortRef { + /// Qualified `node/port` name. + pub fn name(&self) -> String { + self.to_string() + } +} + +impl fmt::Display for PortRef { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}/{}", self.node, self.port) + } +} + +impl TryFrom for PortRef { + type Error = anyhow::Error; + + fn try_from(reference: String) -> Result { + let (node, port) = reference + .split_once('/') + .with_context(|| format!("expected node/port: {reference:?}"))?; + ensure!(!port.contains('/'), "expected node/port: {reference:?}"); + + Ok(Self { + node: node.into(), + port: port.into(), + }) + } +} + +// Value deserialization otherwise accepts null as an empty collection. +fn deserialize_non_null<'de, D: Deserializer<'de>, T: DeserializeOwned>( + deserializer: D, +) -> Result { + let value = MetadataValue::deserialize(deserializer)?; + if value.is_null() { + return Err(serde::de::Error::custom("required field cannot be null")); + } + + serde_path_to_error::deserialize(value).map_err(serde::de::Error::custom) +} + +fn deserialize_list_map<'de, D: Deserializer<'de>, T: DeserializeOwned>( + deserializer: D, +) -> Result { + let value: MetadataValue = deserialize_non_null(deserializer)?; + if let Some(map) = value.as_mapping() { + for (key, item) in map { + if !item.is_sequence() { + return Err(serde::de::Error::custom(format!( + "{key:?}: expected a list" + ))); + } + } + } + + serde_path_to_error::deserialize(value).map_err(serde::de::Error::custom) +} + +fn deserialize_shape<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { + match deserialize_non_null(deserializer)? { + MetadataValue::String(shape_text) => { + serde_json::from_str(&shape_text).map_err(serde::de::Error::custom) + } + value => Vec::deserialize(value).map_err(serde::de::Error::custom), + } +} diff --git a/packages/leapp-runtime-rust/src/validate.rs b/packages/leapp-runtime-rust/src/validate.rs new file mode 100644 index 00000000..59246c61 --- /dev/null +++ b/packages/leapp-runtime-rust/src/validate.rs @@ -0,0 +1,128 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use crate::schema::{BundleDescription, InitialValueSpec, NodeSpec, PortRef}; +use anyhow::{ensure, Context, Result}; +use indexmap::IndexMap; +use safetensors::SafeTensors; +use std::{collections::HashSet, fs, path::Path}; + +pub(crate) fn validate_structure(description: &BundleDescription) -> Result<()> { + let mut inputs = IndexMap::new(); + let mut outputs = IndexMap::new(); + for (name, node) in &description.models { + // Python input buffers keep the last duplicate; edge sources use the first. + for spec in &node.inputs { + inputs.insert( + PortRef { + node: name.clone(), + port: spec.name.clone(), + }, + spec, + ); + } + for spec in &node.outputs { + outputs + .entry(PortRef { + node: name.clone(), + port: spec.name.clone(), + }) + .or_insert(spec); + } + } + + for flow in [ + &description.pipeline.data_flow, + &description.pipeline.feedback_flow, + ] { + for (source, targets) in flow { + let source_spec = outputs + .get(source) + .with_context(|| format!("unknown output port {source}"))?; + for target in targets { + let target_spec = inputs + .get(target) + .with_context(|| format!("unknown input port {target}"))?; + ensure!( + source_spec.shape == target_spec.shape + && source_spec.dtype == target_spec.dtype, + "incompatible tensors: {source} -> {target}" + ); + } + } + } + + // Python inspects published outputs during loading, but not external inputs. + let mut routed = HashSet::new(); + for (node, names) in &description.pipeline.outputs { + ensure!( + description.models.contains_key(node), + "boundary: unknown node {node}" + ); + for name in names { + let key = PortRef { + node: node.clone(), + port: name.clone(), + }; + ensure!(outputs.contains_key(&key), "unknown output port {key}"); + routed.insert(key); + } + } + routed.extend(description.pipeline.data_flow.keys().cloned()); + routed.extend(description.pipeline.feedback_flow.keys().cloned()); + for output in outputs.keys() { + ensure!(routed.contains(output), "output is not routed: {output}"); + } + Ok(()) +} + +pub(crate) fn read_initial_values( + path: &Path, + models: &IndexMap, +) -> Result> { + let bytes = fs::read(path)?; + // Safetensors owns JSON header and payload-layout validation. + let (_, metadata) = SafeTensors::read_metadata(&bytes)?; + + let mut initial_values = IndexMap::new(); + for name in metadata.offset_keys() { + let key = PortRef::try_from(name.clone())?; + models + .get(&key.node) + .and_then(|node| node.inputs.iter().find(|port| port.name == key.port)) + .with_context(|| format!("unknown input port {key}"))?; + + let tensor = metadata + .info(&name) + .context("missing initial tensor metadata")?; + let shape = tensor + .shape + .iter() + .copied() + .map(i64::try_from) + .collect::, _>>()?; + initial_values.insert( + key, + InitialValueSpec { + dtype: tensor.dtype, + shape, + }, + ); + } + + // Return file metadata in YAML node/input order. + let mut ordered = IndexMap::new(); + for (node, model) in models { + for port in &model.inputs { + let key = PortRef { + node: node.clone(), + port: port.name.clone(), + }; + if let Some(spec) = initial_values.shift_remove(&key) { + ordered.insert(key, spec); + } + } + } + + Ok(ordered) +} diff --git a/packages/leapp-runtime-rust/tests/bundle.rs b/packages/leapp-runtime-rust/tests/bundle.rs new file mode 100644 index 00000000..c23705be --- /dev/null +++ b/packages/leapp-runtime-rust/tests/bundle.rs @@ -0,0 +1,682 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use leapp_runtime::{Bundle, DataType, ErrorCode, MetadataValue as Value, PortRef}; +use safetensors::{tensor::TensorView, Dtype}; +use std::{fs, path::PathBuf}; +use tempfile::TempDir; + +const YAML: &str = include_str!("fixtures/pipeline.yaml"); + +struct Fixture { + directory: TempDir, + path: PathBuf, + value: Value, +} + +impl Fixture { + fn new() -> Self { + let directory = tempfile::tempdir().unwrap(); + fs::write(directory.path().join("model.bin"), b"abc").unwrap(); + let path = directory.path().join("pipeline.yaml"); + Self { + directory, + path, + value: serde_yaml_ng::from_str(YAML).unwrap(), + } + } + + fn write(&self) { + fs::write(&self.path, serde_yaml_ng::to_string(&self.value).unwrap()).unwrap(); + } + + fn open(&self) -> Bundle { + self.write(); + Bundle::open(&self.path).unwrap() + } + + fn fail(&self, phase: &str, message: &str) { + self.write(); + let error = Bundle::open(&self.path).unwrap_err(); + assert_eq!(error.phase(), phase, "{error}"); + assert!( + error.message().contains(message), + "{error}: missing {message:?}" + ); + assert_eq!(error.path(), self.path); + } + + fn initial(&mut self, name: &str, dtype: Dtype, shape: Vec, bytes: &[u8]) { + let tensor = TensorView::new(dtype, shape, bytes).unwrap(); + safetensors::serialize_to_file( + [(name, tensor)], + None, + &self.directory.path().join("initial.safetensors"), + ) + .unwrap(); + self.value["pipeline"]["initial_values"] = Value::String("initial.safetensors".into()); + } +} + +fn value(text: &str) -> Value { + serde_yaml_ng::from_str(text).unwrap() +} + +#[test] +fn inspection_preserves_order_and_metadata() { + let f = Fixture::new(); + let b = f.open(); + assert_eq!(b.model_names(), ["z_observation", "a_policy"]); + assert_eq!(b.yaml_path(), f.path); + assert_eq!(b.pipeline_inputs()[0].name(), "z_observation/joint_pos"); + assert_eq!( + b.pipeline_outputs() + .iter() + .map(PortRef::name) + .collect::>(), + [ + "a_policy/next_state", + "a_policy/actions", + "z_observation/features" + ] + ); + let port = b.input_info("a_policy", "features").unwrap(); + assert_eq!(port.shape, [3]); + assert_eq!(port.dtype, DataType::Float32); + assert_eq!(port.tensor_type, "tensor"); + assert_eq!( + b.output_info("a_policy", "actions").unwrap().metadata["temporal_period_ms"], + value("100") + ); + assert_eq!(b.pipeline_configs()["frequency"], value("50")); + assert_eq!(b.system_information()["leapp config version"], "1.3"); + let parameters = b.model_parameters("z_observation").unwrap(); + assert_eq!(parameters["backend"], "jit"); + assert_eq!(parameters["model_path"], "model.bin"); + assert_eq!( + parameters["custom_options"], + value("{device: cpu, flags: [one, two]}") + ); +} + +#[test] +fn getters_preserve_unknown_fields_at_each_level() { + let mut f = Fixture::new(); + f.value["project"] = + value("{name: example, axes: [x, y, z], mixed: [null, true, -2, 2.5, '3']}"); + f.value["pipeline"]["description"] = value("feedback"); + f.value["models"]["a_policy"]["purpose"] = value("control"); + f.value["models"]["a_policy"]["parameters"][""] = value("custom"); + let b = f.open(); + assert_eq!(b.metadata()["project"], f.value["project"]); + assert_eq!(b.pipeline_metadata()["description"], "feedback"); + assert_eq!(b.model_metadata("a_policy").unwrap()["purpose"], "control"); + assert_eq!(b.model_parameters("a_policy").unwrap()[""], "custom"); +} + +#[test] +fn absent_lookups_return_none() { + let b = Fixture::new().open(); + assert!(b.model_parameters("absent").is_none()); + assert!(b.model_metadata("absent").is_none()); + assert!(b.input_info("absent", "features").is_none()); + assert!(b.input_info("a_policy", "absent").is_none()); + assert!(b.output_info("a_policy", "state").is_none()); +} + +#[test] +fn clones_share_a_snapshot_and_survive_file_removal() { + fn assert_send_sync() {} + assert_send_sync::(); + let f = Fixture::new(); + let original = f.open(); + let copy = original.clone(); + assert!(std::ptr::eq( + original.input_info("a_policy", "state").unwrap(), + copy.input_info("a_policy", "state").unwrap() + )); + drop(original); + fs::remove_file(&f.path).unwrap(); + fs::remove_file(f.directory.path().join("model.bin")).unwrap(); + std::thread::spawn(move || { + assert_eq!(copy.input_info("a_policy", "state").unwrap().shape, [3]) + }) + .join() + .unwrap(); +} + +#[test] +fn failed_open_does_not_replace_an_existing_bundle() { + let f = Fixture::new(); + let original = f.open(); + let error = Bundle::open(f.directory.path().join("missing.yaml")).unwrap_err(); + assert_eq!(error.code(), ErrorCode::NotFound); + assert_eq!(error.phase(), "open"); + assert_eq!(original.model_names(), ["z_observation", "a_policy"]); + use std::error::Error; + let mut cause = error.source().unwrap(); + while let Some(next) = cause.source() { + cause = next; + } + assert!(cause.downcast_ref::().is_some()); +} + +#[test] +fn malformed_documents_are_rejected() { + let f = Fixture::new(); + for text in [ + "[", + "null", + "[]", + "{}", + "models: {}\nmodels: {}", + "models: {}\n---\nmodels: {}", + "models: !custom {}", + ] { + fs::write(&f.path, text).unwrap(); + let error = Bundle::open(&f.path).unwrap_err(); + assert_eq!(error.phase(), "parse", "{error}"); + } + assert_eq!( + Bundle::open("").unwrap_err().code(), + ErrorCode::InvalidArgument + ); +} + +#[test] +fn collections_cannot_be_null_or_scalar() { + for section in ["inputs", "outputs", "data_flow", "feedback_flow", "configs"] { + for text in ["null", "invalid"] { + let mut f = Fixture::new(); + f.value["pipeline"][section] = value(text); + f.fail("parse", section); + } + } + for field in ["inputs", "outputs", "parameters"] { + let mut f = Fixture::new(); + f.value["models"]["a_policy"][field] = Value::Null; + f.fail("parse", field); + } +} + +#[test] +fn accepts_empty_graph_and_unspecified_version() { + let mut f = Fixture::new(); + f.value = value("models: {}\npipeline: {inputs: {}, outputs: {}}\nsystem information: {}\n"); + let b = f.open(); + assert!(b.model_names().is_empty()); + assert!(b.pipeline_inputs().is_empty()); + assert!(b.pipeline_outputs().is_empty()); + assert!(b.pipeline_configs().is_empty()); + f.value["system information"]["leapp config version"] = value("future"); + assert_eq!( + f.open().system_information()["leapp config version"], + "future" + ); +} + +#[test] +fn preserves_supported_backend_aliases_in_parameters() { + for backend in [ + "onnx", + "onnx-dynamo", + "onnx-torchscript", + "jit", + "jit-script", + "jit-trace", + "pt2", + "exported-program", + "trt", + "tensorrt", + ] { + let mut f = Fixture::new(); + f.value["models"]["a_policy"]["parameters"]["backend"] = value(backend); + assert_eq!( + f.open().model_parameters("a_policy").unwrap()["backend"], + backend + ); + } + let mut f = Fixture::new(); + f.value["models"]["a_policy"]["parameters"]["backend"] = value("custom"); + assert_eq!( + f.open().model_parameters("a_policy").unwrap()["backend"], + "custom" + ); + f.value["models"]["a_policy"]["parameters"]["backend"] = value(" "); + f.write(); + assert_eq!( + Bundle::open(&f.path).unwrap_err().code(), + ErrorCode::InvalidArgument + ); +} + +#[test] +fn preserves_shapes_and_all_canonical_dtypes() { + for shape in ["[]", "[0]", "[-1]", "[2, 3]", "'[2, 3]'"] { + for (name, dtype) in [ + ("bool", DataType::Bool), + ("uint8", DataType::Uint8), + ("uint16", DataType::Uint16), + ("uint32", DataType::Uint32), + ("uint64", DataType::Uint64), + ("int8", DataType::Int8), + ("int16", DataType::Int16), + ("int32", DataType::Int32), + ("int64", DataType::Int64), + ("float16", DataType::Float16), + ("bfloat16", DataType::Bfloat16), + ("float32", DataType::Float32), + ("float64", DataType::Float64), + ] { + let mut f = Fixture::new(); + for node in f.value["models"].as_mapping_mut().unwrap().values_mut() { + for direction in ["inputs", "outputs"] { + for port in node[direction].as_sequence_mut().unwrap() { + port["shape"] = value(shape); + port["dtype"] = value(name); + } + } + } + let b = f.open(); + assert_eq!(b.input_info("a_policy", "features").unwrap().dtype, dtype); + let expected = if shape.starts_with('\'') { + value("[2, 3]") + } else { + value(shape) + }; + assert_eq!( + b.input_info("a_policy", "features").unwrap().shape, + serde_yaml_ng::from_value::>(expected).unwrap() + ); + } + } +} + +#[test] +fn malformed_shapes_and_unknown_dtypes_fail() { + for shape in [ + "[1.5]", + "[true]", + "[null]", + "[batch]", + "[9223372036854775808]", + "3", + "null", + "'(3,)'", + "'[3,]'", + ] { + let mut f = Fixture::new(); + f.value["models"]["a_policy"]["inputs"][0]["shape"] = value(shape); + f.fail("parse", "shape"); + } + let mut f = Fixture::new(); + f.value["models"]["a_policy"]["inputs"][0]["dtype"] = value("complex64"); + f.fail("parse", "dtype"); +} + +#[test] +fn sha256_is_optional_but_verified_when_present() { + let mut f = Fixture::new(); + fs::write(f.directory.path().join("model.bin"), b"changed").unwrap(); + f.fail("artifact", "SHA256 mismatch"); + for node in f.value["models"].as_mapping_mut().unwrap().values_mut() { + let params = node["parameters"].as_mapping_mut().unwrap(); + params.remove("sha256sum"); + params.remove("md5sum"); + } + f.open(); + fs::remove_file(f.directory.path().join("model.bin")).unwrap(); + f.write(); + assert_eq!( + Bundle::open(&f.path).unwrap_err().code(), + ErrorCode::NotFound + ); +} + +#[test] +fn preserves_md5_metadata_and_accepts_uppercase_sha256() { + let mut f = Fixture::new(); + for node in f.value["models"].as_mapping_mut().unwrap().values_mut() { + node["parameters"]["md5sum"] = value("ignored"); + let hash = node["parameters"]["sha256sum"] + .as_str() + .unwrap() + .to_uppercase(); + node["parameters"]["sha256sum"] = Value::String(hash); + } + assert_eq!( + f.open().model_parameters("a_policy").unwrap()["md5sum"], + "ignored" + ); +} + +#[test] +fn checks_references_and_edge_compatibility() { + for target in ["missing/features", "a_policy/missing", "a_policy/actions"] { + let mut f = Fixture::new(); + f.value["pipeline"]["data_flow"]["z_observation/features"] = + Value::Sequence(vec![Value::String(target.into())]); + f.fail("validate", target); + } + for field in ["dtype", "shape"] { + let mut f = Fixture::new(); + f.value["models"]["a_policy"]["inputs"][0][field] = + value(if field == "dtype" { "float64" } else { "[2]" }); + f.fail("validate", "incompatible tensors"); + } +} + +#[test] +fn accepts_reversed_model_order_and_cycles() { + let mut f = Fixture::new(); + let models = f.value["models"].as_mapping_mut().unwrap(); + let first = models.shift_remove("z_observation").unwrap(); + models.insert(value("z_observation"), first); + assert_eq!(f.open().model_names(), ["a_policy", "z_observation"]); + f.value["pipeline"]["inputs"] = value("{}"); + f.value["pipeline"]["data_flow"]["a_policy/actions"] = value("[z_observation/joint_pos]"); + f.open(); +} + +#[test] +fn permits_multiple_producers_but_requires_output_routing() { + let mut f = Fixture::new(); + f.value["pipeline"]["inputs"]["a_policy"] = value("[features]"); + f.open(); + let mut f = Fixture::new(); + f.value["pipeline"]["outputs"]["a_policy"] = value("[next_state]"); + f.fail("validate", "output is not routed"); + f.value["pipeline"]["data_flow"]["a_policy/actions"] = value("[]"); + f.open(); + f.value["pipeline"]["inputs"] = value("{}"); + f.open(); +} + +#[test] +fn accepts_duplicate_ports_and_boundaries() { + let mut f = Fixture::new(); + let port = f.value["models"]["a_policy"]["inputs"][0].clone(); + f.value["models"]["a_policy"]["inputs"] + .as_sequence_mut() + .unwrap() + .push(port); + f.open(); + let mut f = Fixture::new(); + f.value["pipeline"]["outputs"]["a_policy"] = value("[actions, actions]"); + f.open(); +} + +#[test] +fn initial_values_are_optional_and_can_target_any_input() { + for initial in ["null", "''"] { + let mut f = Fixture::new(); + f.value["pipeline"]["initial_values"] = value(initial); + f.open(); + } + let mut f = Fixture::new(); + f.initial("z_observation/joint_pos", Dtype::I32, vec![1, 3], &[0; 12]); + let b = f.open(); + assert_eq!( + b.input_info("z_observation", "joint_pos").unwrap().dtype, + DataType::Float32 + ); + f.initial("a_policy/state", Dtype::F32, vec![3], &[0; 12]); + f.open(); + f.initial("a_policy/missing", Dtype::F32, vec![3], &[0; 12]); + f.fail("initial_state", "a_policy/missing"); +} + +#[test] +fn malformed_or_truncated_initial_values_fail() { + let mut f = Fixture::new(); + f.initial("a_policy/state", Dtype::F32, vec![3], &[0; 12]); + let path = f.directory.path().join("initial.safetensors"); + let mut bytes = fs::read(&path).unwrap(); + bytes.pop(); + fs::write(&path, bytes).unwrap(); + f.fail("initial_state", "initial.safetensors"); + fs::write(&path, b"invalid").unwrap(); + f.fail("initial_state", "initial.safetensors"); +} + +#[cfg(unix)] +#[test] +fn yaml_symlinks_resolve_artifacts_beside_the_supplied_path() { + use std::os::unix::fs::symlink; + let f = Fixture::new(); + f.write(); + let other = tempfile::tempdir().unwrap(); + let path = other.path().join("alias.yaml"); + symlink(&f.path, &path).unwrap(); + assert_eq!(Bundle::open(&path).unwrap_err().code(), ErrorCode::NotFound); + fs::write(other.path().join("model.bin"), b"abc").unwrap(); + assert_eq!(Bundle::open(&path).unwrap().yaml_path(), path); +} + +#[cfg(unix)] +#[test] +fn yaml_and_artifacts_use_the_same_normalized_path() { + use std::os::unix::fs::symlink; + let mut f = Fixture::new(); + f.value["origin"] = value("normalized"); + f.write(); + let mut other = Fixture::new(); + other.value["origin"] = value("filesystem"); + other.write(); + let child = other.directory.path().join("child"); + fs::create_dir(&child).unwrap(); + symlink(&child, f.directory.path().join("link")).unwrap(); + let path = f.directory.path().join("link/../pipeline.yaml"); + + // The filesystem follows the symlink before processing the parent segment. + assert_eq!( + fs::read_to_string(&path).unwrap(), + fs::read_to_string(&other.path).unwrap() + ); + let bundle = Bundle::open(&path).unwrap(); + assert_eq!(bundle.yaml_path(), f.path); + assert_eq!(bundle.metadata()["origin"], "normalized"); + + fs::remove_file(f.directory.path().join("model.bin")).unwrap(); + assert_eq!(Bundle::open(&path).unwrap_err().code(), ErrorCode::NotFound); +} + +#[test] +fn cpp_bundle_fixtures_load_through_the_new_api() { + let f = Fixture::new(); + fs::create_dir(f.directory.path().join("artifacts")).unwrap(); + fs::write( + f.directory.path().join("artifacts/policy.pt2"), + b"model bytes\n", + ) + .unwrap(); + fs::write( + f.directory.path().join("artifacts/preprocess.onnx"), + b"model bytes\n", + ) + .unwrap(); + fs::write(&f.path, include_str!("fixtures/policy.yaml")).unwrap(); + let bundle = Bundle::open(&f.path).unwrap(); + assert_eq!(bundle.model_names(), ["policy", "preprocess"]); + assert_eq!(bundle.pipeline_inputs()[0].name(), "preprocess/observation"); + let text = + include_str!("fixtures/feedback.yaml").replace(" initial_values: state.safetensors\n", ""); + fs::write(&f.path, text).unwrap(); + let bundle = Bundle::open(&f.path).unwrap(); + assert_eq!(bundle.input_info("policy", "state").unwrap().shape, [2]); +} + +#[test] +fn retains_yaml12_scalar_resolution_from_the_prototype() { + let f = Fixture::new(); + fs::write( + &f.path, + format!("{YAML}\nlabels: [x, y, z, yes, on, true, 'true']\n"), + ) + .unwrap(); + let b = Bundle::open(&f.path).unwrap(); + assert_eq!( + b.metadata()["labels"], + value("[x, y, z, 'yes', 'on', true, 'true']") + ); +} + +#[test] +fn metadata_has_no_additional_depth_limit() { + let f = Fixture::new(); + fs::write( + &f.path, + format!("{YAML}\ndeep: {}0{}\n", "[".repeat(65), "]".repeat(65)), + ) + .unwrap(); + Bundle::open(&f.path).unwrap(); +} + +#[test] +fn duplicate_mapping_keys_and_yaml_merges_follow_python() { + let f = Fixture::new(); + fs::write(&f.path, format!("{YAML}\ncustom: {{a: 1, a: 2}}\nshared: &base {{label: inherited}}\nmerged: {{<<: *base, own: true}}\n")).unwrap(); + let b = Bundle::open(&f.path).unwrap(); + assert_eq!(b.metadata()["custom"]["a"], value("2")); + assert_eq!(b.metadata()["merged"]["label"], "inherited"); + assert_eq!(b.metadata()["merged"]["own"], value("true")); + + fs::write( + &f.path, + format!("{YAML}\nmodels: {{}}\npipeline: {{inputs: {{}}, outputs: {{}}}}\n"), + ) + .unwrap(); + assert!(Bundle::open(&f.path).unwrap().model_names().is_empty()); +} + +#[test] +fn external_input_references_are_descriptive() { + let mut f = Fixture::new(); + f.value["pipeline"]["inputs"] = value("{unknown: [missing, missing], unused: []}"); + let b = f.open(); + assert_eq!(b.pipeline_inputs()[0].name(), "unknown/missing"); + assert_eq!(b.pipeline_inputs().len(), 2); +} + +#[test] +fn duplicate_edge_ports_match_python_lookup_order() { + let mut f = Fixture::new(); + let mut input = f.value["models"]["a_policy"]["inputs"][0].clone(); + input["shape"] = value("[7]"); + f.value["models"]["a_policy"]["inputs"] + .as_sequence_mut() + .unwrap() + .insert(0, input); + let mut output = f.value["models"]["z_observation"]["outputs"][0].clone(); + output["shape"] = value("[7]"); + f.value["models"]["z_observation"]["outputs"] + .as_sequence_mut() + .unwrap() + .push(output); + f.open(); +} + +#[test] +fn names_and_tensor_labels_have_no_extra_validation() { + let mut f = Fixture::new(); + let port = &mut f.value["models"]["z_observation"]["inputs"][0]; + port["name"] = Value::String("/\0".into()); + port["type"] = value("''"); + f.open(); + f.value["models"]["z_observation"]["inputs"][0] + .as_mapping_mut() + .unwrap() + .remove("type"); + assert_eq!( + f.open() + .input_info("z_observation", "/\0") + .unwrap() + .tensor_type, + "" + ); + assert!(PortRef::try_from("/".to_string()).is_ok()); + assert!(PortRef::try_from("node/\0".to_string()).is_ok()); + assert!(PortRef::try_from("node/port/extra".to_string()).is_err()); +} + +#[test] +fn nested_metadata_keys_need_not_be_strings() { + let mut f = Fixture::new(); + f.value["custom"] = value("{1: number, false: boolean}"); + assert_eq!(f.open().metadata()["custom"][value("1")], "number"); +} + +#[test] +fn artifacts_fall_back_to_the_working_directory() { + let working_directory = std::env::current_dir().unwrap(); + let artifact_directory = tempfile::tempdir_in(&working_directory).unwrap(); + let artifact_path = artifact_directory.path().join("model.bin"); + fs::write(&artifact_path, b"abc").unwrap(); + let relative_path = artifact_path.strip_prefix(&working_directory).unwrap(); + let mut f = Fixture::new(); + for model in f.value["models"].as_mapping_mut().unwrap().values_mut() { + model["parameters"]["model_path"] = Value::String(relative_path.to_str().unwrap().into()); + } + f.open(); + + // An existing path beside the YAML takes precedence over the fallback. + let adjacent = f.directory.path().join(relative_path); + fs::create_dir_all(adjacent.parent().unwrap()).unwrap(); + fs::write(adjacent, b"changed").unwrap(); + f.fail("artifact", "SHA256 mismatch"); +} + +#[test] +fn nested_yaml_merges_expand_before_overrides() { + let f = Fixture::new(); + fs::write(&f.path, format!("{YAML}\nbase: &base {{frequency: 50, label: inherited}}\nmiddle: &middle {{<<: *base, frequency: 100}}\nlast: {{<<: *middle}}\ncombined: {{<<: [*middle, *base], label: explicit}}\n")).unwrap(); + let b = Bundle::open(&f.path).unwrap(); + assert_eq!(b.metadata()["last"]["frequency"], value("100")); + assert_eq!(b.metadata()["last"]["label"], "inherited"); + assert!(b.metadata()["last"].get("<<").is_none()); + assert_eq!(b.metadata()["combined"]["frequency"], value("100")); + assert_eq!(b.metadata()["combined"]["label"], "explicit"); +} + +#[test] +fn tagged_metadata_preserves_nested_values_duplicates_and_merges() { + let f = Fixture::new(); + fs::write(&f.path, format!("{YAML}\nbase: &base {{label: inherited}}\ntagged: !sensor {{<<: *base, label: old, label: current, samples: [!units radians, null, 1.25]}}\n")).unwrap(); + let b = Bundle::open(&f.path).unwrap(); + let Value::Tagged(tagged) = &b.metadata()["tagged"] else { + panic!("metadata tag was lost"); + }; + assert_eq!(tagged.tag, "sensor"); + assert_eq!(tagged.value["label"], "current"); + assert_eq!( + tagged.value["samples"], + value("[!units radians, null, 1.25]") + ); +} + +#[test] +fn tagged_metadata_is_preserved_at_every_inspection_level() { + let mut f = Fixture::new(); + let tag = value("!description {units: !unit radians}"); + f.value["description"] = tag.clone(); + f.value["pipeline"]["description"] = tag.clone(); + f.value["pipeline"]["configs"]["description"] = tag.clone(); + f.value["system information"]["description"] = tag.clone(); + f.value["models"]["a_policy"]["description"] = tag.clone(); + f.value["models"]["a_policy"]["parameters"]["description"] = tag.clone(); + f.value["models"]["a_policy"]["inputs"][0]["description"] = tag.clone(); + f.value["models"]["a_policy"]["outputs"][0]["description"] = tag.clone(); + let b = f.open(); + for metadata in [ + b.metadata(), + b.pipeline_metadata(), + b.pipeline_configs(), + b.system_information(), + b.model_metadata("a_policy").unwrap(), + b.model_parameters("a_policy").unwrap(), + &b.input_info("a_policy", "features").unwrap().metadata, + &b.output_info("a_policy", "actions").unwrap().metadata, + ] { + assert_eq!(metadata["description"], tag); + } +} diff --git a/packages/leapp-runtime-rust/tests/fixtures/feedback.yaml b/packages/leapp-runtime-rust/tests/fixtures/feedback.yaml new file mode 100644 index 00000000..b56c25dc --- /dev/null +++ b/packages/leapp-runtime-rust/tests/fixtures/feedback.yaml @@ -0,0 +1,19 @@ +models: + policy: + inputs: + - {name: state, dtype: float32, shape: [2], type: tensor} + outputs: + - {name: next_state, dtype: float32, shape: [2], type: tensor} + parameters: + backend: pt2 + model_path: artifacts/policy.pt2 +pipeline: + data_flow: {} + feedback_flow: + policy/next_state: [policy/state] + inputs: {} + outputs: + policy: [next_state] + initial_values: state.safetensors +system information: + leapp config version: '1.3' diff --git a/packages/leapp-runtime-rust/tests/fixtures/pipeline.yaml b/packages/leapp-runtime-rust/tests/fixtures/pipeline.yaml new file mode 100644 index 00000000..77abc14c --- /dev/null +++ b/packages/leapp-runtime-rust/tests/fixtures/pipeline.yaml @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# Schema-1.3 example adapted from generated_configs.rst and the C++ bundle test. +# Artifacts are deliberately opaque test bytes, never executable models. +models: + z_observation: + inputs: + - name: joint_pos + dtype: float32 + shape: [3] + type: tensor + kind: state/joint/position + element_names: [[hip, knee, ankle]] + coordinate_frame: robot + outputs: + - {name: features, dtype: float32, shape: [3], type: tensor} + parameters: + backend: jit + model_path: model.bin + md5sum: 900150983cd24fb0d6963f7d28e17f72 + sha256sum: ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad + warp_segments: 0 + custom_options: {device: cpu, flags: [one, two]} + a_policy: + inputs: + - {name: features, dtype: float32, shape: '[3]', type: tensor} + - {name: state, dtype: float32, shape: [3], type: tensor} + outputs: + - name: actions + dtype: float32 + shape: [3] + type: tensor + kind: target/joint/position + element_names: [__temporal_axis__] + temporal_period_ms: 100 + - {name: next_state, dtype: float32, shape: [3], type: tensor} + parameters: + backend: pt2 + model_path: model.bin + md5sum: 900150983cd24fb0d6963f7d28e17f72 + sha256sum: ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad +pipeline: + data_flow: + z_observation/features: [a_policy/features] + feedback_flow: + a_policy/next_state: [a_policy/state] + inputs: {z_observation: [joint_pos]} + outputs: {a_policy: [next_state, actions], z_observation: [features]} + configs: {frequency: 50, project: {label: robot}} +system information: + leapp config version: '1.3' + leapp version: 0.7.1 + torch version: test-fixture diff --git a/packages/leapp-runtime-rust/tests/fixtures/policy.yaml b/packages/leapp-runtime-rust/tests/fixtures/policy.yaml new file mode 100644 index 00000000..5507b397 --- /dev/null +++ b/packages/leapp-runtime-rust/tests/fixtures/policy.yaml @@ -0,0 +1,49 @@ +models: + policy: + inputs: + - {name: z_features, dtype: float32, shape: [4], type: tensor} + - {name: a_bias, dtype: float32, shape: [4], type: tensor} + outputs: + - {name: action, dtype: float32, shape: [4], type: tensor} + parameters: + backend: pt2 + model_path: artifacts/policy.pt2 + sha256sum: "5a3949ca84f738f57e5005ddb999bd80bbcded1eef362fcf31a8ff6e21a63915" + md5sum: "eea00da3830d6875a43cdcd6e6cae801" + torch_version: '2.12.1' + description: {purpose: controller} + preprocess: + inputs: + - name: observation + dtype: float32 + shape: '[4]' + type: tensor + element_names: [first, second, third, fourth] + outputs: + - {name: features, dtype: float32, shape: [4], type: tensor} + parameters: + backend: onnx + model_path: artifacts/preprocess.onnx +pipeline: + data_flow: + preprocess/features: [policy/z_features] + feedback_flow: {} + inputs: + preprocess: [observation] + policy: [a_bias] + outputs: + policy: [action] + configs: + frequency: 50 + enabled: true + gain: 0.5 + label: "42" + optional: null + large: 18446744073709551615 + nested: {entries: [1, false, "3"]} + description: command pipeline +system information: + leapp config version: '1.3' + leapp version: '999.0' + torch version: '2.12.1' +project: {name: robotics} From 4601c5f00287d884db8bfa42a96b8444d954e595 Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Wed, 14 Oct 2026 10:21:42 -0700 Subject: [PATCH 04/21] feat: retain shared runtime lifecycle support Signed-off-by: Marc Alwan --- packages/leapp-runtime-rust/Cargo.toml | 4 + packages/leapp-runtime-rust/README.md | 11 ++ .../bindings/common/README.md | 27 +++ .../bindings/common/src/lib.rs | 7 + .../bindings/common/src/runtime.rs | 155 +++++++++++++++ .../bindings/common/tests/lifecycle.rs | 178 ++++++++++++++++++ 6 files changed, 382 insertions(+) create mode 100644 packages/leapp-runtime-rust/bindings/common/README.md create mode 100644 packages/leapp-runtime-rust/bindings/common/src/lib.rs create mode 100644 packages/leapp-runtime-rust/bindings/common/src/runtime.rs create mode 100644 packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs diff --git a/packages/leapp-runtime-rust/Cargo.toml b/packages/leapp-runtime-rust/Cargo.toml index 10596e06..f18f88e6 100644 --- a/packages/leapp-runtime-rust/Cargo.toml +++ b/packages/leapp-runtime-rust/Cargo.toml @@ -26,3 +26,7 @@ safetensors = "=0.6.2" [dev-dependencies] tempfile = "3" + +[[test]] +name = "binding_lifecycle" +path = "bindings/common/tests/lifecycle.rs" diff --git a/packages/leapp-runtime-rust/README.md b/packages/leapp-runtime-rust/README.md index 34c72050..94ebe1a1 100644 --- a/packages/leapp-runtime-rust/README.md +++ b/packages/leapp-runtime-rust/README.md @@ -57,6 +57,17 @@ cargo fmt --all --check The default workspace member is the core crate. Its inspection tests use opaque artifact bytes and require no inference SDK. +## Common lifecycle sources + +The original shared `RuntimeHandle`, `SessionHandle`, and lifecycle errors are +retained unchanged for integration with either runtime implementation. These +sources depend on concrete execution types and are not a separate compiled +crate in this bundle-only checkout. Three regression tests compile the lifecycle +source against a small test-only runtime API. + +See [common lifecycle sources](bindings/common/README.md) for integration and +testing instructions. + ## Extraction provenance This package is the common foundation extracted from the existing Triton-based diff --git a/packages/leapp-runtime-rust/bindings/common/README.md b/packages/leapp-runtime-rust/bindings/common/README.md new file mode 100644 index 00000000..e4f0eb50 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/common/README.md @@ -0,0 +1,27 @@ + + +# Original shared lifecycle sources + +`src/lib.rs` and `src/runtime.rs` are unchanged copies of the common lifecycle +used by both original runtime implementations. They retain session +ownership, locking, explicit closure, and the original error handling. + +This source depends on the concrete `Runtime`, `Session`, tensor, options, and +error types in `leapp_runtime`. Those execution types are not part of this +bundle-only checkout, so this directory is not a separate Cargo crate here. +The Python and C++ inspection bindings use the core bundle crate directly. +When integrating with either runtime branch, use these common sources with that +branch's existing common-crate manifest and execution feature. + +The root Cargo test target includes the original `runtime.rs` against a small +runtime API defined only in the test binary. Three tests verify that sessions +retain their runtime, explicit close invalidates all handles while preserving +bundle info, and close waits for an active run. These test lifecycle behavior, +not model execution. No adapter traits or cleanup behavior changes are included. + +```sh +cargo test --locked --test binding_lifecycle +``` diff --git a/packages/leapp-runtime-rust/bindings/common/src/lib.rs b/packages/leapp-runtime-rust/bindings/common/src/lib.rs new file mode 100644 index 00000000..4eaec5f5 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/common/src/lib.rs @@ -0,0 +1,7 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Shared synchronous lifecycle for language bindings. +#![forbid(unsafe_code)] +mod runtime; +pub use runtime::{Error, RuntimeHandle, SessionHandle}; diff --git a/packages/leapp-runtime-rust/bindings/common/src/runtime.rs b/packages/leapp-runtime-rust/bindings/common/src/runtime.rs new file mode 100644 index 00000000..865821e9 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/common/src/runtime.rs @@ -0,0 +1,155 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use leapp_runtime::{ + Bundle, ErrorCode, RunOptions, RunResult, Runtime, RuntimeError, RuntimeOptions, Session, + TensorView, +}; +use std::{ + fmt, + sync::{Arc, Mutex, RwLock, Weak}, +}; + +/// Binding lifecycle failures preserve native execution diagnostics. +#[derive(Debug)] +pub enum Error { + Closed, + Lock, + Runtime(RuntimeError), +} + +impl Error { + /// Category shared with the core runtime. + pub fn code(&self) -> ErrorCode { + match self { + Self::Closed => ErrorCode::InvalidArgument, + Self::Lock => ErrorCode::Internal, + Self::Runtime(e) => e.code(), + } + } + + /// Operation that failed. + pub fn phase(&self) -> &str { + match self { + Self::Runtime(e) => e.phase(), + _ => "execute", + } + } +} + +impl fmt::Display for Error { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Closed => f.write_str("runtime is closed"), + Self::Lock => f.write_str("runtime lock poisoned"), + Self::Runtime(e) => e.fmt(f), + } + } +} + +impl std::error::Error for Error {} + +type SessionSlot = Mutex>; + +struct State { + bundle: Bundle, + // Session operations hold a read guard; close waits with the write guard. + runtime: RwLock>, + sessions: Mutex>>, +} + +/// Shared runtime with explicit closure of all associated sessions. +#[derive(Clone)] +pub struct RuntimeHandle(Arc); + +/// Independent session that retains its runtime and serializes operations. +pub struct SessionHandle { + state: Arc, + slot: Arc, +} + +impl RuntimeHandle { + /// Load models using the existing core runtime. + pub fn open(bundle: &Bundle, options: RuntimeOptions) -> Result { + let runtime = Runtime::open(bundle, options).map_err(Error::Runtime)?; + Ok(Self(Arc::new(State { + bundle: bundle.clone(), + runtime: RwLock::new(Some(runtime)), + sessions: Mutex::new(Vec::new()), + }))) + } + + /// Snapshot remains available after close. + pub fn info(&self) -> &Bundle { + &self.0.bundle + } + + /// Check whether explicit close has released the runtime. + pub fn is_open(&self) -> Result { + Ok(self.0.runtime.read().map_err(|_| Error::Lock)?.is_some()) + } + + /// Create independently initialized feedback state. + pub fn create_session(&self) -> Result { + let guard = self.0.runtime.read().map_err(|_| Error::Lock)?; + let runtime = guard.as_ref().ok_or(Error::Closed)?; + let session = runtime.create_session().map_err(Error::Runtime)?; + let slot = Arc::new(Mutex::new(Some(session))); + let mut sessions = self.0.sessions.lock().map_err(|_| Error::Lock)?; + sessions.retain(|s| s.strong_count() != 0); + sessions.push(Arc::downgrade(&slot)); + Ok(SessionHandle { + state: self.0.clone(), + slot, + }) + } + + /// Wait for active native calls and close all sessions. Repeated calls succeed. + pub fn close(&self) -> Result<(), Error> { + let mut runtime = self.0.runtime.write().map_err(|_| Error::Lock)?; + let mut sessions = self.0.sessions.lock().map_err(|_| Error::Lock)?; + for weak in sessions.drain(..) { + if let Some(session) = weak.upgrade() { + session.lock().map_err(|_| Error::Lock)?.take(); + } + } + runtime.take(); + Ok(()) + } +} + +impl SessionHandle { + /// Borrow inputs until synchronous execution releases them. + pub fn run( + &self, + inputs: &[(&str, TensorView<'_>)], + options: &RunOptions, + ) -> Result { + let guard = self.state.runtime.read().map_err(|_| Error::Lock)?; + if guard.is_none() { + return Err(Error::Closed); + } + self.slot + .lock() + .map_err(|_| Error::Lock)? + .as_mut() + .ok_or(Error::Closed)? + .run_with_options(inputs, options) + .map_err(Error::Runtime) + } + + /// Restore the session's explicit feedback state. + pub fn reset(&self) -> Result<(), Error> { + let guard = self.state.runtime.read().map_err(|_| Error::Lock)?; + if guard.is_none() { + return Err(Error::Closed); + } + self.slot + .lock() + .map_err(|_| Error::Lock)? + .as_mut() + .ok_or(Error::Closed)? + .reset() + .map_err(Error::Runtime) + } +} diff --git a/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs b/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs new file mode 100644 index 00000000..dc6b9637 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs @@ -0,0 +1,178 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Compile the original lifecycle source against a test-only runtime API. +// The real dependency still supplies Bundle and ErrorCode; no production shim. +extern crate leapp_runtime as bundle_api; +extern crate self as leapp_runtime; + +pub use bundle_api::{Bundle, ErrorCode}; +use std::{ + fmt, + sync::{mpsc, Arc, Weak}, + thread, + time::Duration, +}; + +#[path = "../src/runtime.rs"] +mod lifecycle; +use lifecycle::{Error, RuntimeHandle}; + +pub type RuntimeOptions = Arc<()>; +pub type TensorView<'a> = &'a [u8]; +pub type RunResult = (); + +#[derive(Default)] +pub struct RunOptions(Option<(mpsc::Sender<()>, mpsc::Receiver<()>)>); + +pub struct Runtime(Arc<()>); + +pub struct Session(Weak<()>); + +#[derive(Debug)] +pub struct RuntimeError; + +impl RuntimeError { + pub fn code(&self) -> ErrorCode { + ErrorCode::Internal + } + + pub fn phase(&self) -> &str { + "execute" + } +} + +impl fmt::Display for RuntimeError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("test runtime error") + } +} + +impl Runtime { + pub fn open(_: &Bundle, owner: RuntimeOptions) -> Result { + Ok(Self(owner)) + } + + pub fn create_session(&self) -> Result { + Ok(Session(Arc::downgrade(&self.0))) + } +} + +impl Session { + pub fn run_with_options( + &mut self, + _: &[(&str, TensorView<'_>)], + options: &RunOptions, + ) -> Result { + assert!( + self.0.upgrade().is_some(), + "runtime dropped during session lifetime" + ); + if let Some((entered, resume)) = &options.0 { + entered.send(()).unwrap(); + resume.recv_timeout(Duration::from_secs(5)).unwrap(); + } + Ok(()) + } + + pub fn reset(&mut self) -> Result<(), RuntimeError> { + assert!( + self.0.upgrade().is_some(), + "runtime dropped during session lifetime" + ); + Ok(()) + } +} + +fn runtime() -> (RuntimeHandle, Weak<()>) { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("empty.yaml"); + std::fs::write( + &path, + "models: {}\npipeline: {inputs: {}, outputs: {}}\nsystem information: {}\n", + ) + .unwrap(); + let owner = Arc::new(()); + let weak = Arc::downgrade(&owner); + ( + RuntimeHandle::open(&Bundle::open(path).unwrap(), owner).unwrap(), + weak, + ) +} + +#[test] +fn sessions_keep_runtime_alive_after_the_last_runtime_handle_is_dropped() { + let (runtime, owner) = runtime(); + let session = runtime.create_session().unwrap(); + drop(runtime); + assert!(owner.upgrade().is_some()); + session.run(&[], &RunOptions::default()).unwrap(); + session.reset().unwrap(); + drop(session); + assert!(owner.upgrade().is_none()); +} + +#[test] +fn close_invalidates_all_handles_and_retains_bundle_info() { + let (runtime, owner) = runtime(); + let clone = runtime.clone(); + let path = runtime.info().yaml_path().to_owned(); + let sessions = [ + runtime.create_session().unwrap(), + runtime.create_session().unwrap(), + ]; + clone.close().unwrap(); + assert!(!runtime.is_open().unwrap()); + assert!(owner.upgrade().is_none()); + assert_eq!(runtime.info().yaml_path(), path); + assert!(matches!(runtime.create_session(), Err(Error::Closed))); + for session in sessions { + assert!(matches!( + session.run(&[], &RunOptions::default()), + Err(Error::Closed) + )); + let error = session.reset().unwrap_err(); + assert!(matches!(error, Error::Closed)); + assert_eq!(error.code(), ErrorCode::InvalidArgument); + assert_eq!(error.phase(), "execute"); + } + runtime.close().unwrap(); +} + +#[test] +fn close_waits_until_an_active_run_finishes() { + let (runtime, owner) = runtime(); + let session = Arc::new(runtime.create_session().unwrap()); + let (entered, running) = mpsc::channel(); + let (resume, paused) = mpsc::channel(); + let run = { + let session = session.clone(); + thread::spawn(move || { + session + .run(&[], &RunOptions(Some((entered, paused)))) + .unwrap() + }) + }; + running.recv_timeout(Duration::from_secs(5)).unwrap(); + let (started, closing) = mpsc::channel(); + let (finished, closed) = mpsc::channel(); + let close = { + let runtime = runtime.clone(); + thread::spawn(move || { + started.send(()).unwrap(); + runtime.close().unwrap(); + finished.send(()).unwrap(); + }) + }; + closing.recv_timeout(Duration::from_secs(5)).unwrap(); + assert!(matches!( + closed.recv_timeout(Duration::from_millis(50)), + Err(mpsc::RecvTimeoutError::Timeout) + )); + assert!(owner.upgrade().is_some()); + resume.send(()).unwrap(); + run.join().unwrap(); + close.join().unwrap(); + assert!(owner.upgrade().is_none()); + assert!(matches!(session.reset(), Err(Error::Closed))); +} From 188f45c4434a3be5924a2fd1fd26560f1a49dea1 Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Wed, 14 Oct 2026 10:22:25 -0700 Subject: [PATCH 05/21] feat: add Python bundle inspection bindings Signed-off-by: Marc Alwan --- packages/leapp-runtime-rust/Cargo.lock | 136 ++++++++++- packages/leapp-runtime-rust/Cargo.toml | 2 +- packages/leapp-runtime-rust/README.md | 12 +- .../bindings/python/Cargo.toml | 16 ++ .../bindings/python/README.md | 30 +++ .../bindings/python/pyproject.toml | 18 ++ .../python/python/leapp_runtime/__init__.py | 35 +++ .../python/python/leapp_runtime/_native.pyi | 75 ++++++ .../python/python/leapp_runtime/py.typed | 0 .../bindings/python/src/bundle.rs | 228 ++++++++++++++++++ .../bindings/python/src/lib.rs | 37 +++ .../bindings/python/tests/test_bundle.py | 121 ++++++++++ 12 files changed, 707 insertions(+), 3 deletions(-) create mode 100644 packages/leapp-runtime-rust/bindings/python/Cargo.toml create mode 100644 packages/leapp-runtime-rust/bindings/python/README.md create mode 100644 packages/leapp-runtime-rust/bindings/python/pyproject.toml create mode 100644 packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/__init__.py create mode 100644 packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi create mode 100644 packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/py.typed create mode 100644 packages/leapp-runtime-rust/bindings/python/src/bundle.rs create mode 100644 packages/leapp-runtime-rust/bindings/python/src/lib.rs create mode 100644 packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py diff --git a/packages/leapp-runtime-rust/Cargo.lock b/packages/leapp-runtime-rust/Cargo.lock index aa8bc8a5..5a64be8e 100644 --- a/packages/leapp-runtime-rust/Cargo.lock +++ b/packages/leapp-runtime-rust/Cargo.lock @@ -8,6 +8,12 @@ version = "1.0.104" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + [[package]] name = "bitflags" version = "2.13.2" @@ -107,6 +113,12 @@ version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + [[package]] name = "indexmap" version = "2.14.2" @@ -119,6 +131,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "indoc" +version = "2.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a37b2691796cffeb8a8cd305ac66e65841559f147f4e63231d0eafa4db5384d1" +dependencies = [ + "rustversion", +] + [[package]] name = "itoa" version = "1.0.18" @@ -140,6 +161,14 @@ dependencies = [ "tempfile", ] +[[package]] +name = "leapp-runtime-python" +version = "0.1.0" +dependencies = [ + "leapp-runtime", + "pyo3", +] + [[package]] name = "libc" version = "0.2.189" @@ -158,12 +187,27 @@ version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + [[package]] name = "once_cell" version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "portable-atomic" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" + [[package]] name = "proc-macro2" version = "1.0.107" @@ -173,6 +217,67 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "pyo3" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab53c047fcd1a1d2a8820fe84f05d6be69e9526be40cb03b73f86b6b03e6d87d" +dependencies = [ + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b455933107de8642b4487ed26d912c2d899dec6114884214a0b3bb3be9261ea6" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c85c9cbfaddf651b1221594209aed57e9e5cff63c4d11d1feead529b872a089" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a5b10c9bf9888125d917fb4d2ca2d25c8df94c7ab5a52e13313a07e050a3b02" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03b51720d314836e53327f5871d4c0cfb4fb37cc2c4a11cc71907a86342c40f9" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn 2.0.119", +] + [[package]] name = "quote" version = "1.0.47" @@ -201,6 +306,12 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + [[package]] name = "ryu" version = "1.0.23" @@ -244,7 +355,7 @@ checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.6", ] [[package]] @@ -295,6 +406,17 @@ dependencies = [ "digest", ] +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "syn" version = "3.0.6" @@ -306,6 +428,12 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + [[package]] name = "tempfile" version = "3.27.0" @@ -331,6 +459,12 @@ version = "1.0.26" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d245f478577f809a851594d02313b640fb437e0bb33866753cff937863096954" +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + [[package]] name = "unsafe-libyaml" version = "0.2.11" diff --git a/packages/leapp-runtime-rust/Cargo.toml b/packages/leapp-runtime-rust/Cargo.toml index f18f88e6..eea27cdd 100644 --- a/packages/leapp-runtime-rust/Cargo.toml +++ b/packages/leapp-runtime-rust/Cargo.toml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 [workspace] -members = [] +members = ["bindings/python"] default-members = ["."] resolver = "2" diff --git a/packages/leapp-runtime-rust/README.md b/packages/leapp-runtime-rust/README.md index 94ebe1a1..887d6306 100644 --- a/packages/leapp-runtime-rust/README.md +++ b/packages/leapp-runtime-rust/README.md @@ -38,7 +38,7 @@ execution. Bundle loading validates graph ports and connections, artifacts, and optional initial-value files. It does not check whether an inference engine can execute the graph. Declared external inputs can remain unresolved, matching the source -implementations' inspection contract. +implementations' inspection contract. Python represents missing input descriptors as `None`. The loader normalizes known backend aliases, accepts custom nonempty backend names, and leaves adapter availability to execution. @@ -68,6 +68,16 @@ source against a small test-only runtime API. See [common lifecycle sources](bindings/common/README.md) for integration and testing instructions. +## Python bindings + +The PyO3 binding exposes bundle inspection, detached metadata and port +descriptors, typed stubs, and structured exceptions. Building it requires a +Rust toolchain and Python 3.10 or newer. See +[Python bindings](bindings/python/README.md) for installation and tests. + +This distribution shares the `leapp_runtime` import name with the full runtime +implementations; install it in a separate environment. + ## Extraction provenance This package is the common foundation extracted from the existing Triton-based diff --git a/packages/leapp-runtime-rust/bindings/python/Cargo.toml b/packages/leapp-runtime-rust/bindings/python/Cargo.toml new file mode 100644 index 00000000..bb1f7608 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/python/Cargo.toml @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "leapp-runtime-python" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +publish = false + +[lib] +name = "_native" +crate-type = ["cdylib"] + +[dependencies] +leapp-runtime = { path = "../.." } +pyo3 = { version = "=0.27.2", features = ["extension-module"] } diff --git a/packages/leapp-runtime-rust/bindings/python/README.md b/packages/leapp-runtime-rust/bindings/python/README.md new file mode 100644 index 00000000..d8e3c9f0 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/python/README.md @@ -0,0 +1,30 @@ +# Python bundle API + +Build from the Rust package directory, using Python 3.10 or newer in a virtual +environment and a Rust toolchain: + +```sh +python -m venv .venv +. .venv/bin/activate +python -m pip install maturin pytest +maturin develop --locked --manifest-path bindings/python/Cargo.toml +python -m pytest bindings/python/tests +``` + +```python +from leapp_runtime import Bundle, BundleError + +bundle = Bundle.open("export/policy.yaml") +for name, spec in bundle.pipeline_outputs().items(): + print(name, spec.dtype, spec.shape) +``` + +Metadata and port descriptors are detached from the immutable Rust snapshot. +Editing returned metadata does not change later inspection results. Unknown +YAML tags become `TaggedValue`; mappings with unhashable or colliding Python +keys become `MetadataMapping` so no entries are lost. + +`BundleError` exposes `code` (an `ErrorCode`), `phase`, `path`, and `message`. +Only bundle inspection is exported. This distribution uses the same +`leapp_runtime` package name as the full runtimes, so use a separate environment. +No inference SDK or NumPy is required. diff --git a/packages/leapp-runtime-rust/bindings/python/pyproject.toml b/packages/leapp-runtime-rust/bindings/python/pyproject.toml new file mode 100644 index 00000000..02c2e9c5 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/python/pyproject.toml @@ -0,0 +1,18 @@ +[build-system] +requires = ["maturin>=1.9,<2"] +build-backend = "maturin" + +[project] +name = "leapp-runtime" +version = "0.1.0" +requires-python = ">=3.10" +description = "Backend-independent LEAPP bundle loading and inspection" +license = "Apache-2.0" +dependencies = [] + +[project.optional-dependencies] +test = ["pytest>=8"] + +[tool.maturin] +python-source = "python" +module-name = "leapp_runtime._native" diff --git a/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/__init__.py b/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/__init__.py new file mode 100644 index 00000000..18f73185 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/__init__.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Backend-independent LEAPP bundle loading and inspection.""" + +from enum import IntEnum + +from ._native import ( + Bundle, + BundleError, + MetadataMapping, + PortRef, + PortSpec, + TaggedValue, +) + + +class ErrorCode(IntEnum): + OK = 0 + INVALID_ARGUMENT = 1 + NOT_FOUND = 2 + UNSUPPORTED = 3 + IO = 4 + INTERNAL = 5 + + +__all__ = [ + "Bundle", + "BundleError", + "ErrorCode", + "MetadataMapping", + "PortRef", + "PortSpec", + "TaggedValue", +] diff --git a/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi b/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi new file mode 100644 index 00000000..567f4b28 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi @@ -0,0 +1,75 @@ +# Copyright 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from os import PathLike +from typing import Any + +class BundleError(Exception): + code: int + phase: str + message: str + path: PathLike[str] + +class PortRef: + @property + def node(self) -> str: ... + + @property + def port(self) -> str: ... + + def name(self) -> str: ... + +class PortSpec: + @property + def name(self) -> str: ... + + @property + def dtype(self) -> str: ... + + @property + def shape(self) -> list[int]: ... + + @property + def tensor_type(self) -> str: ... + + @property + def metadata(self) -> dict[str, Any]: ... + +class TaggedValue: + @property + def tag(self) -> str: ... + + @property + def value(self) -> Any: ... + +class MetadataMapping: + @property + def entries(self) -> list[tuple[Any, Any]]: ... + +class Bundle: + @staticmethod + def open(path: str | PathLike[str]) -> Bundle: ... + + def yaml_path(self) -> PathLike[str]: ... + + def model_names(self) -> list[str]: ... + + def pipeline_inputs(self) -> dict[str, PortSpec | None]: ... + + def pipeline_outputs(self) -> dict[str, PortSpec | None]: ... + + def pipeline_configs(self) -> dict[str, Any]: ... + + def pipeline_metadata(self) -> dict[str, Any]: ... + + def metadata(self) -> dict[str, Any]: ... + + def system_information(self) -> dict[str, Any]: ... + + def model_parameters(self, node: str) -> dict[str, Any] | None: ... + + def model_metadata(self, node: str) -> dict[str, Any] | None: ... + + def input_info(self, node: str, port: str) -> PortSpec | None: ... + + def output_info(self, node: str, port: str) -> PortSpec | None: ... diff --git a/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/py.typed b/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/py.typed new file mode 100644 index 00000000..e69de29b diff --git a/packages/leapp-runtime-rust/bindings/python/src/bundle.rs b/packages/leapp-runtime-rust/bindings/python/src/bundle.rs new file mode 100644 index 00000000..4e388718 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/python/src/bundle.rs @@ -0,0 +1,228 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use leapp_runtime::{Bundle, Metadata, MetadataValue}; +use pyo3::{ + prelude::*, + types::{PyDict, PyList}, +}; +use std::path::PathBuf; + +#[pyclass(name = "PortRef", frozen, get_all, module = "leapp_runtime")] +pub struct PortRef { + node: String, + port: String, +} + +#[pymethods] +impl PortRef { + fn name(&self) -> String { + format!("{}/{}", self.node, self.port) + } + + fn __repr__(&self) -> String { + format!("PortRef({:?}, {:?})", self.node, self.port) + } +} + +#[pyclass(frozen, get_all, module = "leapp_runtime")] +pub struct PortSpec { + name: String, + dtype: String, + shape: Vec, + tensor_type: String, + metadata: Py, +} + +#[pyclass(frozen, get_all, module = "leapp_runtime")] +pub struct TaggedValue { + tag: String, + value: Py, +} + +/// Lossless fallback for YAML mappings containing keys Python cannot hash. +#[pyclass(frozen, get_all, module = "leapp_runtime")] +pub struct MetadataMapping { + entries: Py, +} + +pub fn value(py: Python<'_>, v: &MetadataValue) -> PyResult> { + Ok(match v { + MetadataValue::Null => py.None(), + MetadataValue::Bool(v) => v.into_pyobject(py)?.to_owned().into_any().unbind(), + MetadataValue::Number(v) => { + if let Some(n) = v.as_i64() { + n.into_pyobject(py)?.into_any().unbind() + } else if let Some(n) = v.as_u64() { + n.into_pyobject(py)?.into_any().unbind() + } else { + v.as_f64() + .expect("YAML number") + .into_pyobject(py)? + .into_any() + .unbind() + } + } + MetadataValue::String(v) => v.into_pyobject(py)?.into_any().unbind(), + MetadataValue::Sequence(v) => { + let items = v + .iter() + .map(|v| value(py, v)) + .collect::>>()?; + PyList::new(py, items)?.into_any().unbind() + } + MetadataValue::Mapping(v) => { + let entries = v + .iter() + .map(|(k, v)| Ok((value(py, k)?, value(py, v)?))) + .collect::>>()?; + // Do not stringify or discard YAML keys that have no Python dict equivalent. + if entries.iter().any(|(k, _)| k.bind(py).hash().is_err()) { + let entries = PyList::new(py, entries)?.into_any().unbind(); + Py::new(py, MetadataMapping { entries })?.into_any() + } else { + let dict = PyDict::new(py); + let mut collision = false; + for (key, val) in &entries { + collision |= dict.contains(key)?; + dict.set_item(key, val)?; + } + if collision { + let entries = PyList::new(py, entries)?.into_any().unbind(); + Py::new(py, MetadataMapping { entries })?.into_any() + } else { + dict.into_any().unbind() + } + } + } + MetadataValue::Tagged(v) => Py::new( + py, + TaggedValue { + tag: v.tag.to_string(), + value: value(py, &v.value)?, + }, + )? + .into_any(), + }) +} + +pub fn metadata(py: Python<'_>, data: &Metadata) -> PyResult> { + let result = PyDict::new(py); + for (key, val) in data { + result.set_item(key, value(py, val)?)?; + } + Ok(result.into_any().unbind()) +} + +fn spec(py: Python<'_>, v: Option<&leapp_runtime::PortSpec>) -> PyResult> { + v.map(|v| { + Ok(PortSpec { + name: v.name.clone(), + dtype: format!("{:?}", v.dtype).to_lowercase(), + shape: v.shape.clone(), + tensor_type: v.tensor_type.clone(), + metadata: metadata(py, &v.metadata)?, + }) + }) + .transpose() +} + +#[pyclass(name = "Bundle", frozen, module = "leapp_runtime")] +pub struct PyBundle { + pub inner: Bundle, +} + +#[pymethods] +impl PyBundle { + #[staticmethod] + /// Parse and validate a bundle without loading its model backends. + fn open(py: Python<'_>, path: PathBuf) -> PyResult { + let inner = py.detach(|| Bundle::open(&path)).map_err(|e| { + let error = super::diagnostic( + py, + super::BundleError::new_err(e.message()), + e.code(), + e.phase(), + ); + if let Err(e) = error.value(py).setattr("path", e.path().to_path_buf()) { + return e; + } + error + })?; + Ok(Self { inner }) + } + + fn yaml_path(&self) -> PathBuf { + self.inner.yaml_path().into() + } + + fn model_names(&self) -> Vec { + self.inner + .model_names() + .into_iter() + .map(String::from) + .collect() + } + + /// Ordered mapping of qualified input names to detached port descriptors. + fn pipeline_inputs<'py>(&self, py: Python<'py>) -> PyResult> { + let result = PyDict::new(py); + for port in self.inner.pipeline_inputs() { + result.set_item( + port.name(), + spec(py, self.inner.input_info(&port.node, &port.port))?, + )?; + } + Ok(result) + } + + /// Ordered mapping of qualified output names to detached port descriptors. + fn pipeline_outputs<'py>(&self, py: Python<'py>) -> PyResult> { + let result = PyDict::new(py); + for port in self.inner.pipeline_outputs() { + result.set_item( + port.name(), + spec(py, self.inner.output_info(&port.node, &port.port))?, + )?; + } + Ok(result) + } + + fn pipeline_configs(&self, py: Python<'_>) -> PyResult> { + metadata(py, self.inner.pipeline_configs()) + } + + fn pipeline_metadata(&self, py: Python<'_>) -> PyResult> { + metadata(py, self.inner.pipeline_metadata()) + } + + fn metadata(&self, py: Python<'_>) -> PyResult> { + metadata(py, self.inner.metadata()) + } + + fn system_information(&self, py: Python<'_>) -> PyResult> { + metadata(py, self.inner.system_information()) + } + + fn model_parameters(&self, py: Python<'_>, node: &str) -> PyResult>> { + self.inner + .model_parameters(node) + .map(|v| metadata(py, v)) + .transpose() + } + + fn model_metadata(&self, py: Python<'_>, node: &str) -> PyResult>> { + self.inner + .model_metadata(node) + .map(|v| metadata(py, v)) + .transpose() + } + + fn input_info(&self, py: Python<'_>, node: &str, port: &str) -> PyResult> { + spec(py, self.inner.input_info(node, port)) + } + + fn output_info(&self, py: Python<'_>, node: &str, port: &str) -> PyResult> { + spec(py, self.inner.output_info(node, port)) + } +} diff --git a/packages/leapp-runtime-rust/bindings/python/src/lib.rs b/packages/leapp-runtime-rust/bindings/python/src/lib.rs new file mode 100644 index 00000000..7f03cf8b --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/python/src/lib.rs @@ -0,0 +1,37 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![deny(unsafe_op_in_unsafe_fn)] + +mod bundle; + +use leapp_runtime::ErrorCode; +use pyo3::{create_exception, exceptions::PyException, prelude::*}; + +create_exception!(leapp_runtime, BundleError, PyException); + +fn diagnostic(py: Python<'_>, error: PyErr, code: ErrorCode, phase: &str) -> PyErr { + let value = error.value(py); + let result = (|| -> PyResult<()> { + let code = py + .import("leapp_runtime")? + .getattr("ErrorCode")? + .call1((code as i32,))?; + value.setattr("code", code)?; + value.setattr("phase", phase)?; + value.setattr("message", value.str()?)?; + Ok(()) + })(); + result.err().unwrap_or(error) +} + +#[pymodule] +fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add("BundleError", m.py().get_type::())?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) +} diff --git a/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py b/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py new file mode 100644 index 00000000..2ebb9de3 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bundle inspection needs only YAML and opaque artifacts, with no model SDK.""" + +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import leapp_runtime as lr +import pytest + +YAML = (Path(__file__).parents[3] / "tests/fixtures/pipeline.yaml").read_text() + + +@pytest.fixture +def bundle_path(tmp_path): + (tmp_path / "model.bin").write_bytes(b"abc") + path = tmp_path / "pipeline.yaml" + path.write_text(YAML) + return path + + +def test_ordered_inspection_and_detached_values(bundle_path): + bundle = lr.Bundle.open(bundle_path) + assert bundle.yaml_path() == bundle_path + assert bundle.model_names() == ["z_observation", "a_policy"] + assert list(bundle.pipeline_inputs()) == ["z_observation/joint_pos"] + assert list(bundle.pipeline_outputs()) == [ + "a_policy/next_state", + "a_policy/actions", + "z_observation/features", + ] + spec = bundle.output_info("a_policy", "actions") + assert spec.dtype == "float32" + assert spec.shape == [3] + assert spec.tensor_type == "tensor" + assert spec.metadata["temporal_period_ms"] == 100 + assert bundle.system_information()["leapp config version"] == "1.3" + assert bundle.pipeline_configs()["frequency"] == 50 + parameters = bundle.model_parameters("z_observation") + assert parameters["model_path"] == "model.bin" + parameters["custom_options"]["flags"].clear() + spec.metadata["temporal_period_ms"] = 999 + assert bundle.model_parameters("z_observation")["custom_options"]["flags"] == [ + "one", + "two", + ] + assert bundle.output_info("a_policy", "actions").metadata["temporal_period_ms"] == 100 + assert bundle.input_info("absent", "port") is None + assert bundle.model_metadata("absent") is None + assert bundle.model_parameters("absent") is None + with pytest.raises(AttributeError): + spec.dtype = "int8" + del bundle + assert spec.shape == [3] + + +def test_custom_backend_and_unresolved_input(bundle_path): + bundle_path.write_text( + YAML.replace("backend: pt2", "backend: custom").replace( + "inputs: {z_observation: [joint_pos]}", "inputs: {absent: [port]}" + ) + ) + bundle = lr.Bundle.open(bundle_path) + assert bundle.model_parameters("a_policy")["backend"] == "custom" + assert bundle.pipeline_inputs() == {"absent/port": None} + + +def test_lossless_metadata(bundle_path): + bundle_path.write_text( + YAML + + """ +tagged: !sensor {axes: [x, y], enabled: true} +sequence_keys: {[x, y]: coordinates} +colliding_keys: {true: boolean, 1: integer} +unsigned: 18446744073709551615 +""" + ) + metadata = lr.Bundle.open(bundle_path).metadata() + assert isinstance(metadata["tagged"], lr.TaggedValue) + assert metadata["tagged"].tag == "!sensor" + assert metadata["tagged"].value == {"axes": ["x", "y"], "enabled": True} + assert isinstance(metadata["sequence_keys"], lr.MetadataMapping) + assert metadata["sequence_keys"].entries == [(["x", "y"], "coordinates")] + assert isinstance(metadata["colliding_keys"], lr.MetadataMapping) + assert metadata["colliding_keys"].entries == [(True, "boolean"), (1, "integer")] + assert metadata["unsigned"] == 18446744073709551615 + + +@pytest.mark.parametrize( + "failure,code,phase", + [ + ("missing", lr.ErrorCode.NOT_FOUND, "open"), + ("parse", lr.ErrorCode.INVALID_ARGUMENT, "parse"), + ("artifact", lr.ErrorCode.INVALID_ARGUMENT, "artifact"), + ], +) +def test_diagnostics(bundle_path, failure, code, phase): + if failure == "missing": + bundle_path.unlink() + elif failure == "parse": + bundle_path.write_text("models: [") + else: + (bundle_path.parent / "model.bin").write_bytes(b"wrong hash") + with pytest.raises(lr.BundleError) as caught: + lr.Bundle.open(bundle_path) + assert caught.value.code is code + assert caught.value.phase == phase + assert caught.value.path == bundle_path + assert caught.value.message + + +def test_snapshot_and_concurrent_reads(bundle_path): + bundle = lr.Bundle.open(bundle_path) + bundle_path.unlink() + with ThreadPoolExecutor(max_workers=4) as pool: + results = list( + pool.map(lambda _: bundle.input_info("a_policy", "features").shape, range(24)) + ) + assert results == [[3]] * 24 + assert bundle.model_names() == ["z_observation", "a_policy"] From 2b281a0798ef2bd833e8b3e0c9ff5bf9a0ba4eac Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Wed, 14 Oct 2026 10:22:45 -0700 Subject: [PATCH 06/21] feat: add C++ bundle inspection bindings Signed-off-by: Marc Alwan --- packages/leapp-runtime-rust/Cargo.lock | 187 +++++++++++ packages/leapp-runtime-rust/Cargo.toml | 2 +- packages/leapp-runtime-rust/README.md | 96 +++--- .../bindings/cxx/.clang-format | 4 + .../bindings/cxx/CMakeLists.txt | 70 +++++ .../bindings/cxx/CPPLINT.cfg | 2 + .../bindings/cxx/Cargo.toml | 16 + .../leapp-runtime-rust/bindings/cxx/README.md | 39 +++ .../leapp-runtime-rust/bindings/cxx/build.rs | 9 + .../cxx/cmake/leapp_runtimeConfig.cmake.in | 4 + .../bindings/cxx/include/leapp/bundle.h | 79 +++++ .../bindings/cxx/include/leapp/data_type.h | 30 ++ .../bindings/cxx/include/leapp/metadata.h | 46 +++ .../bindings/cxx/include/leapp/status.h | 46 +++ .../bindings/cxx/src/bridge.rs | 95 ++++++ .../bindings/cxx/src/bundle.cc | 297 ++++++++++++++++++ .../bindings/cxx/src/bundle.rs | 192 +++++++++++ .../bindings/cxx/src/dtype.rs | 24 ++ .../bindings/cxx/src/lib.rs | 32 ++ .../bindings/cxx/tests/bundle_test.cc | 107 +++++++ .../cxx/tests/consumer/CMakeLists.txt | 7 + .../bindings/cxx/tests/consumer/main.cc | 13 + 22 files changed, 1358 insertions(+), 39 deletions(-) create mode 100644 packages/leapp-runtime-rust/bindings/cxx/.clang-format create mode 100644 packages/leapp-runtime-rust/bindings/cxx/CMakeLists.txt create mode 100644 packages/leapp-runtime-rust/bindings/cxx/CPPLINT.cfg create mode 100644 packages/leapp-runtime-rust/bindings/cxx/Cargo.toml create mode 100644 packages/leapp-runtime-rust/bindings/cxx/README.md create mode 100644 packages/leapp-runtime-rust/bindings/cxx/build.rs create mode 100644 packages/leapp-runtime-rust/bindings/cxx/cmake/leapp_runtimeConfig.cmake.in create mode 100644 packages/leapp-runtime-rust/bindings/cxx/include/leapp/bundle.h create mode 100644 packages/leapp-runtime-rust/bindings/cxx/include/leapp/data_type.h create mode 100644 packages/leapp-runtime-rust/bindings/cxx/include/leapp/metadata.h create mode 100644 packages/leapp-runtime-rust/bindings/cxx/include/leapp/status.h create mode 100644 packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs create mode 100644 packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc create mode 100644 packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs create mode 100644 packages/leapp-runtime-rust/bindings/cxx/src/dtype.rs create mode 100644 packages/leapp-runtime-rust/bindings/cxx/src/lib.rs create mode 100644 packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc create mode 100644 packages/leapp-runtime-rust/bindings/cxx/tests/consumer/CMakeLists.txt create mode 100644 packages/leapp-runtime-rust/bindings/cxx/tests/consumer/main.cc diff --git a/packages/leapp-runtime-rust/Cargo.lock b/packages/leapp-runtime-rust/Cargo.lock index 5a64be8e..fd5be50c 100644 --- a/packages/leapp-runtime-rust/Cargo.lock +++ b/packages/leapp-runtime-rust/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + [[package]] name = "anyhow" version = "1.0.104" @@ -29,12 +35,59 @@ dependencies = [ "generic-array", ] +[[package]] +name = "cc" +version = "1.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6651c9ed80effdc7db0ff72512157f901af5e3549e341e24b1dd4887d836d838" +dependencies = [ + "find-msvc-tools", + "shlex", +] + [[package]] name = "cfg-if" version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4e7648175b45a9a48536d676f68d918270699102aa8dab5496df06904c914600" +[[package]] +name = "clap" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa8876b300ab35ba921adea3dfd70157a46249b33f95c9084ae5709785478946" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0797fb7aeb1406c84efac526901f7ec3ead2124f946b494e72879d4b54704d" +dependencies = [ + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_lex" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c133bc6a41be0d194c306b5506d15e6feeea7b1d6604bd3f8310dfb2ca96486" + +[[package]] +name = "codespan-reporting" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" +dependencies = [ + "serde", + "termcolor", + "unicode-width", +] + [[package]] name = "cpufeatures" version = "0.2.17" @@ -54,6 +107,68 @@ dependencies = [ "typenum", ] +[[package]] +name = "cxx" +version = "1.0.205" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8a02bfeba9074b2bbb39492c9b0908ec854f455a3cccdcde7a1de55e2315e40" +dependencies = [ + "cc", + "cxx-build", + "cxxbridge-cmd", + "cxxbridge-flags", + "cxxbridge-macro", + "foldhash", + "link-cplusplus", +] + +[[package]] +name = "cxx-build" +version = "1.0.205" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbc2ac76997619d2d888b58ad79488f1c835de41ca4dd3dd7082acf4ecb0a50f" +dependencies = [ + "cc", + "codespan-reporting", + "indexmap", + "proc-macro2", + "quote", + "scratch", + "syn 3.0.6", +] + +[[package]] +name = "cxxbridge-cmd" +version = "1.0.205" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54a486ed16a3f2a053fe7f5b98db4c1a59846a23e8cb9e49f9d1bd48830d51a2" +dependencies = [ + "clap", + "codespan-reporting", + "indexmap", + "proc-macro2", + "quote", + "syn 3.0.6", +] + +[[package]] +name = "cxxbridge-flags" +version = "1.0.205" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c72569be94e2283764d334cd285980b747dbe221e27e315484eb6bcd93fa808a" + +[[package]] +name = "cxxbridge-macro" +version = "1.0.205" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "423e795284a7f7b3f843c956a5322c4bc497b60c2dd1a6339da10313cb18d575" +dependencies = [ + "indexmap", + "proc-macro2", + "quote", + "syn 3.0.6", +] + [[package]] name = "digest" version = "0.10.7" @@ -86,6 +201,18 @@ version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" +[[package]] +name = "find-msvc-tools" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aedcfb3409746eddb02b9e19ebda1c3394f759a152e48ee875a0844d1b955484" + +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "generic-array" version = "0.14.7" @@ -161,6 +288,15 @@ dependencies = [ "tempfile", ] +[[package]] +name = "leapp-runtime-cxx" +version = "0.1.0" +dependencies = [ + "cxx", + "cxx-build", + "leapp-runtime", +] + [[package]] name = "leapp-runtime-python" version = "0.1.0" @@ -175,6 +311,15 @@ version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" +[[package]] +name = "link-cplusplus" +version = "1.0.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f78c730aaa7d0b9336a299029ea49f9ee53b0ed06e9202e8cb7db9bae7b8c82" +dependencies = [ + "cc", +] + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -328,6 +473,12 @@ dependencies = [ "serde_json", ] +[[package]] +name = "scratch" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d68f2ec51b097e4c1a75b681a8bec621909b5e91f15bb7b840c4f2f7b01148b2" + [[package]] name = "serde" version = "1.0.229" @@ -406,6 +557,18 @@ dependencies = [ "digest", ] +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "syn" version = "2.0.119" @@ -447,6 +610,15 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "termcolor" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" +dependencies = [ + "winapi-util", +] + [[package]] name = "typenum" version = "1.20.1" @@ -459,6 +631,12 @@ version = "1.0.26" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d245f478577f809a851594d02313b640fb437e0bb33866753cff937863096954" +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + [[package]] name = "unindent" version = "0.2.4" @@ -477,6 +655,15 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys", +] + [[package]] name = "windows-link" version = "0.2.1" diff --git a/packages/leapp-runtime-rust/Cargo.toml b/packages/leapp-runtime-rust/Cargo.toml index eea27cdd..aaa45100 100644 --- a/packages/leapp-runtime-rust/Cargo.toml +++ b/packages/leapp-runtime-rust/Cargo.toml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 [workspace] -members = ["bindings/python"] +members = ["bindings/python", "bindings/cxx"] default-members = ["."] resolver = "2" diff --git a/packages/leapp-runtime-rust/README.md b/packages/leapp-runtime-rust/README.md index 887d6306..b7653bd9 100644 --- a/packages/leapp-runtime-rust/README.md +++ b/packages/leapp-runtime-rust/README.md @@ -3,16 +3,40 @@ SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All SPDX-License-Identifier: Apache-2.0 --> -# LEAPP Rust bundle API - -This package loads and validates exported LEAPP bundles and exposes their -public inspection API in Rust. Bundle loading does not initialize an inference -backend. It provides `Bundle`, `PortRef`, `PortSpec`, `DataType`, `Metadata`, -`MetadataValue`, `BundleError`, and `ErrorCode`. - -The loader validates graph connections, artifact paths, optional SHA-256 hashes, -and safetensors initial state. It preserves ordered metadata, nested YAML tags, -duplicate mapping keys, and YAML merges. +# LEAPP Rust bundle API and shared lifecycle sources + +This is the standalone common foundation extracted from the Triton-based +and native runtime implementations. It loads and validates +exported LEAPP bundles and exposes their public inspection API in Rust, Python, +and C++. The original common runtime/session lifecycle is retained unchanged +for integration with either engine. Bundle loading does not initialize an +inference backend. + +The package path and library names match the runtime branches to make this a +base for either implementation. This worktree does not refactor those branches +to depend on a new crate. The two Python distributions share the `leapp_runtime` +import name and should be installed in separate environments. + +## Included code + +| Layer | Contents | +| --- | --- | +| Rust core | `Bundle`, YAML schema parsing, graph and artifact validation, SHA-256 checks, safetensors initial-state validation | +| Public inspection types | `PortRef`, `PortSpec`, `DataType`, `Metadata`, `MetadataValue`, `BundleError`, `ErrorCode` | +| Common lifecycle sources | Original `RuntimeHandle`, `SessionHandle`, and lifecycle errors; three focused regression tests | +| Python | PyO3 bundle bindings, detached metadata and port descriptors, typed stubs and structured exceptions | +| C++ | C++20 bundle facade, typed metadata, status and data-type headers, private cxx transport, installable CMake target | +| Tests | Existing Rust bundle regression suite and independent Python/C++ inspection and lifetime tests | + +Concrete engine runtimes, sessions, results, tensor/DLPack code, graph planners, +the Triton converter, and backend adapters remain in the runtime branches. +Python and C++ share the core bundle implementation. The `bindings/common` +lifecycle source depends on concrete execution types, so it is not a separate +compiled crate in this bundle-only checkout. See the +[common lifecycle sources](bindings/common/README.md) for integration and testing. +The extraction has no Triton, CUDA, ONNX Runtime, PyTorch, or NumPy dependency. +The C++ binding currently supports Unix platforms (Linux and macOS), using native +path bytes. ## Rust @@ -38,7 +62,8 @@ execution. Bundle loading validates graph ports and connections, artifacts, and optional initial-value files. It does not check whether an inference engine can execute the graph. Declared external inputs can remain unresolved, matching the source -implementations' inspection contract. Python represents missing input descriptors as `None`. +implementations' inspection contract. Python and C++ represent missing input +descriptors as `None` and `std::nullopt`, respectively. The loader normalizes known backend aliases, accepts custom nonempty backend names, and leaves adapter availability to execution. @@ -54,35 +79,30 @@ cargo clippy --locked --workspace --all-targets -- -D warnings cargo fmt --all --check ``` -The default workspace member is the core crate. Its inspection tests use opaque -artifact bytes and require no inference SDK. - -## Common lifecycle sources - -The original shared `RuntimeHandle`, `SessionHandle`, and lifecycle errors are -retained unchanged for integration with either runtime implementation. These -sources depend on concrete execution types and are not a separate compiled -crate in this bundle-only checkout. Three regression tests compile the lifecycle -source against a small test-only runtime API. +The default workspace member is the core crate, including the three common +lifecycle tests. These tests require no engine or language-binding SDK. Building +the complete workspace +also requires a C++20 compiler and Python 3.10 or newer. See +[Python bindings](bindings/python/README.md) and +[C++ bindings](bindings/cxx/README.md) for installation and binding tests. +All inspection tests create opaque artifact bytes; no model generation is needed. -See [common lifecycle sources](bindings/common/README.md) for integration and -testing instructions. - -## Python bindings +## Extraction provenance -The PyO3 binding exposes bundle inspection, detached metadata and port -descriptors, typed stubs, and structured exceptions. Building it requires a -Rust toolchain and Python 3.10 or newer. See -[Python bindings](bindings/python/README.md) for installation and tests. +The shared sources, loader, and tests were extracted from the existing runtime +implementations for backend-independent inspection. -This distribution shares the `leapp_runtime` import name with the full runtime -implementations; install it in a separate environment. +Execution feature gates and dependencies were removed. The C++ bundle facade +was split out of `runtime.cc`, and `DataType` moved from `tensor.h` into +`data_type.h`. Binding manifests, Python exports and stubs, and CMake packaging +now expose only implemented inspection operations. -## Extraction provenance +The shared `bindings/common/src/lib.rs` and `runtime.rs` from the same source +commit are retained byte-for-byte. The lifecycle regression target compiles the +original implementation against a small test-only runtime API. No generic engine +interface or lifecycle cleanup changes are introduced by this extraction. -This package is the common foundation extracted from the existing Triton-based -and native runtime implementations. Execution engines, sessions, tensor/DLPack -code, graph planners, the Triton converter, and backend adapters remain in those -implementations. The package path and library names match the runtime branches -to support later integration. The extraction has no Triton, CUDA, ONNX Runtime, -PyTorch, or NumPy dependency. +Binding regression tests also exposed a pre-existing tagged-metadata parsing +failure. The loader now reads YAML tags directly and separates unknown metadata +from serde's schema-field buffer before restoring it to the snapshot. Tests cover +tags at every inspection level, nested values, duplicate keys, and YAML merges. diff --git a/packages/leapp-runtime-rust/bindings/cxx/.clang-format b/packages/leapp-runtime-rust/bindings/cxx/.clang-format new file mode 100644 index 00000000..84c46fa0 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/.clang-format @@ -0,0 +1,4 @@ +BasedOnStyle: Google +Standard: c++20 +CommentPragmas: '^ (IWYU pragma:|SPDX-)' +InsertBraces: true diff --git a/packages/leapp-runtime-rust/bindings/cxx/CMakeLists.txt b/packages/leapp-runtime-rust/bindings/cxx/CMakeLists.txt new file mode 100644 index 00000000..44afe34a --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/CMakeLists.txt @@ -0,0 +1,70 @@ +cmake_minimum_required(VERSION 3.22) +project(leapp_runtime VERSION 0.1.0 LANGUAGES CXX) + +include(GNUInstallDirs) +include(CMakePackageConfigHelpers) +include(CTest) + +find_package(Threads REQUIRED) +find_program(CARGO_EXECUTABLE cargo REQUIRED) + +set(RUST_TARGET_DIR "${CMAKE_CURRENT_BINARY_DIR}/cargo") +set(RUST_LIBRARY "${RUST_TARGET_DIR}/release/libleapp_runtime_cxx.a") + +# Cargo tracks Rust sources and dependencies; this also refreshes generated headers. +add_custom_target(leapp_runtime_rust + COMMAND "${CMAKE_COMMAND}" -E env --unset=MAKEFLAGS --unset=MFLAGS + "${CARGO_EXECUTABLE}" build --locked --release + --manifest-path "${CMAKE_CURRENT_SOURCE_DIR}/Cargo.toml" + --target-dir "${RUST_TARGET_DIR}" + BYPRODUCTS "${RUST_LIBRARY}" + "${RUST_TARGET_DIR}/cxxbridge/leapp-runtime-cxx/src/bridge.rs.h" + VERBATIM) + +add_library(leapp_runtime SHARED src/bundle.cc) +add_library(leapp::runtime ALIAS leapp_runtime) +add_dependencies(leapp_runtime leapp_runtime_rust) +set_target_properties(leapp_runtime PROPERTIES EXPORT_NAME runtime + VERSION ${PROJECT_VERSION} SOVERSION 0) +target_compile_features(leapp_runtime PUBLIC cxx_std_20) +target_include_directories(leapp_runtime + PUBLIC "$" + "$" + PRIVATE "${RUST_TARGET_DIR}/cxxbridge") + +target_link_libraries(leapp_runtime PRIVATE "${RUST_LIBRARY}" + Threads::Threads ${CMAKE_DL_LIBS} m) +if(APPLE) + target_link_libraries(leapp_runtime PRIVATE iconv) +elseif(UNIX) + target_link_libraries(leapp_runtime PRIVATE rt util) + target_link_options(leapp_runtime PRIVATE "-Wl,--exclude-libs,ALL" "-Wl,--no-undefined") +endif() + +target_compile_options(leapp_runtime PRIVATE -Wall -Wextra -Wpedantic -Werror -fno-exceptions) + +install(TARGETS leapp_runtime EXPORT leapp_runtimeTargets + LIBRARY DESTINATION "${CMAKE_INSTALL_LIBDIR}") +install(DIRECTORY include/leapp DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}") +install(EXPORT leapp_runtimeTargets NAMESPACE leapp:: + DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/leapp_runtime") + +configure_package_config_file(cmake/leapp_runtimeConfig.cmake.in + "${CMAKE_CURRENT_BINARY_DIR}/leapp_runtimeConfig.cmake" + INSTALL_DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/leapp_runtime") +write_basic_package_version_file( + "${CMAKE_CURRENT_BINARY_DIR}/leapp_runtimeConfigVersion.cmake" + VERSION ${PROJECT_VERSION} COMPATIBILITY SameMajorVersion) + +install(FILES "${CMAKE_CURRENT_BINARY_DIR}/leapp_runtimeConfig.cmake" + "${CMAKE_CURRENT_BINARY_DIR}/leapp_runtimeConfigVersion.cmake" + DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/leapp_runtime") + +if(BUILD_TESTING) + add_executable(leapp_bundle_test tests/bundle_test.cc) + target_link_libraries(leapp_bundle_test PRIVATE leapp::runtime Threads::Threads) + target_compile_options(leapp_bundle_test PRIVATE + -Wall -Wextra -Wpedantic -Werror -fno-exceptions) + add_test(NAME leapp_bundle COMMAND leapp_bundle_test + "${CMAKE_CURRENT_SOURCE_DIR}/../../tests/fixtures/pipeline.yaml") +endif() diff --git a/packages/leapp-runtime-rust/bindings/cxx/CPPLINT.cfg b/packages/leapp-runtime-rust/bindings/cxx/CPPLINT.cfg new file mode 100644 index 00000000..9f68ee9c --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/CPPLINT.cfg @@ -0,0 +1,2 @@ +# The project targets C++20; filesystem is part of the supported standard library. +filter=-build/c++17 diff --git a/packages/leapp-runtime-rust/bindings/cxx/Cargo.toml b/packages/leapp-runtime-rust/bindings/cxx/Cargo.toml new file mode 100644 index 00000000..e442be1d --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "leapp-runtime-cxx" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +publish = false + +[lib] +crate-type = ["staticlib"] + +[dependencies] +leapp-runtime = { path = "../.." } +cxx = "=1.0.205" + +[build-dependencies] +cxx-build = "=1.0.205" diff --git a/packages/leapp-runtime-rust/bindings/cxx/README.md b/packages/leapp-runtime-rust/bindings/cxx/README.md new file mode 100644 index 00000000..0d0439c4 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/README.md @@ -0,0 +1,39 @@ +# C++ bundle API + +The C++20 facade uses a private cxx bridge to the Rust loader. Public headers +contain no Rust transport types. Rust, CMake 3.22+, and a C++20 compiler are +required on Linux or macOS; inference SDKs are unnecessary. + +From the Rust package directory: + +```sh +cmake -S bindings/cxx -B target/cxx -DCMAKE_BUILD_TYPE=Release +cmake --build target/cxx --parallel +ctest --test-dir target/cxx --output-on-failure +cmake --install target/cxx --prefix "$PWD/target/install" +cmake -S bindings/cxx/tests/consumer -B target/consumer \ + -DCMAKE_PREFIX_PATH="$PWD/target/install" +cmake --build target/consumer --parallel +./target/consumer/consumer +``` + +Consumers use `find_package(leapp_runtime CONFIG REQUIRED)` and link +`leapp::runtime`. Include `leapp/bundle.h` for inspection or `leapp/data_type.h` +for scalar type declarations. + +```cpp +leapp::Bundle bundle; +auto status = leapp::Bundle::Open("export/policy.yaml", &bundle); +if (!status.ok()) { + // status.code(), status.phase(), and status.message() describe the failure. + return 1; +} +for (const auto& [name, spec] : bundle.pipeline_outputs()) { + // Inspect the optional PortSpec and its typed Metadata. +} +``` + +`Bundle` copies retain a shared immutable snapshot. Getter references remain +valid while any copy owns the snapshot. Concurrent reads are safe; assignment +or destruction of the same handle requires synchronization. `Open` requires an +empty destination and leaves it unchanged on failure. diff --git a/packages/leapp-runtime-rust/bindings/cxx/build.rs b/packages/leapp-runtime-rust/bindings/cxx/build.rs new file mode 100644 index 00000000..d4c608d5 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/build.rs @@ -0,0 +1,9 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +fn main() { + cxx_build::bridge("src/bridge.rs") + .std("c++20") + .compile("leapp-cxx-bridge"); + println!("cargo:rerun-if-changed=src/bridge.rs"); +} diff --git a/packages/leapp-runtime-rust/bindings/cxx/cmake/leapp_runtimeConfig.cmake.in b/packages/leapp-runtime-rust/bindings/cxx/cmake/leapp_runtimeConfig.cmake.in new file mode 100644 index 00000000..9c0d20c4 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/cmake/leapp_runtimeConfig.cmake.in @@ -0,0 +1,4 @@ +@PACKAGE_INIT@ + +include("${CMAKE_CURRENT_LIST_DIR}/leapp_runtimeTargets.cmake") +check_required_components(leapp_runtime) diff --git a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/bundle.h b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/bundle.h new file mode 100644 index 00000000..74e0a9c6 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/bundle.h @@ -0,0 +1,79 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#ifndef LEAPP_BUNDLE_H_ +#define LEAPP_BUNDLE_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include "leapp/data_type.h" +#include "leapp/metadata.h" +#include "leapp/status.h" + +namespace leapp { + +struct PortRef { + std::string node; + std::string port; + + std::string name() const { return node + "/" + port; } +}; + +struct PortSpec { + std::string name; + DataType dtype = DataType::kFloat32; + std::vector shape; + std::string tensor_type; + Metadata metadata; +}; + +/// @brief Ordered public ports. A missing descriptor preserves an unresolved +/// declaration. +using PortMap = std::vector>>; + +/// @brief Shared immutable bundle snapshot and retained C++ inspection values. +/// Getter references survive while any copy retains the snapshot. Concurrent +/// reads are safe; assignment or destruction of the same handle needs +/// synchronization. +class Bundle { + public: + Bundle() = default; + + /// @brief Open into an empty destination. Failure preserves *bundle. + static Status Open(const std::filesystem::path& path, Bundle* bundle); + + bool is_valid() const { return data_ != nullptr; } + + const std::filesystem::path& yaml_path() const; + const std::vector& model_names() const; + const PortMap& pipeline_inputs() const; + const PortMap& pipeline_outputs() const; + const Metadata& pipeline_configs() const; + const Metadata& system_information() const; + const Metadata& metadata() const; + const Metadata& pipeline_metadata() const; + + /// @brief Borrow model metadata or parameters; absent models return nullptr. + const Metadata* model_parameters(std::string_view node) const; + const Metadata* model_metadata(std::string_view node) const; + + /// @brief Borrow the first matching port descriptor, or nullptr if absent. + const PortSpec* input_info(std::string_view node, + std::string_view port) const; + const PortSpec* output_info(std::string_view node, + std::string_view port) const; + + private: + struct Data; + std::shared_ptr data_; +}; + +} // namespace leapp + +#endif // LEAPP_BUNDLE_H_ diff --git a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/data_type.h b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/data_type.h new file mode 100644 index 00000000..65930339 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/data_type.h @@ -0,0 +1,30 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#ifndef LEAPP_DATA_TYPE_H_ +#define LEAPP_DATA_TYPE_H_ + +#include + +namespace leapp { + +/// @brief Fixed-width scalar types; half and Boolean values use packed bytes. +enum class DataType : uint8_t { + kBool, + kUint8, + kUint16, + kUint32, + kUint64, + kInt8, + kInt16, + kInt32, + kInt64, + kFloat16, + kBfloat16, + kFloat32, + kFloat64 +}; + +} // namespace leapp + +#endif // LEAPP_DATA_TYPE_H_ diff --git a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/metadata.h b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/metadata.h new file mode 100644 index 00000000..342262e9 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/metadata.h @@ -0,0 +1,46 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#ifndef LEAPP_METADATA_H_ +#define LEAPP_METADATA_H_ + +#include +#include +#include +#include +#include +#include + +namespace leapp { + +/// @brief Detached typed metadata, preserving ordered mappings and non-string +/// keys. +struct Metadata { + struct Entry; + struct Tagged; + + using Sequence = std::vector; + using Mapping = std::vector; + using Value = + std::variant>; + + Value value; ///< Defaults to null. + + /// @brief Borrow a string-keyed mapping value; absent keys return nullptr. + const Metadata* Find(std::string_view key) const; +}; + +struct Metadata::Entry { + Metadata key; + Metadata value; +}; + +struct Metadata::Tagged { + std::string tag; + Metadata value; +}; + +} // namespace leapp + +#endif // LEAPP_METADATA_H_ diff --git a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/status.h b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/status.h new file mode 100644 index 00000000..42565841 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/status.h @@ -0,0 +1,46 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#ifndef LEAPP_STATUS_H_ +#define LEAPP_STATUS_H_ + +#include +#include + +namespace leapp { + +/// @brief Failure categories shared with the Rust runtime. +enum class ErrorCode { + kOk, + kInvalidArgument, + kNotFound, + kUnsupported, + kIoError, + kInternal +}; + +/// @brief Owned operation status; default construction represents success. +class [[nodiscard]] Status { + public: + Status() = default; + + Status(ErrorCode code, std::string phase, std::string message) + : code_(code), phase_(std::move(phase)), message_(std::move(message)) {} + + bool ok() const { return code_ == ErrorCode::kOk; } + + ErrorCode code() const { return code_; } + + const std::string& phase() const { return phase_; } + + const std::string& message() const { return message_; } + + private: + ErrorCode code_ = ErrorCode::kOk; + std::string phase_; + std::string message_; +}; + +} // namespace leapp + +#endif // LEAPP_STATUS_H_ diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs b/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs new file mode 100644 index 00000000..a4e9a7a6 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs @@ -0,0 +1,95 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use crate::bundle::*; + +#[cxx::bridge(namespace = "leapp::detail")] +pub mod ffi { + struct StatusData { + code: i32, + phase: String, + message: String, + } + + enum Dtype { + Bool, + Uint8, + Uint16, + Uint32, + Uint64, + Int8, + Int16, + Int32, + Int64, + Float16, + Bfloat16, + Float32, + Float64, + } + + enum MetadataKind { + Null, + Bool, + Int, + Uint, + Float, + String, + Sequence, + Mapping, + Tagged, + } + + // Preorder nodes carry immediate child counts; mappings alternate key/value nodes. + struct MetadataNode { + kind: MetadataKind, + boolean: bool, + integer: i64, + unsigned_integer: u64, + number: f64, + text: String, + children: usize, + } + + struct PortData { + present: bool, + name: String, + dtype: Dtype, + shape: Vec, + tensor_type: String, + metadata: Vec, + } + + struct PortName { + node: String, + port: String, + } + + enum MetadataField { + Root, + System, + Configs, + Pipeline, + Parameters, + Model, + } + + struct MetadataData { + present: bool, + nodes: Vec, + } + + extern "Rust" { + type BundleHandle; + fn bundle_new() -> Box; + fn bundle_open(path: &[u8], output: &mut BundleHandle) -> StatusData; + fn bundle_path(bundle: &BundleHandle) -> Vec; + fn bundle_models(bundle: &BundleHandle) -> Vec; + fn bundle_ports(bundle: &BundleHandle, inputs: bool) -> Vec; + fn bundle_port(bundle: &BundleHandle, node: &[u8], port: &[u8], input: bool) -> PortData; + fn bundle_metadata( + bundle: &BundleHandle, + field: MetadataField, + node: &[u8], + ) -> MetadataData; + } +} diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc new file mode 100644 index 00000000..bd2262f1 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc @@ -0,0 +1,297 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "leapp/bundle.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "leapp-runtime-cxx/src/bridge.rs.h" + +namespace leapp { +namespace { +namespace wire = detail; + +// Keep casts across the private bridge checked when either enum changes. +static_assert(static_cast(DataType::kBool) == + static_cast(wire::Dtype::Bool)); +static_assert(static_cast(DataType::kUint8) == + static_cast(wire::Dtype::Uint8)); +static_assert(static_cast(DataType::kUint16) == + static_cast(wire::Dtype::Uint16)); +static_assert(static_cast(DataType::kUint32) == + static_cast(wire::Dtype::Uint32)); +static_assert(static_cast(DataType::kUint64) == + static_cast(wire::Dtype::Uint64)); +static_assert(static_cast(DataType::kInt8) == + static_cast(wire::Dtype::Int8)); +static_assert(static_cast(DataType::kInt16) == + static_cast(wire::Dtype::Int16)); +static_assert(static_cast(DataType::kInt32) == + static_cast(wire::Dtype::Int32)); +static_assert(static_cast(DataType::kInt64) == + static_cast(wire::Dtype::Int64)); +static_assert(static_cast(DataType::kFloat16) == + static_cast(wire::Dtype::Float16)); +static_assert(static_cast(DataType::kBfloat16) == + static_cast(wire::Dtype::Bfloat16)); +static_assert(static_cast(DataType::kFloat32) == + static_cast(wire::Dtype::Float32)); +static_assert(static_cast(DataType::kFloat64) == + static_cast(wire::Dtype::Float64)); + +std::string String(const rust::String& value) { + return {value.data(), value.size()}; +} + +rust::Slice Bytes(std::string_view value) { + return {reinterpret_cast(value.data()), value.size()}; +} + +Status Convert(const wire::StatusData& status) { + return {static_cast(status.code), String(status.phase), + String(status.message)}; +} + +Status Invalid(std::string message) { + return {ErrorCode::kInvalidArgument, "binding", std::move(message)}; +} + +const Metadata& EmptyMetadata() { + static const Metadata empty; + return empty; +} + +// Decode the typed preorder traversal without reparsing YAML or serializing +// text. +Metadata ReadMetadata(const rust::Vec& nodes, + size_t* index) { + const auto& node = nodes[(*index)++]; + using Kind = wire::MetadataKind; + switch (node.kind) { + case Kind::Null: + return {}; + case Kind::Bool: + return {node.boolean}; + case Kind::Int: + return {node.integer}; + case Kind::Uint: + return {node.unsigned_integer}; + case Kind::Float: + return {node.number}; + case Kind::String: + return {String(node.text)}; + case Kind::Sequence: { + Metadata::Sequence items; + items.reserve(node.children); + for (size_t i = 0; i < node.children; ++i) { + items.push_back(ReadMetadata(nodes, index)); + } + return {std::move(items)}; + } + case Kind::Mapping: { + Metadata::Mapping items; + items.reserve(node.children); + for (size_t i = 0; i < node.children; ++i) { + auto key = ReadMetadata(nodes, index); + auto value = ReadMetadata(nodes, index); + items.push_back({std::move(key), std::move(value)}); + } + return {std::move(items)}; + } + case Kind::Tagged: + return {std::make_shared( + Metadata::Tagged{String(node.text), ReadMetadata(nodes, index)})}; + } + // An unknown kind violates the private generated bridge contract. + std::abort(); +} + +Metadata ReadMetadata(const rust::Vec& nodes) { + size_t index = 0; + return nodes.empty() ? Metadata{} : ReadMetadata(nodes, &index); +} + +Metadata Field(const wire::BundleHandle& bundle, wire::MetadataField field, + std::string_view name = {}) { + return ReadMetadata(wire::bundle_metadata(bundle, field, Bytes(name)).nodes); +} + +} // namespace + +const Metadata* Metadata::Find(std::string_view key) const { + const auto* mapping = std::get_if(&value); + if (mapping == nullptr) { + return nullptr; + } + for (const auto& entry : *mapping) { + const auto* name = std::get_if(&entry.key.value); + if (name != nullptr && *name == key) { + return &entry.value; + } + } + return nullptr; +} + +struct Bundle::Data { + explicit Data(rust::Box value) + : handle(std::move(value)) { + const auto bytes = wire::bundle_path(*handle); + path = std::string(bytes.begin(), bytes.end()); + for (const auto& name : wire::bundle_models(*handle)) { + names.push_back(String(name)); + } + metadata = Field(*handle, wire::MetadataField::Root); + system = Field(*handle, wire::MetadataField::System); + configs = Field(*handle, wire::MetadataField::Configs); + pipeline = Field(*handle, wire::MetadataField::Pipeline); + for (const auto& name : names) { + parameters.emplace(name, + Field(*handle, wire::MetadataField::Parameters, name)); + models.emplace(name, Field(*handle, wire::MetadataField::Model, name)); + } + inputs = Ports(true); + outputs = Ports(false); + } + + const PortSpec* Port(std::string_view node, std::string_view port, + bool input) const { + std::lock_guard lock(mutex); + auto key = std::make_tuple(std::string(node), std::string(port), input); + if (auto it = ports.find(key); it != ports.end()) { + return it->second ? &*it->second : nullptr; + } + const auto value = + wire::bundle_port(*handle, Bytes(node), Bytes(port), input); + std::optional spec; + if (value.present) { + spec = PortSpec{String(value.name), + static_cast(value.dtype), + {value.shape.begin(), value.shape.end()}, + String(value.tensor_type), + ReadMetadata(value.metadata)}; + } + const auto it = ports.emplace(std::move(key), std::move(spec)).first; + return it->second ? &*it->second : nullptr; + } + + PortMap Ports(bool input) const { + PortMap result; + for (const auto& p : wire::bundle_ports(*handle, input)) { + const auto node = String(p.node); + const auto port = String(p.port); + const auto name = node + "/" + port; + if (std::any_of(result.begin(), result.end(), + [&](const auto& entry) { return entry.first == name; })) { + continue; + } + const auto* info = Port(node, port, input); + result.emplace_back(name, info == nullptr + ? std::nullopt + : std::optional(*info)); + } + return result; + } + + rust::Box handle; + std::filesystem::path path; + std::vector names; + PortMap inputs; + PortMap outputs; + Metadata metadata; + Metadata system; + Metadata configs; + Metadata pipeline; + std::map> parameters; + std::map> models; + // Map nodes never move; returned descriptor addresses survive later lookups. + mutable std::mutex mutex; + mutable std::map, + std::optional> + ports; +}; + +Status Bundle::Open(const std::filesystem::path& path, Bundle* bundle) { + if (bundle == nullptr || bundle->is_valid()) { + return Invalid("bundle destination must be non-null and empty"); + } + auto handle = wire::bundle_new(); + auto status = Convert(wire::bundle_open(Bytes(path.native()), *handle)); + if (!status.ok()) { + return status; + } + bundle->data_ = std::make_shared(std::move(handle)); + return {}; +} + +const std::filesystem::path& Bundle::yaml_path() const { + static const std::filesystem::path empty; + return data_ ? data_->path : empty; +} + +const std::vector& Bundle::model_names() const { + static const std::vector empty; + return data_ ? data_->names : empty; +} + +const PortMap& Bundle::pipeline_inputs() const { + static const PortMap empty; + return data_ ? data_->inputs : empty; +} + +const PortMap& Bundle::pipeline_outputs() const { + static const PortMap empty; + return data_ ? data_->outputs : empty; +} + +const Metadata& Bundle::metadata() const { + return data_ ? data_->metadata : EmptyMetadata(); +} + +const Metadata& Bundle::system_information() const { + return data_ ? data_->system : EmptyMetadata(); +} + +const Metadata& Bundle::pipeline_configs() const { + return data_ ? data_->configs : EmptyMetadata(); +} + +const Metadata& Bundle::pipeline_metadata() const { + return data_ ? data_->pipeline : EmptyMetadata(); +} + +const Metadata* Bundle::model_parameters(std::string_view node) const { + if (!data_) { + return nullptr; + } + const auto it = data_->parameters.find(node); + return it == data_->parameters.end() ? nullptr : &it->second; +} + +const Metadata* Bundle::model_metadata(std::string_view node) const { + if (!data_) { + return nullptr; + } + const auto it = data_->models.find(node); + return it == data_->models.end() ? nullptr : &it->second; +} + +const PortSpec* Bundle::input_info(std::string_view node, + std::string_view port) const { + return data_ ? data_->Port(node, port, true) : nullptr; +} + +const PortSpec* Bundle::output_info(std::string_view node, + std::string_view port) const { + return data_ ? data_->Port(node, port, false) : nullptr; +} + +} // namespace leapp diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs new file mode 100644 index 00000000..7e2b2280 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs @@ -0,0 +1,192 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use crate::{dtype::wire_dtype, ffi}; +use leapp_runtime::{Bundle, Metadata, MetadataValue}; +use std::{ffi::OsStr, os::unix::ffi::OsStrExt, path::Path}; + +pub struct BundleHandle(pub Option); + +pub fn bundle_new() -> Box { + Box::new(BundleHandle(None)) +} + +pub fn bundle_open(path: &[u8], output: &mut BundleHandle) -> ffi::StatusData { + if output.0.is_some() { + return crate::invalid("bundle destination is not empty"); + } + match Bundle::open(Path::new(OsStr::from_bytes(path))) { + Ok(bundle) => { + output.0 = Some(bundle); + crate::success() + } + Err(e) => crate::error(e.code(), e.phase(), e.message()), + } +} + +pub fn bundle_path(bundle: &BundleHandle) -> Vec { + bundle + .0 + .as_ref() + .map(|b| b.yaml_path().as_os_str().as_bytes().to_vec()) + .unwrap_or_default() +} + +pub fn bundle_models(bundle: &BundleHandle) -> Vec { + bundle + .0 + .as_ref() + .map(|b| b.model_names().into_iter().map(String::from).collect()) + .unwrap_or_default() +} + +pub fn bundle_ports(bundle: &BundleHandle, inputs: bool) -> Vec { + let Some(b) = &bundle.0 else { + return Vec::new(); + }; + let ports = if inputs { + b.pipeline_inputs() + } else { + b.pipeline_outputs() + }; + ports + .iter() + .map(|p| ffi::PortName { + node: p.node.clone(), + port: p.port.clone(), + }) + .collect() +} + +fn node(kind: ffi::MetadataKind) -> ffi::MetadataNode { + ffi::MetadataNode { + kind, + boolean: false, + integer: 0, + unsigned_integer: 0, + number: 0., + text: String::new(), + children: 0, + } +} + +fn append(value: &MetadataValue, nodes: &mut Vec) { + use ffi::MetadataKind as K; + let mut n = node(K::Null); + match value { + MetadataValue::Null => {} + MetadataValue::Bool(v) => { + n.kind = K::Bool; + n.boolean = *v; + } + MetadataValue::Number(v) => { + if let Some(i) = v.as_i64() { + n.kind = K::Int; + n.integer = i; + } else if let Some(u) = v.as_u64() { + n.kind = K::Uint; + n.unsigned_integer = u; + } else { + n.kind = K::Float; + n.number = v.as_f64().expect("YAML number"); + } + } + MetadataValue::String(v) => { + n.kind = K::String; + n.text = v.clone(); + } + MetadataValue::Sequence(v) => { + n.kind = K::Sequence; + n.children = v.len(); + } + MetadataValue::Mapping(v) => { + n.kind = K::Mapping; + n.children = v.len(); + } + MetadataValue::Tagged(v) => { + n.kind = K::Tagged; + n.text = v.tag.to_string(); + n.children = 1; + } + } + nodes.push(n); + match value { + MetadataValue::Sequence(v) => { + for v in v { + append(v, nodes); + } + } + MetadataValue::Mapping(v) => { + for (k, v) in v { + append(k, nodes); + append(v, nodes); + } + } + MetadataValue::Tagged(v) => append(&v.value, nodes), + _ => {} + } +} + +fn metadata(data: &Metadata) -> Vec { + let mut root = node(ffi::MetadataKind::Mapping); + root.children = data.len(); + let mut nodes = vec![root]; + for (key, value) in data { + let mut key_node = node(ffi::MetadataKind::String); + key_node.text = key.clone(); + nodes.push(key_node); + append(value, &mut nodes); + } + nodes +} + +pub fn bundle_metadata( + b: &BundleHandle, + field: ffi::MetadataField, + name: &[u8], +) -> ffi::MetadataData { + let name = std::str::from_utf8(name).ok(); + let data = b.0.as_ref().and_then(|b| match field { + ffi::MetadataField::Root => Some(b.metadata()), + ffi::MetadataField::System => Some(b.system_information()), + ffi::MetadataField::Configs => Some(b.pipeline_configs()), + ffi::MetadataField::Pipeline => Some(b.pipeline_metadata()), + ffi::MetadataField::Parameters => name.and_then(|name| b.model_parameters(name)), + ffi::MetadataField::Model => name.and_then(|name| b.model_metadata(name)), + _ => None, + }); + ffi::MetadataData { + present: data.is_some(), + nodes: data.map(metadata).unwrap_or_default(), + } +} + +pub fn bundle_port(b: &BundleHandle, node: &[u8], port: &[u8], input: bool) -> ffi::PortData { + let spec = b.0.as_ref().and_then(|b| { + let node = std::str::from_utf8(node).ok()?; + let port = std::str::from_utf8(port).ok()?; + if input { + b.input_info(node, port) + } else { + b.output_info(node, port) + } + }); + match spec { + Some(s) => ffi::PortData { + present: true, + name: s.name.clone(), + dtype: wire_dtype(s.dtype), + shape: s.shape.clone(), + tensor_type: s.tensor_type.clone(), + metadata: metadata(&s.metadata), + }, + None => ffi::PortData { + present: false, + name: String::new(), + dtype: ffi::Dtype::Float32, + shape: Vec::new(), + tensor_type: String::new(), + metadata: Vec::new(), + }, + } +} diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/dtype.rs b/packages/leapp-runtime-rust/bindings/cxx/src/dtype.rs new file mode 100644 index 00000000..bf9f111c --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/src/dtype.rs @@ -0,0 +1,24 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use crate::ffi; +use leapp_runtime::DataType; + +pub fn wire_dtype(value: DataType) -> ffi::Dtype { + use DataType as D; + match value { + D::Bool => ffi::Dtype::Bool, + D::Uint8 => ffi::Dtype::Uint8, + D::Uint16 => ffi::Dtype::Uint16, + D::Uint32 => ffi::Dtype::Uint32, + D::Uint64 => ffi::Dtype::Uint64, + D::Int8 => ffi::Dtype::Int8, + D::Int16 => ffi::Dtype::Int16, + D::Int32 => ffi::Dtype::Int32, + D::Int64 => ffi::Dtype::Int64, + D::Float16 => ffi::Dtype::Float16, + D::Bfloat16 => ffi::Dtype::Bfloat16, + D::Float32 => ffi::Dtype::Float32, + D::Float64 => ffi::Dtype::Float64, + } +} diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/lib.rs b/packages/leapp-runtime-rust/bindings/cxx/src/lib.rs new file mode 100644 index 00000000..046953cb --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/src/lib.rs @@ -0,0 +1,32 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Private cxx transport for the public C++ facade. +#![deny(unsafe_op_in_unsafe_fn)] + +mod bridge; +mod bundle; +mod dtype; + +use bridge::ffi; +use leapp_runtime::ErrorCode; + +fn success() -> ffi::StatusData { + ffi::StatusData { + code: 0, + phase: String::new(), + message: String::new(), + } +} + +fn error(code: ErrorCode, phase: &str, message: impl ToString) -> ffi::StatusData { + ffi::StatusData { + code: code as i32, + phase: phase.into(), + message: message.to_string(), + } +} + +fn invalid(message: &str) -> ffi::StatusData { + error(ErrorCode::InvalidArgument, "binding", message) +} diff --git a/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc b/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc new file mode 100644 index 00000000..288b11b7 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc @@ -0,0 +1,107 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "leapp/bundle.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +void Check(bool condition, const char* message) { + if (!condition) { + std::cerr << message << '\n'; + std::exit(1); + } +} +} // namespace + +int main(int argc, char** argv) { + Check(argc == 2, "expected the pipeline fixture path"); + const auto nonce = + std::chrono::steady_clock::now().time_since_epoch().count(); + const auto directory = std::filesystem::temp_directory_path() / + ("leapp-bundle-" + std::to_string(nonce)); + Check(std::filesystem::create_directory(directory), + "create fixture directory"); + const auto path = directory / "pipeline.yaml"; + std::filesystem::copy_file(argv[1], path); + std::ofstream(directory / "model.bin", std::ios::binary) << "abc"; + std::ofstream(path, std::ios::app) + << "tagged: !sensor {axes: [x, y]}\n" + "sequence_keys: {[x, y]: coordinates}\n" + "unsigned: 18446744073709551615\n"; + + leapp::Bundle empty; + Check(!empty.is_valid() && empty.model_names().empty(), "empty handle"); + Check(empty.input_info("missing", "port") == nullptr, "empty descriptor"); + Check(!leapp::Bundle::Open(path, nullptr).ok(), "null destination"); + auto status = leapp::Bundle::Open(directory / "absent.yaml", &empty); + Check(status.code() == leapp::ErrorCode::kNotFound, "missing YAML category"); + Check(status.phase() == "open" && !status.message().empty(), + "error diagnostics"); + Check(!empty.is_valid(), "failure preserves empty destination"); + + leapp::Bundle bundle; + Check(leapp::Bundle::Open(path, &bundle).ok(), "load opaque artifacts"); + Check(bundle.yaml_path() == path, "absolute path"); + Check(bundle.model_names() == + std::vector({"z_observation", "a_policy"}), + "model declaration order"); + Check(bundle.pipeline_inputs().size() == 1, "pipeline inputs"); + Check(bundle.pipeline_outputs().size() == 3, "pipeline outputs"); + Check(bundle.pipeline_outputs()[0].first == "a_policy/next_state", + "output order"); + const auto* port = bundle.output_info("a_policy", "actions"); + Check(port != nullptr && port->dtype == leapp::DataType::kFloat32, + "output dtype"); + Check(port->shape == std::vector({3}), "output shape"); + Check(bundle.input_info("absent", "port") == nullptr, "absent descriptor"); + Check(bundle.model_parameters("absent") == nullptr, "absent parameters"); + const auto* frequency = bundle.pipeline_configs().Find("frequency"); + Check(frequency && std::get(frequency->value) == 50, "typed config"); + const auto* tag = bundle.metadata().Find("tagged"); + Check(tag != nullptr, "tag metadata"); + const auto& tagged = + std::get>(tag->value); + Check(tagged->tag == "!sensor" && tagged->value.Find("axes"), + "tag preserved"); + const auto* mapping = bundle.metadata().Find("sequence_keys"); + Check(mapping != nullptr, "mapping metadata"); + const auto& entries = std::get(mapping->value); + Check(std::get(entries[0].key.value).size() == 2, + "non-string mapping key"); + const auto* unsigned_value = bundle.metadata().Find("unsigned"); + Check(unsigned_value && + std::get(unsigned_value->value) == UINT64_MAX, + "unsigned metadata"); + + Check(!leapp::Bundle::Open(path, &bundle).ok(), + "reject occupied destination"); + Check(bundle.output_info("a_policy", "actions") == port, + "failed open preserves snapshot"); + auto retained = bundle; + bundle = {}; + std::filesystem::remove(path); + std::vector readers; + for (int i = 0; i < 4; ++i) { + readers.emplace_back([retained, port] { + for (int j = 0; j < 100; ++j) { + Check(retained.output_info("a_policy", "actions") == port, + "shared snapshot lifetime"); + Check(retained.input_info("z_observation", "joint_pos") != nullptr, + "concurrent descriptor lookup"); + } + }); + } + for (auto& reader : readers) { + reader.join(); + } + std::filesystem::remove_all(directory); +} diff --git a/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/CMakeLists.txt b/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/CMakeLists.txt new file mode 100644 index 00000000..15f51a8a --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/CMakeLists.txt @@ -0,0 +1,7 @@ +cmake_minimum_required(VERSION 3.22) +project(leapp_bundle_consumer LANGUAGES CXX) + +find_package(leapp_runtime CONFIG REQUIRED) + +add_executable(consumer main.cc) +target_link_libraries(consumer PRIVATE leapp::runtime) diff --git a/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/main.cc b/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/main.cc new file mode 100644 index 00000000..2c2f2e68 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/main.cc @@ -0,0 +1,13 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "leapp/bundle.h" + +int main() { + leapp::Bundle bundle; + const auto status = leapp::Bundle::Open("", &bundle); + return !status.ok() && !bundle.is_valid() && + status.code() == leapp::ErrorCode::kInvalidArgument + ? 0 + : 1; +} From c6e70cb899c385d60e9a6e5d43103e7f8d54227d Mon Sep 17 00:00:00 2001 From: Frank Lai NV Date: Fri, 16 Oct 2026 15:23:32 -0700 Subject: [PATCH 07/21] Functionalize safe in-place tensor methods from ATen schemas (#28) replace the hand-maintained in-place allowlist with cached ATen schema discovery preserve eager receiver identity while recording deterministic functional graph operations support metadata-mutating views such as squeeze_, unsqueeze_, and transpose_ cover masked_fill_, index_copy_, cumsum_, and other schema-compatible methods reject nondeterministic or incompatible methods before mutation consolidate overlapping traced-tensor regression coverage while retaining distinct semantics --- .../datatypes/torch/traced_tensor.py | 226 ++++++- tests/unit_tests/test_traced_tensor.py | 582 ++++++++++-------- 2 files changed, 506 insertions(+), 302 deletions(-) diff --git a/leapp/leapp_graph/datatypes/torch/traced_tensor.py b/leapp/leapp_graph/datatypes/torch/traced_tensor.py index 5804ef15..94d1fe62 100644 --- a/leapp/leapp_graph/datatypes/torch/traced_tensor.py +++ b/leapp/leapp_graph/datatypes/torch/traced_tensor.py @@ -10,6 +10,7 @@ and graph recording. """ +import functools import operator from abc import ABCMeta @@ -78,18 +79,6 @@ class TracedTensor(TracedData, torch.Tensor, metaclass=_TracedTensorMeta): _EQUIVALENT_COPY_NAMES = frozenset( {"clone", "detach", "contiguous", "cpu", "cuda"} ) - # In-place methods whose allocating torch.* counterpart has the same - # argument order and describes the receiver's complete post-write value. - # This lets one mutation path preserve receiver identity eagerly while - # recording the corresponding functional operation in FX. - _FUNCTIONAL_INPLACE_NAMES = frozenset({ - "add_", "sub_", "mul_", "div_", "pow_", - "subtract_", "multiply_", "divide_", "true_divide_", - "floor_divide_", "remainder_", "fmod_", - "clamp_", "clamp_min_", "clamp_max_", "clip_", "lerp_", - "addcmul_", "addcdiv_", - "addmm_", "addmv_", "addr_", "addbmm_", "baddbmm_", - }) _NATIVE_TYPE = torch.Tensor @staticmethod @@ -172,6 +161,145 @@ def _new_alias(self, tensor: torch.Tensor) -> "TracedTensor": # Static Methods # ========================================================================= + @staticmethod + def _schema_arguments_match(inplace_arg, functional_arg) -> bool: + """Whether two schema arguments expose the same Python call contract.""" + if ( + inplace_arg.name != functional_arg.name + or str(inplace_arg.type) != str(functional_arg.type) + or inplace_arg.kwarg_only != functional_arg.kwarg_only + or inplace_arg.has_default_value() + != functional_arg.has_default_value() + ): + return False + if not inplace_arg.has_default_value(): + return True + return inplace_arg.default_value == functional_arg.default_value + + @classmethod + def _schemas_are_functional_twins( + cls, inplace_schema, functional_schema + ) -> bool: + """Check that a functional schema describes one receiver-only write.""" + if ( + not inplace_schema.is_mutable + or functional_schema.is_mutable + or len(inplace_schema.arguments) != len(functional_schema.arguments) + or len(inplace_schema.returns) != 1 + or len(functional_schema.returns) != 1 + or str(inplace_schema.returns[0].type) != "Tensor" + or str(functional_schema.returns[0].type) != "Tensor" + ): + return False + + if not all( + cls._schema_arguments_match(inplace_arg, functional_arg) + for inplace_arg, functional_arg in zip( + inplace_schema.arguments, functional_schema.arguments + ) + ): + return False + + writes = [ + index + for index, argument in enumerate(inplace_schema.arguments) + if argument.alias_info is not None + and argument.alias_info.is_write + ] + if writes != [0] or any( + argument.alias_info is not None + and argument.alias_info.is_write + for argument in functional_schema.arguments + ): + return False + + receiver_alias = inplace_schema.arguments[0].alias_info + return_alias = inplace_schema.returns[0].alias_info + if receiver_alias is None or return_alias is None: + return False + return bool( + receiver_alias.after_set & return_alias.after_set + or receiver_alias.before_set & return_alias.before_set + ) + + @staticmethod + def _schema_tags(packet, overload_name): + overload = getattr(packet, overload_name or "default") + return set(overload.tags) + + @classmethod + @functools.lru_cache(maxsize=None) + def _functional_inplace_spec(cls, func_name): + """Return ``(torch functional, mutates metadata)`` when schemas prove it.""" + if ( + not func_name.endswith("_") + or func_name.startswith("_") + or len(func_name) == 1 + ): + return None + + functional_name = func_name[:-1] + functional = getattr(torch, functional_name, None) + if not callable(functional): + return None + + try: + inplace_packet = getattr(torch.ops.aten, func_name) + functional_packet = getattr(torch.ops.aten, functional_name) + inplace_schemas = inplace_packet._schemas + functional_schemas = functional_packet._schemas + inplace_view_tag = torch.Tag.inplace_view + nondeterministic_tags = { + torch.Tag.nondeterministic_bitwise, + torch.Tag.nondeterministic_seeded, + } + except (AttributeError, RuntimeError): + return None + + tensor_inplace_schemas = [ + (overload_name, schema) + for overload_name, schema in inplace_schemas.items() + if schema.arguments + and str(schema.arguments[0].type) == "Tensor" + and len(schema.returns) == 1 + and str(schema.returns[0].type) == "Tensor" + ] + metadata_flags = [] + for overload_name, inplace_schema in tensor_inplace_schemas: + try: + inplace_tags = cls._schema_tags( + inplace_packet, overload_name + ) + except AttributeError: + return None + if inplace_tags & nondeterministic_tags: + return None + + matched = False + for functional_overload, functional_schema in ( + functional_schemas.items() + ): + try: + functional_tags = cls._schema_tags( + functional_packet, functional_overload + ) + except AttributeError: + continue + if functional_tags & nondeterministic_tags: + continue + if cls._schemas_are_functional_twins( + inplace_schema, functional_schema + ): + matched = True + break + if not matched: + return None + metadata_flags.append(inplace_view_tag in inplace_tags) + + if not metadata_flags or len(set(metadata_flags)) != 1: + return None + return functional, metadata_flags[0] + @staticmethod def _convert_method_descriptor(func, context=None): """Convert torch.Tensor method descriptors to torch.* function equivalents. @@ -193,14 +321,6 @@ def _convert_method_descriptor(func, context=None): if torch_func is not None and callable(torch_func): return torch_func - # Handle in-place operations (e.g., add_ -> torch.add) - # In-place methods end with '_' but torch.* functions don't - if func_name.endswith('_') and len(func_name) > 1: - base_name = func_name[:-1] # Remove trailing underscore - torch_func = getattr(torch, base_name, None) - if torch_func is not None and callable(torch_func): - return torch_func - if context is not None: _get_logger().debug( f"while tracing for {context} detected method_descriptor or builtin_function_or_method {func_name} " @@ -451,10 +571,29 @@ def _handle_functional_inplace(cls, func, args=(), kwargs=None): """ kwargs = kwargs or {} func_name = getattr(func, "__name__", "") - if func_name not in cls._FUNCTIONAL_INPLACE_NAMES or not args: + if not args: return False, None target = args[0] + if type(target) is not torch.Tensor and not isinstance(target, cls): + return False, None + + spec = cls._functional_inplace_spec(func_name) + if spec is None: + exact_inplace = getattr(torch, func_name, None) + functional = getattr(torch, func_name[:-1], None) + if callable(exact_inplace) and not callable(functional): + return False, None + if func_name.endswith("_") and not func_name.startswith("__"): + _get_logger().fatal( + f"TracedTensor cannot safely lower in-place method " + f"{func_name}: no deterministic functional ATen schema " + f"with the same call contract was found.", + error_type=NotImplementedError, + ) + return False, None + + functional, mutates_metadata = spec operands = args[1:] if type(target) is torch.Tensor: anchor = cls.find_traced_tensor((operands, kwargs)) @@ -467,12 +606,13 @@ def _handle_functional_inplace(cls, func, args=(), kwargs=None): cls._promote_plain_tensor( target, anchor.name, anchor.context_obj, dest_proxy ) - elif not isinstance(target, cls): - return False, None - functional = getattr(torch, func_name[:-1]) return True, target._apply_inplace( - functional, func_name, *operands, **kwargs + functional, + func_name, + *operands, + _mutates_metadata=mutates_metadata, + **kwargs, ) @classmethod @@ -562,6 +702,12 @@ def shed_replaced_graph(value): if traced_tensor is not None: cls._close_warp_segment() + # A carrier is an alias, which Torch refuses to detach in place, and the + # autograd state detach_ edits is never read while tracing. + if (getattr(func, "__name__", "") == "detach_" + and args and isinstance(args[0], TracedTensor)): + return args[0] + # A plain receiver must be promoted before method-descriptor conversion # turns an in-place method into an allocating function. handled, result = cls._handle_plain_assignment(func, args, kwargs) @@ -581,12 +727,6 @@ def shed_replaced_graph(value): # Fallback to default behavior if no TracedTensor found return NotImplemented - # A carrier is an alias, which Torch refuses to detach in place, and the - # autograd state detach_ edits is never read while tracing. - if (getattr(func, "__name__", "") == "detach_" - and args and isinstance(args[0], TracedTensor)): - return args[0] - # Extract real tensors for actual computation real_args = tuple(TracedTensor.unwrap_traced_tensor(arg) for arg in args) real_kwargs = {k: TracedTensor.unwrap_traced_tensor(v) for k, v in kwargs.items()} @@ -886,7 +1026,14 @@ def __rmatmul__(self, other): # In-place Arithmetic Operators (for +=, -=, etc.) # ========================================================================= - def _apply_inplace(self, functional, inplace_name, *operands, **kwargs): + def _apply_inplace( + self, + functional, + inplace_name, + *operands, + _mutates_metadata=False, + **kwargs, + ): """Mutate this tensor in place, recording the operation functionally. Native eager execution preserves PyTorch's identity, casting, and error @@ -910,6 +1057,13 @@ def _apply_inplace(self, functional, inplace_name, *operands, **kwargs): underlying.copy_( functional(underlying, *real_operands, **real_kwargs) ) + elif _mutates_metadata: + # Shape/view mutations change the receiver's tensor metadata, not + # its storage, so mutating an ``as_subclass`` alias is insufficient. + with torch._C.DisableTorchFunction(): + getattr(torch.Tensor, inplace_name)( + self, *real_operands, **real_kwargs + ) else: getattr(underlying, inplace_name)( *real_operands, **real_kwargs @@ -993,6 +1147,14 @@ def index_add_(self, dim, index, source, *, alpha=1): # Descriptor conversion would otherwise discard the receiver mutation. return self.copy_(torch.index_add(self, dim, index, source, alpha=alpha)) + def polygamma_(self, n): + """Handle the functional form's reversed ``n, input`` arguments.""" + return self._apply_inplace( + lambda tensor, order: torch.polygamma(order, tensor), + "polygamma_", + n, + ) + def copy_(self, src, non_blocking=False): """Copy ``src`` into this tensor and record a full functional write.""" self._close_warp_segment() diff --git a/tests/unit_tests/test_traced_tensor.py b/tests/unit_tests/test_traced_tensor.py index e19f8b87..40544067 100644 --- a/tests/unit_tests/test_traced_tensor.py +++ b/tests/unit_tests/test_traced_tensor.py @@ -6,6 +6,7 @@ """Tests for basic TrackedTensor operations.""" import inspect +import operator import pathlib import tempfile import unittest @@ -595,11 +596,28 @@ def validate_export( verify_exported_program( self, graph_module, inputs, expected, test_name=test_name) - def test_shape(self): - """Test shape property.""" + def test_tensor_metadata_properties(self): + """Tensor metadata access matches native Tensor behavior.""" ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.randn(3, 4, 5), name="x") - self.assertEqual(x.shape, torch.Size([3, 4, 5])) + x_float = ctx.create_input(torch.randn(3, 4, 5), name="x_float") + x_int = ctx.create_input( + torch.tensor([1, 2, 3], dtype=torch.int32), name="x_int" + ) + + self.assertEqual(x_float.shape, torch.Size([3, 4, 5])) + self.assertEqual(x_float.dim(), 3) + self.assertEqual(x_float.ndim, 3) + self.assertEqual(len(x_float), 3) + self.assertEqual(x_float.numel(), 60) + self.assertEqual(x_float.size(), torch.Size([3, 4, 5])) + self.assertEqual(x_float.size(0), 3) + self.assertEqual(x_float.size(1), 4) + self.assertEqual(x_float.size(-1), 5) + self.assertEqual(x_float.device.type, "cpu") + self.assertTrue(x_float.is_contiguous()) + self.assertTrue(x_float.is_floating_point()) + self.assertEqual(x_int.dtype, torch.int32) + self.assertFalse(x_int.is_floating_point()) def test_reshape_accepts_list_shape(self): """Match Tensor.reshape semantics used by libraries such as einops.""" @@ -615,68 +633,6 @@ def test_reshape_accepts_tuple_shape(self): reshaped = x.reshape((2, 6)) self.assertEqual(reshaped.shape, torch.Size([2, 6])) - def test_dim(self): - """Test dim property.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.randn(3, 4, 5), name="x") - self.assertEqual(x.dim(), 3) - - def test_dtype(self): - """Test dtype property.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor( - [1, 2, 3], dtype=torch.int32), name="x") - self.assertEqual(x.dtype, torch.int32) - - def test_device(self): - """Test device property.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0]), name="x") - self.assertEqual(x.device.type, "cpu") - - def test_len(self): - """Test len() function.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.randn(10, 5), name="x") - self.assertEqual(len(x), 10) - - def test_ndim(self): - """Test ndim property.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.randn(3, 4, 5), name="x") - self.assertEqual(x.ndim, 3) - - def test_numel(self): - """Test numel() method.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.randn(3, 4, 5), name="x") - self.assertEqual(x.numel(), 60) - - def test_size_method(self): - """Test size() method with and without dimension.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.randn(3, 4, 5), name="x") - self.assertEqual(x.size(), torch.Size([3, 4, 5])) - self.assertEqual(x.size(0), 3) - self.assertEqual(x.size(1), 4) - self.assertEqual(x.size(-1), 5) - - def test_is_contiguous(self): - """Test is_contiguous() method.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.randn(3, 4), name="x") - self.assertTrue(x.is_contiguous()) - - def test_is_floating_point(self): - """Test is_floating_point() method.""" - ctx = TracedTensorNode(name="test", node_index=0) - x_float = ctx.create_input(torch.randn(3), name="x_float") - self.assertTrue(x_float.is_floating_point()) - - ctx2 = TracedTensorNode(name="test2", node_index=1) - x_int = ctx2.create_input(torch.tensor([1, 2, 3]), name="x_int") - self.assertFalse(x_int.is_floating_point()) - def test_T_property(self): """Test .T transpose property.""" ctx = TracedTensorNode(name="test", node_index=0) @@ -731,40 +687,6 @@ def test_tensor_norm_method_rewrites_torch_tensor_function_target(self): self.assertTrue(torch.allclose( traced(torch.tensor([[3.0, 4.0], [5.0, 12.0]])), expected)) - def test_regular_tensor_mul_traced_tensor(self): - """Test regular_tensor * traced_tensor.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - regular = torch.tensor([2.0, 2.0, 2.0]) - - result = regular * x - - self.assertIsInstance(result, TracedTensor) - expected = torch.tensor([2.0, 4.0, 6.0]) - self.assertTrue(torch.allclose(result.tensor, expected)) - - def test_regular_tensor_sub_traced_tensor(self): - """Test regular_tensor - traced_tensor.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - regular = torch.tensor([10.0, 10.0, 10.0]) - - result = regular - x - - expected = torch.tensor([9.0, 8.0, 7.0]) - self.assertTrue(torch.allclose(result.tensor, expected)) - - def test_regular_tensor_add_traced_tensor(self): - """Test regular_tensor + traced_tensor.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - regular = torch.tensor([10.0, 10.0, 10.0]) - - result = regular + x - - expected = torch.tensor([11.0, 12.0, 13.0]) - self.assertTrue(torch.allclose(result.tensor, expected)) - # ==================== is_tracing=False Tests ==================== def test_operations_after_compile_return_native_tensors(self): @@ -817,14 +739,21 @@ def test_to_after_compile_returns_native_tensor(self): # ==================== Clone and Contiguous Tests ==================== - def test_clone(self): - """Test clone operation.""" + def test_clone_preserves_type_and_compiles(self): + """Clone stays traced and replays through the compiled graph.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - y = x.clone() + cloned = x.clone() - self.assertIsInstance(y, TracedTensor) - self.assertTrue(torch.allclose(y.tensor, x.tensor)) + self.assertIsInstance(cloned, TracedTensor) + self.assertTrue(torch.allclose(cloned.tensor, x.tensor)) + + output = cloned * 2 + ctx.compile_trace({'output': output}) + runtime_input = torch.tensor([4.0, 5.0, 6.0]) + self.assertTrue(torch.allclose( + ctx.m(runtime_input), runtime_input.clone() * 2 + )) def test_contiguous(self): """Test contiguous operation.""" @@ -835,16 +764,24 @@ def test_contiguous(self): self.assertIsInstance(y, TracedTensor) self.assertEqual(y.shape, torch.Size([4, 3])) - def test_detach(self): - """Test detach operation returns TracedTensor and is recorded in graph.""" + def test_detach_preserves_type_and_compiles(self): + """Detach stays traced, clears autograd state, and replays.""" ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0], requires_grad=True), name="x") - y = x.detach() + x = ctx.create_input( + torch.tensor([1.0, 2.0, 3.0], requires_grad=True), name="x" + ) + detached = x.detach() - self.assertIsInstance(y, TracedTensor) - self.assertTrue(torch.allclose(y.tensor, x.tensor)) - # Detached tensor should not require grad - self.assertFalse(y.requires_grad) + self.assertIsInstance(detached, TracedTensor) + self.assertTrue(torch.allclose(detached.tensor, x.tensor)) + self.assertFalse(detached.requires_grad) + + output = detached * 2 + ctx.compile_trace({'output': output}) + runtime_input = torch.tensor([4.0, 5.0, 6.0]) + self.assertTrue(torch.allclose( + ctx.m(runtime_input), runtime_input.detach() * 2 + )) def test_detach_in_place_returns_the_receiver(self): """``detach_`` keeps its receiver so chained in-place idioms survive. @@ -862,20 +799,6 @@ def test_detach_in_place_returns_the_receiver(self): self.assertIs(result, x) self.assertTrue(torch.equal(result.tensor, torch.zeros(3))) - def test_detach_compiled(self): - """Test detach operation compiles correctly.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - y = x.detach() * 2 # Use detached tensor in computation - - ctx.compile_trace({'output': y}) - - # Test compiled graph works - input_tensor = torch.tensor([4.0, 5.0, 6.0]) - expected = input_tensor.detach() * 2 - output = ctx.m(input_tensor) - self.assertTrue(torch.allclose(output, expected)) - def test_cpu(self): """Test cpu operation returns TracedTensor.""" ctx = TracedTensorNode(name="test", node_index=0) @@ -886,20 +809,6 @@ def test_cpu(self): self.assertTrue(torch.allclose(y.tensor, x.tensor)) self.assertEqual(y.device.type, 'cpu') - def test_clone_compiled(self): - """Test clone operation compiles correctly.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - y = x.clone() * 2 # Use cloned tensor in computation - - ctx.compile_trace({'output': y}) - - # Test compiled graph works - input_tensor = torch.tensor([4.0, 5.0, 6.0]) - expected = input_tensor.clone() * 2 - output = ctx.m(input_tensor) - self.assertTrue(torch.allclose(output, expected)) - # ==================== Additional Math Operations ==================== def test_abs(self): @@ -1131,90 +1040,32 @@ def test_comparison_operators(self): # ==================== In-Place Operations Tests ==================== - def test_inplace_add(self): - """Test in-place addition (+=).""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - original_id = id(x) - - x += 2.0 - - # Should return the same object - self.assertEqual(id(x), original_id) - # Should have updated value - expected = torch.tensor([3.0, 4.0, 5.0]) - self.assertTrue(torch.allclose(x.tensor, expected)) - # Should still be TracedTensor - self.assertIsInstance(x, TracedTensor) - - # Should compile successfully - ctx.compile_trace({'x': x}) - result = ctx.m(torch.tensor([1.0, 2.0, 3.0])) - self.assertTrue(torch.allclose(result, expected)) - - def test_inplace_sub(self): - """Test in-place subtraction (-=).""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([5.0, 6.0, 7.0]), name="x") - original_id = id(x) - - x -= 2.0 - - self.assertEqual(id(x), original_id) - expected = torch.tensor([3.0, 4.0, 5.0]) - self.assertTrue(torch.allclose(x.tensor, expected)) - - ctx.compile_trace({'x': x}) - result = ctx.m(torch.tensor([5.0, 6.0, 7.0])) - self.assertTrue(torch.allclose(result, expected)) - - def test_inplace_mul(self): - """Test in-place multiplication (*=).""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - original_id = id(x) - - x *= 3.0 - - self.assertEqual(id(x), original_id) - expected = torch.tensor([3.0, 6.0, 9.0]) - self.assertTrue(torch.allclose(x.tensor, expected)) - - ctx.compile_trace({'x': x}) - result = ctx.m(torch.tensor([1.0, 2.0, 3.0])) - self.assertTrue(torch.allclose(result, expected)) - - def test_inplace_div(self): - """Test in-place division (/=).""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([4.0, 8.0, 12.0]), name="x") - original_id = id(x) - - x /= 2.0 - - self.assertEqual(id(x), original_id) - expected = torch.tensor([2.0, 4.0, 6.0]) - self.assertTrue(torch.allclose(x.tensor, expected)) - - ctx.compile_trace({'x': x}) - result = ctx.m(torch.tensor([4.0, 8.0, 12.0])) - self.assertTrue(torch.allclose(result, expected)) + def test_inplace_operators_preserve_receiver_and_compile(self): + """In-place operators mutate one receiver and replay functionally.""" + cases = ( + ("add", operator.iadd, torch.tensor([1., 2., 3.]), 2.0), + ("sub", operator.isub, torch.tensor([5., 6., 7.]), 2.0), + ("mul", operator.imul, torch.tensor([1., 2., 3.]), 3.0), + ("div", operator.itruediv, torch.tensor([4., 8., 12.]), 2.0), + ("pow", operator.ipow, torch.tensor([2., 3., 4.]), 2.0), + ) - def test_inplace_pow(self): - """Test in-place power (**=).""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([2.0, 3.0, 4.0]), name="x") - original_id = id(x) + for name, operation, input_tensor, operand in cases: + with self.subTest(operation=name): + ctx = TracedTensorNode(name="test", node_index=0) + receiver = ctx.create_input(input_tensor.clone(), name="x") + expected = operation(input_tensor.clone(), operand) - x **= 2.0 + result = operation(receiver, operand) - self.assertEqual(id(x), original_id) - expected = torch.tensor([4.0, 9.0, 16.0]) - self.assertTrue(torch.allclose(x.tensor, expected)) + self.assertIs(result, receiver) + self.assertIsInstance(receiver, TracedTensor) + self.assertTrue(torch.allclose(receiver.tensor, expected)) - ctx.compile_trace({'x': x}) - result = ctx.m(torch.tensor([2.0, 3.0, 4.0])) - self.assertTrue(torch.allclose(result, expected)) + ctx.compile_trace({'x': receiver}) + self.assertTrue(torch.allclose( + ctx.m(input_tensor.clone()), expected + )) def test_inplace_method_add_(self): """Test in-place method add_().""" @@ -2267,20 +2118,6 @@ def test_all_arithmetic_functions(self): # ==================== Setitem Tests ==================== - def test_setitem_full_slice_with_regular_tensor(self): - """Test full slice assignment with a regular tensor value.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - - # Assign a new value via full slice - x[:] = torch.tensor([10.0, 20.0, 30.0]) - - # Should still be a TracedTensor - self.assertIsInstance(x, TracedTensor) - # Should have updated values - expected = torch.tensor([10.0, 20.0, 30.0]) - self.assertTrue(torch.allclose(x.tensor, expected)) - def test_setitem_full_slice_with_traced_tensor(self): """Test full slice assignment where value is another TracedTensor. @@ -2322,32 +2159,6 @@ def test_setitem_full_slice_with_traced_tensor(self): expected_different = torch.tensor([20.0, 40.0, 60.0]) self.assertTrue(torch.allclose(result_different, expected_different)) - def test_setitem_partial_slice(self): - """Test partial slice assignment.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]), name="x") - - # Assign to a partial slice - x[1:3] = torch.tensor([20.0, 30.0]) - - # Should still be a TracedTensor - self.assertIsInstance(x, TracedTensor) - # Should have updated values - expected = torch.tensor([1.0, 20.0, 30.0, 4.0, 5.0]) - self.assertTrue(torch.allclose(x.tensor, expected)) - - def test_setitem_single_index(self): - """Test single index assignment.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - - # Assign to a single index - x[1] = 99.0 - - self.assertIsInstance(x, TracedTensor) - expected = torch.tensor([1.0, 99.0, 3.0]) - self.assertTrue(torch.allclose(x.tensor, expected)) - def test_setitem_2d_slice(self): """Test 2D slice assignment.""" ctx = TracedTensorNode(name="test", node_index=0) @@ -2403,12 +2214,18 @@ def test_setitem_single_index_fx_torchscript_onnx(self): """Test: x[0] = 10 exports correctly to FX, TorchScript, and ONNX.""" ctx = TracedTensorNode(name="test_setitem_export", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]), name="x") - + receiver = x + x[0] = 10.0 + self.assertIs(x, receiver) + self.assertIsInstance(x, TracedTensor) + self.assertTrue(torch.equal( + x.tensor, torch.tensor([10.0, 2.0, 3.0, 4.0, 5.0]) + )) y = x * 2 - + ctx.compile_trace({'y': y}) - + input_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]) expected = torch.tensor([20.0, 4.0, 6.0, 8.0, 10.0]) self.validate_export(ctx.m, (input_tensor,), expected, "setitem_single") @@ -2417,12 +2234,18 @@ def test_setitem_slice_fx_torchscript_onnx(self): """Test: x[1:3] = [10, 20] exports correctly to FX, TorchScript, and ONNX.""" ctx = TracedTensorNode(name="test_setitem_slice_export", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]), name="x") - + receiver = x + x[1:3] = torch.tensor([10.0, 20.0]) + self.assertIs(x, receiver) + self.assertIsInstance(x, TracedTensor) + self.assertTrue(torch.equal( + x.tensor, torch.tensor([1.0, 10.0, 20.0, 4.0, 5.0]) + )) y = x * 2 - + ctx.compile_trace({'y': y}) - + input_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]) expected = torch.tensor([2.0, 20.0, 40.0, 8.0, 10.0]) self.validate_export(ctx.m, (input_tensor,), expected, "setitem_slice") @@ -2431,12 +2254,18 @@ def test_setitem_full_slice_constant_fx_torchscript_onnx(self): """Test: x[:] = constant exports correctly to FX, TorchScript, and ONNX.""" ctx = TracedTensorNode(name="test_setitem_full_export", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - + receiver = x + x[:] = torch.tensor([10.0, 20.0, 30.0]) + self.assertIs(x, receiver) + self.assertIsInstance(x, TracedTensor) + self.assertTrue(torch.equal( + x.tensor, torch.tensor([10.0, 20.0, 30.0]) + )) y = x * 2 - + ctx.compile_trace({'y': y}) - + input_tensor = torch.tensor([1.0, 2.0, 3.0]) expected = torch.tensor([20.0, 40.0, 60.0]) self.validate_export(ctx.m, (input_tensor,), expected, "setitem_full") @@ -3065,6 +2894,219 @@ def test_registered_inplace_method_mutates_traced_receiver(self): "traced_receiver_remainder", ) + def test_schema_discovered_view_methods_mutate_traced_receiver(self): + """View mutations preserve identity and export their functional form.""" + cases = ( + ("squeeze_", (2, 3, 1), (-1,)), + ("unsqueeze_", (2, 3), (-1,)), + ("transpose_", (2, 3), (0, 1)), + ("t_", (2, 3), ()), + ("swapaxes_", (2, 3), (0, 1)), + ("swapdims_", (2, 3), (0, 1)), + ) + + for method, shape, args in cases: + with self.subTest(method=method): + trace_input = torch.arange( + torch.tensor(shape).prod(), dtype=torch.float32 + ).reshape(shape) + ctx = TracedTensorNode(name="test", node_index=0) + receiver = ctx.create_input(trace_input.clone(), name="x") + + result = getattr(receiver, method)(*args) + + expected = trace_input.clone() + getattr(expected, method)(*args) + self.assertIs(result, receiver) + self.assertEqual(receiver.shape, expected.shape) + self.assertTrue(torch.equal(receiver.tensor, expected)) + + ctx.compile_trace({"out": receiver}) + runtime_input = trace_input + 10 + runtime_expected = runtime_input.clone() + getattr(runtime_expected, method)(*args) + self.validate_export( + ctx.m, + runtime_input, + runtime_expected, + f"traced_receiver_{method.rstrip('_')}", + ) + + def test_inplace_squeeze_transit_invalidates_port(self): + """A post-capture squeeze mutates the carrier and drops provenance.""" + ctx = TracedTensorNode(name="test", node_index=0) + receiver = ctx.create_input(torch.ones(2, 3, 1), name="x") + ctx.compile_trace({"out": receiver}) + ctx.publish_output_port(receiver, "out") + + result = receiver.squeeze_(-1) + + self.assertIs(result, receiver) + self.assertEqual(receiver.shape, torch.Size([2, 3])) + self.assertIsNone(receiver.output_port) + + def test_inplace_masked_fill_exports_functionally(self): + """A discarded masked_fill_ result still updates receiver and graph.""" + ctx = TracedTensorNode(name="test", node_index=0) + trace_input = torch.tensor([1.0, 2.0, 3.0, 4.0]) + trace_mask = torch.tensor([True, False, True, False]) + receiver = ctx.create_input(trace_input.clone(), name="x") + mask = ctx.create_input(trace_mask, name="mask") + + result = receiver.masked_fill_(mask, -1.0) + + self.assertIs(result, receiver) + self.assertTrue( + torch.equal(receiver.tensor, torch.tensor([-1., 2., -1., 4.])) + ) + ctx.compile_trace({"out": receiver}) + runtime_input = torch.tensor([10.0, 20.0, 30.0, 40.0]) + runtime_mask = torch.tensor([False, True, True, False]) + expected = runtime_input.masked_fill(runtime_mask, -1.0) + self.validate_export( + ctx.m, + (runtime_input, runtime_mask), + expected, + "traced_receiver_masked_fill", + ) + + def test_inplace_index_copy_exports_functionally(self): + """index_copy_ preserves identity and replays runtime index and source.""" + ctx = TracedTensorNode(name="test", node_index=0) + trace_input = torch.zeros(4, 2) + trace_index = torch.tensor([0, 2]) + trace_source = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + receiver = ctx.create_input(trace_input.clone(), name="x") + index = ctx.create_input(trace_index, name="index") + source = ctx.create_input(trace_source, name="source") + + result = receiver.index_copy_(0, index, source) + + expected = trace_input.index_copy(0, trace_index, trace_source) + self.assertIs(result, receiver) + self.assertTrue(torch.equal(receiver.tensor, expected)) + ctx.compile_trace({"out": receiver}) + runtime_input = torch.full((4, 2), 5.0) + runtime_index = torch.tensor([1, 3]) + runtime_source = torch.tensor([[10.0, 20.0], [30.0, 40.0]]) + runtime_expected = runtime_input.index_copy( + 0, runtime_index, runtime_source + ) + self.validate_export( + ctx.m, + (runtime_input, runtime_index, runtime_source), + runtime_expected, + "traced_receiver_index_copy", + verify_legacy_onnx=False, + verify_dynamo_onnx=True, + ) + + def test_inplace_cumsum_exports_functionally(self): + """cumsum_ preserves identity and records torch.cumsum for replay.""" + ctx = TracedTensorNode(name="test", node_index=0) + trace_input = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + receiver = ctx.create_input(trace_input.clone(), name="x") + + result = receiver.cumsum_(1) + + self.assertIs(result, receiver) + self.assertTrue(torch.equal(receiver.tensor, trace_input.cumsum(1))) + ctx.compile_trace({"out": receiver}) + runtime_input = torch.tensor([[2.0, 1.0, 4.0], [3.0, 6.0, 5.0]]) + self.validate_export( + ctx.m, + runtime_input, + runtime_input.cumsum(1), + "traced_receiver_cumsum", + ) + + def test_inplace_schema_discovery_is_conservative(self): + """ATen schemas identify safe functions, views, and exclusions.""" + expected = { + "add_": ("add", False), + "eq_": ("eq", False), + "index_put_": ("index_put", False), + "squeeze_": ("squeeze", True), + } + for method, expected_spec in expected.items(): + with self.subTest(method=method): + functional, mutates_metadata = ( + TracedTensor._functional_inplace_spec(method) + ) + self.assertEqual( + (functional.__name__, mutates_metadata), expected_spec + ) + + for method in ("normal_", "polygamma_", "resize_"): + with self.subTest(method=method): + self.assertIsNone( + TracedTensor._functional_inplace_spec(method) + ) + + def test_schema_discovered_comparison_preserves_receiver_dtype(self): + """Functional comparison output is cast to the receiver's dtype.""" + ctx = TracedTensorNode(name="test", node_index=0) + trace_input = torch.tensor([1.0, 4.0, 9.0]) + trace_other = torch.tensor([1.0, 3.0, 3.0]) + receiver = ctx.create_input(trace_input.clone(), name="x") + other = ctx.create_input(trace_other, name="other") + + self.assertIs(receiver.eq_(other), receiver) + expected = trace_input.eq(trace_other).to(trace_input.dtype) + self.assertTrue(torch.equal(receiver.tensor, expected)) + + ctx.compile_trace({"out": receiver}) + runtime_input = torch.tensor([4.0, 16.0, 25.0]) + runtime_other = torch.tensor([4.0, 5.0, 25.0]) + runtime_expected = runtime_input.eq(runtime_other).to( + runtime_input.dtype + ) + self.validate_export( + ctx.m, + (runtime_input, runtime_other), + runtime_expected, + "schema_inplace_comparison", + ) + + def test_inplace_polygamma_uses_argument_adapter(self): + """polygamma_ handles the functional API's reversed arguments.""" + ctx = TracedTensorNode(name="test", node_index=0) + trace_input = torch.tensor([1.0, 2.0, 3.0]) + receiver = ctx.create_input(trace_input.clone(), name="x") + + result = receiver.polygamma_(1) + + self.assertIs(result, receiver) + self.assertTrue( + torch.allclose(receiver.tensor, torch.polygamma(1, trace_input)) + ) + ctx.compile_trace({"out": receiver}) + runtime_input = torch.tensor([1.5, 2.5, 3.5]) + self.validate_export( + ctx.m, + runtime_input, + torch.polygamma(1, runtime_input), + "inplace_polygamma", + verify_legacy_onnx=False, + ) + + def test_nondeterministic_inplace_methods_fail_before_mutation(self): + """Random mutations without safe functional semantics fail clearly.""" + cases = (("normal_", (0.0, 1.0)), ("bernoulli_", (0.5,))) + for method, args in cases: + with self.subTest(method=method): + ctx = TracedTensorNode(name="test", node_index=0) + receiver = ctx.create_input(torch.ones(3), name="x") + before = receiver.tensor.clone() + + with self.assertRaisesRegex( + NotImplementedError, + f"cannot safely lower in-place method {method}", + ): + getattr(receiver, method)(*args) + + self.assertTrue(torch.equal(receiver.tensor, before)) + def test_registered_inplace_method_preserves_destination_dtype(self): """Functional lowering casts back to the in-place receiver dtype.""" ctx = TracedTensorNode(name="test", node_index=0) From 88bce07f8f5d5295514e5a1a41f99f4615a7a3df Mon Sep 17 00:00:00 2001 From: Frank Lai NV Date: Fri, 16 Oct 2026 16:25:33 -0700 Subject: [PATCH 08/21] Adopt Ruff formatting and enforce it in CI (#30) Signed-off-by: Frank Lai --- .github/workflows/ci.yml | 13 + AGENTS.md | 28 + CONTRIBUTING.md | 25 +- README.md | 2 +- docs/source/conf.py | 87 +- leapp/__init__.py | 1 + leapp/backends/export_backend.py | 159 ++-- leapp/backends/exported_program_backend.py | 15 +- leapp/backends/module_builder.py | 212 +++-- leapp/backends/onnx_export_backend.py | 126 +-- leapp/backends/torch_export_backend.py | 21 +- leapp/buffer_tracker.py | 11 +- leapp/export_manager.py | 346 ++++---- leapp/inference_manager.py | 264 +++--- leapp/leapp.py | 63 +- .../warp_operator/__init__.py | 12 +- .../warp_operator/bundle.py | 1 - leapp/leapp_graph/datatypes/__init__.py | 32 +- .../datatypes/numpy/traced_np_array.py | 151 ++-- leapp/leapp_graph/datatypes/patching.py | 13 +- leapp/leapp_graph/datatypes/proxy_view.py | 8 +- .../datatypes/torch/traced_tensor.py | 168 ++-- leapp/leapp_graph/datatypes/traced_data.py | 135 +-- leapp/leapp_graph/datatypes/traced_tensor.py | 185 +++-- leapp/leapp_graph/datatypes/warp/patching.py | 31 +- leapp/leapp_graph/datatypes/warp/session.py | 2 +- .../datatypes/warp/traced_wp_array.py | 76 +- .../datatypes/warp/warp_segment.py | 10 +- leapp/leapp_graph/function_decorator_node.py | 169 ++-- leapp/leapp_graph/leapp_graph.py | 168 ++-- leapp/leapp_graph/leapp_node.py | 306 ++++--- leapp/leapp_graph/traced_node.py | 247 +++--- leapp/leapp_graph/visualization_adapter.py | 73 +- leapp/leapp_graph/warp_op.py | 23 +- leapp/utils/caller_identity.py | 55 +- leapp/utils/dtype.py | 24 +- leapp/utils/enums.py | 62 +- leapp/utils/logging.py | 48 +- leapp/utils/tensor_description.py | 114 ++- leapp/utils/tracing_lock.py | 19 +- leapp/utils/utils.py | 151 ++-- leapp/warp_runtime.py | 12 +- .../src/leapp_visualization/geometry.py | 109 ++- .../src/leapp_visualization/layout.py | 19 +- .../src/leapp_visualization/model.py | 4 +- .../src/leapp_visualization/png_renderer.py | 145 +++- .../tests/test_graph_visualization_builder.py | 26 +- .../test_graph_visualization_compatibility.py | 6 +- .../test_graph_visualization_geometry.py | 708 ++++++++++++++-- .../tests/test_graph_visualization_imports.py | 1 + .../test_graph_visualization_integration.py | 6 +- .../tests/test_graph_visualization_layout.py | 52 +- .../tests/test_graph_visualization_png.py | 78 +- pyproject.toml | 12 +- scripts/leapp-build-warp-runtime | 25 +- tests/functional_tests/base.py | 325 +++++--- tests/functional_tests/test_annotate.py | 766 ++++++++++-------- .../test_config_generation.py | 254 +++--- .../functional_tests/test_connection_case.py | 128 +-- tests/functional_tests/test_export_backend.py | 627 +++++++------- .../functional_tests/test_export_situation.py | 282 ++++--- tests/functional_tests/test_inference.py | 103 ++- tests/functional_tests/test_leapp_tools.py | 37 +- tests/functional_tests/test_legacy_method.py | 89 +- tests/functional_tests/test_logging.py | 4 +- tests/functional_tests/test_passthrough.py | 176 ++-- .../test_shared_memory_arrays.py | 335 +++++--- tests/functional_tests/test_state_tracking.py | 494 +++++++---- .../functional_tests/test_unsupported_fail.py | 428 ++++++---- tests/functional_tests/test_warp_tracing.py | 72 +- tests/test_examples/base_example_test.py | 37 +- tests/test_examples/test_function_patch.py | 14 +- .../test_examples/test_stateful_gru_export.py | 13 +- tests/test_examples/test_wbc_obj.py | 50 +- tests/test_examples/test_wbc_plain.py | 52 +- tests/unit_tests/export_format_validation.py | 11 +- tests/unit_tests/test_conversion.py | 141 ++-- .../unit_tests/test_export_inference_mode.py | 33 +- tests/unit_tests/test_module_builder.py | 169 ++-- tests/unit_tests/test_numpy_compatibility.py | 107 +-- .../test_torch_builtin_compatibility.py | 259 +++--- tests/unit_tests/test_traced_np_array.py | 408 ++++++---- tests/unit_tests/test_traced_tensor.py | 408 +++++----- tests/unit_tests/test_two_pass_planning.py | 8 +- tests/unit_tests/test_user_patching.py | 35 +- tests/unit_tests/test_warp_runtime.py | 61 +- .../unit_tests/test_warp_torch_boundaries.py | 57 +- tests/warp_support.py | 12 +- uv.lock | 130 +-- 89 files changed, 6749 insertions(+), 4235 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f8644b43..277aa8bf 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -11,6 +11,19 @@ concurrency: cancel-in-progress: true jobs: + format: + name: Ruff formatting + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Install Ruff + run: python -m pip install ruff==0.17.1 + - name: Check formatting + run: ruff format --check . + test: runs-on: [self-hosted, gpu] container: diff --git a/AGENTS.md b/AGENTS.md index 540456fd..561368df 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -74,6 +74,34 @@ Narrower (use these while iterating): | Visualization | `pytest packages/leapp-visualization/tests/ -v` (Python 3.11+) | | Warp tracing | see `warp_support.md` | +### Python formatting + +Use Ruff for Python formatting. The pinned development dependency and +`[tool.ruff]` configuration in `pyproject.toml` set an 88-character line width +and Python 3.10 compatibility. Run these commands from the repository root +after editing Python code and before handing off changes: + +```bash +ruff format . +ruff format --check . +``` + +This covers the main package, visualization package, tests, Python +docs configuration, and `scripts/leapp-build-warp-runtime`. The separate +`Ruff formatting` CI job runs `ruff format --check .` on pull requests and +fails if formatting changes are needed. No local Git hook is required. + +The entire `examples/` directory is excluded from Ruff formatting, including +when example files are passed explicitly. Preserve their existing layout. +When editing examples, check documentation excerpts and source links; prefer +function names or text boundaries over fixed line ranges. + +Use the pinned Ruff version instead of Black. Formatting does not include +lint fixes or import sorting; keep those changes out of formatting-only work. +When upgrading Ruff, update the dependency and `required-version` in +`pyproject.toml`, the install pin in `.github/workflows/ci.yml`, and `uv.lock` +together. + ### Warp runtime build Pip distributions ship the native sources under diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 0efad270..9fd46bad 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -2,6 +2,29 @@ SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 --> +#### Formatting + +LEAPP uses Ruff for Python formatting, with an 88-character line width. Install +the development dependencies and run the formatter from the repository root: + +```bash +python -m pip install -e ".[dev]" +ruff format . +ruff format --check . +``` + +The separate `Ruff formatting` CI job checks pull requests and fails if files +need formatting. Run `ruff format .` locally and include the resulting changes +in your PR. No local Git hook is required. You can also enable format on save +with the Ruff editor extension. + +The `examples/` directory is excluded from formatting to preserve the examples' +existing layout, including when files are passed explicitly to Ruff. + +Keep the Ruff version in `pyproject.toml` and `.github/workflows/ci.yml` in sync +and regenerate `uv.lock` when upgrading. Formatting does not run lint fixes or +sort imports. + #### Signing Your Work * We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license. @@ -53,4 +76,4 @@ SPDX-License-Identifier: Apache-2.0 are public and that a record of the contribution (including all personal information I submit with it, including my sign-off) is maintained indefinitely and may be redistributed consistent with - this project or the open source license(s) involved. \ No newline at end of file + this project or the open source license(s) involved. diff --git a/README.md b/README.md index 18f29ad7..6585c639 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ leapp.stop() leapp.compile_graph() ``` -Running this produces a `my_pipeline/` directory containing the exported models, a YAML pipeline spec, and, on Python 3.11 or later, a PNG graph visualization. On Python 3.10, LEAPP warns and skips visualization while completing the rest of the export. +Running this produces a `my_pipeline/` directory containing the exported models and a YAML pipeline spec. ## Documentation diff --git a/docs/source/conf.py b/docs/source/conf.py index abee0867..20870655 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -7,63 +7,64 @@ from pathlib import Path from zipfile import ZIP_DEFLATED, ZipFile -sys.path.insert(0, os.path.abspath('../..')) +sys.path.insert(0, os.path.abspath("../..")) # -- Project information ----------------------------------------------------- -project = 'LEAPP' -author = 'NVIDIA Corporation' -copyright = f'{datetime.now().year}, NVIDIA Corporation' +project = "LEAPP" +author = "NVIDIA Corporation" +copyright = f"{datetime.now().year}, NVIDIA Corporation" try: from importlib.metadata import version as _pkg_version - release = _pkg_version('leapp') + + release = _pkg_version("leapp") except Exception: - release = '0.7.2' + release = "0.7.2" version = release # -- General configuration --------------------------------------------------- extensions = [ - 'nvidia_sphinx_theme', - 'sphinx.ext.autodoc', - 'sphinx.ext.autosectionlabel', - 'sphinx.ext.intersphinx', - 'sphinx.ext.mathjax', - 'sphinx.ext.napoleon', - 'sphinx.ext.todo', - 'sphinx.ext.viewcode', - 'sphinx_copybutton', - 'sphinx_design', - 'sphinx_tabs.tabs', - 'sphinxcontrib.spelling', + "nvidia_sphinx_theme", + "sphinx.ext.autodoc", + "sphinx.ext.autosectionlabel", + "sphinx.ext.intersphinx", + "sphinx.ext.mathjax", + "sphinx.ext.napoleon", + "sphinx.ext.todo", + "sphinx.ext.viewcode", + "sphinx_copybutton", + "sphinx_design", + "sphinx_tabs.tabs", + "sphinxcontrib.spelling", ] -templates_path = ['_templates'] -exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store', '_redirects'] +templates_path = ["_templates"] +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", "_redirects"] autosectionlabel_prefix_document = True autosectionlabel_maxdepth = 2 sphinx_tabs_disable_tab_closing = True -copybutton_prompt_text = r'>>> |\$ ' +copybutton_prompt_text = r">>> |\$ " copybutton_prompt_is_regexp = True copybutton_copy_empty_lines = False todo_include_todos = True intersphinx_mapping = { - 'python': ('https://docs.python.org/3', None), - 'torch': ('https://pytorch.org/docs/stable', None), - 'numpy': ('https://numpy.org/doc/stable', None), + "python": ("https://docs.python.org/3", None), + "torch": ("https://pytorch.org/docs/stable", None), + "numpy": ("https://numpy.org/doc/stable", None), } # -- Spell checking ---------------------------------------------------------- -spelling_lang = 'en_US' -tokenizer_lang = 'en_US' -spelling_word_list_filename = 'spelling_wordlist.txt' +spelling_lang = "en_US" +tokenizer_lang = "en_US" +spelling_word_list_filename = "spelling_wordlist.txt" spelling_show_suggestions = True spelling_warning = True @@ -71,7 +72,7 @@ linkcheck_ignore = [ # local file references in examples that aren't published URLs - r'.*\.gitlab-master\.nvidia\.com.*', + r".*\.gitlab-master\.nvidia\.com.*", ] linkcheck_anchors = False linkcheck_workers = 10 @@ -80,37 +81,37 @@ # -- HTML output ------------------------------------------------------------- -html_theme = 'nvidia_sphinx_theme' +html_theme = "nvidia_sphinx_theme" html_theme_options = { - 'copyright_override': { - 'start': 2025, + "copyright_override": { + "start": 2025, }, - 'collapse_navigation': False, - 'navigation_depth': -1, + "collapse_navigation": False, + "navigation_depth": -1, } -html_static_path = ['_static'] -html_extra_path = ['_redirects'] -html_css_files = ['css/leapp.css'] +html_static_path = ["_static"] +html_extra_path = ["_redirects"] +html_css_files = ["css/leapp.css"] -html_last_updated_fmt = '%b %d, %Y' +html_last_updated_fmt = "%b %d, %Y" def package_agent_skill(app, exception): """Bundle the repository skill and its references with the published HTML.""" - if exception is not None or app.builder.format != 'html': + if exception is not None or app.builder.format != "html": return repo_root = Path(__file__).resolve().parents[2] - download_dir = Path(app.outdir) / '_downloads' + download_dir = Path(app.outdir) / "_downloads" download_dir.mkdir(parents=True, exist_ok=True) - files = [repo_root / 'SKILL.md', *sorted((repo_root / 'references').rglob('*.md'))] - with ZipFile(download_dir / 'leapp-integrate.zip', 'w', ZIP_DEFLATED) as bundle: + files = [repo_root / "SKILL.md", *sorted((repo_root / "references").rglob("*.md"))] + with ZipFile(download_dir / "leapp-integrate.zip", "w", ZIP_DEFLATED) as bundle: for source in files: - name = (Path('leapp-integrate') / source.relative_to(repo_root)).as_posix() + name = (Path("leapp-integrate") / source.relative_to(repo_root)).as_posix() bundle.writestr(name, source.read_bytes()) def setup(app): - app.connect('build-finished', package_agent_skill) + app.connect("build-finished", package_agent_skill) diff --git a/leapp/__init__.py b/leapp/__init__.py index 8a79bafc..b2b6820d 100644 --- a/leapp/__init__.py +++ b/leapp/__init__.py @@ -26,6 +26,7 @@ from .export_manager import ExportManager from .inference_manager import InferenceManager + # Importing this module registers the ``leapp::warp_runner`` custom op as a # side effect, so it is available to tracing/export without explicit setup. from .leapp_graph.custom_operator_registry import warp_operator # noqa: F401 diff --git a/leapp/backends/export_backend.py b/leapp/backends/export_backend.py index 0fcfffbf..958ca0a8 100644 --- a/leapp/backends/export_backend.py +++ b/leapp/backends/export_backend.py @@ -110,7 +110,7 @@ def _output_shape_is_static(shape): return all(isinstance(dim, int) and dim >= 0 for dim in shape) def _can_use_cuda_iobinding(self, args): - if self._active_provider != 'CUDAExecutionProvider': + if self._active_provider != "CUDAExecutionProvider": return False if not args: @@ -164,7 +164,7 @@ def _run_with_cuda_iobinding(self, args): ) binding.bind_input( name=name, - device_type='cuda', + device_type="cuda", device_id=device_id, element_type=np_dtype, shape=tuple(prepared.shape), @@ -188,7 +188,7 @@ def _run_with_cuda_iobinding(self, args): ).contiguous() binding.bind_output( name=output_meta.name, - device_type='cuda', + device_type="cuda", device_id=output_device.index or 0, element_type=self._torch_dtype_to_numpy_dtype(torch_dtype), shape=tuple(output_tensor.shape), @@ -201,7 +201,7 @@ def _run_with_cuda_iobinding(self, args): def __del__(self): """Clean up temp directory if we own it.""" - if getattr(self, '_temp_dir', None) is not None: + if getattr(self, "_temp_dir", None) is not None: try: if os.path.exists(self._temp_dir): shutil.rmtree(self._temp_dir) @@ -212,13 +212,13 @@ def _get_providers(self): """Get execution providers, always preferring CUDA when available.""" available_providers = ort.get_available_providers() - if 'CUDAExecutionProvider' in available_providers: - return ['CUDAExecutionProvider', 'CPUExecutionProvider'] + if "CUDAExecutionProvider" in available_providers: + return ["CUDAExecutionProvider", "CPUExecutionProvider"] _get_logger().warning( "CUDA execution provider not available. Falling back to CPU." ) - return ['CPUExecutionProvider'] + return ["CPUExecutionProvider"] def _get_source_size(self): """Get total size of all files in the source directory.""" @@ -233,7 +233,9 @@ def _get_source_size(self): return total_size - def save(self, destination, include_initializers=True, keep_initializers_as_inputs=False): + def save( + self, destination, include_initializers=True, keep_initializers_as_inputs=False + ): """Save the ONNX model to disk. For large models (>2GB), re-saves with external data format. @@ -265,7 +267,8 @@ def save(self, destination, include_initializers=True, keep_initializers_as_inpu os.remove(data_path) _get_logger().info( - f"Saving model to {destination} with external data in {data_filename}") + f"Saving model to {destination} with external data in {data_filename}" + ) onnx.save( model_proto, destination, @@ -279,11 +282,11 @@ def save(self, destination, include_initializers=True, keep_initializers_as_inpu # Verify external data file was created expected_data_path = os.path.join(dest_dir, data_filename) if os.path.exists(expected_data_path): - _get_logger().info( - f"External data file created: {expected_data_path}") + _get_logger().info(f"External data file created: {expected_data_path}") else: _get_logger().warning( - f"External data file NOT found: {expected_data_path}") + f"External data file NOT found: {expected_data_path}" + ) else: # Small model - just copy files shutil.copy2(src_path, destination) @@ -307,14 +310,18 @@ def _ensure_session(self): # patterns (e.g. Gemm chains produced by FX make_fx decomposition). # ORT_ENABLE_BASIC is safe and still applies constant folding. sess_options.graph_optimization_level = ( - ort.GraphOptimizationLevel.ORT_ENABLE_BASIC) + ort.GraphOptimizationLevel.ORT_ENABLE_BASIC + ) self._session = ort.InferenceSession( - model_path, sess_options, providers=providers) - self._input_names = [ - inp.name for inp in self._session.get_inputs()] + model_path, sess_options, providers=providers + ) + self._input_names = [inp.name for inp in self._session.get_inputs()] self._output_metas = list(self._session.get_outputs()) - self._active_provider = self._session.get_providers( - )[0] if self._session.get_providers() else 'CPUExecutionProvider' + self._active_provider = ( + self._session.get_providers()[0] + if self._session.get_providers() + else "CPUExecutionProvider" + ) def __call__(self, *args): """Run inference on the ONNX model. @@ -335,7 +342,7 @@ def __call__(self, *args): ) # Determine output device from input tensors - output_device = 'cpu' + output_device = "cpu" for tensor in args: if isinstance(tensor, torch.Tensor) and tensor.is_cuda: output_device = tensor.device @@ -368,7 +375,7 @@ def prepare_tensors_for_export(tensors): with torch.inference_mode(False): for t in tensors: if isinstance(t, torch.Tensor): - if hasattr(t, 'original_clone'): + if hasattr(t, "original_clone"): prepared.append(t.original_clone()) else: prepared.append(t.clone()) @@ -390,7 +397,7 @@ def __init__(self, node_context, backend_params=None): self.compiled_model = None self.compiled_module = None self.runtime_device = None - + def override_module_builder(self, module_builder: Callable): self.module_builder = module_builder @@ -398,7 +405,7 @@ def _verify_model_location_and_get_hash(self, model_path): if not os.path.exists(model_path): _get_logger().error(f"Model file not found at {model_path}") return None, None - with open(model_path, 'rb') as f: + with open(model_path, "rb") as f: file_data = f.read() md5sum = hashlib.md5(file_data).hexdigest() sha256sum = hashlib.sha256(file_data).hexdigest() @@ -437,54 +444,56 @@ def get_backend_metadata(self) -> dict: @abc.abstractmethod def compile(self, m: torch.nn.Module = None) -> Any: - ''' + """ Compiles the model. - This function should return the compiled model. The resulting compiled model + This function should return the compiled model. The resulting compiled model can be used in recombination with other models. - ''' + """ raise NotImplementedError @abc.abstractmethod def save(self, save_path: str) -> Tuple[str, str, str]: - ''' + """ Save the compiled model to the given path This function should apply all necessary optimizations - ''' + """ raise NotImplementedError @abc.abstractmethod def load(self, model_path: str, sha256sum: str): raise NotImplementedError - + def _select_runtime_device(self): - return 'cuda' if torch.cuda.is_available() else 'cpu' + return "cuda" if torch.cuda.is_available() else "cpu" def _load_onnx(self, model_path: str, sha256sum: str): - _, actual_sha256sum = self._verify_model_location_and_get_hash( - model_path) + _, actual_sha256sum = self._verify_model_location_and_get_hash(model_path) if actual_sha256sum != sha256sum: _get_logger().fatal( f"SHA256 checksum mismatch for {model_path}: " f"expected {sha256sum}, got {actual_sha256sum}", - error_type=ValueError) + error_type=ValueError, + ) model = SimplifiedONNXProgram(model_path) self.compiled_model = model - has_ort_cuda = 'CUDAExecutionProvider' in ort.get_available_providers() - self.runtime_device = 'cuda' if ( - torch.cuda.is_available() and has_ort_cuda - ) else 'cpu' - self.compiled_module = None # ONNX models cannot be represented as a module for reexport. - + has_ort_cuda = "CUDAExecutionProvider" in ort.get_available_providers() + self.runtime_device = ( + "cuda" if (torch.cuda.is_available() and has_ort_cuda) else "cpu" + ) + self.compiled_module = ( + None # ONNX models cannot be represented as a module for reexport. + ) + def _load_torchscript(self, model_path: str, sha256sum: str): - _, actual_sha256sum = self._verify_model_location_and_get_hash( - model_path) + _, actual_sha256sum = self._verify_model_location_and_get_hash(model_path) if actual_sha256sum != sha256sum: _get_logger().fatal( f"SHA256 checksum mismatch for {model_path}: " f"expected {sha256sum}, got {actual_sha256sum}", - error_type=ValueError) + error_type=ValueError, + ) device = self._select_runtime_device() model = torch.jit.load(model_path, map_location=device) model = model.to(device) @@ -506,30 +515,48 @@ def get_backend_metadata(self): return {} def compile(self, m: torch.nn.Module = None): - if "model_path" not in self.backend_params or self.backend_params['model_path'] is None: + if ( + "model_path" not in self.backend_params + or self.backend_params["model_path"] is None + ): _get_logger().warning( - f"No model path provided for {self.node_context.name}") - _get_logger().warning("if this is intentional, please provide a path to the correct model " - "in the generated yaml file. Otherwise, please manually fill in the backend parameters.") + f"No model path provided for {self.node_context.name}" + ) + _get_logger().warning( + "if this is intentional, please provide a path to the correct model " + "in the generated yaml file. Otherwise, please manually fill in the backend parameters." + ) else: - _, sha256sum = self._verify_model_location_and_get_hash(self.backend_params['model_path']) - self.load(self.backend_params['model_path'], sha256sum, self.get_backend_model_type()) - + _, sha256sum = self._verify_model_location_and_get_hash( + self.backend_params["model_path"] + ) + self.load( + self.backend_params["model_path"], + sha256sum, + self.get_backend_model_type(), + ) def save(self, save_path: str) -> Tuple[str, str, str]: - if "model_path" not in self.backend_params or self.backend_params['model_path'] is None: + if ( + "model_path" not in self.backend_params + or self.backend_params["model_path"] is None + ): return None, None, None md5sum, sha256sum = self._verify_model_location_and_get_hash( - self.backend_params['model_path']) - model_path = self.backend_params['model_path'] + self.backend_params["model_path"] + ) + model_path = self.backend_params["model_path"] - if "copy_original_model" in self.backend_params and self.backend_params['copy_original_model'] is True: + if ( + "copy_original_model" in self.backend_params + and self.backend_params["copy_original_model"] is True + ): model_path = self._copy_model_to_path(model_path, save_path) return model_path, md5sum, sha256sum def load(self, model_path: str, sha256sum: str, model_type=None): - if model_type is None: + if model_type is None: return if model_type == "onnx": self._load_onnx(model_path, sha256sum) @@ -537,28 +564,28 @@ def load(self, model_path: str, sha256sum: str, model_type=None): self._load_torchscript(model_path, sha256sum) else: _get_logger().fatal( - f"Unsupported model type: {model_type}", - error_type=ValueError) + f"Unsupported model type: {model_type}", error_type=ValueError + ) def get_backend_model_type(self): if "model_path" not in self.backend_params: return None - path = self.backend_params['model_path'] - suffix = path.split('.')[-1] - if suffix == 'pt': + path = self.backend_params["model_path"] + suffix = path.split(".")[-1] + if suffix == "pt": return "jit" - elif suffix == 'pt2': + elif suffix == "pt2": return "pt2" - elif suffix == 'onnx': + elif suffix == "onnx": return "onnx" - elif suffix == 'cpp' or suffix == "cc": + elif suffix == "cpp" or suffix == "cc": return "cpp" - elif suffix == 'py': + elif suffix == "py": return "py" - elif suffix == 'engine' or suffix == 'plan': - return 'trt' + elif suffix == "engine" or suffix == "plan": + return "trt" else: _get_logger().fatal( - f"Unsupported model file suffix: {suffix}", - error_type=ValueError) + f"Unsupported model file suffix: {suffix}", error_type=ValueError + ) diff --git a/leapp/backends/exported_program_backend.py b/leapp/backends/exported_program_backend.py index 7b707e49..4d87a995 100644 --- a/leapp/backends/exported_program_backend.py +++ b/leapp/backends/exported_program_backend.py @@ -36,8 +36,7 @@ def _load_warp_runtime( torch.ops.load_library(str(library_path)) except OSError as exc: _get_logger().fatal( - "Failed to load the PT2 Warp custom operator library: " - f"{library_path}", + f"Failed to load the PT2 Warp custom operator library: {library_path}", error_type=RuntimeError, cause=exc, ) @@ -53,8 +52,7 @@ def compile(self, m: torch.nn.Module = None): m = self.module_builder() m = m.eval() - input_values = [ - tensor_desc.value for tensor_desc in self.node_context.inputs] + input_values = [tensor_desc.value for tensor_desc in self.node_context.inputs] input_values = prepare_tensors_for_export(input_values) export_args = tuple(input_values) @@ -62,8 +60,7 @@ def compile(self, m: torch.nn.Module = None): if "strict" in self.backend_params: export_kwargs["strict"] = self.backend_params["strict"] - self.exported_program = torch.export.export( - m, export_args, **export_kwargs) + self.exported_program = torch.export.export(m, export_args, **export_kwargs) self._load_warp_runtime(self.exported_program) self.compiled_model = self.exported_program.module() self.compiled_module = m @@ -73,7 +70,8 @@ def save(self, save_path: str) -> Tuple[str, str, str]: if self.exported_program is None: _get_logger().fatal( f"No exported program found for {self.node_context.name}", - error_type=RuntimeError) + error_type=RuntimeError, + ) path = os.path.join(save_path, f"{self.node_context.name}.pt2") torch.export.save(self.exported_program, path) @@ -86,7 +84,8 @@ def load(self, model_path: str, sha256sum: str): _get_logger().fatal( f"SHA256 checksum mismatch for {model_path}: " f"expected {sha256sum}, got {actual_sha256sum}", - error_type=ValueError) + error_type=ValueError, + ) self.exported_program = torch.export.load(model_path) self._load_warp_runtime(self.exported_program) diff --git a/leapp/backends/module_builder.py b/leapp/backends/module_builder.py index 15b715ed..a805e6f0 100644 --- a/leapp/backends/module_builder.py +++ b/leapp/backends/module_builder.py @@ -32,8 +32,8 @@ def __init__(self, *args, **kwargs): self.saved_buffers: List[str] = [] def add_buffer(self, name, value): - if 'self.' in name: - name = name.split('self.')[1] + if "self." in name: + name = name.split("self.")[1] if name in self.saved_buffers: _get_logger().fatal( f"Buffer with name '{name}' was already registered", @@ -60,39 +60,46 @@ def forward(self): name, bases, { - '__init__': __init__, - 'add_buffer': add_buffer, + "__init__": __init__, + "add_buffer": add_buffer, # '_core' : intentionally not included here - it's dynamically bound - 'forward': forward, - '__constant__': list(), - } + "forward": forward, + "__constant__": list(), + }, ) module_instance = ModuleTemplate() return module_instance -class ModuleBuilder: +class ModuleBuilder: def __init__(self, node_context): self.node_context = node_context def __call__(self): # validation - if self.node_context.input_namespace is None or self.node_context.output_namespace is None: + if ( + self.node_context.input_namespace is None + or self.node_context.output_namespace is None + ): _get_logger().fatal( f"Input or output namespace not found for {self.node_context.name}", - error_type=Exception) + error_type=Exception, + ) # if the function or snippet is a class method, we need to create a module that inherits # from the class to preserve class functions input_namespace = self.node_context.input_namespace - if 'self' in input_namespace: - parent_class = type(input_namespace['self']) + if "self" in input_namespace: + parent_class = type(input_namespace["self"]) else: parent_class = None self.module_instance = get_module_template( - self.node_context.name, parent_class, self.node_context.cached_constant_values) + self.node_context.name, + parent_class, + self.node_context.cached_constant_values, + ) # setup the module with required attributes from the environment self._register_explicit_values() @@ -100,14 +107,19 @@ def __call__(self): # extract the body of the function body = self._extract_source_code() - core_header, prepended_conditioning, return_statement, outputs_str = self._create_core_header_and_return() - forward_header, forward_body, forward_return_statement = self._create_forward_header_and_body() - - #TODO: this can be removed once we are sure that the new method works with self. returns + core_header, prepended_conditioning, return_statement, outputs_str = ( + self._create_core_header_and_return() + ) + forward_header, forward_body, forward_return_statement = ( + self._create_forward_header_and_body() + ) + + # TODO: this can be removed once we are sure that the new method works with self. returns if self.node_context.from_function: # For functions, append outputs_str to all return statements - modified_source, replacement_count = self._append_outputs_to_return_statements( - body, outputs_str) + modified_source, replacement_count = ( + self._append_outputs_to_return_statements(body, outputs_str) + ) body = modified_source if replacement_count != 0 or len(self.node_context._declared_outputs) == 0: return_statement = "" @@ -115,36 +127,41 @@ def __call__(self): # stitch the _core function together: final_core_body = prepended_conditioning + body + return_statement - indented_core_body = textwrap.indent(final_core_body, ' ') + indented_core_body = textwrap.indent(final_core_body, " ") core_function_string = core_header + indented_core_body final_forward_body = forward_body + forward_return_statement - indented_forward_body = textwrap.indent(final_forward_body, ' ') + indented_forward_body = textwrap.indent(final_forward_body, " ") forward_function_string = forward_header + indented_forward_body # combine _core and forward functions function_string = core_function_string + "\n" + forward_function_string - _get_logger().debug("\n"+function_string) + _get_logger().debug("\n" + function_string) self._compile_forward_function(function_string) return self.module_instance def _register_explicit_values(self): - for constant_name, constant_value in self.node_context.cached_constant_values.items(): - if 'self.' in constant_name: - constant_name = constant_name.split('self.')[1] + for ( + constant_name, + constant_value, + ) in self.node_context.cached_constant_values.items(): + if "self." in constant_name: + constant_name = constant_name.split("self.")[1] setattr(self.module_instance, constant_name, constant_value) self.module_instance.__constant__.append(constant_name) _get_logger().debug( - f"Set instance attribute (constant value): {constant_name} = {constant_value}") + f"Set instance attribute (constant value): {constant_name} = {constant_value}" + ) # Set default values as instance attributes (from stored dictionary) for default_name, default_value in self.node_context.default_kwargs.items(): # Set as instance attribute - accessible as self.default_name setattr(self.module_instance, default_name, default_value) _get_logger().debug( - f"Set instance attribute (default value): {default_name} = {default_value}") + f"Set instance attribute (default value): {default_name} = {default_value}" + ) # add buffers to the module for buffer_name in self.node_context.register_buffers: @@ -155,32 +172,49 @@ def _register_explicit_values(self): _get_logger().info("Created the following buffers:") for buffer_name in self.module_instance.saved_buffers: _get_logger().info( - f" - self.{buffer_name}: initialized as {getattr(self.module_instance, buffer_name)}") + f" - self.{buffer_name}: initialized as {getattr(self.module_instance, buffer_name)}" + ) def _duplicate_attributes(self): # copy all the attributes from the original object to the module input_namespace = self.node_context.input_namespace - if 'self' in input_namespace: - original_obj = input_namespace['self'] + if "self" in input_namespace: + original_obj = input_namespace["self"] # Warning #1: Lazy modules may not transfer correctly - if isinstance(original_obj, LazyModuleMixin) and original_obj.has_uninitialized_params(): + if ( + isinstance(original_obj, LazyModuleMixin) + and original_obj.has_uninitialized_params() + ): _get_logger().error( "Parent class contains uninitialized lazy parameters. " - "Export may fail or produce incorrect results.") + "Export may fail or produce incorrect results." + ) # Warning #3: Check for registered hooks that may not transfer correctly - hook_attrs = ['_forward_hooks', '_backward_hooks', '_forward_pre_hooks', '_backward_pre_hooks'] + hook_attrs = [ + "_forward_hooks", + "_backward_hooks", + "_forward_pre_hooks", + "_backward_pre_hooks", + ] for hook_attr in hook_attrs: - if hasattr(original_obj, hook_attr) and len(getattr(original_obj, hook_attr)) > 0: + if ( + hasattr(original_obj, hook_attr) + and len(getattr(original_obj, hook_attr)) > 0 + ): _get_logger().error( - f"Parent class has registered {hook_attr}. These may not transfer correctly.") + f"Parent class has registered {hook_attr}. These may not transfer correctly." + ) for attr_name in dir(original_obj): # we don't copy any private attributes or register buffers or registered constants - if attr_name.startswith('__') or attr_name.endswith('__') or \ - f"self.{attr_name}" in self.node_context.register_buffers or \ - f"{attr_name}" in self.module_instance.__constant__: + if ( + attr_name.startswith("__") + or attr_name.endswith("__") + or f"self.{attr_name}" in self.node_context.register_buffers + or f"{attr_name}" in self.module_instance.__constant__ + ): continue try: @@ -209,44 +243,45 @@ def _duplicate_attributes(self): bound_method = types.MethodType(value, self.module_instance) setattr(self.module_instance, attr_name, bound_method) except Exception as e: - _get_logger().error( - f"Error setting attribute {attr_name}: {e}") + _get_logger().error(f"Error setting attribute {attr_name}: {e}") # Fix #2: Also copy slotted attributes (not in __dict__) for cls in type(original_obj).__mro__: - if hasattr(cls, '__slots__'): + if hasattr(cls, "__slots__"): for slot in cls.__slots__: - if slot.startswith('__') or slot.endswith('__'): + if slot.startswith("__") or slot.endswith("__"): continue - if hasattr(original_obj, slot) and not hasattr(self.module_instance, slot): + if hasattr(original_obj, slot) and not hasattr( + self.module_instance, slot + ): try: value = getattr(original_obj, slot) setattr(self.module_instance, slot, value) _get_logger().debug(f"Copied slotted attribute: {slot}") except Exception as e: - _get_logger().error(f"Error copying slotted attribute {slot}: {e}") + _get_logger().error( + f"Error copying slotted attribute {slot}: {e}" + ) def _extract_source_code(self): """Extract the source code body for the traced function or block. - + Uses the unified extract_source_from_line_range function which handles both function decorators and block contexts, and properly captures multiline statements. """ executed_lines = self.node_context.executed_lines - + source_code, message = extract_source_from_line_range( executed_lines, self.node_context.name, - is_function=self.node_context.from_function + is_function=self.node_context.from_function, ) - + if source_code: _get_logger().info(message) else: - _get_logger().fatal( - message, - error_type=Exception) + _get_logger().fatal(message, error_type=Exception) return source_code @@ -254,39 +289,52 @@ def _create_core_header_and_return(self): """Create header, body prepend (data patching), and return statement for _core function.""" # Output type annotation for _core output_types = [ - parameter_description.dtype for parameter_description in self.node_context.output_formats] + parameter_description.dtype + for parameter_description in self.node_context.output_formats + ] if len(output_types) == 0: output_types_str = "None" elif len(output_types) == 1: output_types_str = output_types[0] else: output_types_str = f"Tuple[{', '.join(output_types)}]" - + # Core function needs type annotations for TorchScript - core_inputs = ', '.join([f"{input_val.name_str}: {input_val.dtype}" - for input_val in self.node_context.input_formats]) + core_inputs = ", ".join( + [ + f"{input_val.name_str}: {input_val.dtype}" + for input_val in self.node_context.input_formats + ] + ) header = f"def _core(self, {core_inputs}) -> {output_types_str}:\n" # Build data patching string (prepended to core body) prepended_conditioning = "" - + # Handle input name transformations (e.g., "name_str" -> "name_raw") for input in self.node_context.input_formats: if input.name_raw != input.name_str: prepended_conditioning += f"{input.name_raw} = {input.name_str}\n" # Handle constants, buffers, and default kwargs - targets = self.node_context.environment_constants | self.node_context.register_buffers | set( - self.node_context.default_kwargs.keys()) + targets = ( + self.node_context.environment_constants + | self.node_context.register_buffers + | set(self.node_context.default_kwargs.keys()) + ) for const_name in targets: if "self." in const_name or const_name == "self": continue # special case for any class variables, no rename required prepended_conditioning += f"{const_name} = self.{const_name}\n" # only include outputs that are declared by user. - outputs = ", ".join([output.name_raw - for output in self.node_context.output_formats - if output.name_raw in self.node_context._declared_outputs]) + outputs = ", ".join( + [ + output.name_raw + for output in self.node_context.output_formats + if output.name_raw in self.node_context._declared_outputs + ] + ) return_statement = f"\nreturn {outputs}\n" if outputs else "" return header, prepended_conditioning, return_statement, outputs @@ -295,8 +343,9 @@ def _create_forward_header_and_body(self): """Create header and body for forward function.""" # Forward function accepts flat tensor inputs (no type annotations) forward_input_names = ", ".join( - [tensor_desc.name for tensor_desc in self.node_context.inputs]) - + [tensor_desc.name for tensor_desc in self.node_context.inputs] + ) + header = f"def forward(self, {forward_input_names}):\n" prepended_conditioning = "" @@ -305,15 +354,17 @@ def _create_forward_header_and_body(self): packing_expr = param_format.packed_tensor_expr if packing_expr: prepended_conditioning += packing_expr + "\n" - + appended_conditioning = "" for param_format in self.node_context.output_formats: appended_conditioning += param_format.unpacked_tensor_expr + "\n" # Build forward function body with explicit unpacking - core_output_names = [output.name_str for output in self.node_context.output_formats] + core_output_names = [ + output.name_str for output in self.node_context.output_formats + ] core_input_names = [input.name_str for input in self.node_context.input_formats] - core_input_names_str = ', '.join(core_input_names) + core_input_names_str = ", ".join(core_input_names) forward_output_names = [output.name for output in self.node_context.outputs] if len(core_output_names) == 0: @@ -325,7 +376,7 @@ def _create_forward_header_and_body(self): body = f"({output_names_str}) = self._core({core_input_names_str})\n" body = prepended_conditioning + body + appended_conditioning - + # Build return statement based on the actual number of forward outputs (after unpacking) if len(forward_output_names) == 0: return_statement = "" @@ -334,7 +385,7 @@ def _create_forward_header_and_body(self): else: forward_output_names_str = ", ".join(forward_output_names) return_statement = f"return {forward_output_names_str}\n" - + return header, body, return_statement def _append_outputs_to_return_statements(self, source_code, outputs_str): @@ -360,20 +411,21 @@ def replace_return(match): # Match return statements: capture indentation and the return value # Negative lookahead to avoid matching 'return' in strings - pattern = r'^(\s*)return(\s+[^\n]*)?$' - modified_code = re.sub(pattern, replace_return, - source_code, flags=re.MULTILINE) + pattern = r"^(\s*)return(\s+[^\n]*)?$" + modified_code = re.sub(pattern, replace_return, source_code, flags=re.MULTILINE) return modified_code, replacement_count[0] def _compile_forward_function(self, function_string): - filename = f'generated_{self.node_context.name}_{id(self.module_instance)}_function' + filename = ( + f"generated_{self.node_context.name}_{id(self.module_instance)}_function" + ) try: code = compile(function_string, filename, "exec") # recreate the environment of the function using stored namespace namespace = { - 'Tensor': torch.Tensor, + "Tensor": torch.Tensor, "Tuple": Tuple, "List": List, "Dict": Dict, @@ -387,15 +439,15 @@ def _compile_forward_function(self, function_string): _get_logger().fatal( f"Error compiling function: {e}\n{function_string}", error_type=type(e), - cause=e) + cause=e, + ) + + _core = namespace["_core"] + forward = namespace["forward"] - _core = namespace['_core'] - forward = namespace['forward'] + lines = [line + "\n" for line in function_string.splitlines()] + linecache.cache[filename] = (len(function_string), None, lines, filename) - lines = [line + '\n' for line in function_string.splitlines()] - linecache.cache[filename] = ( - len(function_string), None, lines, filename) - # Set methods on the class (not instance) so TorchScript can recognize them module_class = type(self.module_instance) module_class._core = _core diff --git a/leapp/backends/onnx_export_backend.py b/leapp/backends/onnx_export_backend.py index 89ce18c2..0275a3e9 100644 --- a/leapp/backends/onnx_export_backend.py +++ b/leapp/backends/onnx_export_backend.py @@ -5,7 +5,11 @@ import torch from typing import Tuple -from leapp.backends.export_backend import ExportBackend, prepare_tensors_for_export, SimplifiedONNXProgram +from leapp.backends.export_backend import ( + ExportBackend, + prepare_tensors_for_export, + SimplifiedONNXProgram, +) from leapp.utils.logging import _get_logger import os import onnx @@ -16,32 +20,41 @@ from torch.onnx import _constants from leapp.leapp_graph.custom_operator_registry import warp_operator + class ONNXExportBackend(ExportBackend): def get_backend_metadata(self): metadata = {} - metadata['opset_version'] = getattr(self, 'opset_version', _constants.ONNX_DEFAULT_OPSET) + metadata["opset_version"] = getattr( + self, "opset_version", _constants.ONNX_DEFAULT_OPSET + ) return metadata def get_backend_model_type(self): return "onnx" - + def _get_onnx_model(self, onnx_program): """Get the ONNX model proto from an ONNX program.""" - if hasattr(onnx_program, 'model_proto'): + if hasattr(onnx_program, "model_proto"): # dynamo export returns ONNXProgram with model_proto return onnx_program.model_proto else: # SimplifiedONNXProgram - load from file - model_path = os.path.join(onnx_program._source_dir, onnx_program._source_filename) + model_path = os.path.join( + onnx_program._source_dir, onnx_program._source_filename + ) return onnx.load(model_path) def _get_actual_onnx_inputs(self, onnx_model): """Get actual input names from ONNX model, excluding initializers (constants).""" # Initializers are constants baked into the model - exclude them initializer_names = {init.name for init in onnx_model.graph.initializer} - + # Return only true dynamic inputs - return [inp.name for inp in onnx_model.graph.input if inp.name not in initializer_names] + return [ + inp.name + for inp in onnx_model.graph.input + if inp.name not in initializer_names + ] def _handle_duplicate_io_names(self): """Rename overlapping input/output names for ONNX's flat I/O namespace.""" @@ -72,25 +85,31 @@ def _handle_duplicate_io_names(self): def _sync_inputs_with_onnx(self, onnx_program): """Sync inputs with what ONNX actually exported. - + ONNX export can remove unused inputs (constant folding, dead code elimination). This is allowed and we update node_context.inputs accordingly. - + However, if ONNX renames inputs (which happens with identity/passthrough functions when using onnx-dynamo), we raise an error and suggest using onnx-torchscript instead. """ onnx_model = self._get_onnx_model(onnx_program) - + expected_input_names = set(td.name for td in self.node_context.inputs) actual_input_names = self._get_actual_onnx_inputs(onnx_model) actual_input_set = set(actual_input_names) - + # Check for removed inputs (this is OK - ONNX optimized them away) - removed_inputs = [td.name for td in self.node_context.inputs if td.name not in actual_input_set] - + removed_inputs = [ + td.name + for td in self.node_context.inputs + if td.name not in actual_input_set + ] + # Check for renamed/new inputs (this is NOT OK - we can't track them) - renamed_inputs = [name for name in actual_input_names if name not in expected_input_names] - + renamed_inputs = [ + name for name in actual_input_names if name not in expected_input_names + ] + if renamed_inputs: _get_logger().fatal( f"[{self.node_context.name}] ONNX export renamed inputs: {renamed_inputs}. " @@ -101,7 +120,7 @@ def _sync_inputs_with_onnx(self, onnx_program): "3. remove the inputs that are not used in the computation or directly returned as output", error_type=ValueError, ) - + if removed_inputs: _get_logger().warning( f"[{self.node_context.name}] ONNX export optimized away {len(removed_inputs)} inputs: {removed_inputs}. " @@ -109,13 +128,12 @@ def _sync_inputs_with_onnx(self, onnx_program): ) # Track trimmed inputs self.node_context.trimmed_inputs.update(removed_inputs) - + # Filter inputs to only keep actual ONNX inputs (preserving order) self.node_context.inputs = [ td for td in self.node_context.inputs if td.name in actual_input_set ] - def load(self, model_path: str, sha256sum: str): self._load_onnx(model_path, sha256sum) @@ -145,7 +163,11 @@ def _fix_scalar_slice_inputs(model: onnx.ModelProto) -> int: if inp_name in init_map: arr = numpy_helper.to_array(init_map[inp_name]) if arr.ndim == 0: - role = ["data", "starts", "ends", "axes", "steps"][i] if i < 5 else f"input_{i}" + role = ( + ["data", "starts", "ends", "axes", "steps"][i] + if i < 5 + else f"input_{i}" + ) new_name = f"{inp_name}_1d_{node.name}_{role}" new_tensor = numpy_helper.from_array( arr.reshape(1), name=new_name @@ -157,7 +179,8 @@ def _fix_scalar_slice_inputs(model: onnx.ModelProto) -> int: if num_fixed > 0: _get_logger().debug( f"Fixed {num_fixed} scalar Slice initializer(s) " - f"(ONNX exporter bug: shared scalar initializers between Gather and Slice nodes)") + f"(ONNX exporter bug: shared scalar initializers between Gather and Slice nodes)" + ) @staticmethod def _set_onnx_string_attr(node: onnx.NodeProto, name: str, value: str) -> None: @@ -244,26 +267,28 @@ def save(self, save_path: str) -> Tuple[str, str, str]: ) # Validate model (skip for large models as it can fail with protobuf limits) - skip_validation = self.backend_params.get('skip_validation', False) + skip_validation = self.backend_params.get("skip_validation", False) if not skip_validation: try: # Try to check model, but don't fail the export if validation fails onnx.checker.check_model(onnx_path) except Exception as e: _get_logger().error( - f"ONNX model validation warning (model still saved): {e}") + f"ONNX model validation warning (model still saved): {e}" + ) md5sum, sha256sum = self._verify_model_location_and_get_hash(onnx_path) return onnx_path, md5sum, sha256sum - def compile(self, m: torch.nn.Module = None) -> SimplifiedONNXProgram: + def compile(self, m: torch.nn.Module = None) -> SimplifiedONNXProgram: raise NotImplementedError( - "TorchExportBackend does not support compilation, please use torch-script or torch-trace instead") + "TorchExportBackend does not support compilation, please use torch-script or torch-trace instead" + ) return None -class ONNXTorchScriptExportBackend(ONNXExportBackend): +class ONNXTorchScriptExportBackend(ONNXExportBackend): def compile(self, m: torch.nn.Module = None): if m is None: m = self.module_builder() @@ -272,32 +297,34 @@ def compile(self, m: torch.nn.Module = None): # Optionally pre-script the module before ONNX export # This is useful when using traced models as environment constants - if self.backend_params.get('prescript', False): + if self.backend_params.get("prescript", False): m = torch.jit.script(m) - + # Get flat tensor values directly from inputs (not input_formats which preserves nested structure) input_values = tuple( - [tensor_desc.value for tensor_desc in self.node_context.inputs]) + [tensor_desc.value for tensor_desc in self.node_context.inputs] + ) # Clone tensors to escape inference mode (inference tensors can't participate in autograd) input_values = prepare_tensors_for_export(input_values) # Create temp directory manually (not context manager) so we can control its lifetime tmpdir = tempfile.mkdtemp() save_path = os.path.join(tmpdir, "model.onnx") - self.opset_version = self.backend_params.get('opset_version', _constants.ONNX_DEFAULT_OPSET) + self.opset_version = self.backend_params.get( + "opset_version", _constants.ONNX_DEFAULT_OPSET + ) torch.onnx.export( m, input_values, save_path, dynamo=False, - input_names=[ - tensor_desc.name for tensor_desc in self.node_context.inputs], + input_names=[tensor_desc.name for tensor_desc in self.node_context.inputs], output_names=[ - tensor_desc.name for tensor_desc in self.node_context.outputs], - + tensor_desc.name for tensor_desc in self.node_context.outputs + ], verbose=_get_logger().is_verbose(), - report=self.backend_params.get('report', False), + report=self.backend_params.get("report", False), opset_version=self.opset_version, ) @@ -312,6 +339,7 @@ def compile(self, m: torch.nn.Module = None): self.compiled_model = onnx_program self.compiled_module = m + class ONNXDynamoExportBackend(ONNXExportBackend): @staticmethod def _fix_gathernd_int32_indices(model: onnx.ModelProto) -> int: @@ -370,28 +398,29 @@ def compile(self, m: torch.nn.Module = None): self._handle_duplicate_io_names() # Get flat tensor values directly from inputs (not input_formats which preserves nested structure) input_values = tuple( - [tensor_desc.value for tensor_desc in self.node_context.inputs]) + [tensor_desc.value for tensor_desc in self.node_context.inputs] + ) # Clone tensors to escape inference mode (inference tensors can't participate in autograd) input_values = prepare_tensors_for_export(input_values) - self.opset_version = self.backend_params.get('opset_version', _constants.ONNX_DEFAULT_OPSET) + self.opset_version = self.backend_params.get( + "opset_version", _constants.ONNX_DEFAULT_OPSET + ) onnx_program = torch.onnx.export( m, input_values, None, # no need to save the model dynamo=True, - input_names=[ - tensor_desc.name for tensor_desc in self.node_context.inputs], + input_names=[tensor_desc.name for tensor_desc in self.node_context.inputs], output_names=[ - tensor_desc.name for tensor_desc in self.node_context.outputs], - - verify=self.backend_params.get('verify', False), - optimize=self.backend_params.get('optimize', True), - + tensor_desc.name for tensor_desc in self.node_context.outputs + ], + verify=self.backend_params.get("verify", False), + optimize=self.backend_params.get("optimize", True), verbose=_get_logger().is_verbose(), - report=self.backend_params.get('report', False), + report=self.backend_params.get("report", False), opset_version=self.opset_version, ) - + # Sync inputs with what ONNX actually exported # (ONNX can remove unused inputs during optimization) self._sync_inputs_with_onnx(onnx_program) @@ -411,10 +440,11 @@ def compile(self, m: torch.nn.Module = None): tmpdir = tempfile.mkdtemp() tmp_path = os.path.join(tmpdir, f"{self.node_context.name}.onnx") - PROTOBUF_LIMIT = 1.5 * 1024 * 1024 * 1024 # 1.5 GB — headroom below 2 GB protobuf limit + PROTOBUF_LIMIT = ( + 1.5 * 1024 * 1024 * 1024 + ) # 1.5 GB — headroom below 2 GB protobuf limit estimated_size = sum( - numpy_helper.to_array(init).nbytes - for init in model_proto.graph.initializer + numpy_helper.to_array(init).nbytes for init in model_proto.graph.initializer ) if estimated_size > PROTOBUF_LIMIT: data_filename = f"{self.node_context.name}.onnx.data" diff --git a/leapp/backends/torch_export_backend.py b/leapp/backends/torch_export_backend.py index 95f7096a..0eca8a19 100644 --- a/leapp/backends/torch_export_backend.py +++ b/leapp/backends/torch_export_backend.py @@ -26,7 +26,7 @@ class TorchExportBackend(ExportBackend): def get_backend_metadata(self): return {} - + def get_backend_model_type(self): return "jit" @@ -35,14 +35,16 @@ def save(self, save_path: str) -> Tuple[str, str, str]: # Freeze the model before saving for optimization if compiled_model is not None: preserved_attrs = [] - if hasattr(self, 'node_context') and hasattr(self.node_context, 'saved_buffers'): + if hasattr(self, "node_context") and hasattr( + self.node_context, "saved_buffers" + ): preserved_attrs = self.node_context.saved_buffers compiled_model = torch.jit.freeze( - compiled_model.eval(), preserved_attrs=preserved_attrs) + compiled_model.eval(), preserved_attrs=preserved_attrs + ) else: - _get_logger().error( - f"No compiled model found for {self.node_context.name}") + _get_logger().error(f"No compiled model found for {self.node_context.name}") path = os.path.join(save_path, f"{self.node_context.name}.pt") compiled_model.save(path) @@ -54,7 +56,8 @@ def load(self, model_path: str, sha256sum: str): def compile(self): raise NotImplementedError( - "TorchExportBackend does not support compilation, please use torch-script or torch-trace instead") + "TorchExportBackend does not support compilation, please use torch-script or torch-trace instead" + ) return None @@ -70,13 +73,11 @@ def compile(self, m: torch.nn.Module = None): m = self.module_builder() m = m.eval() # Get flat tensor values directly from inputs (not input_formats which preserves nested structure) - input_values = [ - tensor_desc.value for tensor_desc in self.node_context.inputs] + input_values = [tensor_desc.value for tensor_desc in self.node_context.inputs] # Clone tensors to escape inference mode (inference tensors can't participate in autograd) input_values = prepare_tensors_for_export(input_values) - self.compiled_model = torch.jit.trace( - m, input_values, **self.backend_params) + self.compiled_model = torch.jit.trace(m, input_values, **self.backend_params) self.compiled_module = m # Freezing moved to save() method to allow node combination diff --git a/leapp/buffer_tracker.py b/leapp/buffer_tracker.py index da4cba46..2bbbaff0 100644 --- a/leapp/buffer_tracker.py +++ b/leapp/buffer_tracker.py @@ -59,6 +59,7 @@ @dataclass class _BufferInfo: """Tracking info for a single injected buffer.""" + clean_name: str original_traced: TracedTensor original_buffer: torch.Tensor @@ -160,7 +161,8 @@ def inject(self) -> None: _get_logger().info( f"BufferTracker: injected {len(self._injected)} buffer(s) " - f"into node '{self._node_name}'") + f"into node '{self._node_name}'" + ) def collect(self) -> None: """Detect mutated buffers and register as state outputs. @@ -196,7 +198,8 @@ def collect(self) -> None: self._export_manager.update_state(self._node_name, mutated) _get_logger().info( f"BufferTracker: detected {len(mutated)} mutated buffer(s): " - f"{list(mutated.keys())}") + f"{list(mutated.keys())}" + ) # Bake non-mutated buffers as constants in the FX graph. # They should keep their trained values (e.g. normalizer mean/var) @@ -230,7 +233,8 @@ def collect(self) -> None: _get_logger().info( f"BufferTracker: baked {len(non_mutated)} non-mutated buffer(s) " - f"as constants: {non_mutated}") + f"as constants: {non_mutated}" + ) self._collected = True @@ -249,4 +253,3 @@ def restore(self) -> None: info.owning_module.register_buffer(info.attr_name, raw) self._injected.clear() - diff --git a/leapp/export_manager.py b/leapp/export_manager.py index 07519215..49f79505 100644 --- a/leapp/export_manager.py +++ b/leapp/export_manager.py @@ -31,20 +31,30 @@ is_traced_type, is_tracable_tensor_type, ) -from leapp.leapp_graph.datatypes.patching import TracingPatcher, warn_if_script_functions_in_scope +from leapp.leapp_graph.datatypes.patching import ( + TracingPatcher, + warn_if_script_functions_in_scope, +) from leapp.utils.tensor_description import TensorSemantics -from leapp.utils.tensor_description import (verify_data_exact_match, - flatten_io_structure, - unwrap_tensor_semantics) -from leapp.utils.caller_identity import (get_caller_stack_identity, - caller_identity_has_same_anchor, - format_caller_identity) -from leapp.utils.utils import (get_relative_path, - mirror_all_traced_state, - extract_return_names, - frame_to_namespace) +from leapp.utils.tensor_description import ( + verify_data_exact_match, + flatten_io_structure, + unwrap_tensor_semantics, +) +from leapp.utils.caller_identity import ( + get_caller_stack_identity, + caller_identity_has_same_anchor, + format_caller_identity, +) +from leapp.utils.utils import ( + get_relative_path, + mirror_all_traced_state, + extract_return_names, + frame_to_namespace, +) from contextlib import nullcontext + class ExportManager: _instance = None _initialized = False # True after singleton __init__ completes @@ -80,8 +90,6 @@ def __init__(self): ExportManager._initialized = True - - ######################################################### # state accessors (for runtime API) ######################################################### @@ -139,12 +147,12 @@ def reset_tracing_lock(self): def restore_pending_buffer_trackers(self): for node in self.nodes.values(): - if hasattr(node, '_buffer_tracker') and node._buffer_tracker is not None: + if hasattr(node, "_buffer_tracker") and node._buffer_tracker is not None: node._buffer_tracker.restore() def set_detected_graph(self, models: dict, pipeline: dict): - self.detected_nodes = models['models'] - self.detected_pipeline = pipeline['pipeline'] + self.detected_nodes = models["models"] + self.detected_pipeline = pipeline["pipeline"] ######################################################### # node setup @@ -158,12 +166,12 @@ def _assign_completion_index(self, node: LeappNode): def validate_nodes_ready_for_compile(self): incomplete_nodes = [ - name for name, node in self.nodes.items() + name + for name, node in self.nodes.items() if node.node_index == LeappNode.UNSET_NODE_INDEX ] pending_warp_nodes = [ - name for name, node in self.nodes.items() - if node.has_pending_warp_segments + name for name, node in self.nodes.items() if node.has_pending_warp_segments ] if incomplete_nodes: incomplete_nodes.sort() @@ -172,7 +180,8 @@ def validate_nodes_ready_for_compile(self): "The following nodes were created but never completed: " f"{formatted}. Did you forget to call output_tensors() " "or finish the annotated function?", - error_type=Exception) + error_type=Exception, + ) if pending_warp_nodes: pending_warp_nodes.sort() formatted = ", ".join(pending_warp_nodes) @@ -188,23 +197,25 @@ def _setup_new_node(self, name, node_class: LeappNode, **kwargs): _get_logger().fatal( f"Error: node '{name}' already exists. " f"Cannot create a new node with the same name.", - error_type=Exception) + error_type=Exception, + ) if self.is_dry_run(name): - kwargs['export_with'] = None - kwargs.setdefault('backend_params', {}) + kwargs["export_with"] = None + kwargs.setdefault("backend_params", {}) if node_class is FunctionDecoratorNode: # kept for backward compatibility we still suppport old style initialization # TODO: this needs to be removed in the future - node = node_class(name, - backend=kwargs.get("export_with", None), - backend_params=kwargs.get("backend_params", None), - inputs=kwargs.get("inputs", None), - outputs=kwargs.get("outputs", None), - environment_constants=kwargs.get( - "environment_constants", None), - register_buffers=kwargs.get("register_buffers", None)) + node = node_class( + name, + backend=kwargs.get("export_with", None), + backend_params=kwargs.get("backend_params", None), + inputs=kwargs.get("inputs", None), + outputs=kwargs.get("outputs", None), + environment_constants=kwargs.get("environment_constants", None), + register_buffers=kwargs.get("register_buffers", None), + ) else: node = node_class(name, **kwargs) @@ -222,8 +233,8 @@ def _passthrough_dict_values(tensors: dict): def _normalize_named_tensor_payload(api_name: str, node_name: str, tensors): metadata = {} if isinstance(tensors, (TensorSemantics, list, tuple)) and ( - isinstance(tensors, TensorSemantics) or - any(isinstance(t, TensorSemantics) for t in tensors) + isinstance(tensors, TensorSemantics) + or any(isinstance(t, TensorSemantics) for t in tensors) ): tensors, metadata = unwrap_tensor_semantics(tensors) @@ -231,7 +242,8 @@ def _normalize_named_tensor_payload(api_name: str, node_name: str, tensors): _get_logger().fatal( f"{api_name}() for node '{node_name}' does not accept a bare tensor. " "Pass a dict of named tensors or a TensorSemantics/list of TensorSemantics.", - error_type=TypeError) + error_type=TypeError, + ) if isinstance(tensors, dict): return tensors, metadata @@ -239,7 +251,8 @@ def _normalize_named_tensor_payload(api_name: str, node_name: str, tensors): _get_logger().fatal( f"{api_name}() for node '{node_name}' expects either a dict of named tensors " f"or a TensorSemantics/list of TensorSemantics. Received {type(tensors).__name__}.", - error_type=TypeError) + error_type=TypeError, + ) ######################################################### # annotation APIs @@ -249,10 +262,12 @@ def input_tensors(self, node_name: str, tensors): _get_logger().fatal( "Cannot call input_tensors() while a _method()-traced function " "is executing. Mixing active contexts is not allowed.", - error_type=Exception) + error_type=Exception, + ) tensors, metadata = self._normalize_named_tensor_payload( - "input_tensors", node_name, tensors) + "input_tensors", node_name, tensors + ) if not ExportManager._interpret_graph: return self._passthrough_dict_values(tensors) @@ -261,8 +276,7 @@ def input_tensors(self, node_name: str, tensors): if node_name in self.nodes: traced_tensors_node = self.nodes[node_name] else: - traced_tensors_node = self._setup_new_node( - node_name, TracedTensorNode) + traced_tensors_node = self._setup_new_node(node_name, TracedTensorNode) _caller_identity = get_caller_stack_identity() @@ -275,7 +289,8 @@ def input_tensors(self, node_name: str, tensors): if not traced_tensors_node.is_tracing: matching_origin = next( ( - identity for identity in traced_tensors_node._caller_identities + identity + for identity in traced_tensors_node._caller_identities if caller_identity_has_same_anchor(identity, _caller_identity) ), None, @@ -285,14 +300,16 @@ def input_tensors(self, node_name: str, tensors): f"Error: node '{node_name}' is being called from a different annotation origin " f"than the first trace. Cannot reuse a node name from a different call site.\n" f"New call site:\n{format_caller_identity(_caller_identity)}", - error_type=Exception) + error_type=Exception, + ) if _caller_identity not in traced_tensors_node._caller_identities: _get_logger().warning( f"Warning: node '{node_name}' is being re-entered from a new caller context, " f"but the normalized annotation origin matches a previously seen site. " f"Allowing re-entry.\n" f"Original origin:\n{format_caller_identity(matching_origin)}\n" - f"New call site:\n{format_caller_identity(_caller_identity)}") + f"New call site:\n{format_caller_identity(_caller_identity)}" + ) traced_tensors_node._caller_identities.add(_caller_identity) traced_tensors_node.reentry_validate_inputs(tensors) rebound_tensors = { @@ -318,7 +335,8 @@ def input_tensors(self, node_name: str, tensors): traced_tensors = [] for tensor_name, tensor in tensors.items(): traced_tensor = traced_tensors_node.create_input( - tensor, tensor_name, semantics=metadata.get(tensor_name)) + tensor, tensor_name, semantics=metadata.get(tensor_name) + ) traced_tensors.append(traced_tensor) # if node is tracing we return the traced tensors @@ -326,11 +344,11 @@ def input_tensors(self, node_name: str, tensors): def output_tensors(self, node_name: str, tensors, static_outputs=None, **kwargs): tensors, metadata = self._normalize_named_tensor_payload( - "output_tensors", node_name, tensors) + "output_tensors", node_name, tensors + ) if not ExportManager._interpret_graph: return self._passthrough_dict_values(tensors) - if node_name in self.nodes: traced_tensors_node = self.nodes[node_name] @@ -338,46 +356,58 @@ def output_tensors(self, node_name: str, tensors, static_outputs=None, **kwargs) _get_logger().fatal( f"output_tensors() called for node '{node_name}' but input_tensors() was never called for it. " "Call annotate.input_tensors() before annotate.output_tensors() for the same node name.", - error_type=Exception) + error_type=Exception, + ) # force the warp segment to close if any. if self.patcher.warp is not None: self.patcher.warp.close_warp_segment() # process outputs - flattened_tensors = flatten_io_structure(tensors, '') + flattened_tensors = flatten_io_structure(tensors, "") static_outputs_metadata = {} normalized_static_outputs = None if static_outputs is not None: - normalized_static_outputs, static_outputs_metadata = self._normalize_named_tensor_payload( - "output_tensors static_outputs", node_name, static_outputs) + normalized_static_outputs, static_outputs_metadata = ( + self._normalize_named_tensor_payload( + "output_tensors static_outputs", node_name, static_outputs + ) + ) if not traced_tensors_node.is_tracing: flattened_static = {} if normalized_static_outputs is not None: - flattened_static = flatten_io_structure(normalized_static_outputs, '') + flattened_static = flatten_io_structure(normalized_static_outputs, "") traced_tensors_node.reentry_validate_outputs( - flattened_tensors, flattened_static) + flattened_tensors, flattened_static + ) return self._passthrough_dict_values(tensors) # Warn if pre-compiled ScriptFunctions are visible in the caller's scope warn_if_script_functions_in_scope() self._validate_initial_traced_payload( - "output_tensors", node_name, traced_tensors_node, flattened_tensors) + "output_tensors", node_name, traced_tensors_node, flattened_tensors + ) # process static outputs (constant tensors that should be returned but aren't derived from inputs) flattened_static_outputs = None if normalized_static_outputs is not None: - flattened_static_outputs = flatten_io_structure(normalized_static_outputs, '') - - export_with = None if self.is_dry_run(node_name) else kwargs.get("export_with", None) - traced_tensors_node.compile_trace(flattened_tensors, - backend=export_with, - backend_params=kwargs.get("backend_params", {}), - static_tensors=flattened_static_outputs, - semantics_map=metadata, - static_semantics_map=static_outputs_metadata) + flattened_static_outputs = flatten_io_structure( + normalized_static_outputs, "" + ) + + export_with = ( + None if self.is_dry_run(node_name) else kwargs.get("export_with", None) + ) + traced_tensors_node.compile_trace( + flattened_tensors, + backend=export_with, + backend_params=kwargs.get("backend_params", {}), + static_tensors=flattened_static_outputs, + semantics_map=metadata, + static_semantics_map=static_outputs_metadata, + ) self._assign_completion_index(traced_tensors_node) return self._passthrough_dict_values(tensors) @@ -425,7 +455,8 @@ def register_buffer(self, node_name: str, tensors): if node_name not in self.nodes: msg = ( f"register_buffer() called for node '{node_name}' but node not found. " - "Call annotate.input_tensors() first to create the node.") + "Call annotate.input_tensors() first to create the node." + ) _get_logger().fatal(msg, error_type=Exception) # Normalize input to a dict with auto-generated names if needed @@ -436,7 +467,8 @@ def register_buffer(self, node_name: str, tensors): if not isinstance(traced_node, TracedTensorNode): msg = ( f"register_buffer() is not supported for node '{node_name}' — " - "it was created with the legacy method annotation.") + "it was created with the legacy method annotation." + ) _get_logger().fatal(msg, error_type=Exception) if not traced_node.is_tracing: @@ -479,12 +511,16 @@ def _validate_flat_state_payload(api_name: str, node_name: str, tensors): f"State '{state_name}' has unsupported top-level type {type(value).__name__}. " "Please either explicitly list out each state as its own named tensor " "or use input_tensors() and rely on LEAPP feedback detection.", - error_type=TypeError) + error_type=TypeError, + ) @staticmethod - def _validate_initial_traced_payload(api_name: str, node_name: str, - traced_tensors_node: TracedTensorNode, - flattened_tensors: dict) -> None: + def _validate_initial_traced_payload( + api_name: str, + node_name: str, + traced_tensors_node: TracedTensorNode, + flattened_tensors: dict, + ) -> None: """Validate first-trace payloads for APIs that require active TracedTensors. This helper is only for the initial tracing path, before a node has been @@ -497,32 +533,35 @@ def _validate_initial_traced_payload(api_name: str, node_name: str, ``output_tensors()`` and ``update_state()`` and avoids later internal FX failures when raw tensors slip into graph outputs/state updates. """ - instances = set(is_traced_type(tensor) - for tensor in flattened_tensors.values()) + instances = set(is_traced_type(tensor) for tensor in flattened_tensors.values()) if not all(instances): - types = set(type(tensor).__name__ - for tensor in flattened_tensors.values()) + types = set(type(tensor).__name__ for tensor in flattened_tensors.values()) msg = ( f"{api_name}() for node '{node_name}' received non-traced tensors: {types}\n" "Possible causes:\n" "1. You did not use the traced tensors returned by leapp functions — " "make sure to replace your original tensors with the return value of annotation functions.\n" "2. An operation in your computation broke tracing (e.g. converting to numpy and back).\n" - f"3. You passed raw tensors instead of the traced ones to {api_name}().") + f"3. You passed raw tensors instead of the traced ones to {api_name}()." + ) _get_logger().fatal(msg, error_type=Exception) - context_names = { - tensor.context for tensor in flattened_tensors.values() - } - if not (len(context_names) == 1 and next(iter(context_names)) == traced_tensors_node.name): + context_names = {tensor.context for tensor in flattened_tensors.values()} + if not ( + len(context_names) == 1 + and next(iter(context_names)) == traced_tensors_node.name + ): msg = ( f"{api_name}() for node '{node_name}' received tensors that belong to a different node: " f"{context_names}.\n Make sure you are passing tensors derived from " - f"annotate.input_tensors('{node_name}', ...) to annotate.{api_name}('{node_name}', ...).") + f"annotate.input_tensors('{node_name}', ...) to annotate.{api_name}('{node_name}', ...)." + ) _get_logger().fatal(msg, error_type=Exception) - def state_tensors(self, node_name: str, tensors: dict[str, torch.Tensor]) -> TracedTensor | tuple[TracedTensor, ...]: + def state_tensors( + self, node_name: str, tensors: dict[str, torch.Tensor] + ) -> TracedTensor | tuple[TracedTensor, ...]: """Register state tensors (both inputs AND outputs) for a traced node. Call input_tensors() first to create the node. Use update_state() to set output values. @@ -531,12 +570,11 @@ def state_tensors(self, node_name: str, tensors: dict[str, torch.Tensor]) -> Tra self._validate_flat_state_payload("state_tensors", node_name, tensors) return self._passthrough_dict_values(tensors) - - if node_name not in self.nodes: msg = ( f"state_tensors() called for node '{node_name}' but node not found. " - "Call annotate.input_tensors() first to create the node.") + "Call annotate.input_tensors() first to create the node." + ) _get_logger().fatal(msg, error_type=Exception) traced_node = self.nodes[node_name] @@ -544,7 +582,8 @@ def state_tensors(self, node_name: str, tensors: dict[str, torch.Tensor]) -> Tra if not isinstance(traced_node, TracedTensorNode): msg = ( f"state_tensors() is not supported for node '{node_name}' — " - "it was created with the legacy method annotation.") + "it was created with the legacy method annotation." + ) _get_logger().fatal(msg, error_type=Exception) self._validate_flat_state_payload("state_tensors", node_name, tensors) @@ -558,7 +597,9 @@ def state_tensors(self, node_name: str, tensors: dict[str, torch.Tensor]) -> Tra values = list(state_dict.values()) return values[0] if len(values) == 1 else tuple(values) - def update_state(self, node_name: str, tensors: dict[str, TracedTensor]) -> TracedTensor | tuple[TracedTensor, ...]: + def update_state( + self, node_name: str, tensors: dict[str, TracedTensor] + ) -> TracedTensor | tuple[TracedTensor, ...]: """Set output values for state tensors and return passthrough values.""" if not ExportManager._interpret_graph: self._validate_flat_state_payload("update_state", node_name, tensors) @@ -567,7 +608,8 @@ def update_state(self, node_name: str, tensors: dict[str, TracedTensor]) -> Trac if node_name not in self.nodes: msg = ( f"update_state() called for node '{node_name}' but node not found. " - "Call annotate.input_tensors() first to create the node.") + "Call annotate.input_tensors() first to create the node." + ) _get_logger().fatal(msg, error_type=Exception) traced_node = self.nodes[node_name] @@ -575,7 +617,8 @@ def update_state(self, node_name: str, tensors: dict[str, TracedTensor]) -> Trac if not isinstance(traced_node, TracedTensorNode): msg = ( f"update_state() is not supported for node '{node_name}' — " - "it was created with the legacy method annotation.") + "it was created with the legacy method annotation." + ) _get_logger().fatal(msg, error_type=Exception) self._validate_flat_state_payload("update_state", node_name, tensors) @@ -585,12 +628,17 @@ def update_state(self, node_name: str, tensors: dict[str, TracedTensor]) -> Trac return self._passthrough_dict_values(tensors) self._validate_initial_traced_payload( - "update_state", node_name, traced_node, tensors) + "update_state", node_name, traced_node, tensors + ) traced_node.update_state_tensors(tensors) return self._passthrough_dict_values(tensors) - def module(self, node_name: str, model: torch.nn.Module, - buffer_names: list[str] | None = None) -> None: + def module( + self, + node_name: str, + model: torch.nn.Module, + buffer_names: list[str] | None = None, + ) -> None: """Register a module for automatic stateful buffer tracking. Replaces registered buffers with TracedTensor inputs so the forward @@ -635,40 +683,40 @@ def module(self, node_name: str, model: torch.nn.Module, _get_logger().fatal(msg, error_type=Exception) from leapp.buffer_tracker import BufferTracker + tracker = BufferTracker(model, node_name, self, buffer_names=buffer_names) tracker.inject() self.nodes[node_name]._buffer_tracker = tracker - def method(self, **params): def decorator(func): name = params.get("node_name", func.__name__) export_with = params.get("export_with", None) backend_params = params.get("backend_params", {}) - + @functools.wraps(func) def wrapper(*args, **kwargs): if not ExportManager._interpret_graph: return func(*args, **kwargs) - + # ~~~~~~~~~~~~~~~~~~~ ensure node exists ~~~~~~~~~~~~~~~~~~~~~~~~ # if name not in self.nodes: self._setup_new_node(name, TracedTensorNode) - + # ~~~~~~~~~~~~~~~~~~~ set up inputs ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # sig = inspect.signature(func) bound_args = sig.bind(*args, **kwargs) bound_args.apply_defaults() - + params_list = list(sig.parameters.items()) new_args = [] new_kwargs = {} for i, arg in enumerate(args): param_name, param = params_list[i] - if i == 0 and param_name in ('self', 'cls'): + if i == 0 and param_name in ("self", "cls"): new_args.append(arg) continue if param.kind == inspect.Parameter.VAR_POSITIONAL: @@ -676,24 +724,23 @@ def wrapper(*args, **kwargs): new_args.append(self.input_tensors(name, {f"arg_{j}": a})) break new_args.append(self.input_tensors(name, {param_name: arg})) - + for key, value in kwargs.items(): new_kwargs[key] = self.input_tensors(name, {key: value}) - + # ~~~~~~~~~~~~~~~~~~~ register default kwargs as buffers ~~~~~~~~ # for param_name, param_value in bound_args.arguments.items(): - if param_name in ('self', 'cls'): + if param_name in ("self", "cls"): continue param = sig.parameters[param_name] - was_provided = ( - param_name in kwargs or - (param.kind != inspect.Parameter.VAR_POSITIONAL and - list(sig.parameters.keys()).index(param_name) < len(args)) + was_provided = param_name in kwargs or ( + param.kind != inspect.Parameter.VAR_POSITIONAL + and list(sig.parameters.keys()).index(param_name) < len(args) ) if not was_provided and is_tracable_tensor_type(param_value): traced = self.register_buffer(name, {param_name: param_value}) new_kwargs[param_name] = traced - + # ~~~~~~~~~~~~~~~~~~~ run the function ~~~~~~~~~~~~~~~~~~~~~~~~~ # result = func(*new_args, **new_kwargs) # ~~~~~~~~~~~~~~~~~~~ set up outputs ~~~~~~~~~~~~~~~~~~~~~~~~~~~ # @@ -702,14 +749,18 @@ def wrapper(*args, **kwargs): if result is None: _get_logger().fatal( f"Error: annotated method {name} returned None, but LEAPP expects a return value", - error_type=Exception) + error_type=Exception, + ) elif isinstance(result, tuple): if len(return_names) != len(result): _get_logger().fatal( f"Error: annotated method {name} returned {len(result)} values, " f"but LEAPP detected the following return names {return_names} from source", - error_type=Exception) - output_dict = {return_names[i]: result[i] for i in range(len(result))} + error_type=Exception, + ) + output_dict = { + return_names[i]: result[i] for i in range(len(result)) + } self.output_tensors( name, output_dict, @@ -725,10 +776,18 @@ def wrapper(*args, **kwargs): ) # ~~~~~~~~~~~~~~~~~~~ set up outputs ~~~~~~~~~~~~~~~~~~~~~~~~~~~ # return result + return wrapper + return decorator - def warp_op(self, node_name: str, inputs: list[str] = None, outputs: list[str] = None, **params): + def warp_op( + self, + node_name: str, + inputs: list[str] = None, + outputs: list[str] = None, + **params, + ): if not ExportManager._interpret_graph: return nullcontext() warp_backend = self.patcher.warp @@ -761,6 +820,7 @@ def _method(self, **params): and is kept for use cases where the newer TracedTensorNode-based method() does not cover all patterns. Not advertised in the public API. """ + def decorator(func): if "node_name" in params: @@ -779,7 +839,8 @@ def wrapper(*args, **kwargs): else: new_node = True node_context = self._setup_new_node( - name, FunctionDecoratorNode, **params) + name, FunctionDecoratorNode, **params + ) caller_namespace = frame_to_namespace(sys._getframe(1)) @@ -790,21 +851,22 @@ def wrapper(*args, **kwargs): bound_args.apply_defaults() node_context.inspect_function_inputs(func, args, kwargs) - input_namespace = { - **caller_namespace, **bound_args.arguments} - if hasattr(func, '__self__'): - input_namespace['self'] = func.__self__ + input_namespace = {**caller_namespace, **bound_args.arguments} + if hasattr(func, "__self__"): + input_namespace["self"] = func.__self__ node_context.capture_inputs_from_namespace(input_namespace) node_context.compile_trace(func) node_context.snapshot_buffer_values(input_namespace) trace_fn = node_context.create_trace_function( - os.path.basename(__file__), entry_hook=None) + os.path.basename(__file__), entry_hook=None + ) else: node_context.validate_function_boundaries(func) node_context.validate_function_inputs(func, args, kwargs) trace_fn = node_context.create_trace_function( - os.path.basename(__file__), entry_hook=None) + os.path.basename(__file__), entry_hook=None + ) sys._getframe(1).f_trace = trace_fn @@ -819,21 +881,26 @@ def wrapper(*args, **kwargs): if new_node: _get_logger().info( - f"****Tracing stopped for {node_context.name}****\n\n") + f"****Tracing stopped for {node_context.name}****\n\n" + ) node_context.inspect_function_outputs(func, result) if node_context.output_namespace is not None: node_context.capture_outputs_from_namespace( - node_context.output_namespace) + node_context.output_namespace + ) self._assign_completion_index(node_context) else: node_context.validate_function_outputs(func, result) if node_context.output_namespace is not None: node_context.validate_outputs_from_namespace( - node_context.output_namespace) + node_context.output_namespace + ) node_context.increment_cache_idx() return result + return wrapper + return decorator ######################################################### @@ -851,20 +918,22 @@ def mirror_leapp_tags(self, source, target): return mirror_all_traced_state(source, target) def get_io_descriptions(self): - _get_logger().section( - f"Compiling graph parameters for {len(self.nodes)} nodes") + _get_logger().section(f"Compiling graph parameters for {len(self.nodes)} nodes") models = {"models": {}} - nodes_to_describe = sorted( - self.nodes.values(), key=lambda x: x.node_index) + nodes_to_describe = sorted(self.nodes.values(), key=lambda x: x.node_index) for node in nodes_to_describe: _get_logger().info(f"Compiling parameters for {node.name}") description = node.get_description() - if 'parameters' in description and 'model_path' in description['parameters']: + if ( + "parameters" in description + and "model_path" in description["parameters"] + ): # Convert model path to be relative to YAML file location - model_path = description['parameters']['model_path'] + model_path = description["parameters"]["model_path"] if model_path: - description['parameters']['model_path'] = get_relative_path( - model_path, self.SAVE_PATH) + description["parameters"]["model_path"] = get_relative_path( + model_path, self.SAVE_PATH + ) models["models"][node.name] = description return models @@ -879,18 +948,19 @@ def save_models(self): if self.SAVE_PATH is None: _get_logger().fatal( "Error: No save path provided, please provide a save path to export the graph", - error_type=Exception) - _get_logger().section( - f"Saving {len(self.nodes)} models to {self.SAVE_PATH}") + error_type=Exception, + ) + _get_logger().section(f"Saving {len(self.nodes)} models to {self.SAVE_PATH}") if not os.path.exists(self.SAVE_PATH): os.makedirs(self.SAVE_PATH) for node_context in self.nodes.values(): - _get_logger().info( - f"Saving {node_context.name}") + _get_logger().info(f"Saving {node_context.name}") node_context.save_model(self.SAVE_PATH) _get_logger().info("Success\n") - def validate_all_models(self, rtol: float = 1e-3, atol: float = 1e-5, strict: bool = True): + def validate_all_models( + self, rtol: float = 1e-3, atol: float = 1e-5, strict: bool = True + ): """Validate all exported models by comparing computed outputs against captured outputs. This method runs each compiled model with its captured input values and verifies @@ -928,19 +998,18 @@ def validate_all_models(self, rtol: float = 1e-3, atol: float = 1e-5, strict: bo _get_logger().section(f"Validating {len(self.nodes)} models") for node_name, node in self.nodes.items(): - _get_logger().info(f"Validating {node_name}...") try: - passed, error_hint = node.validate_compiled_model( - rtol=rtol, atol=atol) + passed, error_hint = node.validate_compiled_model(rtol=rtol, atol=atol) results[node_name] = passed if error_hint is not None: error_hints.append(f"{node_name}: {error_hint}") except Exception as e: _get_logger().error( - f"{node_name}: Validation failed with exception: {e}") + f"{node_name}: Validation failed with exception: {e}" + ) results[node_name] = e finally: @@ -964,13 +1033,10 @@ def validate_all_models(self, rtol: float = 1e-3, atol: float = 1e-5, strict: bo _get_logger().error(f" Errors: {errors}/{len(results)}") if strict and (failed > 0 or errors > 0): - failed_nodes = [name for name, - v in results.items() if v is not True] - message = ( - f"Model validation failed for {len(failed_nodes)} node(s): {failed_nodes}" - ) + failed_nodes = [name for name, v in results.items() if v is not True] + message = f"Model validation failed for {len(failed_nodes)} node(s): {failed_nodes}" if error_hints: message += "\n\n" + "\n".join(error_hints) _get_logger().fatal(message, error_type=Exception) - return results \ No newline at end of file + return results diff --git a/leapp/inference_manager.py b/leapp/inference_manager.py index 79a6b043..1f8efb86 100644 --- a/leapp/inference_manager.py +++ b/leapp/inference_manager.py @@ -12,7 +12,10 @@ from safetensors.torch import load_file -from leapp.utils.tensor_description import map_to_torch_dtype, validate_connection_compatibility +from leapp.utils.tensor_description import ( + map_to_torch_dtype, + validate_connection_compatibility, +) from leapp.utils.logging import _get_logger from leapp.backends.torch_export_backend import TorchExportBackend @@ -27,47 +30,50 @@ def __init__(self, name, model_path, inputs, outputs, parameters): self.input_descriptions = inputs self.output_descriptions = outputs - - backend = self._create_backend(parameters['backend']) - backend.load(model_path, parameters['sha256sum']) + backend = self._create_backend(parameters["backend"]) + backend.load(model_path, parameters["sha256sum"]) self.model = backend.compiled_model self.device = backend.runtime_device @property def input_names(self): - return [input['name'] for input in self.input_descriptions] + return [input["name"] for input in self.input_descriptions] @property def input_shapes(self): - return [input['shape'] for input in self.input_descriptions] + return [input["shape"] for input in self.input_descriptions] @property def input_dtypes(self): - return [input['dtype'] for input in self.input_descriptions] + return [input["dtype"] for input in self.input_descriptions] @property def output_names(self): - return [output['name'] for output in self.output_descriptions] + return [output["name"] for output in self.output_descriptions] @property def output_shapes(self): - return [output['shape'] for output in self.output_descriptions] + return [output["shape"] for output in self.output_descriptions] @property def output_dtypes(self): - return [output['dtype'] for output in self.output_descriptions] + return [output["dtype"] for output in self.output_descriptions] @property def mock_input(self): input_values_to_populate = {} for input_val in self.input_descriptions: - input_name = input_val['name'] - input_shape = json.loads(input_val['shape']) if isinstance( - input_val['shape'], str) else input_val['shape'] - input_dtype = input_val['dtype'] + input_name = input_val["name"] + input_shape = ( + json.loads(input_val["shape"]) + if isinstance(input_val["shape"], str) + else input_val["shape"] + ) + input_dtype = input_val["dtype"] input_value = torch.zeros( - tuple(input_shape), dtype=map_to_torch_dtype(input_dtype)) + tuple(input_shape), dtype=map_to_torch_dtype(input_dtype) + ) input_values_to_populate[input_name] = input_value.to(self.device) return input_values_to_populate @@ -76,14 +82,17 @@ def mock_input(self): def mock_output(self): output_values_to_populate = {} for output_val in self.output_descriptions: - output_name = output_val['name'] - output_shape = json.loads(output_val['shape']) if isinstance( - output_val['shape'], str) else output_val['shape'] - output_dtype = output_val['dtype'] + output_name = output_val["name"] + output_shape = ( + json.loads(output_val["shape"]) + if isinstance(output_val["shape"], str) + else output_val["shape"] + ) + output_dtype = output_val["dtype"] output_value = torch.zeros( - tuple(output_shape), dtype=map_to_torch_dtype(output_dtype)) - output_values_to_populate[output_name] = output_value.to( - self.device) + tuple(output_shape), dtype=map_to_torch_dtype(output_dtype) + ) + output_values_to_populate[output_name] = output_value.to(self.device) return output_values_to_populate @@ -96,13 +105,16 @@ def _create_backend(self, backend): elif backend == "pt2": return ExportedProgramExportBackend(None) else: - _get_logger().fatal(f"Unsupported backend: {backend}", error_type=ValueError) + _get_logger().fatal( + f"Unsupported backend: {backend}", error_type=ValueError + ) def __call__(self, *inputs): if len(inputs) != len(self.input_descriptions): _get_logger().fatal( f"Expected {len(self.input_descriptions)} inputs, got {len(inputs)}", - error_type=ValueError) + error_type=ValueError, + ) outputs = self.model(*inputs) @@ -117,7 +129,8 @@ def __call__(self, *inputs): else: _get_logger().fatal( f"Expected {len(self.output_descriptions)} outputs, got {len(outputs)}", - error_type=ValueError) + error_type=ValueError, + ) class InferenceManager: @@ -132,12 +145,14 @@ def __init__(self, model_path): if not os.path.exists(model_path): _get_logger().fatal( f"Leapp description file not found at {model_path}", - error_type=FileNotFoundError) + error_type=FileNotFoundError, + ) if not model_path.endswith(".yaml"): _get_logger().fatal( f"Leapp description file must end with .yaml, got {model_path}", - error_type=ValueError) + error_type=ValueError, + ) self.model_path = model_path @@ -153,10 +168,14 @@ def _load_description(self): with open(self.model_path, "r") as f: self.leapp_description = yaml.safe_load(f) - if any(key not in self.leapp_description for key in ["models", "pipeline", "system information"]): + if any( + key not in self.leapp_description + for key in ["models", "pipeline", "system information"] + ): _get_logger().fatal( f"Leapp description file must contain models, pipeline, and system_info, got {self.leapp_description.keys()}", - error_type=ValueError) + error_type=ValueError, + ) self.models = self.leapp_description["models"] self.pipeline = self.leapp_description["pipeline"] @@ -165,18 +184,25 @@ def _load_description(self): def _create_nodes(self): nodes = {} for name, description in self.models.items(): - if any(key not in description for key in ["inputs", "outputs", "parameters"]): + if any( + key not in description for key in ["inputs", "outputs", "parameters"] + ): _get_logger().fatal( f"Model description must contain inputs, outputs, and parameters, got {description.keys()}", - error_type=ValueError) + error_type=ValueError, + ) # load the model - parameters = description['parameters'] - if any(key not in parameters for key in ["model_path", "md5sum", "sha256sum", "backend"]): + parameters = description["parameters"] + if any( + key not in parameters + for key in ["model_path", "md5sum", "sha256sum", "backend"] + ): _get_logger().fatal( f"Model description must contain model_path, md5sum, sha256sum, and backend, got {parameters.keys()}", - error_type=ValueError) + error_type=ValueError, + ) - model_path = parameters['model_path'] + model_path = parameters["model_path"] base_path = os.path.dirname(self.model_path) relative_model_path = os.path.join(base_path, model_path) @@ -184,7 +210,12 @@ def _create_nodes(self): model_path = relative_model_path node = NodeManager( - name, model_path, description['inputs'], description['outputs'], description['parameters']) + name, + model_path, + description["inputs"], + description["outputs"], + description["parameters"], + ) nodes[name] = node return nodes @@ -199,7 +230,10 @@ def _validate_and_create_inference_graph(self): self.organized_pipeline_connections = {} # Merge data_flow and feedback_flow, combining target lists for shared keys all_flows = {} - for flow_dict in [self.pipeline.get('data_flow', {}), self.pipeline.get('feedback_flow', {})]: + for flow_dict in [ + self.pipeline.get("data_flow", {}), + self.pipeline.get("feedback_flow", {}), + ]: for source, targets in flow_dict.items(): if source in all_flows: all_flows[source] = all_flows[source] + targets @@ -207,44 +241,49 @@ def _validate_and_create_inference_graph(self): all_flows[source] = targets # organize the pipeline connections for source, targets in all_flows.items(): - source_node_name, source_output_name = source.split('/') + source_node_name, source_output_name = source.split("/") if source_node_name not in self.organized_pipeline_connections: self.organized_pipeline_connections[source_node_name] = {} - self.organized_pipeline_connections[source_node_name][source_output_name] = [ - ] + self.organized_pipeline_connections[source_node_name][ + source_output_name + ] = [] for target in targets: - target_node_name, target_input_name = target.split('/') - self.organized_pipeline_connections[source_node_name][source_output_name].append( - (target_node_name, target_input_name)) + target_node_name, target_input_name = target.split("/") + self.organized_pipeline_connections[source_node_name][ + source_output_name + ].append((target_node_name, target_input_name)) # Create '==out==' slot and route final outputs to it output_cache = {} - for node_name, output_names in self.pipeline['outputs'].items(): + for node_name, output_names in self.pipeline["outputs"].items(): node = self.nodes[node_name] # Get output descriptions for this node - output_descs = {desc['name'] - : desc for desc in node.output_descriptions} + output_descs = {desc["name"]: desc for desc in node.output_descriptions} for output_name in output_names: desc = output_descs[output_name] - output_shape = json.loads(desc['shape']) if isinstance( - desc['shape'], str) else desc['shape'] - output_dtype = map_to_torch_dtype(desc['dtype']) + output_shape = ( + json.loads(desc["shape"]) + if isinstance(desc["shape"], str) + else desc["shape"] + ) + output_dtype = map_to_torch_dtype(desc["dtype"]) # Create cache tensor with unique key: node_name/output_name cache_key = f"{node_name}/{output_name}" output_cache[cache_key] = torch.zeros( - tuple(output_shape), dtype=output_dtype, device=node.device) + tuple(output_shape), dtype=output_dtype, device=node.device + ) # Add connection from this output to ==out== if node_name not in self.organized_pipeline_connections: self.organized_pipeline_connections[node_name] = {} if output_name not in self.organized_pipeline_connections[node_name]: - self.organized_pipeline_connections[node_name][output_name] = [ - ] + self.organized_pipeline_connections[node_name][output_name] = [] self.organized_pipeline_connections[node_name][output_name].append( - ('==out==', cache_key)) + ("==out==", cache_key) + ) - self.value_dict['==out=='] = output_cache + self.value_dict["==out=="] = output_cache self._validate_output_routing() @@ -256,7 +295,8 @@ def _validate_output_routing(self): for node_name, node in self.nodes.items(): pipeline_map = self.organized_pipeline_connections.get(node_name, {}) missing_outputs = [ - output_name for output_name in node.output_names + output_name + for output_name in node.output_names if output_name not in pipeline_map ] if missing_outputs: @@ -265,11 +305,12 @@ def _validate_output_routing(self): "Every model output must appear in pipeline data_flow, " "feedback_flow, or pipeline outputs. This may indicate the " "YAML was edited manually and is inconsistent with the exported model.", - error_type=ValueError) + error_type=ValueError, + ) def _validate_connection_compatibility(self): all_flows = {} - for flow_key in ('data_flow', 'feedback_flow'): + for flow_key in ("data_flow", "feedback_flow"): for source, targets in self.pipeline.get(flow_key, {}).items(): if source in all_flows: all_flows[source] = all_flows[source] + targets @@ -277,51 +318,59 @@ def _validate_connection_compatibility(self): all_flows[source] = list(targets) for source, targets in all_flows.items(): - source_node_name, source_output_name = source.split('/') + source_node_name, source_output_name = source.split("/") # Validate source node exists if source_node_name not in self.nodes: _get_logger().fatal( f"Source node '{source_node_name}' not found in models", - error_type=ValueError) + error_type=ValueError, + ) source_node = self.nodes[source_node_name] # Validate source output exists - source_output_names = [desc['name'] - for desc in source_node.output_descriptions] + source_output_names = [ + desc["name"] for desc in source_node.output_descriptions + ] if source_output_name not in source_output_names: _get_logger().fatal( f"Source output '{source_output_name}' not found in node '{source_node_name}'. " f"Available outputs: {source_output_names}", - error_type=ValueError) + error_type=ValueError, + ) # Get source output description source_desc = next( - desc for desc in source_node.output_descriptions - if desc['name'] == source_output_name + desc + for desc in source_node.output_descriptions + if desc["name"] == source_output_name + ) + source_shape = ( + json.loads(source_desc["shape"]) + if isinstance(source_desc["shape"], str) + else source_desc["shape"] ) - source_shape = json.loads(source_desc['shape']) if isinstance( - source_desc['shape'], str) else source_desc['shape'] - source_dtype = map_to_torch_dtype(source_desc['dtype']) + source_dtype = map_to_torch_dtype(source_desc["dtype"]) for target in targets: - target_node_name, target_input_name = target.split('/') + target_node_name, target_input_name = target.split("/") # Validate target node exists if target_node_name not in self.nodes: _get_logger().fatal( f"Target node '{target_node_name}' not found in models", - error_type=ValueError) + error_type=ValueError, + ) # Validate target input exists if target_input_name not in self.value_dict[target_node_name]: - available_inputs = list( - self.value_dict[target_node_name].keys()) + available_inputs = list(self.value_dict[target_node_name].keys()) _get_logger().fatal( f"Target input '{target_input_name}' not found in node '{target_node_name}'. " f"Available inputs: {available_inputs}", - error_type=ValueError) + error_type=ValueError, + ) target_tensor = self.value_dict[target_node_name][target_input_name] @@ -336,12 +385,12 @@ def _validate_connection_compatibility(self): def _prepopulate_feedback_inputs(self): """Load feedback initial values from safetensors and populate input buffers. - + If the pipeline has an 'initial_values' field pointing to a safetensors file, load it and overwrite the corresponding feedback input buffers so deployers don't need to know what to initialize them as. """ - initial_values_file = self.pipeline.get('initial_values') + initial_values_file = self.pipeline.get("initial_values") if not initial_values_file: return @@ -351,39 +400,46 @@ def _prepopulate_feedback_inputs(self): if not os.path.exists(safetensors_path): _get_logger().fatal( f"Feedback initial values file not found at {safetensors_path}", - error_type=FileNotFoundError) + error_type=FileNotFoundError, + ) initial_values = load_file(safetensors_path) for key, tensor in initial_values.items(): - node_name, input_name = key.split('/') - if node_name in self.value_dict and input_name in self.value_dict[node_name]: + node_name, input_name = key.split("/") + if ( + node_name in self.value_dict + and input_name in self.value_dict[node_name] + ): device = self.value_dict[node_name][input_name].device self.value_dict[node_name][input_name] = tensor.to(device) else: _get_logger().fatal( f"Feedback initial value key '{key}' does not match any node input. " f"Available nodes: {list(self.value_dict.keys())}", - error_type=ValueError) + error_type=ValueError, + ) def run_policy(self, inputs: dict[str, torch.Tensor]): # Update input tensors with provided values (keys are "node_name/input_name") for key, input_value in inputs.items(): try: - node_name, input_name = key.split('/') + node_name, input_name = key.split("/") except ValueError as e: _get_logger().fatal( f"Invalid input key: {key}\n Expected format: node_name/input_name", error_type=ValueError, - cause=e) + cause=e, + ) self.value_dict[node_name][input_name] = input_value for node_name, node in self.nodes.items(): # Run inference via node's __call__ # Use node.input_names to ensure inputs are passed in the correct order - node_inputs = [self.value_dict[node_name][name] - for name in node.input_names] - + node_inputs = [ + self.value_dict[node_name][name] for name in node.input_names + ] + outputs = node(*node_inputs) output_order = node.output_names # Ensure outputs is always a tuple/list for consistent iteration @@ -396,19 +452,22 @@ def run_policy(self, inputs: dict[str, torch.Tensor]): for i, (target_node_name, target_input_name) in enumerate(targets): if i == 0: # Zero-copy: direct assignment for first consumer - self.value_dict[target_node_name][target_input_name] = output_value + self.value_dict[target_node_name][target_input_name] = ( + output_value + ) else: # Clone for 2nd+ consumers to prevent data corruption - self.value_dict[target_node_name][target_input_name] = output_value.clone( + self.value_dict[target_node_name][target_input_name] = ( + output_value.clone() ) - return dict(self.value_dict['==out==']) + return dict(self.value_dict["==out=="]) @property def inputs(self) -> list: """Returns list of expected input keys in 'node_name/input_name' format.""" input_keys = [] - for node_name, input_names in self.pipeline['inputs'].items(): + for node_name, input_names in self.pipeline["inputs"].items(): for input_name in input_names: input_keys.append(f"{node_name}/{input_name}") return input_keys @@ -417,7 +476,7 @@ def inputs(self) -> list: def outputs(self) -> list: """Returns list of output keys in 'node_name/output_name' format.""" output_keys = [] - for node_name, output_names in self.pipeline['outputs'].items(): + for node_name, output_names in self.pipeline["outputs"].items(): for output_name in output_names: output_keys.append(f"{node_name}/{output_name}") return output_keys @@ -425,7 +484,7 @@ def outputs(self) -> list: @property def feedback_inputs(self) -> list: keys = [] - for targets in self.pipeline.get('feedback_flow', {}).values(): + for targets in self.pipeline.get("feedback_flow", {}).values(): keys.extend(targets) return keys @@ -440,33 +499,35 @@ def get_mock_input(self): """ mock_inputs = {} - for node_name, input_names in self.pipeline['inputs'].items(): + for node_name, input_names in self.pipeline["inputs"].items(): node = self.nodes[node_name] # Create a lookup for input descriptions by name - input_descs = {desc['name'] - : desc for desc in node.input_descriptions} + input_descs = {desc["name"]: desc for desc in node.input_descriptions} for input_name in input_names: desc = input_descs[input_name] - shape = json.loads(desc['shape']) if isinstance( - desc['shape'], str) else desc['shape'] - dtype = map_to_torch_dtype(desc['dtype']) + shape = ( + json.loads(desc["shape"]) + if isinstance(desc["shape"], str) + else desc["shape"] + ) + dtype = map_to_torch_dtype(desc["dtype"]) device = node.device # Generate random tensor (randn for floats, randint for integers) if dtype.is_floating_point: - tensor = torch.randn( - tuple(shape), dtype=dtype, device=device) + tensor = torch.randn(tuple(shape), dtype=dtype, device=device) else: - tensor = torch.randint(0, 256, tuple( - shape), dtype=dtype, device=device) + tensor = torch.randint( + 0, 256, tuple(shape), dtype=dtype, device=device + ) mock_inputs[f"{node_name}/{input_name}"] = tensor return mock_inputs - + def reset(self): """resets the inference manager to its initial state - + This method does not reset the internal state of the nodes, only the input and output buffers. It uses default initial values for feedback inputs. """ @@ -477,6 +538,5 @@ def reset(self): self._prepopulate_feedback_inputs() - def __call__(self, inputs: dict[str, torch.Tensor]): return self.run_policy(inputs) diff --git a/leapp/leapp.py b/leapp/leapp.py index 88769edf..29df2309 100644 --- a/leapp/leapp.py +++ b/leapp/leapp.py @@ -34,8 +34,16 @@ _MANAGER = ExportManager() -def start(name, save_path=".", verbose=False, dry_run=False, non_traced=None, - max_cached_io=5, global_patching=True, patching=None): +def start( + name, + save_path=".", + verbose=False, + dry_run=False, + non_traced=None, + max_cached_io=5, + global_patching=True, + patching=None, +): """Initialize and start LEAPP graph interpretation. ``name`` may be a bare graph name (``"my_graph"``), a relative path @@ -49,19 +57,22 @@ def start(name, save_path=".", verbose=False, dry_run=False, non_traced=None, if not isinstance(name, str): _get_logger().fatal( f"leapp.start(name=...) must be a string, got {type(name).__name__}", - error_type=TypeError) + error_type=TypeError, + ) normalized_name = os.path.normpath(os.path.expanduser(name)) if normalized_name in ("", ".", os.sep): _get_logger().fatal( f"leapp.start(name={name!r}) does not contain a usable graph name; " "provide a non-empty basename such as 'my_graph' or '/tmp/my_graph'.", - error_type=ValueError) + error_type=ValueError, + ) name_dir, graph_name = os.path.split(normalized_name) if not graph_name: _get_logger().fatal( f"leapp.start(name={name!r}) resolved to an empty graph name basename; " "provide a name whose basename is non-empty (e.g. 'my_graph').", - error_type=ValueError) + error_type=ValueError, + ) manager = _MANAGER manager.set_graph_name(graph_name) @@ -70,8 +81,10 @@ def start(name, save_path=".", verbose=False, dry_run=False, non_traced=None, manager.configure_logger(verbose=verbose) if ExportManager.is_interpret_graph_enabled(): - _get_logger().warning("LEAPP graph interpretation is already enabled, " - "calling start() again will reset the graph") + _get_logger().warning( + "LEAPP graph interpretation is already enabled, " + "calling start() again will reset the graph" + ) _get_logger().warning("Resetting graph...") manager.reset_tracing_lock() @@ -101,11 +114,13 @@ def stop(): if TracingLock().is_active: _get_logger().fatal( "leapp.stop() was called while a traced function is still executing", - error_type=Exception) + error_type=Exception, + ) if not ExportManager.is_interpret_graph_enabled(): _get_logger().fatal( "leapp.stop() called with no active tracing session — did you call leapp.start()?", - error_type=Exception) + error_type=Exception, + ) ExportManager.set_interpret_graph(False) manager.restore_pending_buffer_trackers() @@ -115,8 +130,15 @@ def stop(): manager.patcher.uninstall() -def compile_graph(visualize=True, verbose=None, validate=True, - rtol=1e-3, atol=1e-5, strict=True, graph_configs=None): +def compile_graph( + visualize=True, + verbose=None, + validate=True, + rtol=1e-3, + atol=1e-5, + strict=True, + graph_configs=None, +): """Compile and save the computational graph from traced nodes. Dry run is declared once via ``leapp.start(dry_run=True)``; when it is @@ -132,7 +154,8 @@ def compile_graph(visualize=True, verbose=None, validate=True, if ExportManager.is_interpret_graph_enabled(): _get_logger().fatal( "LEAPP graph interpretation is enabled. Call leapp.stop() before leapp.compile_graph().", - error_type=Exception) + error_type=Exception, + ) if verbose is not None: _get_logger().set_verbose(verbose) @@ -149,16 +172,18 @@ def compile_graph(visualize=True, verbose=None, validate=True, if not isinstance(graph_configs, GraphConfigs): _get_logger().fatal( "compile_graph(graph_configs=...) expects a GraphConfigs instance", - error_type=TypeError) + error_type=TypeError, + ) pipeline["pipeline"]["configs"] = graph_configs.to_dict() initial_value_filename = None if not manager.is_dry_run(): initial_value_filename = graph.save_feedback_initial_values( - manager.get_save_path(), manager.get_graph_name()) + manager.get_save_path(), manager.get_graph_name() + ) if initial_value_filename is not None: - pipeline['pipeline']['initial_values'] = initial_value_filename + pipeline["pipeline"]["initial_values"] = initial_value_filename if not manager.is_dry_run(): manager.save_models() @@ -178,7 +203,9 @@ def compile_graph(visualize=True, verbose=None, validate=True, _get_logger().info(f"- Total edges: {total_edges}") system_info = get_system_info() - with open(os.path.join(manager.get_save_path(), f"{manager.get_graph_name()}.yaml"), "w") as f: + with open( + os.path.join(manager.get_save_path(), f"{manager.get_graph_name()}.yaml"), "w" + ) as f: yaml.dump(models, f, sort_keys=False) f.write("\n") yaml.dump(pipeline, f) @@ -220,9 +247,7 @@ class AnnotateAPI: def __getattr__(self, name): if name in self._ALLOWED_APIS: return getattr(_MANAGER, name) - raise AttributeError( - f"leapp.annotate has no attribute '{name}'. " - ) + raise AttributeError(f"leapp.annotate has no attribute '{name}'. ") annotate = AnnotateAPI() diff --git a/leapp/leapp_graph/custom_operator_registry/warp_operator/__init__.py b/leapp/leapp_graph/custom_operator_registry/warp_operator/__init__.py index 7e6f0e59..201a10e5 100644 --- a/leapp/leapp_graph/custom_operator_registry/warp_operator/__init__.py +++ b/leapp/leapp_graph/custom_operator_registry/warp_operator/__init__.py @@ -80,9 +80,13 @@ def _onnx_graph_contains_warp_runner(graph) -> bool: if node.domain == ONNX_WRP_DOMAIN and node.op_type == ONNX_WRP_OP_TYPE: return True for attribute in node.attribute: - if attribute.HasField("g") and _onnx_graph_contains_warp_runner(attribute.g): + if attribute.HasField("g") and _onnx_graph_contains_warp_runner( + attribute.g + ): return True - if any(_onnx_graph_contains_warp_runner(graph) for graph in attribute.graphs): + if any( + _onnx_graph_contains_warp_runner(graph) for graph in attribute.graphs + ): return True return False @@ -116,9 +120,7 @@ def _register() -> None: and hasattr(getattr(torch.ops, NAMESPACE), OP_NAME) ): _LIB.define(_SCHEMA) - torch.library.register_fake( - QUALIFIED_NAME, fx_impl.warp_runner_fake, lib=_LIB - ) + torch.library.register_fake(QUALIFIED_NAME, fx_impl.warp_runner_fake, lib=_LIB) _LIB.impl(OP_NAME, fx_impl.warp_runner_eager, "CompositeExplicitAutograd") _get_logger().debug( diff --git a/leapp/leapp_graph/custom_operator_registry/warp_operator/bundle.py b/leapp/leapp_graph/custom_operator_registry/warp_operator/bundle.py index 63397ad8..7ea93661 100644 --- a/leapp/leapp_graph/custom_operator_registry/warp_operator/bundle.py +++ b/leapp/leapp_graph/custom_operator_registry/warp_operator/bundle.py @@ -129,4 +129,3 @@ def iter_warp_segments_from_graph(graph: fx.Graph) -> list: if segment is not None: segments.append(segment) return segments - diff --git a/leapp/leapp_graph/datatypes/__init__.py b/leapp/leapp_graph/datatypes/__init__.py index 1c2da318..5f54cf54 100644 --- a/leapp/leapp_graph/datatypes/__init__.py +++ b/leapp/leapp_graph/datatypes/__init__.py @@ -57,15 +57,16 @@ # Type Checking Functions # ============================================================================= + def is_tracable_tensor_type(obj) -> bool: """Check if an object is a type that can be traced. - + This includes both base types (torch.Tensor, np.ndarray) and their traced counterparts (TracedTensor, TracedNpArray). - + Args: obj: Object to check - + Returns: True if the object is a traceable tensor type """ @@ -74,25 +75,27 @@ def is_tracable_tensor_type(obj) -> bool: def is_traced_type(obj) -> bool: """Check if an object is a TracedData subclass instance. - + Args: obj: Object to check - + Returns: True if the object is a TracedData instance """ return isinstance(obj, TracedData) -def get_traced_class_for(obj_or_type: Union[type, object]) -> Optional[Type[TracedData]]: +def get_traced_class_for( + obj_or_type: Union[type, object], +) -> Optional[Type[TracedData]]: """Get the appropriate TracedData subclass for a given type or object. - + Args: obj_or_type: Either a type (e.g., torch.Tensor) or an instance - + Returns: The corresponding TracedData subclass, or None if not traceable - + Examples: >>> get_traced_class_for(torch.Tensor) @@ -104,7 +107,7 @@ def get_traced_class_for(obj_or_type: Union[type, object]) -> Optional[Type[Trac # If already a traced type, return its class if isinstance(obj_or_type, TracedData): return type(obj_or_type) - + # If it's a type, look it up directly if isinstance(obj_or_type, type): # Check for exact match first @@ -115,12 +118,12 @@ def get_traced_class_for(obj_or_type: Union[type, object]) -> Optional[Type[Trac if issubclass(obj_or_type, base_type): return traced_type return None - + # It's an instance, find matching type for base_type, traced_type in TRACED_TYPE_REGISTRY.items(): if isinstance(obj_or_type, base_type): return traced_type - + return None @@ -170,9 +173,7 @@ def as_traced( if wp is not None and TracedWpArray is not None and isinstance(data, wp.array): if was_traced: return TracedWpArray(data, name, context, proxy, view=view) - return TracedWpArray.make_traced_in_place( - data, name, context, proxy, view=view - ) + return TracedWpArray.make_traced_in_place(data, name, context, proxy, view=view) # Find the appropriate traced class traced_class = get_traced_class_for(data) @@ -238,6 +239,7 @@ def to_export_torch_tensor(data) -> _torch.Tensor: error_type=TypeError, ) + # ============================================================================= # Exports # ============================================================================= diff --git a/leapp/leapp_graph/datatypes/numpy/traced_np_array.py b/leapp/leapp_graph/datatypes/numpy/traced_np_array.py index bb1e08e2..d09d38ea 100644 --- a/leapp/leapp_graph/datatypes/numpy/traced_np_array.py +++ b/leapp/leapp_graph/datatypes/numpy/traced_np_array.py @@ -43,12 +43,14 @@ for dtype in (scalar, np.dtype(scalar)) ) -register_dtype_codec(DtypeCodec( - backend="numpy", - matches=lambda v: isinstance(v, np.ndarray), - value_dtype=lambda v: v.dtype, - dtype_to_name=_NUMPY_DTYPE_TO_NAME, -)) +register_dtype_codec( + DtypeCodec( + backend="numpy", + matches=lambda v: isinstance(v, np.ndarray), + value_dtype=lambda v: v.dtype, + dtype_to_name=_NUMPY_DTYPE_TO_NAME, + ) +) def _torch_dtype_for(dtype): @@ -79,25 +81,21 @@ def _torch_dtype_for(dtype): np.mod: torch.remainder, np.remainder: torch.remainder, np.fmod: torch.fmod, - # Absolute and sign np.absolute: torch.abs, np.abs: torch.abs, np.sign: torch.sign, - # Powers and roots np.sqrt: torch.sqrt, np.square: torch.square, np.exp: torch.exp, np.exp2: lambda x: torch.pow(2, x), np.expm1: torch.expm1, - # Logarithms np.log: torch.log, np.log2: torch.log2, np.log10: torch.log10, np.log1p: torch.log1p, - # Trigonometric functions np.sin: torch.sin, np.cos: torch.cos, @@ -107,7 +105,6 @@ def _torch_dtype_for(dtype): np.arctan: torch.atan, np.arctan2: torch.atan2, np.hypot: torch.hypot, - # Hyperbolic functions np.sinh: torch.sinh, np.cosh: torch.cosh, @@ -115,13 +112,11 @@ def _torch_dtype_for(dtype): np.arcsinh: torch.asinh, np.arccosh: torch.acosh, np.arctanh: torch.atanh, - # Rounding np.floor: torch.floor, np.ceil: torch.ceil, np.trunc: torch.trunc, np.rint: torch.round, - # Comparison (element-wise, return boolean tensor) np.greater: torch.gt, np.greater_equal: torch.ge, @@ -131,17 +126,14 @@ def _torch_dtype_for(dtype): np.not_equal: torch.ne, np.maximum: torch.maximum, np.minimum: torch.minimum, - # Matrix multiplication. ``np.matmul`` is a ufunc, so it never reaches # ``__array_function__``; ``np.dot`` is not, and is mapped there instead. np.matmul: torch.matmul, - # Logical operations np.logical_and: torch.logical_and, np.logical_or: torch.logical_or, np.logical_xor: torch.logical_xor, np.logical_not: torch.logical_not, - # Bitwise operations np.bitwise_and: torch.bitwise_and, np.bitwise_or: torch.bitwise_or, @@ -149,12 +141,10 @@ def _torch_dtype_for(dtype): np.invert: torch.bitwise_not, np.left_shift: torch.bitwise_left_shift, np.right_shift: torch.bitwise_right_shift, - # Special values np.isnan: torch.isnan, np.isinf: torch.isinf, np.isfinite: torch.isfinite, - # Clipping np.clip: torch.clamp, } @@ -177,7 +167,6 @@ def _torch_dtype_for(dtype): np.cumprod: torch.cumprod, np.all: torch.all, np.any: torch.any, - # Array manipulation np.concatenate: torch.cat, np.stack: torch.stack, @@ -196,13 +185,11 @@ def _torch_dtype_for(dtype): np.flip: torch.flip, np.roll: torch.roll, np.rot90: torch.rot90, - # Sorting and searching np.sort: torch.sort, np.argsort: torch.argsort, np.where: torch.where, np.nonzero: torch.nonzero, - # Element-wise (also available as ufuncs) np.clip: torch.clamp, np.abs: torch.abs, @@ -217,7 +204,6 @@ def _torch_dtype_for(dtype): np.tanh: torch.tanh, np.round: torch.round, np.around: torch.round, - # Linear algebra np.dot: torch.matmul, # Note: torch.dot is only for 1D vectors np.tensordot: torch.tensordot, @@ -226,13 +212,11 @@ def _torch_dtype_for(dtype): np.diagonal: torch.diagonal, np.tril: torch.tril, np.triu: torch.triu, - # Creation functions (when operating on TracedTensor) np.zeros_like: torch.zeros_like, np.ones_like: torch.ones_like, np.full_like: torch.full_like, np.empty_like: torch.empty_like, - # Standalone creation (less commonly needed with TracedTensor) np.eye: torch.eye, np.zeros: torch.zeros, @@ -285,6 +269,7 @@ def _torch_dtype_for(dtype): # and np.ndarray's metaclass class _TracedNpArrayMeta(ABCMeta, type(np.ndarray)): """Metaclass combining ABCMeta and np.ndarray's metaclass.""" + pass @@ -344,9 +329,9 @@ def __array_finalize__(self, obj): # ndarray source and a source still mid-construction. bind_new_view( self, - getattr(obj, '_name', 'derived'), - getattr(obj, '_context', None), - getattr(obj, 'proxy', None), + getattr(obj, "_name", "derived"), + getattr(obj, "_context", None), + getattr(obj, "proxy", None), ) # ========================================================================= @@ -423,6 +408,7 @@ def find_all_contexts(obj, contexts=None): def _extract_proxy(self, obj): """Recursively extract proxies and convert constant NumPy arrays.""" + def convert(item): if isinstance(item, TracedData): return item.proxy @@ -435,8 +421,8 @@ def convert(item): return TracedData._map_structure(obj, convert) def _convert_numpy_call_to_torch( - self, proxy_args, proxy_kwargs, torch_func, - original_kwargs=None, args=None): + self, proxy_args, proxy_kwargs, torch_func, original_kwargs=None, args=None + ): """Convert numpy args/kwargs to the matching torch call. Handles: @@ -458,26 +444,30 @@ def _convert_numpy_call_to_torch( # Convert axis → dim has_dim = False - if torch_func in AXIS_TO_DIM_FUNCTIONS and 'axis' in proxy_kwargs: - proxy_kwargs['dim'] = proxy_kwargs.pop('axis') + if torch_func in AXIS_TO_DIM_FUNCTIONS and "axis" in proxy_kwargs: + proxy_kwargs["dim"] = proxy_kwargs.pop("axis") has_dim = True if torch_func in AXIS_TO_DIMS_FUNCTIONS: - if 'axis' in proxy_kwargs: - dims = proxy_kwargs.pop('axis') - if dims is None and torch_func is torch.flip and hasattr(input_arr, 'ndim'): + if "axis" in proxy_kwargs: + dims = proxy_kwargs.pop("axis") + if ( + dims is None + and torch_func is torch.flip + and hasattr(input_arr, "ndim") + ): dims = tuple(range(input_arr.ndim)) elif torch_func is torch.flip and isinstance(dims, int): dims = (dims,) if dims is not None: - proxy_kwargs['dims'] = dims + proxy_kwargs["dims"] = dims elif torch_func is torch.flip: if args is not None and len(args) > 1: dims = proxy_args[1] if isinstance(dims, int): proxy_args = (proxy_args[0], (dims,)) + proxy_args[2:] - elif hasattr(input_arr, 'ndim'): - proxy_kwargs['dims'] = tuple(range(input_arr.ndim)) + elif hasattr(input_arr, "ndim"): + proxy_kwargs["dims"] = tuple(range(input_arr.ndim)) if torch_func in SCALAR_TENSOR_FUNCTIONS: proxy_args = tuple( @@ -487,27 +477,29 @@ def _convert_numpy_call_to_torch( # Convert keepdims → keepdim (numpy uses 's', torch doesn't) # Important: torch.sum/mean only support keepdim when dim is also specified - if 'keepdims' in proxy_kwargs: - keepdims_value = proxy_kwargs.pop('keepdims') + if "keepdims" in proxy_kwargs: + keepdims_value = proxy_kwargs.pop("keepdims") if has_dim: # dim is specified, we can use keepdim - proxy_kwargs['keepdim'] = keepdims_value - elif keepdims_value and input_arr is not None and hasattr(input_arr, 'ndim'): + proxy_kwargs["keepdim"] = keepdims_value + elif ( + keepdims_value and input_arr is not None and hasattr(input_arr, "ndim") + ): # keepdims=True but no axis specified # torch requires dim when using keepdim, so specify all dims - proxy_kwargs['dim'] = tuple(range(input_arr.ndim)) - proxy_kwargs['keepdim'] = True + proxy_kwargs["dim"] = tuple(range(input_arr.ndim)) + proxy_kwargs["keepdim"] = True # If keepdims=False without axis, we can just omit both (default behavior) # np.var/std default to ddof=0 (population), torch defaults to correction=1 (sample) if torch_func in (torch.var, torch.std): - if 'ddof' in proxy_kwargs: - proxy_kwargs['correction'] = proxy_kwargs.pop('ddof') - elif 'correction' not in proxy_kwargs: - proxy_kwargs['correction'] = 0 + if "ddof" in proxy_kwargs: + proxy_kwargs["correction"] = proxy_kwargs.pop("ddof") + elif "correction" not in proxy_kwargs: + proxy_kwargs["correction"] = 0 - if torch_func is torch.tensordot and 'axes' in proxy_kwargs: - proxy_kwargs['dims'] = proxy_kwargs.pop('axes') + if torch_func is torch.tensordot and "axes" in proxy_kwargs: + proxy_kwargs["dims"] = proxy_kwargs.pop("axes") # Handle np.transpose with axes=None or missing # torch.permute requires explicit dims, but numpy reverses all dims when axes is omitted @@ -515,17 +507,21 @@ def _convert_numpy_call_to_torch( # Check if axes was passed as positional arg (2nd argument) axes_in_args = len(args) > 1 if args else False - if 'axes' in proxy_kwargs: - axes = proxy_kwargs.pop('axes') - if axes is None and input_arr is not None and hasattr(input_arr, 'ndim'): + if "axes" in proxy_kwargs: + axes = proxy_kwargs.pop("axes") + if ( + axes is None + and input_arr is not None + and hasattr(input_arr, "ndim") + ): # axes=None means reverse all dimensions - proxy_kwargs['dims'] = tuple(range(input_arr.ndim - 1, -1, -1)) + proxy_kwargs["dims"] = tuple(range(input_arr.ndim - 1, -1, -1)) elif axes is not None: - proxy_kwargs['dims'] = axes - elif not axes_in_args and 'axes' not in original_kwargs: + proxy_kwargs["dims"] = axes + elif not axes_in_args and "axes" not in original_kwargs: # axes not provided at all (neither positional nor kwarg) - numpy reverses all dims - if input_arr is not None and hasattr(input_arr, 'ndim'): - proxy_kwargs['dims'] = tuple(range(input_arr.ndim - 1, -1, -1)) + if input_arr is not None and hasattr(input_arr, "ndim"): + proxy_kwargs["dims"] = tuple(range(input_arr.ndim - 1, -1, -1)) # If axes_in_args is True, it will be passed as positional arg, don't add to kwargs return proxy_args, proxy_kwargs @@ -537,12 +533,9 @@ def _convert_numpy_call_to_torch( def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): """Execute NumPy ufuncs and record them only while the source is active.""" traced_array = TracedNpArray.find_traced_array([inputs, kwargs]) - unwrapped_inputs = tuple( - TracedData.unwrap_traced_data(inp) for inp in inputs - ) + unwrapped_inputs = tuple(TracedData.unwrap_traced_data(inp) for inp in inputs) unwrapped_kwargs = { - key: TracedData.unwrap_traced_data(value) - for key, value in kwargs.items() + key: TracedData.unwrap_traced_data(value) for key, value in kwargs.items() } if method != "__call__": @@ -605,12 +598,12 @@ def _patch_sort(self, torch_func, traced_array, proxy_args, args, kwargs): return None input_array = args[0] - dim = kwargs.get('axis', -1) + dim = kwargs.get("axis", -1) if dim is None: dim = -1 k = input_array.shape[dim] - topk_kwargs = {'dim': dim, 'largest': False, 'sorted': True} + topk_kwargs = {"dim": dim, "largest": False, "sorted": True} proxy_out = traced_array._context.tracer.create_proxy( "call_function", torch.topk, (proxy_args[0], k), topk_kwargs ) @@ -623,12 +616,9 @@ def _patch_sort(self, torch_func, traced_array, proxy_args, args, kwargs): def __array_function__(self, func, types, args, kwargs): """Execute NumPy functions and record them only while the source is active.""" traced_array = TracedNpArray.find_traced_array([args, kwargs]) - unwrapped_args = tuple( - TracedData.unwrap_traced_data(arg) for arg in args - ) + unwrapped_args = tuple(TracedData.unwrap_traced_data(arg) for arg in args) unwrapped_kwargs = { - key: TracedData.unwrap_traced_data(value) - for key, value in kwargs.items() + key: TracedData.unwrap_traced_data(value) for key, value in kwargs.items() } result_array = func(*unwrapped_args, **unwrapped_kwargs) @@ -872,8 +862,7 @@ def __getitem__(self, key): if proxy_out is None: proxy_key = self._extract_proxy(key) proxy_out = self._context.tracer.create_proxy( - "call_function", operator.getitem, - (self.proxy, proxy_key), {} + "call_function", operator.getitem, (self.proxy, proxy_key), {} ) if isinstance(result, np.ndarray): @@ -883,9 +872,7 @@ def __getitem__(self, key): def _record_assignment(self, key, value, real_value): """Record one functional assignment and update this object's proxy.""" value_proxy = value.proxy if isinstance(value, TracedData) else value - if self._update_setitem_proxy( - key, value_proxy, real_value=real_value - ): + if self._update_setitem_proxy(key, value_proxy, real_value=real_value): return True if isinstance(key, tuple): @@ -1029,10 +1016,10 @@ def copy(self): def __array__(self, dtype=None, copy=None): """Convert to numpy array (NumPy array protocol). - + Returns self (or a converted TracedNpArray) for both active and inactive traced values so representation conversion cannot silently demote it. - + Args: dtype: Optional dtype for the resulting array copy: Optional copy flag (NumPy 2.0+) @@ -1042,7 +1029,7 @@ def __array__(self, dtype=None, copy=None): """ # Check if dtype change requires a copy needs_dtype_copy = dtype is not None and dtype != self.dtype - + # copy=False but copy is required → raise error (NumPy 2.0 semantics) if copy is False and needs_dtype_copy: _get_logger().fatal( @@ -1050,7 +1037,7 @@ def __array__(self, dtype=None, copy=None): f"from array with dtype {self.dtype}.", error_type=ValueError, ) - + # Preserve the carrier even after its source context stops tracing. # astype() and copy() already retain the complete tracing state. @@ -1065,7 +1052,7 @@ def __array__(self, dtype=None, copy=None): def tolist(self): """Convert to Python list. - + Warning: During tracing, this extracts concrete values and breaks the trace. If used in a conditional, the graph will only capture one branch. """ @@ -1079,7 +1066,7 @@ def tolist(self): def item(self): """Return scalar value. - + Warning: During tracing, this extracts a concrete scalar and breaks the trace. If used in a conditional, the graph will only capture one branch. """ @@ -1096,7 +1083,9 @@ def item(self): # ========================================================================= def __repr__(self): - return f"TracedNpArray(name={self._name}, shape={self.shape}, dtype={self.dtype})" + return ( + f"TracedNpArray(name={self._name}, shape={self.shape}, dtype={self.dtype})" + ) def __str__(self): return f"TracedNpArray({self.view(np.ndarray)})" diff --git a/leapp/leapp_graph/datatypes/patching.py b/leapp/leapp_graph/datatypes/patching.py index d8ae5d63..45853091 100644 --- a/leapp/leapp_graph/datatypes/patching.py +++ b/leapp/leapp_graph/datatypes/patching.py @@ -68,10 +68,11 @@ def _validate_user_patches( raise TypeError("patching entries must be FunctionPatch instances") if not isinstance(definition.module, ModuleType): raise TypeError("FunctionPatch.module must be a Python module") - if not isinstance(definition.function_name, str) or not definition.function_name: - raise ValueError( - "FunctionPatch.function_name must be a non-empty string" - ) + if ( + not isinstance(definition.function_name, str) + or not definition.function_name + ): + raise ValueError("FunctionPatch.function_name must be a non-empty string") try: target = getattr(definition.module, definition.function_name) except AttributeError: @@ -110,9 +111,7 @@ def _try_create_warp_backend() -> WarpPatchBackend | None: cached_paths = warp_runtime_artifact_paths() runtime_library_available = any( - Path(os.environ.get(env_name) or cached_paths[backend]) - .expanduser() - .is_file() + Path(os.environ.get(env_name) or cached_paths[backend]).expanduser().is_file() for backend, env_name in WARP_RUNTIME_ENVIRONMENTS.items() ) if ( diff --git a/leapp/leapp_graph/datatypes/proxy_view.py b/leapp/leapp_graph/datatypes/proxy_view.py index 701c8dd8..0942a073 100644 --- a/leapp/leapp_graph/datatypes/proxy_view.py +++ b/leapp/leapp_graph/datatypes/proxy_view.py @@ -91,9 +91,7 @@ def bind_shared_view(carrier: Any, name: str, context, view: ProxyView) -> None: carrier._output_port = None -def update_view_proxy( - carrier: Any, name: str, context, proxy: Optional[Proxy] -) -> None: +def update_view_proxy(carrier: Any, name: str, context, proxy: Optional[Proxy]) -> None: """Write a new proxy into ``carrier``'s existing view after an in-place change.""" carrier._name = name carrier._context = context @@ -134,9 +132,7 @@ def layout_key(value): try: pointer = native.data_ptr() element_size = native.element_size() - byte_strides = tuple( - stride * element_size for stride in native.stride() - ) + byte_strides = tuple(stride * element_size for stride in native.stride()) contiguous = native.is_contiguous() except RuntimeError: # Layouts without strides or an addressable buffer, e.g. sparse. diff --git a/leapp/leapp_graph/datatypes/torch/traced_tensor.py b/leapp/leapp_graph/datatypes/torch/traced_tensor.py index 94d1fe62..97f6d75d 100644 --- a/leapp/leapp_graph/datatypes/torch/traced_tensor.py +++ b/leapp/leapp_graph/datatypes/torch/traced_tensor.py @@ -46,13 +46,14 @@ torch.bfloat16: "bfloat16", } -register_dtype_codec(DtypeCodec( - backend="torch", - matches=lambda v: isinstance(v, torch.Tensor), - value_dtype=lambda v: v.dtype, - dtype_to_name=_TORCH_DTYPE_TO_NAME, -)) - +register_dtype_codec( + DtypeCodec( + backend="torch", + matches=lambda v: isinstance(v, torch.Tensor), + value_dtype=lambda v: v.dtype, + dtype_to_name=_TORCH_DTYPE_TO_NAME, + ) +) # Combined metaclass to resolve conflict between ABCMeta (from TracedData) @@ -60,6 +61,7 @@ # the TracedTensor class to not inherit from tracedData and reimplement all the common functions class _TracedTensorMeta(ABCMeta, type(torch.Tensor)): """Metaclass combining ABCMeta and torch.Tensor's metaclass.""" + pass @@ -76,21 +78,19 @@ class TracedTensor(TracedData, torch.Tensor, metaclass=_TracedTensorMeta): # Torch ops whose result holds the same values, shape, and dtype as their # source, so a finished boundary value survives them intact. ``to`` is # handled separately because device moves and dtype casts share one method. - _EQUIVALENT_COPY_NAMES = frozenset( - {"clone", "detach", "contiguous", "cpu", "cuda"} - ) + _EQUIVALENT_COPY_NAMES = frozenset({"clone", "detach", "contiguous", "cpu", "cuda"}) _NATIVE_TYPE = torch.Tensor @staticmethod def __new__(cls, tensor: torch.Tensor, name: str, context, proxy: Proxy): """Create a new TracedTensor instance. - + Args: tensor: The actual torch.Tensor data to wrap name: Name for the tensor (used in ONNX export and graph) context: The TraceContext that owns this tensor proxy: The fx.Proxy for graph recording - + Returns: A new TracedTensor instance sharing storage with the input tensor """ @@ -121,7 +121,7 @@ def __init__(self, tensor: torch.Tensor, name: str, context, proxy: Proxy): def tensor(self) -> torch.Tensor: """Get the underlying torch.Tensor (for compatibility with original API).""" return self.as_subclass(torch.Tensor) - + @property def data(self) -> torch.Tensor: """Get the underlying data.""" @@ -168,8 +168,7 @@ def _schema_arguments_match(inplace_arg, functional_arg) -> bool: inplace_arg.name != functional_arg.name or str(inplace_arg.type) != str(functional_arg.type) or inplace_arg.kwarg_only != functional_arg.kwarg_only - or inplace_arg.has_default_value() - != functional_arg.has_default_value() + or inplace_arg.has_default_value() != functional_arg.has_default_value() ): return False if not inplace_arg.has_default_value(): @@ -177,9 +176,7 @@ def _schema_arguments_match(inplace_arg, functional_arg) -> bool: return inplace_arg.default_value == functional_arg.default_value @classmethod - def _schemas_are_functional_twins( - cls, inplace_schema, functional_schema - ) -> bool: + def _schemas_are_functional_twins(cls, inplace_schema, functional_schema) -> bool: """Check that a functional schema describes one receiver-only write.""" if ( not inplace_schema.is_mutable @@ -203,12 +200,10 @@ def _schemas_are_functional_twins( writes = [ index for index, argument in enumerate(inplace_schema.arguments) - if argument.alias_info is not None - and argument.alias_info.is_write + if argument.alias_info is not None and argument.alias_info.is_write ] if writes != [0] or any( - argument.alias_info is not None - and argument.alias_info.is_write + argument.alias_info is not None and argument.alias_info.is_write for argument in functional_schema.arguments ): return False @@ -267,18 +262,14 @@ def _functional_inplace_spec(cls, func_name): metadata_flags = [] for overload_name, inplace_schema in tensor_inplace_schemas: try: - inplace_tags = cls._schema_tags( - inplace_packet, overload_name - ) + inplace_tags = cls._schema_tags(inplace_packet, overload_name) except AttributeError: return None if inplace_tags & nondeterministic_tags: return None matched = False - for functional_overload, functional_schema in ( - functional_schemas.items() - ): + for functional_overload, functional_schema in functional_schemas.items(): try: functional_tags = cls._schema_tags( functional_packet, functional_overload @@ -287,9 +278,7 @@ def _functional_inplace_spec(cls, func_name): continue if functional_tags & nondeterministic_tags: continue - if cls._schemas_are_functional_twins( - inplace_schema, functional_schema - ): + if cls._schemas_are_functional_twins(inplace_schema, functional_schema): matched = True break if not matched: @@ -312,9 +301,9 @@ def _convert_method_descriptor(func, context=None): """ # Check if it's a method descriptor on torch.Tensor func_type = type(func).__name__ - if func_type in ('method_descriptor', 'builtin_function_or_method'): + if func_type in ("method_descriptor", "builtin_function_or_method"): # Try to get the method name - func_name = getattr(func, '__name__', None) + func_name = getattr(func, "__name__", None) if func_name: # Look for equivalent in torch module torch_func = getattr(torch, func_name, None) @@ -384,9 +373,9 @@ def _handle_TS_decomposition(func, real_args, real_kwargs): # ScriptMethod. func.owner returns the C++ ScriptModule # (not an nn.Module). Save → reload via BytesIO so # torch.jit.load performs full Python-level init. - script_module = getattr(func, '__self__', None) + script_module = getattr(func, "__self__", None) if not isinstance(script_module, torch.nn.Module): - cpp_module = getattr(func, 'owner', None) + cpp_module = getattr(func, "owner", None) if cpp_module is not None: _buf = _io.BytesIO() _tmp = torch.jit.RecursiveScriptModule._construct( @@ -395,15 +384,12 @@ def _handle_TS_decomposition(func, real_args, real_kwargs): torch.jit.save(_tmp, _buf) _buf.seek(0) _dev = next( - (a.device for a in real_args - if isinstance(a, torch.Tensor)), - torch.device('cpu'), - ) - script_module = torch.jit.load( - _buf, map_location=_dev + (a.device for a in real_args if isinstance(a, torch.Tensor)), + torch.device("cpu"), ) + script_module = torch.jit.load(_buf, map_location=_dev) - #validate the script_module + # validate the script_module if not isinstance(script_module, torch.nn.Module): _get_logger().fatal( f"Could not obtain nn.Module from {type(func).__name__}. " @@ -424,9 +410,7 @@ def _handle_TS_decomposition(func, real_args, real_kwargs): k: v.clone().detach() if isinstance(v, torch.Tensor) else v for k, v in (real_kwargs or {}).items() } - ep = TS2EPConverter( - script_module, _clean, _clean_kw - ).convert() + ep = TS2EPConverter(script_module, _clean, _clean_kw).convert() fx_module = ep.module() except Exception as e: @@ -506,9 +490,7 @@ def _record_assignment(self, key, value, real_value): return True value_proxy = value.proxy if isinstance(value, TracedTensor) else value - return self._update_setitem_proxy( - key, value_proxy, real_value=real_value - ) + return self._update_setitem_proxy(key, value_proxy, real_value=real_value) @classmethod def _handle_plain_assignment(cls, func, args=(), kwargs=None): @@ -548,8 +530,7 @@ def _handle_plain_assignment(cls, func, args=(), kwargs=None): dest_proxy = anchor._register_setitem_tensor( target.clone().detach(), "_setitem_destination" ) - cls._promote_plain_tensor( - target, anchor.name, anchor.context_obj, dest_proxy) + cls._promote_plain_tensor(target, anchor.name, anchor.context_obj, dest_proxy) if full_replacement: # The destination now holds exactly the source's data, so it also # presents the source's boundary identity to the next node. @@ -636,15 +617,19 @@ def _handle_scripted_call( kwargs = {} func_type_name = type(func).__name__ - func_module = type(func).__module__ if hasattr(type(func), '__module__') else '' + func_module = type(func).__module__ if hasattr(type(func), "__module__") else "" - is_scripted = ('ScriptMethod' in func_type_name or 'ScriptModule' in func_type_name) + is_scripted = ( + "ScriptMethod" in func_type_name or "ScriptModule" in func_type_name + ) if not is_scripted: - if func_module.startswith('torch.jit') or func_module.startswith('torch._C'): - if hasattr(func, '__self__') and hasattr(func.__self__, '__class__'): + if func_module.startswith("torch.jit") or func_module.startswith( + "torch._C" + ): + if hasattr(func, "__self__") and hasattr(func.__self__, "__class__"): self_class_name = func.__self__.__class__.__name__ - is_scripted = 'Script' in self_class_name + is_scripted = "Script" in self_class_name if not is_scripted: return False, None @@ -704,8 +689,11 @@ def shed_replaced_graph(value): # A carrier is an alias, which Torch refuses to detach in place, and the # autograd state detach_ edits is never read while tracing. - if (getattr(func, "__name__", "") == "detach_" - and args and isinstance(args[0], TracedTensor)): + if ( + getattr(func, "__name__", "") == "detach_" + and args + and isinstance(args[0], TracedTensor) + ): return args[0] # A plain receiver must be promoted before method-descriptor conversion @@ -729,7 +717,9 @@ def shed_replaced_graph(value): # Extract real tensors for actual computation real_args = tuple(TracedTensor.unwrap_traced_tensor(arg) for arg in args) - real_kwargs = {k: TracedTensor.unwrap_traced_tensor(v) for k, v in kwargs.items()} + real_kwargs = { + k: TracedTensor.unwrap_traced_tensor(v) for k, v in kwargs.items() + } # Only args[0]: donor/mutation tracking is receiver-only; out=/foreach stay as today. receiver = args[0] if args and isinstance(args[0], TracedTensor) else None real_receiver = real_args[0] if receiver is not None else None @@ -869,18 +859,22 @@ def __getattr__(self, name: str): return self.__dict__[name] # Handle in-place operations (methods ending with _) - if name.endswith('_') and not name.startswith('_') and len(name) > 1: + if name.endswith("_") and not name.startswith("_") and len(name) > 1: base_name = name[:-1] # Remove trailing underscore torch_func = getattr(torch, base_name, None) if torch_func is not None and callable(torch_func): + def inplace_method(*args, **kwargs): if not self.is_tracing: # Just call the in-place method on the tensor and return raw tensor underlying = self.as_subclass(torch.Tensor) method = getattr(underlying, name, None) if method is not None: - method(*TracedData.unwrap_traced_data(args), **TracedData.unwrap_traced_data(kwargs)) + method( + *TracedData.unwrap_traced_data(args), + **TracedData.unwrap_traced_data(kwargs), + ) self._output_port = None return self @@ -915,6 +909,7 @@ def method(*args, **kwargs): if attr is not None: # If it's a callable method, wrap it to ensure it goes through __torch_function__ if callable(attr): + def wrapped_method(*args, **kwargs): # Call the method on the underlying tensor result = attr(*args, **kwargs) @@ -925,15 +920,16 @@ def wrapped_method(*args, **kwargs): # Create a proxy node for this operation self_proxy = self.proxy proxy = self_proxy._tracer.create_proxy( - 'call_method', name, (self_proxy,) + args, kwargs) + "call_method", name, (self_proxy,) + args, kwargs + ) return self._new(result, proxy) return result + return wrapped_method return attr # Attribute not found - raise AttributeError( - f"'TracedTensor' object has no attribute '{name}'") + raise AttributeError(f"'TracedTensor' object has no attribute '{name}'") # ========================================================================= # Special Methods: reshape and permute need custom handling @@ -1054,20 +1050,14 @@ def _apply_inplace( before = underlying.clone() if should_trace else None if inplace_name is None: - underlying.copy_( - functional(underlying, *real_operands, **real_kwargs) - ) + underlying.copy_(functional(underlying, *real_operands, **real_kwargs)) elif _mutates_metadata: # Shape/view mutations change the receiver's tensor metadata, not # its storage, so mutating an ``as_subclass`` alias is insufficient. with torch._C.DisableTorchFunction(): - getattr(torch.Tensor, inplace_name)( - self, *real_operands, **real_kwargs - ) + getattr(torch.Tensor, inplace_name)(self, *real_operands, **real_kwargs) else: - getattr(underlying, inplace_name)( - *real_operands, **real_kwargs - ) + getattr(underlying, inplace_name)(*real_operands, **real_kwargs) if not should_trace: self._output_port = None @@ -1174,8 +1164,10 @@ def copy_(self, src, non_blocking=False): ) value_proxy = src.proxy if isinstance(src, TracedTensor) else src self._proxy_view.proxy = self._context.tracer.create_proxy( - "call_method", "__setitem__", - (self.proxy, Ellipsis, value_proxy), {}, + "call_method", + "__setitem__", + (self.proxy, Ellipsis, value_proxy), + {}, ) return self @@ -1251,7 +1243,10 @@ def __getitem__(self, key): "call_method", "reshape", (key.proxy, (-1,)), {} ) proxy_out = self._context.tracer.create_proxy( - "call_function", torch.index_select, (self.proxy, 0, flat_index_proxy), {} + "call_function", + torch.index_select, + (self.proxy, 0, flat_index_proxy), + {}, ) if key.ndim > 1: @@ -1372,12 +1367,12 @@ def to(self, *args, **kwargs): def numpy(self): """Convert to numpy array, preserving tracing as TracedNpArray. - + If tracing is active, returns a TracedNpArray that shares the proxy and context, allowing the trace to continue across the conversion. - + If not tracing, returns a regular numpy array. - + Returns: TracedNpArray if tracing, else np.ndarray """ @@ -1391,6 +1386,7 @@ def numpy(self): # Conversion is representation-only, so both active and inactive # sources forward their complete tracing state. from leapp.leapp_graph.datatypes import as_traced + if may_adopt_view(self, np_data): view, proxy = self.proxy_view, None else: @@ -1401,24 +1397,24 @@ def numpy(self): def __array__(self, dtype=None, copy=None): """NumPy array protocol - called by np.array(), np.asarray(), etc. - + This allows seamless conversion to numpy while preserving tracing. - + Args: dtype: Optional dtype for the resulting array copy: Optional copy flag (NumPy 2.0+) - None: Copy only if necessary (default) - True: Always copy - False: Never copy; raise ValueError if copy required - + Returns: TracedNpArray if tracing, else np.ndarray """ result = self.numpy() - + # Check if dtype change requires a copy needs_dtype_copy = dtype is not None and dtype != result.dtype - + # copy=False but copy is required → raise error (NumPy 2.0 semantics) if copy is False and needs_dtype_copy: _get_logger().fatal( @@ -1426,11 +1422,11 @@ def __array__(self, dtype=None, copy=None): f"from array with dtype {result.dtype}.", error_type=ValueError, ) - + if needs_dtype_copy: result = result.astype(dtype) - + if copy is True: result = result.copy() - + return result diff --git a/leapp/leapp_graph/datatypes/traced_data.py b/leapp/leapp_graph/datatypes/traced_data.py index 1c7b700d..f2fc7d3d 100644 --- a/leapp/leapp_graph/datatypes/traced_data.py +++ b/leapp/leapp_graph/datatypes/traced_data.py @@ -20,21 +20,21 @@ class TracedData(ABC): """Abstract base class for traced data types. - + This class provides the common infrastructure for tracing operations on different data types (torch.Tensor, numpy.ndarray, etc.) using torch.fx. - + Child classes implement type-specific operation interception: - TracedTensor: Uses __torch_function__ for torch operations - TracedArray (future): Uses __array_ufunc__/__array_function__ for numpy - + The FX graph always records torch operations, so child classes must map their native operations to torch equivalents for proxy recording. """ - + def __init__(self, value: Any, name: str, context, proxy: Proxy): """Initialize a TracedData instance. - + Args: value: The underlying data (torch.Tensor, numpy.ndarray, etc.) name: Name for the data (used in ONNX export and graph) @@ -43,7 +43,7 @@ def __init__(self, value: Any, name: str, context, proxy: Proxy): """ self._value = value bind_new_view(self, name, context, proxy) - + # ========================================================================= # Common Properties # ========================================================================= @@ -54,7 +54,7 @@ def _name_from_proxy(proxy: Proxy) -> str: if proxy is not None: return str(proxy.node.name) return "untraced" - + @property def proxy(self) -> Proxy: """Get the fx.Proxy for graph recording.""" @@ -69,19 +69,19 @@ def proxy_view(self) -> ProxyView: a mutation in between go unseen. """ return self._proxy_view - + @property def name(self) -> str: """Get the name of this traced data.""" return self._name - + @property def context(self) -> str: """Get the name of the context that owns this data.""" if self._context is None: return "untraced" return self._context.name - + @property def context_obj(self): """Get the TraceContext that owns this data.""" @@ -100,7 +100,7 @@ def output_port(self) -> Optional[str]: @output_port.setter def output_port(self, port: Optional[str]) -> None: self._output_port = port - + @property def is_tracing(self) -> bool: """Get the tracing status of the context that owns this data.""" @@ -124,7 +124,7 @@ def describes_replaced_graph(self) -> bool: # no provenance to shed. graph = getattr(self._context, "graph", None) return graph is not None and proxy.node.graph is not graph - + @property @abstractmethod def tensor(self) -> torch.Tensor: @@ -144,18 +144,19 @@ def get_dtype_name(self) -> str: mapping lives with that backend rather than in leapp core. """ return dtype_to_name(self.data.dtype) + # ========================================================================= # Abstract Methods - Must be implemented by child classes # ========================================================================= - + @abstractmethod def _new(self, value: Any, proxy: Proxy = None) -> "TracedData": """Create a new TracedData of the same type in the same context. - + Args: value: The new underlying value proxy: The proxy for the new data (can be None if not tracing) - + Returns: A new TracedData instance of the same type """ @@ -164,14 +165,14 @@ def _new(self, value: Any, proxy: Proxy = None) -> "TracedData": # ========================================================================= # Common Static Methods # ========================================================================= - + @staticmethod def find_traced_data(obj): """Find the first TracedData in a supported nested structure. - + Args: obj: Object to search (can be TracedData, list, tuple, dict, or other) - + Returns: The first TracedData found, or None if not found """ @@ -185,7 +186,7 @@ def visit(item): TracedData._map_structure(obj, visit) return found - + @staticmethod def _map_structure(obj, leaf_fn): """Recursively apply leaf_fn while preserving common containers.""" @@ -196,9 +197,7 @@ def _map_structure(obj, leaf_fn): TracedData._map_structure(obj.step, leaf_fn), ) if isinstance(obj, (list, tuple)): - mapped = tuple( - TracedData._map_structure(item, leaf_fn) for item in obj - ) + mapped = tuple(TracedData._map_structure(item, leaf_fn) for item in obj) if hasattr(obj, "_fields"): # namedtuples take their fields positionally. Plain sequences and # torch's structseq return types instead take a single sequence. @@ -214,10 +213,10 @@ def _map_structure(obj, leaf_fn): @staticmethod def unwrap_traced_data(obj): """Recursively unwrap TracedData to get raw values. - + Args: obj: Object to unwrap (can be TracedData, list, tuple, dict, or other) - + Returns: The unwrapped value with all TracedData replaced by their raw values """ @@ -225,15 +224,15 @@ def unwrap_traced_data(obj): obj, lambda item: item.data if isinstance(item, TracedData) else item, ) - + @staticmethod def find_all_contexts(obj, contexts: Optional[Set[str]] = None) -> Set[str]: """Recursively find all unique context names. - + Args: obj: Object to search contexts: Set to accumulate context names (created if None) - + Returns: Set of context names found """ @@ -247,14 +246,14 @@ def collect(item): TracedData._map_structure(obj, collect) return contexts - + @staticmethod def extract_proxy(obj): """Recursively extract proxies from TracedData objects. - + Args: obj: Object to extract proxies from - + Returns: Object with TracedData replaced by their proxies """ @@ -262,26 +261,26 @@ def extract_proxy(obj): obj, lambda item: item.proxy if isinstance(item, TracedData) else item, ) - + # ========================================================================= # Common Validation # ========================================================================= - + def validate_status(self, args=None, kwargs=None) -> bool: """Validate that this TracedData can be used in the current context. - + Checks: 1. If not tracing, returns False (operation should not be recorded) 2. If tracing inside a traced function, raises an error 3. If mixing TracedData from different contexts, raises an error - + Args: args: Positional arguments to check for other TracedData kwargs: Keyword arguments to check for other TracedData - + Returns: True if tracing should proceed, False if not tracing - + Raises: Exception: If TracedData is used inside a traced function Exception: If TracedData from multiple contexts are mixed @@ -296,7 +295,7 @@ def validate_status(self, args=None, kwargs=None) -> bool: if kwargs is not None: for kwarg in kwargs.values(): contexts = TracedData.find_all_contexts(kwarg, contexts) - + if len(contexts) > 1: cls_name = self.__class__.__name__ _get_logger().fatal( @@ -307,7 +306,8 @@ def validate_status(self, args=None, kwargs=None) -> bool: "\n" "1. call output_tensors() to finalize one of the nodes first\n" "2. combine both nodes into a single node by calling input_tensors() with the same node name", - error_type=Exception) + error_type=Exception, + ) return True # ========================================================================= @@ -392,9 +392,7 @@ def overwrite_port(self, key, value) -> None: and tuple(value.shape) == tuple(self.shape) and value.dtype == self.dtype ): - update_view_proxy( - self, value.name, value.context_obj, value.proxy - ) + update_view_proxy(self, value.name, value.context_obj, value.proxy) self._output_port = value.output_port return self._output_port = None @@ -402,14 +400,14 @@ def overwrite_port(self, key, value) -> None: # ========================================================================= # Common Magic Methods # ========================================================================= - + def __len__(self) -> int: """Length operator.""" return len(self._value) - + def __bool__(self) -> bool: """Boolean conversion for TracedNpArray. - + During tracing, this indicates tensor-dependent control flow which produces a silently incorrect graph. Logs an error to alert the user. """ @@ -434,15 +432,15 @@ def __bool__(self) -> bool: f" then start a new input_tensors() after it" ) return bool(self.data) - + def __str__(self) -> str: """String representation.""" return f"{type(self).__name__}({self._value})" - + def __repr__(self) -> str: """String representation.""" return f"{type(self).__name__}({self._value})" - + def __format__(self, format_spec: str) -> str: """Format the TracedData by delegating to the underlying value.""" return self._value.__format__(format_spec) @@ -450,7 +448,7 @@ def __format__(self, format_spec: str) -> str: # ========================================================================= # Setitem Proxy Generation # ========================================================================= - + def _register_setitem_tensor(self, value: torch.Tensor, prefix: str): """Register a tensor constant and return its get_attr proxy.""" attr_base = f"{prefix}_{id(value)}" @@ -459,12 +457,8 @@ def _register_setitem_tensor(self, value: torch.Tensor, prefix: str): while hasattr(self._context.tracer.root, attr_name): attr_name = f"{attr_base}_{suffix}" suffix += 1 - self._context.tracer.root.register_buffer( - attr_name, value.clone().detach() - ) - return self._context.tracer.create_proxy( - "get_attr", attr_name, (), {} - ) + self._context.tracer.root.register_buffer(attr_name, value.clone().detach()) + return self._context.tracer.create_proxy("get_attr", attr_name, (), {}) @staticmethod def _tensor_index_key(key): @@ -476,6 +470,7 @@ def _tensor_index_key(key): def _proxy_index_key(self, key): """Replace tensor indices with graph values or registered constants.""" + def convert(item): if isinstance(item, TracedData): return item.proxy @@ -577,8 +572,10 @@ def _lower_index_key(self, proxy, key): return proxy residual_key = tuple(residual) if isinstance(key, tuple) else residual[0] return self._context.tracer.create_proxy( - "call_function", operator.getitem, - (proxy, self._proxy_index_key(residual_key)), {} + "call_function", + operator.getitem, + (proxy, self._proxy_index_key(residual_key)), + {}, ) def _create_getitem_proxy(self, key): @@ -626,17 +623,25 @@ def _create_setitem_proxy(self, key, value_proxy, real_value=None): dtype=destination.dtype, device=destination.device, ) - if source_tensor.numel() != 1 and tuple(source_tensor.shape) != target_shape: + if ( + source_tensor.numel() != 1 + and tuple(source_tensor.shape) != target_shape + ): source_tensor = torch.broadcast_to(source_tensor, target_shape).clone() source_proxy = self._register_setitem_tensor( source_tensor, "_setitem_value" ) else: source_numel = getattr(real_value, "numel", lambda: 1)() - if source_numel != 1 and tuple(getattr(real_value, "shape", ())) != target_shape: + if ( + source_numel != 1 + and tuple(getattr(real_value, "shape", ())) != target_shape + ): source_proxy = self._context.tracer.create_proxy( - "call_function", torch.broadcast_to, - (source_proxy, target_shape), {} + "call_function", + torch.broadcast_to, + (source_proxy, target_shape), + {}, ) destination_proxy = self.proxy @@ -650,8 +655,10 @@ def _create_setitem_proxy(self, key, value_proxy, real_value=None): "call_method", "reshape", (destination_proxy, (-1,)), {} ) updated_proxy = self._context.tracer.create_proxy( - "call_function", torch.index_put, - (flat_destination_proxy, (indices_proxy,), flat_source_proxy), {} + "call_function", + torch.index_put, + (flat_destination_proxy, (indices_proxy,), flat_source_proxy), + {}, ) return self._context.tracer.create_proxy( "call_method", "reshape", (updated_proxy, tensor_shape), {} @@ -659,9 +666,7 @@ def _create_setitem_proxy(self, key, value_proxy, real_value=None): def _update_setitem_proxy(self, key, value_proxy, real_value=None): """Record a functional assignment and install its resulting proxy.""" - proxy_out = self._create_setitem_proxy( - key, value_proxy, real_value=real_value - ) + proxy_out = self._create_setitem_proxy(key, value_proxy, real_value=real_value) if proxy_out is None: return False self._proxy_view.proxy = proxy_out diff --git a/leapp/leapp_graph/datatypes/traced_tensor.py b/leapp/leapp_graph/datatypes/traced_tensor.py index 6372ffd3..d200e1fe 100644 --- a/leapp/leapp_graph/datatypes/traced_tensor.py +++ b/leapp/leapp_graph/datatypes/traced_tensor.py @@ -23,12 +23,12 @@ from .traced_data import TracedData - # Combined metaclass to resolve conflict between ABCMeta (from TracedData) # and torch.Tensor's metaclass. In the future if this breaks, we need to convert # the TracedTensor class to not inherit from tracedData and reimplement all the common functions class _TracedTensorMeta(ABCMeta, type(torch.Tensor)): """Metaclass combining ABCMeta and torch.Tensor's metaclass.""" + pass @@ -45,13 +45,13 @@ class TracedTensor(TracedData, torch.Tensor, metaclass=_TracedTensorMeta): @staticmethod def __new__(cls, tensor: torch.Tensor, name: str, context, proxy: Proxy): """Create a new TracedTensor instance. - + Args: tensor: The actual torch.Tensor data to wrap name: Name for the tensor (used in ONNX export and graph) context: The TraceContext that owns this tensor proxy: The fx.Proxy for graph recording - + Returns: A new TracedTensor instance sharing storage with the input tensor """ @@ -82,7 +82,7 @@ def __init__(self, tensor: torch.Tensor, name: str, context, proxy: Proxy): def tensor(self) -> torch.Tensor: """Get the underlying torch.Tensor (for compatibility with original API).""" return self.as_subclass(torch.Tensor) - + @property def data(self) -> torch.Tensor: """Get the underlying data.""" @@ -122,9 +122,9 @@ def _convert_method_descriptor(func, context=None): """ # Check if it's a method descriptor on torch.Tensor func_type = type(func).__name__ - if func_type in ('method_descriptor', 'builtin_function_or_method'): + if func_type in ("method_descriptor", "builtin_function_or_method"): # Try to get the method name - func_name = getattr(func, '__name__', None) + func_name = getattr(func, "__name__", None) if func_name: # Look for equivalent in torch module torch_func = getattr(torch, func_name, None) @@ -133,7 +133,7 @@ def _convert_method_descriptor(func, context=None): # Handle in-place operations (e.g., add_ -> torch.add) # In-place methods end with '_' but torch.* functions don't - if func_name.endswith('_') and len(func_name) > 1: + if func_name.endswith("_") and len(func_name) > 1: base_name = func_name[:-1] # Remove trailing underscore torch_func = getattr(torch, base_name, None) if torch_func is not None and callable(torch_func): @@ -223,13 +223,8 @@ def copy_into(target: torch.Tensor, source: "TracedTensor") -> "TracedTensor": # Return a new TracedTensor that wraps target but uses source's proxy # This effectively makes target "become" traced - return TracedTensor( - target, - source.name, - source._context, - source.proxy - ) - + return TracedTensor(target, source.name, source._context, source.proxy) + @staticmethod def _handle_TS_decomposition(func, real_args, real_kwargs): """Handle TorchScript decomposition.""" @@ -247,9 +242,9 @@ def _handle_TS_decomposition(func, real_args, real_kwargs): # ScriptMethod. func.owner returns the C++ ScriptModule # (not an nn.Module). Save → reload via BytesIO so # torch.jit.load performs full Python-level init. - script_module = getattr(func, '__self__', None) + script_module = getattr(func, "__self__", None) if not isinstance(script_module, torch.nn.Module): - cpp_module = getattr(func, 'owner', None) + cpp_module = getattr(func, "owner", None) if cpp_module is not None: _buf = _io.BytesIO() _tmp = torch.jit.RecursiveScriptModule._construct( @@ -258,15 +253,12 @@ def _handle_TS_decomposition(func, real_args, real_kwargs): torch.jit.save(_tmp, _buf) _buf.seek(0) _dev = next( - (a.device for a in real_args - if isinstance(a, torch.Tensor)), - torch.device('cpu'), - ) - script_module = torch.jit.load( - _buf, map_location=_dev + (a.device for a in real_args if isinstance(a, torch.Tensor)), + torch.device("cpu"), ) + script_module = torch.jit.load(_buf, map_location=_dev) - #validate the script_module + # validate the script_module if not isinstance(script_module, torch.nn.Module): raise RuntimeError( f"Could not obtain nn.Module from {type(func).__name__}. " @@ -286,9 +278,7 @@ def _handle_TS_decomposition(func, real_args, real_kwargs): k: v.clone().detach() if isinstance(v, torch.Tensor) else v for k, v in (real_kwargs or {}).items() } - ep = TS2EPConverter( - script_module, _clean, _clean_kw - ).convert() + ep = TS2EPConverter(script_module, _clean, _clean_kw).convert() fx_module = ep.module() except Exception as e: @@ -325,18 +315,19 @@ def _handle_class_swap(cls, func, args=(), kwargs=None): if kwargs is None: kwargs = {} - original_func_name = getattr(func, '__name__', '') + original_func_name = getattr(func, "__name__", "") - if original_func_name == '__setitem__' and len(args) >= 3: + if original_func_name == "__setitem__" and len(args) >= 3: target, key, value = args[0], args[1], args[2] if type(target) is torch.Tensor and isinstance(value, TracedTensor): is_full_copy = ( - (isinstance(key, slice) and key == slice(None)) - or key is Ellipsis - ) + isinstance(key, slice) and key == slice(None) + ) or key is Ellipsis if is_full_copy: target.__class__ = TracedTensor - target._init_tracing_state(value.name, value.context_obj, value.proxy) + target._init_tracing_state( + value.name, value.context_obj, value.proxy + ) return True, None _get_logger().warning( @@ -347,11 +338,13 @@ def _handle_class_swap(cls, func, args=(), kwargs=None): "reassign the variable (target = source) instead." ) - elif original_func_name == 'copy_' and len(args) >= 2: + elif original_func_name == "copy_" and len(args) >= 2: target, source = args[0], args[1] if type(target) is torch.Tensor and isinstance(source, TracedTensor): target.__class__ = TracedTensor - target._init_tracing_state(source.name, source.context_obj, source.proxy) + target._init_tracing_state( + source.name, source.context_obj, source.proxy + ) return True, target return False, None @@ -377,15 +370,19 @@ def _handle_scripted_call( kwargs = {} func_type_name = type(func).__name__ - func_module = type(func).__module__ if hasattr(type(func), '__module__') else '' + func_module = type(func).__module__ if hasattr(type(func), "__module__") else "" - is_scripted = ('ScriptMethod' in func_type_name or 'ScriptModule' in func_type_name) + is_scripted = ( + "ScriptMethod" in func_type_name or "ScriptModule" in func_type_name + ) if not is_scripted: - if func_module.startswith('torch.jit') or func_module.startswith('torch._C'): - if hasattr(func, '__self__') and hasattr(func.__self__, '__class__'): + if func_module.startswith("torch.jit") or func_module.startswith( + "torch._C" + ): + if hasattr(func, "__self__") and hasattr(func.__self__, "__class__"): self_class_name = func.__self__.__class__.__name__ - is_scripted = 'Script' in self_class_name + is_scripted = "Script" in self_class_name if not is_scripted: return False, None @@ -435,11 +432,13 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): # Extract real tensors for actual computation real_args = tuple(TracedTensor.unwrap_traced_tensor(arg) for arg in args) - real_kwargs = {k: TracedTensor.unwrap_traced_tensor(v) for k, v in kwargs.items()} + real_kwargs = { + k: TracedTensor.unwrap_traced_tensor(v) for k, v in kwargs.items() + } # used to detect in-place mutations # specific to fill_ copy_ things that mutate the 0th index # more general case is not supported yet - receiver = args[0] if args and isinstance(args[0], TracedTensor) else None + receiver = args[0] if args and isinstance(args[0], TracedTensor) else None real_receiver = real_args[0] if receiver is not None else None receiver_version = cls._safe_tensor_version(real_receiver) @@ -550,18 +549,22 @@ def __getattr__(self, name: str): return self.__dict__[name] # Handle in-place operations (methods ending with _) - if name.endswith('_') and not name.startswith('_') and len(name) > 1: + if name.endswith("_") and not name.startswith("_") and len(name) > 1: base_name = name[:-1] # Remove trailing underscore torch_func = getattr(torch, base_name, None) if torch_func is not None and callable(torch_func): + def inplace_method(*args, **kwargs): if not self.is_tracing: # Just call the in-place method on the tensor and return raw tensor underlying = self.as_subclass(torch.Tensor) method = getattr(underlying, name, None) if method is not None: - method(*TracedData.unwrap_traced_data(args), **TracedData.unwrap_traced_data(kwargs)) + method( + *TracedData.unwrap_traced_data(args), + **TracedData.unwrap_traced_data(kwargs), + ) return underlying # Record as functional operation @@ -595,6 +598,7 @@ def method(*args, **kwargs): if attr is not None: # If it's a callable method, wrap it to ensure it goes through __torch_function__ if callable(attr): + def wrapped_method(*args, **kwargs): # Call the method on the underlying tensor result = attr(*args, **kwargs) @@ -605,15 +609,16 @@ def wrapped_method(*args, **kwargs): return result # Create a proxy node for this operation proxy = self._proxy._tracer.create_proxy( - 'call_method', name, (self._proxy,) + args, kwargs) + "call_method", name, (self._proxy,) + args, kwargs + ) return self._new(result, proxy) return result + return wrapped_method return attr # Attribute not found - raise AttributeError( - f"'TracedTensor' object has no attribute '{name}'") + raise AttributeError(f"'TracedTensor' object has no attribute '{name}'") # ========================================================================= # Special Methods: reshape and permute need custom handling @@ -721,7 +726,9 @@ def __iadd__(self, other): # Update internal state to maintain in-place semantics with torch.no_grad(): - torch.Tensor.copy_(self, result.tensor if isinstance(result, TracedTensor) else result) + torch.Tensor.copy_( + self, result.tensor if isinstance(result, TracedTensor) else result + ) if isinstance(result, TracedTensor): self._proxy = result.proxy @@ -737,7 +744,9 @@ def __isub__(self, other): result = torch.sub(self, other) with torch.no_grad(): - torch.Tensor.copy_(self, result.tensor if isinstance(result, TracedTensor) else result) + torch.Tensor.copy_( + self, result.tensor if isinstance(result, TracedTensor) else result + ) if isinstance(result, TracedTensor): self._proxy = result.proxy return self @@ -752,7 +761,9 @@ def __imul__(self, other): result = torch.mul(self, other) with torch.no_grad(): - torch.Tensor.copy_(self, result.tensor if isinstance(result, TracedTensor) else result) + torch.Tensor.copy_( + self, result.tensor if isinstance(result, TracedTensor) else result + ) if isinstance(result, TracedTensor): self._proxy = result.proxy return self @@ -767,7 +778,9 @@ def __itruediv__(self, other): result = torch.div(self, other) with torch.no_grad(): - torch.Tensor.copy_(self, result.tensor if isinstance(result, TracedTensor) else result) + torch.Tensor.copy_( + self, result.tensor if isinstance(result, TracedTensor) else result + ) if isinstance(result, TracedTensor): self._proxy = result.proxy return self @@ -782,7 +795,9 @@ def __ipow__(self, other): result = torch.pow(self, other) with torch.no_grad(): - torch.Tensor.copy_(self, result.tensor if isinstance(result, TracedTensor) else result) + torch.Tensor.copy_( + self, result.tensor if isinstance(result, TracedTensor) else result + ) if isinstance(result, TracedTensor): self._proxy = result.proxy return self @@ -799,7 +814,9 @@ def __imatmul__(self, other): result = torch.matmul(self, other) with torch.no_grad(): - torch.Tensor.copy_(self, result.tensor if isinstance(result, TracedTensor) else result) + torch.Tensor.copy_( + self, result.tensor if isinstance(result, TracedTensor) else result + ) if isinstance(result, TracedTensor): self._proxy = result.proxy return self @@ -818,7 +835,9 @@ def add_(self, other, *, alpha=1): # Record as functional operation result = torch.add(self, other, alpha=alpha) with torch.no_grad(): - torch.Tensor.copy_(self, result.tensor if isinstance(result, TracedTensor) else result) + torch.Tensor.copy_( + self, result.tensor if isinstance(result, TracedTensor) else result + ) if isinstance(result, TracedTensor): self._proxy = result.proxy return self @@ -833,7 +852,9 @@ def sub_(self, other, *, alpha=1): result = torch.sub(self, other, alpha=alpha) with torch.no_grad(): - torch.Tensor.copy_(self, result.tensor if isinstance(result, TracedTensor) else result) + torch.Tensor.copy_( + self, result.tensor if isinstance(result, TracedTensor) else result + ) if isinstance(result, TracedTensor): self._proxy = result.proxy return self @@ -848,7 +869,9 @@ def mul_(self, other): result = torch.mul(self, other) with torch.no_grad(): - torch.Tensor.copy_(self, result.tensor if isinstance(result, TracedTensor) else result) + torch.Tensor.copy_( + self, result.tensor if isinstance(result, TracedTensor) else result + ) if isinstance(result, TracedTensor): self._proxy = result.proxy return self @@ -863,7 +886,9 @@ def div_(self, other, *, rounding_mode=None): result = torch.div(self, other, rounding_mode=rounding_mode) with torch.no_grad(): - torch.Tensor.copy_(self, result.tensor if isinstance(result, TracedTensor) else result) + torch.Tensor.copy_( + self, result.tensor if isinstance(result, TracedTensor) else result + ) if isinstance(result, TracedTensor): self._proxy = result.proxy return self @@ -878,7 +903,9 @@ def pow_(self, exponent): result = torch.pow(self, exponent) with torch.no_grad(): - torch.Tensor.copy_(self, result.tensor if isinstance(result, TracedTensor) else result) + torch.Tensor.copy_( + self, result.tensor if isinstance(result, TracedTensor) else result + ) if isinstance(result, TracedTensor): self._proxy = result.proxy return self @@ -967,7 +994,10 @@ def __getitem__(self, key): "call_method", "reshape", (key.proxy, (-1,)), {} ) proxy_out = self._context.tracer.create_proxy( - "call_function", torch.index_select, (self._proxy, 0, flat_index_proxy), {} + "call_function", + torch.index_select, + (self._proxy, 0, flat_index_proxy), + {}, ) if key.ndim > 1: @@ -1009,7 +1039,7 @@ def __getitem__(self, key): def __setitem__(self, key, value): """Indexed assignment using functional operations for graph compatibility. - + Uses the shared _create_setitem_proxy helper from TracedData to convert __setitem__ to torch.index_put for FX/TorchScript/ONNX compatibility. """ @@ -1028,11 +1058,11 @@ def __setitem__(self, key, value): # Use shared helper to create the proxy proxy_out = self._create_setitem_proxy(key, value_proxy) - + if proxy_out is not None: self._proxy = proxy_out return - + # Handle unsupported cases with warnings has_slice = isinstance(key, tuple) and any(isinstance(k, slice) for k in key) if has_slice: @@ -1061,7 +1091,7 @@ def __setitem__(self, key, value): "Consider using torch.index_put or torch.masked_scatter directly.", stacklevel=2, ) - + # Fallback: record __setitem__ directly (may not export) self._proxy = self._context.tracer.create_proxy( "call_method", "__setitem__", (self._proxy, key, value_proxy), {} @@ -1111,56 +1141,57 @@ def to(self, *args, **kwargs): def numpy(self): """Convert to numpy array, preserving tracing as TracedNpArray. - + If tracing is active, returns a TracedNpArray that shares the proxy and context, allowing the trace to continue across the conversion. - + If not tracing, returns a regular numpy array. - + Returns: TracedNpArray if tracing, else np.ndarray """ np_data = self.as_subclass(torch.Tensor).detach().cpu().numpy() - + # Skip tracing if context is not tracing - return raw numpy array if not self.validate_status(): return np_data - + # Return TracedNpArray with inherited proxy/context - from . import as_traced #lazy import to avoid circular dependency + from . import as_traced # lazy import to avoid circular dependency + return as_traced(np_data, self.name, self.context_obj, self.proxy) def __array__(self, dtype=None, copy=None): """NumPy array protocol - called by np.array(), np.asarray(), etc. - + This allows seamless conversion to numpy while preserving tracing. - + Args: dtype: Optional dtype for the resulting array copy: Optional copy flag (NumPy 2.0+) - None: Copy only if necessary (default) - True: Always copy - False: Never copy; raise ValueError if copy required - + Returns: TracedNpArray if tracing, else np.ndarray """ result = self.numpy() - + # Check if dtype change requires a copy needs_dtype_copy = dtype is not None and dtype != result.dtype - + # copy=False but copy is required → raise error (NumPy 2.0 semantics) if copy is False and needs_dtype_copy: raise ValueError( f"Unable to avoid copy while creating an array with dtype {dtype} " f"from array with dtype {result.dtype}." ) - + if needs_dtype_copy: result = result.astype(dtype) - + if copy is True: result = result.copy() - - return result \ No newline at end of file + + return result diff --git a/leapp/leapp_graph/datatypes/warp/patching.py b/leapp/leapp_graph/datatypes/warp/patching.py index 819b252c..ec15d05a 100644 --- a/leapp/leapp_graph/datatypes/warp/patching.py +++ b/leapp/leapp_graph/datatypes/warp/patching.py @@ -61,7 +61,10 @@ import warp as wp from warp._src.context import Function as WarpKernelLanguageFunction -from leapp.utils.caller_identity import caller_identity_has_same_anchor, get_caller_stack_identity +from leapp.utils.caller_identity import ( + caller_identity_has_same_anchor, + get_caller_stack_identity, +) from leapp.utils.logging import _get_logger from .._attribute_patching import AttributePatchRegistry @@ -191,9 +194,7 @@ def is_full_copy_function(self, func: Callable) -> bool: return id(func) == self._full_copy_function_id def function_qualname(self, func: Callable) -> str: - return self._qualnames_by_original_id.get( - id(func), f"warp.{func.__qualname__}" - ) + return self._qualnames_by_original_id.get(id(func), f"warp.{func.__qualname__}") def uninstall(self) -> None: """Restore every attribute patched by this detector.""" @@ -293,7 +294,9 @@ def _patch_namespace( # this is not a namespace; maybe it is a callable we should wrap, otherwise ignore it. self._patch_callable_attr(owner, attr_name, value, attr_qualname) - def _should_scan_class(self, cls: Any, owner_module_name: str, class_depth: int) -> bool: + def _should_scan_class( + self, cls: Any, owner_module_name: str, class_depth: int + ) -> bool: # Class scan criteria: # - The depth is not too deep. # - Only scan actual class objects. @@ -318,7 +321,9 @@ def _should_scan_class(self, cls: Any, owner_module_name: str, class_depth: int) return True - def _patch_callable_attr(self, owner: Any, attr_name: str, raw_value: Any, qualname: str) -> None: + def _patch_callable_attr( + self, owner: Any, attr_name: str, raw_value: Any, qualname: str + ) -> None: # Patch only safe callable attributes: # - Preserve staticmethod/classmethod binding by unwrapping, wrapping, then # restoring the original descriptor type. @@ -336,9 +341,13 @@ def _patch_callable_attr(self, owner: Any, attr_name: str, raw_value: Any, qualn if attr_name.startswith("__") and attr_name not in _ALLOWED_DUNDER_METHODS: return - if callable_original is None or getattr(callable_original, _WRAPPER_MARKER, False): + if callable_original is None or getattr( + callable_original, _WRAPPER_MARKER, False + ): return - if inspect.isclass(callable_original) or isinstance(callable_original, property): + if inspect.isclass(callable_original) or isinstance( + callable_original, property + ): return if isinstance(callable_original, WarpKernelLanguageFunction): """In Warp, symbols like these often look callable from Python: @@ -380,7 +389,11 @@ def _patch_attr( return wrapper_func = self._get_or_make_wrapper(qualname, callable_original) - wrapper = descriptor_type(wrapper_func) if descriptor_type is not None else wrapper_func + wrapper = ( + descriptor_type(wrapper_func) + if descriptor_type is not None + else wrapper_func + ) try: self._patches.install( diff --git a/leapp/leapp_graph/datatypes/warp/session.py b/leapp/leapp_graph/datatypes/warp/session.py index 26dfc12f..89ad9af1 100644 --- a/leapp/leapp_graph/datatypes/warp/session.py +++ b/leapp/leapp_graph/datatypes/warp/session.py @@ -19,7 +19,7 @@ def __init__(self) -> None: self.active_warp_op: Any | None = None self.owner_token: Any | None = None # TODO: revisit if we need 2 different pause states for cupti and patching - self._pause_depth = 0 # used to detect nested pauses. + self._pause_depth = 0 # used to detect nested pauses. @property def paused(self) -> bool: diff --git a/leapp/leapp_graph/datatypes/warp/traced_wp_array.py b/leapp/leapp_graph/datatypes/warp/traced_wp_array.py index 03d46208..43581af9 100644 --- a/leapp/leapp_graph/datatypes/warp/traced_wp_array.py +++ b/leapp/leapp_graph/datatypes/warp/traced_wp_array.py @@ -25,25 +25,27 @@ # Register the Warp dtype codec so leapp core can convert warp dtypes to common # name strings without importing warp itself. Mirrors the torch/numpy entries # and only runs when this (warp-gated) module is imported. -register_dtype_codec(DtypeCodec( - backend="warp", - matches=lambda v: isinstance(v, wp.array), - value_dtype=lambda v: v.dtype, - dtype_to_name={ - wp.float16: "float16", - wp.float32: "float32", - wp.float64: "float64", - wp.int8: "int8", - wp.int16: "int16", - wp.int32: "int32", - wp.int64: "int64", - wp.uint8: "uint8", - wp.uint16: "uint16", - wp.uint32: "uint32", - wp.uint64: "uint64", - wp.bool: "bool", - }, -)) +register_dtype_codec( + DtypeCodec( + backend="warp", + matches=lambda v: isinstance(v, wp.array), + value_dtype=lambda v: v.dtype, + dtype_to_name={ + wp.float16: "float16", + wp.float32: "float32", + wp.float64: "float64", + wp.int8: "int8", + wp.int16: "int16", + wp.int32: "int32", + wp.int64: "int64", + wp.uint8: "uint8", + wp.uint16: "uint16", + wp.uint32: "uint32", + wp.uint64: "uint64", + wp.bool: "bool", + }, + ) +) class TracedWpArray(wp.array, TracedData): @@ -121,8 +123,13 @@ def clear_tracing_state(cls, array: "TracedWpArray") -> wp.array: the pass that promoted it has to shed that provenance before the next pass can describe it. """ - for attribute in ("_name", "_context", "_proxy_view", "_output_port", - "_leapp_warp_segment"): + for attribute in ( + "_name", + "_context", + "_proxy_view", + "_output_port", + "_leapp_warp_segment", + ): array.__dict__.pop(attribute, None) array.__class__ = wp.array return array @@ -156,15 +163,11 @@ def __warp_function__( qualname = backend.function_qualname(func) if backend.is_boundary_function(func): - handled, result = cls._handle_boundary_call( - backend, func, args, kwargs - ) + handled, result = cls._handle_boundary_call(backend, func, args, kwargs) if handled: return result - traced_inputs, call_args, call_kwargs = cls._normalize_and_collect( - args, kwargs - ) + traced_inputs, call_args, call_kwargs = cls._normalize_and_collect(args, kwargs) trace_source = cls._select_trace_source(qualname, traced_inputs) segment = backend.resolve_or_begin_warp_segment(trace_source) @@ -174,9 +177,7 @@ def __warp_function__( with backend.pause_context(): result = func(*call_args, **call_kwargs) - cls._process_post_call_arrays( - func, segment, args, kwargs, result, trace_source - ) + cls._process_post_call_arrays(func, segment, args, kwargs, result, trace_source) if segment is not None and getattr(func, "__name__", None) == "launch": outputs = kwargs.get("outputs", args[3] if len(args) > 3 else ()) if any(getattr(array, "_tensor", None) is not None for array in outputs): @@ -212,16 +213,17 @@ def _handle_boundary_call( # Keep the caller's Torch object as the owner of the Warp view so # writes through that view can update the caller's tracing state. alias = getattr(raw, "_tensor", None) - if (isinstance(src, torch.Tensor) and type(alias) is torch.Tensor - and layout_key(alias) == layout_key(src)): + if ( + isinstance(src, torch.Tensor) + and type(alias) is torch.Tensor + and layout_key(alias) == layout_key(src) + ): raw._tensor = src if may_adopt_view(src, raw): view, proxy = src.proxy_view, None else: view, proxy = None, src.proxy - return True, as_traced( - raw, src.name, src.context_obj, proxy, view=view - ) + return True, as_traced(raw, src.name, src.context_obj, proxy, view=view) return True, raw @staticmethod @@ -280,8 +282,7 @@ def _normalize_node( } if type(obj) in (list, tuple, set, frozenset): return type(obj)( - cls._normalize_node(item, traced, depth=depth + 1) - for item in obj + cls._normalize_node(item, traced, depth=depth + 1) for item in obj ) return obj @@ -485,7 +486,6 @@ def data(self) -> wp.array: raw._ref = self return raw - @property def tensor(self) -> torch.Tensor: return wp.to_torch(self) diff --git a/leapp/leapp_graph/datatypes/warp/warp_segment.py b/leapp/leapp_graph/datatypes/warp/warp_segment.py index c6f8cdd6..f688e9df 100644 --- a/leapp/leapp_graph/datatypes/warp/warp_segment.py +++ b/leapp/leapp_graph/datatypes/warp/warp_segment.py @@ -60,8 +60,8 @@ def from_value( storage_dtype, storage_shape = value_to_name_and_shape(value) storage_shape = tuple(storage_shape) component_shape = () - if isinstance(shape, tuple) and storage_shape[:len(shape)] == shape: - component_shape = storage_shape[len(shape):] + if isinstance(shape, tuple) and storage_shape[: len(shape)] == shape: + component_shape = storage_shape[len(shape) :] return cls( name=name, @@ -107,11 +107,7 @@ def is_open(self) -> bool: @property def is_empty(self) -> bool: - return not ( - self.events - or self.input_refs - or self.output_refs - ) + return not (self.events or self.input_refs or self.output_refs) def add_event(self, event: Any) -> None: self.events.append(event) diff --git a/leapp/leapp_graph/function_decorator_node.py b/leapp/leapp_graph/function_decorator_node.py index 164ddbcf..7fb878c2 100644 --- a/leapp/leapp_graph/function_decorator_node.py +++ b/leapp/leapp_graph/function_decorator_node.py @@ -41,9 +41,16 @@ class FunctionDecoratorNode(LeappNode): an nn.Module via ModuleBuilder for export. """ - def __init__(self, name, backend=None, - backend_params=None, inputs=None, outputs=None, - environment_constants=None, register_buffers=None): + def __init__( + self, + name, + backend=None, + backend_params=None, + inputs=None, + outputs=None, + environment_constants=None, + register_buffers=None, + ): super().__init__(name) if inputs is not None: @@ -79,12 +86,12 @@ def __init__(self, name, backend=None, self.setup_backend(backend, backend_params) self.executed_lines = { - 'filename': None, - 'function_name': None, - 'min_line': None, - 'max_line': None, - 'lines': set(), - 'source_code': None, + "filename": None, + "function_name": None, + "min_line": None, + "max_line": None, + "lines": set(), + "source_code": None, } self.input_namespace = None @@ -104,8 +111,8 @@ def compile_trace(self, func, *args): ) func_code = unwrapped_func.__code__ - self.executed_lines['filename'] = func_code.co_filename - self.executed_lines['function_name'] = func_code.co_name + self.executed_lines["filename"] = func_code.co_filename + self.executed_lines["function_name"] = func_code.co_name try: func_lines, start_line = inspect.getsourcelines(unwrapped_func) @@ -115,25 +122,32 @@ def compile_trace(self, func, *args): f"If this function is wrapped by another decorator, try reordering so @annotate._method() " f"is the innermost decorator.", error_type=OSError, - cause=e) + cause=e, + ) - self.executed_lines['min_line'] = start_line - self.executed_lines['max_line'] = start_line + len(func_lines) - 1 + self.executed_lines["min_line"] = start_line + self.executed_lines["max_line"] = start_line + len(func_lines) - 1 def create_trace_function(self, skip_file, entry_hook=None): """Create a trace function for use with sys.settrace.""" + def trace_function(frame, event, arg): - if event == 'call': + if event == "call": code = frame.f_code if os.path.basename(code.co_filename) == skip_file: return trace_function - if (code.co_filename == self.executed_lines['filename'] and - self.executed_lines['min_line'] <= frame.f_lineno <= self.executed_lines['max_line']): + if ( + code.co_filename == self.executed_lines["filename"] + and self.executed_lines["min_line"] + <= frame.f_lineno + <= self.executed_lines["max_line"] + ): namespace = frame_to_namespace(frame) self.output_namespace = namespace return trace_function + return trace_function # ── function input / output inspection ─────────────────────────────── @@ -147,12 +161,11 @@ def inspect_function_inputs(self, func, args, kwargs): for param_name, param_value in bound_args.arguments.items(): self._check_for_active_traced_tensors(param_value, param_name) - if param_name == 'self': + if param_name == "self": continue param = sig.parameters[param_name] - was_provided = ( - param_name in kwargs or - param_names.index(param_name) < len(args) + was_provided = param_name in kwargs or param_names.index(param_name) < len( + args ) if was_provided: self.add_input(param_name, param_name, param_value) @@ -168,7 +181,8 @@ def publish_output_port(self, value, port: str): if not is_tracable_tensor_type(value): _get_logger().warning( f"\033[93mWarning: output '{port}' of node '{self.name}' has " - f"unconnectable type {type(value)}\033[0m") + f"unconnectable type {type(value)}\033[0m" + ) return value promoted = promote_in_place(value, port, self, None) return super().publish_output_port(promoted, port) @@ -183,11 +197,9 @@ def inspect_function_outputs(self, func, result): if i < len(return_names): output_name = return_names[i] else: - output_name = f"output{i+1}" - descriptions = self.add_output( - output_name, output_name, result[i]) - self.publish_output_ports( - result[i], output_name, descriptions) + output_name = f"output{i + 1}" + descriptions = self.add_output(output_name, output_name, result[i]) + self.publish_output_ports(result[i], output_name, descriptions) else: if not len(return_names) == 1: _get_logger().fatal( @@ -195,8 +207,7 @@ def inspect_function_outputs(self, func, result): " outputs, but only one output is detected", error_type=Exception, ) - descriptions = self.add_output( - return_names[0], return_names[0], result) + descriptions = self.add_output(return_names[0], return_names[0], result) self.publish_output_ports(result, return_names[0], descriptions) # ── validation (re-entry) ──────────────────────────────────────────── @@ -208,24 +219,28 @@ def validate_function_boundaries(self, func): func_lines, start_line = inspect.getsourcelines(unwrapped_func) new_boundaries = { - 'filename': func_code.co_filename, - 'function_name': func_code.co_name, - 'min_line': start_line, - 'max_line': start_line + len(func_lines) - 1 + "filename": func_code.co_filename, + "function_name": func_code.co_name, + "min_line": start_line, + "max_line": start_line + len(func_lines) - 1, } - for key in ('filename', 'function_name', 'min_line', 'max_line'): + for key in ("filename", "function_name", "min_line", "max_line"): if self.executed_lines.get(key) != new_boundaries[key]: - original_data = {k: self.executed_lines.get(k) for k in ( - 'filename', 'function_name', 'min_line', 'max_line')} + original_data = { + k: self.executed_lines.get(k) + for k in ("filename", "function_name", "min_line", "max_line") + } _get_logger().fatal( f"Error: {self.name} seen twice but function boundaries do not match\n" f"Original: {original_data}\n" f"New: {new_boundaries}", - error_type=Exception) + error_type=Exception, + ) - self.executed_lines['lines'] = set( - range(start_line, start_line + len(func_lines))) + self.executed_lines["lines"] = set( + range(start_line, start_line + len(func_lines)) + ) def validate_function_inputs(self, func, args, kwargs): """Validate function inputs on re-entry against previously captured inputs.""" @@ -237,16 +252,16 @@ def validate_function_inputs(self, func, args, kwargs): for param_name, param_value in bound_args.arguments.items(): self._check_for_active_traced_tensors(param_value, param_name) - if param_name == 'self': + if param_name == "self": continue param = sig.parameters[param_name] - was_provided = ( - param_name in kwargs or - param_names.index(param_name) < len(args) + was_provided = param_name in kwargs or param_names.index(param_name) < len( + args ) if was_provided: self.validate_input_and_update_sources( - param_name, param_name, param_value) + param_name, param_name, param_value + ) def validate_function_outputs(self, func, result): """Validate function outputs on re-entry against previously captured outputs.""" @@ -259,9 +274,10 @@ def validate_function_outputs(self, func, result): if i < len(return_names): output_name = return_names[i] else: - output_name = f"output{i+1}" + output_name = f"output{i + 1}" self.validate_output_and_update_sources( - output_name, output_name, result[i]) + output_name, output_name, result[i] + ) else: if not len(return_names) == 1: _get_logger().fatal( @@ -270,7 +286,8 @@ def validate_function_outputs(self, func, result): error_type=Exception, ) self.validate_output_and_update_sources( - return_names[0], return_names[0], result) + return_names[0], return_names[0], result + ) # ── namespace capture / validation ─────────────────────────────────── @@ -278,15 +295,16 @@ def capture_inputs_from_namespace(self, namespace): """Capture declared inputs from a namespace dictionary.""" try: for input_name in self._declared_inputs: - final_input_name, final_input_value = self._capture_specified_value_from_namespace( - input_name, namespace) - self.add_input(final_input_name, input_name, - final_input_value) + final_input_name, final_input_value = ( + self._capture_specified_value_from_namespace(input_name, namespace) + ) + self.add_input(final_input_name, input_name, final_input_value) except Exception as e: _get_logger().fatal( f"Error capturing inputs from namespace: {e}", error_type=type(e), - cause=e) + cause=e, + ) self.input_namespace = namespace def capture_outputs_from_namespace(self, namespace): @@ -296,46 +314,55 @@ def capture_outputs_from_namespace(self, namespace): # The namespace still holds the object the caller will use, while # the captured value is a snapshot, so the port goes on the former. obj, _ = get_attribute_value_from_namespace(namespace, output_name) - final_output_name, final_output_value = self._capture_specified_value_from_namespace( - output_name, namespace) - descriptions = self.add_output(final_output_name, - output_name, final_output_value) + final_output_name, final_output_value = ( + self._capture_specified_value_from_namespace(output_name, namespace) + ) + descriptions = self.add_output( + final_output_name, output_name, final_output_value + ) self.publish_output_ports(obj, final_output_name, descriptions) except Exception as e: _get_logger().fatal( f"Error capturing outputs from namespace: {e}", error_type=type(e), - cause=e) + cause=e, + ) self.output_namespace = namespace def validate_inputs_from_namespace(self, namespace): """Validate declared inputs against a namespace dictionary.""" try: for input_name in self._declared_inputs: - final_input_name, final_input_value = self._capture_specified_value_from_namespace( - input_name, namespace) + final_input_name, final_input_value = ( + self._capture_specified_value_from_namespace(input_name, namespace) + ) self.validate_input_and_update_sources( - final_input_name, input_name, final_input_value) + final_input_name, input_name, final_input_value + ) except Exception as e: _get_logger().fatal( f"Error validating inputs from namespace: {e}", error_type=type(e), - cause=e) + cause=e, + ) self.input_namespace = namespace def validate_outputs_from_namespace(self, namespace): """Validate declared outputs against a namespace dictionary.""" try: for output_name in self._declared_outputs: - final_output_name, final_output_value = self._capture_specified_value_from_namespace( - output_name, namespace) + final_output_name, final_output_value = ( + self._capture_specified_value_from_namespace(output_name, namespace) + ) self.validate_output_and_update_sources( - final_output_name, output_name, final_output_value) + final_output_name, output_name, final_output_value + ) except Exception as e: _get_logger().fatal( f"Error validating outputs from namespace: {e}", error_type=type(e), - cause=e) + cause=e, + ) self.output_namespace = namespace # ── buffer / constant snapshot ─────────────────────────────────────── @@ -357,7 +384,8 @@ def snapshot_buffer_values(self, namespace): def _capture_specified_value_from_namespace(self, variable_name, namespace): """Extract a variable value from a namespace dictionary.""" obj, final_variable_name = get_attribute_value_from_namespace( - namespace, variable_name) + namespace, variable_name + ) self._check_for_active_traced_tensors(obj, variable_name) return final_variable_name, safe_deepcopy(obj) @@ -374,12 +402,15 @@ def _check_for_active_traced_tensors(self, data, variable_name, path=None): f"into code traced by _method().\n" f"\n" f"Call output_tensors() on the source node first, or use .tensor to get the underlying tensor.", - error_type=Exception) + error_type=Exception, + ) elif isinstance(data, (list, tuple)): for i, item in enumerate(data): self._check_for_active_traced_tensors( - item, f"{path}[{i}]" if path else f"[{i}]") + item, f"{path}[{i}]" if path else f"[{i}]" + ) elif isinstance(data, dict): for key, value in data.items(): self._check_for_active_traced_tensors( - value, f"{path}['{key}']" if path else f"['{key}']") + value, f"{path}['{key}']" if path else f"['{key}']" + ) diff --git a/leapp/leapp_graph/leapp_graph.py b/leapp/leapp_graph/leapp_graph.py index 007a50b7..d74f414c 100644 --- a/leapp/leapp_graph/leapp_graph.py +++ b/leapp/leapp_graph/leapp_graph.py @@ -21,7 +21,10 @@ from safetensors.torch import save_file -from leapp.utils.tensor_description import CompactYamlList, validate_connection_compatibility +from leapp.utils.tensor_description import ( + CompactYamlList, + validate_connection_compatibility, +) from leapp.utils.logging import _get_logger @@ -58,14 +61,15 @@ def __init__(self, nodes, graph_name="combined_node"): self.graph_name = graph_name # process graph connections - _get_logger().section( - "Processing Node Connections Using Traced Output Ports") + _get_logger().section("Processing Node Connections Using Traced Output Ports") self.connections, self.feedback_connections = self._build_connections( - self.nodes) + self.nodes + ) _get_logger().section("Discovering graph inputs and outputs") self.graph_inputs, self.graph_outputs = self._compile_graph_io( - self.nodes, self.connections, self.feedback_connections) + self.nodes, self.connections, self.feedback_connections + ) _get_logger().section("Validating connection compatibility") self._validate_connection_compatibility() @@ -73,9 +77,9 @@ def __init__(self, nodes, graph_name="combined_node"): def get_feedback_initial_values(self): feedback_initial_values = {} for connection in self.feedback_connections: - for target in connection['targets']: - node = target['node'] - input_desc = node.inputs[target['idx']] + for target in connection["targets"]: + node = target["node"] + input_desc = node.inputs[target["idx"]] key = f"{node.name}/{input_desc.name_str}" value = input_desc.value feedback_initial_values[key] = value @@ -95,25 +99,30 @@ def save_feedback_initial_values(self, save_path, name): def get_full_pipeline_description(self): processed_connections = self._finalize_connections(self.connections) processed_feedback_connections = self._finalize_connections( - self.feedback_connections) + self.feedback_connections + ) graph_inputs = {} for input in self.graph_inputs: - if '/' in input: - node_name, field_name = input.split('/', 1) + if "/" in input: + node_name, field_name = input.split("/", 1) if node_name not in graph_inputs: graph_inputs[node_name] = CompactYamlList() graph_inputs[node_name].append(field_name) graph_outputs = {} for output in self.graph_outputs: - if '/' in output: - node_name, field_name = output.split('/', 1) + if "/" in output: + node_name, field_name = output.split("/", 1) if node_name not in graph_outputs: graph_outputs[node_name] = CompactYamlList() graph_outputs[node_name].append(field_name) - pipeline = {'pipeline': {'data_flow': processed_connections, - 'feedback_flow': processed_feedback_connections, - 'inputs': graph_inputs, - 'outputs': graph_outputs}} + pipeline = { + "pipeline": { + "data_flow": processed_connections, + "feedback_flow": processed_feedback_connections, + "inputs": graph_inputs, + "outputs": graph_outputs, + } + } return pipeline @@ -141,25 +150,32 @@ def visualize(self, save_path, graph_name): def get_graph_statistics(self): internal_connections = 0 for connection in self.connections: - internal_connections += len(connection['targets']) + internal_connections += len(connection["targets"]) for connection in self.feedback_connections: - internal_connections += len(connection['targets']) + internal_connections += len(connection["targets"]) - total_edges = internal_connections + \ - len(self.graph_inputs) + len(self.graph_outputs) + total_edges = ( + internal_connections + len(self.graph_inputs) + len(self.graph_outputs) + ) return internal_connections, total_edges def _finalize_connections(self, connections_dict): processed_connections = {} for connection in connections_dict: - source = connection['source'] - targets = connection['targets'] - source_port = source['node'].name + '/' + \ - source['node'].outputs[source['idx']].name_str + source = connection["source"] + targets = connection["targets"] + source_port = ( + source["node"].name + + "/" + + source["node"].outputs[source["idx"]].name_str + ) target_ports = CompactYamlList() for target in targets: - target_ports.append(target['node'].name + '/' + - target['node'].inputs[target['idx']].name_str) + target_ports.append( + target["node"].name + + "/" + + target["node"].inputs[target["idx"]].name_str + ) processed_connections[source_port] = target_ports return processed_connections @@ -175,7 +191,8 @@ def _resolve_source_output(self, node, input_desc): _get_logger().error( f"Error: input {node.name}/{input_desc.name_str} came from node " f"'{source_node.name}', which is not part of this graph.\n" - f"Treating it as a dangling input.") + f"Treating it as a dangling input." + ) return None for out_idx, output in enumerate(source_node.outputs): @@ -187,7 +204,8 @@ def _resolve_source_output(self, node, input_desc): f"'{input_desc.port}' from node '{source_node.name}', " f"which does not produce it.\n" f"Available outputs: {[o.port for o in source_node.outputs]}\n" - f"Treating it as a dangling input.") + f"Treating it as a dangling input." + ) return None def _build_state_connections(self, nodes, feedback_connections): @@ -197,6 +215,7 @@ def _build_state_connections(self, nodes, feedback_connections): its own inputs, so it is always a feedback edge and never depends on what value happened to cross the boundary. """ + def index_of(descriptions, description): for idx, candidate in enumerate(descriptions): if candidate is description: @@ -212,19 +231,20 @@ def index_of(descriptions, description): _get_logger().error( f"Error: state pair '{input_desc.name_str}' of node " f"'{node.name}' is not registered on the node, so its " - f"feedback edge cannot be built.") + f"feedback edge cannot be built." + ) continue claimed.add(id(input_desc)) key = (node, output_desc.port) - bucket = feedback_connections.setdefault(key, { - 'source': {'node': node, 'idx': out_idx}, - 'targets': [] - }) - bucket['targets'].append({'node': node, 'idx': in_idx}) + bucket = feedback_connections.setdefault( + key, {"source": {"node": node, "idx": out_idx}, "targets": []} + ) + bucket["targets"].append({"node": node, "idx": in_idx}) _get_logger().info( f"Found state feedback connection: " f"{node.name}/{output_desc.name_str} " - f" -> {node.name}/{input_desc.name_str}") + f" -> {node.name}/{input_desc.name_str}" + ) return claimed def _build_connections(self, nodes): @@ -234,18 +254,23 @@ def _build_connections(self, nodes): for node in nodes.values(): # A node cannot receive the same source output on two of its inputs, # because the runtime has one value to deliver and two places to put it. - sources = [(input.source_node, input.port) - for input in node.inputs - if input.has_source and id(input) not in state_inputs] - duplicates = {source for source, - count in Counter(sources).items() if count > 1} + sources = [ + (input.source_node, input.port) + for input in node.inputs + if input.has_source and id(input) not in state_inputs + ] + duplicates = { + source for source, count in Counter(sources).items() if count > 1 + } if duplicates: duplicate_list = ", ".join( - sorted(f"{source.name}/{port}" for source, port in duplicates)) + sorted(f"{source.name}/{port}" for source, port in duplicates) + ) _get_logger().fatal( f"Error: unsupported use of sending the same tensor multiple times to the same node. " f"Duplicate input sources in node {node.name}: {duplicate_list}", - error_type=Exception) + error_type=Exception, + ) for in_idx, input in enumerate(node.inputs): if id(input) in state_inputs: # already wired from its declaration @@ -265,14 +290,15 @@ def _build_connections(self, nodes): key = (source_node, input.port) if key not in bucket: bucket[key] = { - 'source': {'node': source_node, 'idx': out_idx}, - 'targets': [] + "source": {"node": source_node, "idx": out_idx}, + "targets": [], } - bucket[key]['targets'].append({'node': node, 'idx': in_idx}) + bucket[key]["targets"].append({"node": node, "idx": in_idx}) _get_logger().info( f"Found {'connection' if is_forward else 'feedback connection'}: " f"{source_node.name}/{source_node.outputs[out_idx].name_str} " - f" -> {node.name}/{node.inputs[in_idx].name_str}") + f" -> {node.name}/{node.inputs[in_idx].name_str}" + ) return list(connections.values()), list(feedback_connections.values()) @@ -284,29 +310,41 @@ def _compile_graph_io(self, nodes, connections, feedback_connections): # Collect all target ports (destinations) from both forward and feedback connections all_target_ports = set() for connection in connections: - for target in connection['targets']: - target_port = target['node'].name + '/' + \ - target['node'].inputs[target['idx']].name_str + for target in connection["targets"]: + target_port = ( + target["node"].name + + "/" + + target["node"].inputs[target["idx"]].name_str + ) all_target_ports.add(target_port) for connection in feedback_connections: - for target in connection['targets']: - target_port = target['node'].name + '/' + \ - target['node'].inputs[target['idx']].name_str + for target in connection["targets"]: + target_port = ( + target["node"].name + + "/" + + target["node"].inputs[target["idx"]].name_str + ) all_target_ports.add(target_port) # Collect all source ports from both forward and feedback connections all_source_ports = set() for connection in connections: - source = connection['source'] - source_port = source['node'].name + '/' + \ - source['node'].outputs[source['idx']].name_str + source = connection["source"] + source_port = ( + source["node"].name + + "/" + + source["node"].outputs[source["idx"]].name_str + ) all_source_ports.add(source_port) for connection in feedback_connections: - source = connection['source'] - source_port = source['node'].name + '/' + \ - source['node'].outputs[source['idx']].name_str + source = connection["source"] + source_port = ( + source["node"].name + + "/" + + source["node"].outputs[source["idx"]].name_str + ) all_source_ports.add(source_port) # Find dangling inputs and outputs @@ -314,14 +352,14 @@ def _compile_graph_io(self, nodes, connections, feedback_connections): # An input is dangling if it's not the target of any internal connection for input_desc in node.inputs: input_name = input_desc.name_str - node_input = node.name + '/' + input_name + node_input = node.name + "/" + input_name if node_input not in all_target_ports: graph_inputs.append(node_input) # An output is dangling if it's not the source of any internal connection for output_desc in node.outputs: output_name = output_desc.name_str - node_output = node.name + '/' + output_name + node_output = node.name + "/" + output_name if node_output not in all_source_ports: graph_outputs.append(node_output) @@ -331,12 +369,12 @@ def _compile_graph_io(self, nodes, connections, feedback_connections): def _validate_connection_compatibility(self): for connection in self.connections + self.feedback_connections: - source = connection['source'] - source_desc = source['node'].outputs[source['idx']] + source = connection["source"] + source_desc = source["node"].outputs[source["idx"]] source_name = f"{source['node'].name}/{source_desc.name_str}" - for target in connection['targets']: - target_desc = target['node'].inputs[target['idx']] + for target in connection["targets"]: + target_desc = target["node"].inputs[target["idx"]] target_name = f"{target['node'].name}/{target_desc.name_str}" validate_connection_compatibility( source_name=source_name, diff --git a/leapp/leapp_graph/leapp_node.py b/leapp/leapp_graph/leapp_node.py index 7b457bfa..78fbb26e 100644 --- a/leapp/leapp_graph/leapp_node.py +++ b/leapp/leapp_graph/leapp_node.py @@ -30,12 +30,12 @@ from leapp.utils.logging import _get_logger -class LeappNode(): +class LeappNode: def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) - if '__init__' in cls.__dict__: + if "__init__" in cls.__dict__: original_init = cls.__init__ - if getattr(original_init, '_leapp_wrapped', False): + if getattr(original_init, "_leapp_wrapped", False): return @functools.wraps(original_init) @@ -44,13 +44,13 @@ def wrapped_init(self, *args, **kwargs): # Only log if this is the actual class being instantiated # (not a parent's __init__ being called via super()) if type(self) is cls: - _get_logger().info( - f"Node context initialized: {self.name}") + _get_logger().info(f"Node context initialized: {self.name}") + wrapped_init._leapp_wrapped = True cls.__init__ = wrapped_init # Wrap compile_trace to set _model_captured = True after execution - if 'compile_trace' in cls.__dict__: + if "compile_trace" in cls.__dict__: original_compile_trace = cls.compile_trace @functools.wraps(original_compile_trace) @@ -58,6 +58,7 @@ def wrapped_compile_trace(self, *args, **kwargs): result = original_compile_trace(self, *args, **kwargs) self._model_captured = True return result + cls.compile_trace = wrapped_compile_trace UNSET_NODE_INDEX = -1 @@ -65,7 +66,7 @@ def wrapped_compile_trace(self, *args, **kwargs): def __init__(self, name): self.name = name self._node_index = self.UNSET_NODE_INDEX - + # Attributes expected by export backends (subclasses may override) self.register_buffers = set() self.environment_constants = set() @@ -105,7 +106,8 @@ def node_index(self, value): _get_logger().fatal( f"Node index for '{self.name}' is already set to {self._node_index} " "and cannot be reassigned.", - error_type=Exception) + error_type=Exception, + ) self._node_index = value @property @@ -126,7 +128,7 @@ def exports_model(self) -> bool: Work that only feeds export can be skipped when this is ``False``. """ return self.backend not in (None, "None") - + @property def compiled_model(self): if self.export_backend is None: @@ -134,7 +136,7 @@ def compiled_model(self): if self.export_backend.compiled_model is None: return None return self.export_backend.compiled_model - + @property def compiled_module(self): if self.export_backend is None: @@ -142,7 +144,7 @@ def compiled_module(self): if self.export_backend.compiled_module is None: return None return self.export_backend.compiled_module - + def delete_compiled_model(self): if self.export_backend is None: return @@ -160,12 +162,13 @@ def get_description(self): output_descriptions = [output.dict() for output in self.outputs] parameters = { - 'model_path': self.model_path, - 'md5sum': self.md5sum, - 'sha256sum': self.sha256sum, - 'backend': self.get_backend(), - 'warp_segments': len(iter_warp_segments_from_graph(self.graph)) - if hasattr(self, 'graph') else 0, + "model_path": self.model_path, + "md5sum": self.md5sum, + "sha256sum": self.sha256sum, + "backend": self.get_backend(), + "warp_segments": len(iter_warp_segments_from_graph(self.graph)) + if hasattr(self, "graph") + else 0, } backend_metadata = self.export_backend.get_backend_metadata() @@ -173,9 +176,9 @@ def get_description(self): parameters.update(backend_metadata) description = {} - description['inputs'] = input_descriptions - description['outputs'] = output_descriptions - description['parameters'] = parameters + description["inputs"] = input_descriptions + description["outputs"] = output_descriptions + description["parameters"] = parameters return description @@ -184,38 +187,44 @@ def setup_backend(self, backend, backend_params): def _create_backend(self, backend, backend_params): available_backends = [ - "jit-script", "jit-trace", "onnx-dynamo", "onnx-torchscript", + "jit-script", + "jit-trace", + "onnx-dynamo", + "onnx-torchscript", "exported-program", ] if backend is None: self.backend = "None" - self.export_backend = NoneExportBackend( - self, backend_params) - elif backend == "jit-script" or backend == "jit": # default jit export method + self.export_backend = NoneExportBackend(self, backend_params) + elif backend == "jit-script" or backend == "jit": # default jit export method self.backend = "jit-script" from leapp.backends.torch_export_backend import TorchScriptExportBackend - self.export_backend = TorchScriptExportBackend( - self, backend_params) + + self.export_backend = TorchScriptExportBackend(self, backend_params) elif backend == "jit-trace": self.backend = "jit-trace" from leapp.backends.torch_export_backend import TorchTraceExportBackend - self.export_backend = TorchTraceExportBackend( - self, backend_params) - elif backend == "onnx-dynamo" or backend == "onnx": #default onnx export method + + self.export_backend = TorchTraceExportBackend(self, backend_params) + elif ( + backend == "onnx-dynamo" or backend == "onnx" + ): # default onnx export method self.backend = "onnx-dynamo" from leapp.backends.onnx_export_backend import ONNXDynamoExportBackend - self.export_backend = ONNXDynamoExportBackend( - self, backend_params) + + self.export_backend = ONNXDynamoExportBackend(self, backend_params) elif backend == "onnx-torchscript": self.backend = "onnx-torchscript" from leapp.backends.onnx_export_backend import ONNXTorchScriptExportBackend - self.export_backend = ONNXTorchScriptExportBackend( - self, backend_params) + + self.export_backend = ONNXTorchScriptExportBackend(self, backend_params) elif backend == "exported-program" or backend == "pt2": self.backend = "exported-program" - from leapp.backends.exported_program_backend import ExportedProgramExportBackend - self.export_backend = ExportedProgramExportBackend( - self, backend_params) + from leapp.backends.exported_program_backend import ( + ExportedProgramExportBackend, + ) + + self.export_backend = ExportedProgramExportBackend(self, backend_params) elif backend == "cpp": _get_logger().fatal("C++ backend not implemented", error_type=Exception) elif backend == "py": @@ -224,10 +233,13 @@ def _create_backend(self, backend, backend_params): _get_logger().fatal( f"{self.name} Unexpected backend: {backend}, \n" f"please use one of the following: {available_backends}", - error_type=Exception) + error_type=Exception, + ) def save_model(self, save_path: str): - self.model_path, self.md5sum, self.sha256sum = self.export_backend.save(save_path) + self.model_path, self.md5sum, self.sha256sum = self.export_backend.save( + save_path + ) def compile_model(self): try: @@ -241,10 +253,13 @@ def compile_model(self): seen = set() while error is not None and id(error) not in seen: seen.add(id(error)) - if any(message in str(error) for message in ( - "Inference tensors cannot be saved for backward", - "Setting requires_grad=True on inference tensor outside InferenceMode", - )): + if any( + message in str(error) + for message in ( + "Inference tensors cannot be saved for backward", + "Setting requires_grad=True on inference tensor outside InferenceMode", + ) + ): _get_logger().fatal( f"LEAPP could not export node '{self.name}' with backend " f"'{self.backend}' because PyTorch rejected an inference tensor. " @@ -253,20 +268,19 @@ def compile_model(self): "Use torch.no_grad() instead in those paths. Recreate affected " "tensors outside inference mode; leaving inference mode does " "not convert existing tensors.", - cause=e) + cause=e, + ) error = error.__cause__ or error.__context__ _get_logger().fatal( - f"Error compiling model {self.name}: {e}", - error_type=type(e), - cause=e) + f"Error compiling model {self.name}: {e}", error_type=type(e), cause=e + ) def get_backend(self): return self.export_backend.get_backend_model_type() @property def is_tracing(self) -> bool: - """Whether this node is currently recording operations. - """ + """Whether this node is currently recording operations.""" return False def publish_output_port(self, value, port: str): @@ -289,7 +303,8 @@ def publish_output_port(self, value, port: str): _get_logger().debug( f"Node '{self.name}' publishes one value as both output " f"'{existing}' and output '{port}'; consumers of it will " - f"connect through '{existing}'.") + f"connect through '{existing}'." + ) return value value.output_port = port return value @@ -306,7 +321,8 @@ def publish_output_ports(self, value, name: str, descriptions): _get_logger().error( f"Error: output '{name}' of node '{self.name}' flattened to " f"{len(leaves)} values but {len(descriptions)} descriptions; " - "skipping output-port registration for it.") + "skipping output-port registration for it." + ) return for leaf, description in zip(leaves, descriptions): self.publish_output_port(leaf, description.port) @@ -320,7 +336,8 @@ def _validate_and_add_to_list(descriptions, current_io_list, node_name): f"Duplicate i/o name '{description.name_str}' in node '{node_name}'. \n" f"Currently existing names: {existing_names}\n" f"Each input/output in the same node must have a unique name.", - error_type=Exception) + error_type=Exception, + ) existing_names.add(description.name_str) current_io_list.append(description) @@ -334,20 +351,24 @@ def _apply_semantics_to_descriptions(io_descriptions, semantics): semantic_fields = semantics.to_dict() for desc in io_descriptions: - desc.init_semantics(TensorSemantics( - name=desc.name, - ref=desc.value, - **semantic_fields, - )) + desc.init_semantics( + TensorSemantics( + name=desc.name, + ref=desc.value, + **semantic_fields, + ) + ) def add_output(self, output_name, raw_output_name, output_value, semantics=None): io_descriptions, output_format = describe_io( - output_name, raw_output_name, output_value) + output_name, raw_output_name, output_value + ) self._apply_semantics_to_descriptions(io_descriptions, semantics) for desc in io_descriptions: - desc.cached_values = [torch.zeros_like(desc.value) for _ in range(self._max_cached_io)] - self._validate_and_add_to_list( - io_descriptions, self.outputs, self.name) + desc.cached_values = [ + torch.zeros_like(desc.value) for _ in range(self._max_cached_io) + ] + self._validate_and_add_to_list(io_descriptions, self.outputs, self.name) for desc in io_descriptions: # An output is its own source. The name was just proven unique on # this node, which is what makes it usable as the port identity. @@ -359,10 +380,13 @@ def add_output(self, output_name, raw_output_name, output_value, semantics=None) def add_input(self, input_name, raw_input_name, input_value, semantics=None): io_descriptions, input_format = describe_io( - input_name, raw_input_name, input_value) + input_name, raw_input_name, input_value + ) self._apply_semantics_to_descriptions(io_descriptions, semantics) for desc in io_descriptions: - desc.cached_values = [torch.zeros_like(desc.value) for _ in range(self._max_cached_io)] + desc.cached_values = [ + torch.zeros_like(desc.value) for _ in range(self._max_cached_io) + ] self._validate_and_add_to_list(io_descriptions, self.inputs, self.name) # used to rebuild the nested i/o self.input_formats.append(input_format) @@ -384,26 +408,33 @@ def get_io_description_by_name(name, io_list): return io_description return None - def validate_io_and_update_sources(self, io_name, raw_io_name, io_value, current_io_list): - ''' + def validate_io_and_update_sources( + self, io_name, raw_io_name, io_value, current_io_list + ): + """ this is used for rerunning the the tracing. each time we run it we we want to validate that the inputs are consistent with the previous run and also resolve sources that only appear on a later pass, which is how feedback edges are discovered. - ''' + """ io_descriptions, _ = describe_io(io_name, raw_io_name, io_value) for io_description in io_descriptions: if io_description.name_str in self.trimmed_inputs: # we skip validation for inputs that are not used in the model continue existing_io_description = LeappNode.get_io_description_by_name( - io_description.name_str, current_io_list) + io_description.name_str, current_io_list + ) if existing_io_description is None: _get_logger().fatal( f"Error: Reentering {self.name} with new i/o {io_description.name_str} but failed to find it in the current i/o list.\n" f"available i/o names: {[io.name_str for io in current_io_list]}", - error_type=Exception) - elif current_io_list is self.inputs and not existing_io_description.has_source: + error_type=Exception, + ) + elif ( + current_io_list is self.inputs + and not existing_io_description.has_source + ): # A previously dangling input may acquire its source on a later # pass; this is how feedback edges become visible. Only an input # can gain a source this way, because an output is its own @@ -415,25 +446,37 @@ def validate_io_and_update_sources(self, io_name, raw_io_name, io_value, current current_io_description_dict = io_description.dict() current_io_description_dict["name"] = existing_io_description.name - common_keys = existing_io_description_dict.keys() & current_io_description_dict.keys() - if not all(existing_io_description_dict[key] == current_io_description_dict[key] for key in common_keys): + common_keys = ( + existing_io_description_dict.keys() & current_io_description_dict.keys() + ) + if not all( + existing_io_description_dict[key] == current_io_description_dict[key] + for key in common_keys + ): _get_logger().fatal( f"Error: Reentering {self.name} with new i/o {io_description.name_str} \n" f"but the description has changed from {existing_io_description_dict} to {current_io_description_dict}.\n" f"This can happen if some dynamic behavior is not captured by the annotations", - error_type=Exception) + error_type=Exception, + ) if self._cache_write_idx < self._max_cached_io: - existing_io_description.cached_values[self._cache_write_idx] = io_description.value + existing_io_description.cached_values[self._cache_write_idx] = ( + io_description.value + ) def increment_cache_idx(self): if self._cache_write_idx < self._max_cached_io: self._cache_write_idx += 1 def _build_validation_examples(self): - examples = [(0, - [td.value for td in self.inputs], - tuple(td.value for td in self.outputs))] + examples = [ + ( + 0, + [td.value for td in self.inputs], + tuple(td.value for td in self.outputs), + ) + ] for cache_idx in range(self._cache_write_idx): cached_inputs = [td.cached_values[cache_idx] for td in self.inputs] @@ -457,10 +500,13 @@ def _reentry_validation_hint(self, sample_passed: dict) -> str | None: "it with annotate.input_tensors() for this node." ) - def validate_compiled_model(self, rtol: float = 1e-3, atol: float = 1e-5) -> tuple[bool, str | None]: + def validate_compiled_model( + self, rtol: float = 1e-3, atol: float = 1e-5 + ) -> tuple[bool, str | None]: if self.compiled_model is None: _get_logger().warning( - f"Model {self.name} does not have a compiled model. Skipping validation.") + f"Model {self.name} does not have a compiled model. Skipping validation." + ) return True, None examples = self._build_validation_examples() @@ -477,14 +523,20 @@ def validate_compiled_model(self, rtol: float = 1e-3, atol: float = 1e-5) -> tup if len(exported_outputs) != len(source_outputs): _get_logger().error( f"{self.name} sample {sample_idx}: Output count mismatch - " - f"got {len(exported_outputs)}, expected {len(source_outputs)}") + f"got {len(exported_outputs)}, expected {len(source_outputs)}" + ) all_match = False sample_passed[sample_idx] = False continue - for idx, (exported, source) in enumerate(zip(exported_outputs, source_outputs)): - output_name = self.outputs[idx].name if idx < len( - self.outputs) else f"output_{idx}" + for idx, (exported, source) in enumerate( + zip(exported_outputs, source_outputs) + ): + output_name = ( + self.outputs[idx].name + if idx < len(self.outputs) + else f"output_{idx}" + ) if exported.device != source.device: exported = exported.to(source.device) @@ -494,24 +546,34 @@ def validate_compiled_model(self, rtol: float = 1e-3, atol: float = 1e-5) -> tup source_nan = torch.isnan(source).sum().item() source_inf = torch.isinf(source).sum().item() - if exported_nan > 0 or exported_inf > 0 or source_nan > 0 or source_inf > 0: + if ( + exported_nan > 0 + or exported_inf > 0 + or source_nan > 0 + or source_inf > 0 + ): all_match = False sample_passed[sample_idx] = False num_elements = exported.numel() _get_logger().error( - f"{self.name}/{output_name} sample {sample_idx}: NaN/Inf detected!") + f"{self.name}/{output_name} sample {sample_idx}: NaN/Inf detected!" + ) if exported_nan > 0: _get_logger().error( - f" Exported has {exported_nan}/{num_elements} NaN values ({100*exported_nan/num_elements:.3f}%)") + f" Exported has {exported_nan}/{num_elements} NaN values ({100 * exported_nan / num_elements:.3f}%)" + ) if exported_inf > 0: _get_logger().error( - f" Exported has {exported_inf}/{num_elements} Inf values ({100*exported_inf/num_elements:.3f}%)") + f" Exported has {exported_inf}/{num_elements} Inf values ({100 * exported_inf / num_elements:.3f}%)" + ) if source_nan > 0: _get_logger().warning( - f" Source has {source_nan}/{num_elements} NaN values ({100*source_nan/num_elements:.3f}%)") + f" Source has {source_nan}/{num_elements} NaN values ({100 * source_nan / num_elements:.3f}%)" + ) if source_inf > 0: _get_logger().warning( - f" Source has {source_inf}/{num_elements} Inf values ({100*source_inf/num_elements:.3f}%)") + f" Source has {source_inf}/{num_elements} Inf values ({100 * source_inf / num_elements:.3f}%)" + ) continue if not torch.allclose(exported, source, rtol=rtol, atol=atol): @@ -523,28 +585,40 @@ def validate_compiled_model(self, rtol: float = 1e-3, atol: float = 1e-5) -> tup max_diff = diff.max().item() mean_diff = diff.mean().item() - percentiles = torch.tensor([0.50, 0.75, 0.90, 0.99, 0.995], device=diff_flat.device) + percentiles = torch.tensor( + [0.50, 0.75, 0.90, 0.99, 0.995], device=diff_flat.device + ) pct_values = torch.quantile(diff_flat, percentiles) p50, p75, p90, p99, p995 = pct_values.tolist() source_min, source_max = source.min().item(), source.max().item() - exported_min, exported_max = exported.min().item(), exported.max().item() + exported_min, exported_max = ( + exported.min().item(), + exported.max().item(), + ) log_path = _get_logger().path _get_logger().error( - f"{self.name}/{output_name} sample {sample_idx}: Mismatch detected (rtol={rtol}, atol={atol}). Please check {log_path} for more details.") + f"{self.name}/{output_name} sample {sample_idx}: Mismatch detected (rtol={rtol}, atol={atol}). Please check {log_path} for more details." + ) _get_logger().info( - f" Source shape: {source.shape}, dtype: {source.dtype}") + f" Source shape: {source.shape}, dtype: {source.dtype}" + ) _get_logger().info( - f" Exported shape: {exported.shape}, dtype: {exported.dtype}") + f" Exported shape: {exported.shape}, dtype: {exported.dtype}" + ) _get_logger().info( - f" Source range: [{source_min:.6e}, {source_max:.6e}]") + f" Source range: [{source_min:.6e}, {source_max:.6e}]" + ) _get_logger().info( - f" Exported range: [{exported_min:.6e}, {exported_max:.6e}]") + f" Exported range: [{exported_min:.6e}, {exported_max:.6e}]" + ) _get_logger().info( - f" Diff stats: max={max_diff:.6e}, mean={mean_diff:.6e}") + f" Diff stats: max={max_diff:.6e}, mean={mean_diff:.6e}" + ) _get_logger().info( - f" Diff percentiles: p50={p50:.6e}, p75={p75:.6e}, p90={p90:.6e}, p99={p99:.6e}, p995={p995:.6e}") + f" Diff percentiles: p50={p50:.6e}, p75={p75:.6e}, p90={p90:.6e}, p99={p99:.6e}, p995={p995:.6e}" + ) error_hint = self._reentry_validation_hint(sample_passed) if error_hint is not None: @@ -553,25 +627,32 @@ def validate_compiled_model(self, rtol: float = 1e-3, atol: float = 1e-5) -> tup if all_match: num_examples = len(examples) _get_logger().info( - f" [PASS] {self.name} passed validation ({num_examples} example{'s' if num_examples > 1 else ''})") + f" [PASS] {self.name} passed validation ({num_examples} example{'s' if num_examples > 1 else ''})" + ) return all_match, error_hint - def validate_input_and_update_sources(self, input_name, raw_input_name, input_value): + def validate_input_and_update_sources( + self, input_name, raw_input_name, input_value + ): self.validate_io_and_update_sources( - input_name, raw_input_name, input_value, self.inputs) + input_name, raw_input_name, input_value, self.inputs + ) - def validate_output_and_update_sources(self, output_name, raw_output_name, output_value): + def validate_output_and_update_sources( + self, output_name, raw_output_name, output_value + ): self.validate_io_and_update_sources( - output_name, raw_output_name, output_value, self.outputs) + output_name, raw_output_name, output_value, self.outputs + ) def reentry_validate_inputs(self, tensors: dict): """Validate input tensors on re-entry (node already compiled).""" for tensor_name, tensor in tensors.items(): - self.validate_input_and_update_sources( - tensor_name, tensor_name, tensor) + self.validate_input_and_update_sources(tensor_name, tensor_name, tensor) - def reentry_validate_outputs(self, tensors: dict, - static_tensors: dict | None = None): + def reentry_validate_outputs( + self, tensors: dict, static_tensors: dict | None = None + ): """Validate output tensors on re-entry, then advance the cache index. Ports are not re-published here. The descriptions registered on the @@ -582,20 +663,19 @@ def reentry_validate_outputs(self, tensors: dict, """ all_tensors = {**tensors, **(static_tensors or {})} for tensor_name, tensor in all_tensors.items(): - self.validate_output_and_update_sources( - tensor_name, tensor_name, tensor) + self.validate_output_and_update_sources(tensor_name, tensor_name, tensor) self.increment_cache_idx() def reentry_validate_state_update(self, tensors: dict): """Validate state tensor updates on re-entry.""" for tensor_name, tensor in tensors.items(): - self.validate_output_and_update_sources( - tensor_name, tensor_name, tensor) + self.validate_output_and_update_sources(tensor_name, tensor_name, tensor) def change_input_name(self, old_name, new_name): _get_logger().debug( - f"changing input name from {old_name} to {new_name} for model {self.name}") + f"changing input name from {old_name} to {new_name} for model {self.name}" + ) if old_name == new_name: return current_input_names = [input.name_str for input in self.inputs] @@ -611,7 +691,8 @@ def change_input_name(self, old_name, new_name): def change_output_name(self, old_name, new_name): _get_logger().debug( - f"changing output name from {old_name} to {new_name} for model {self.name}") + f"changing output name from {old_name} to {new_name} for model {self.name}" + ) if old_name == new_name: return current_output_names = [output.name_str for output in self.outputs] @@ -627,4 +708,5 @@ def change_output_name(self, old_name, new_name): def compile_trace(self, *args): raise NotImplementedError( - f"compile_trace not implemented for {self.__class__.__name__}") + f"compile_trace not implemented for {self.__class__.__name__}" + ) diff --git a/leapp/leapp_graph/traced_node.py b/leapp/leapp_graph/traced_node.py index b7844364..03300e95 100644 --- a/leapp/leapp_graph/traced_node.py +++ b/leapp/leapp_graph/traced_node.py @@ -31,6 +31,7 @@ resolve_tensor_descriptions_to_names, flatten_io_structure, ) + # Importing registers the ``leapp::warp_runner`` custom op (import side effect) # and exposes ``get_op``/``QUALIFIED_NAME`` used when emitting segment nodes. from leapp.leapp_graph.custom_operator_registry import warp_operator @@ -40,6 +41,7 @@ if TYPE_CHECKING: from leapp.leapp_graph.datatypes.warp import WarpSegment + # fx graph overloads to create custom behavior when interacting with the graph. # keep local so that only the tracedTensorNode can create it. class _LeappFXGraph(fx.Graph): @@ -55,7 +57,9 @@ def create_node(self, *args, **kwargs): class TracedTensorNode(LeappNode): def __init__(self, name, *args, **kwargs): if args or kwargs: - _get_logger().warning(f"{name} received unexpected arguments on initialization. these arguments will be ignored.") + _get_logger().warning( + f"{name} received unexpected arguments on initialization. these arguments will be ignored." + ) super().__init__(name) self.graph = _LeappFXGraph() self.tracer = fx.Tracer() @@ -120,7 +124,10 @@ def add_warp_segment(self, segment: "WarpSegment") -> None: self._pending_warp_segments.append(segment) def complete_warp_segment(self, segment: "WarpSegment") -> None: - if not self._pending_warp_segments or self._pending_warp_segments[0] is not segment: + if ( + not self._pending_warp_segments + or self._pending_warp_segments[0] is not segment + ): _get_logger().fatal( f"[{self.name}] Captured Warp regions out of discovery order.", error_type=RuntimeError, @@ -195,10 +202,15 @@ def _validate_declared_input_aliases(self, outputs: dict) -> None: error_type=Exception, ) - def compile_trace(self, tensors: dict[str, "TracedTensor"], backend=None, backend_params={}, - static_tensors: dict[str, torch.Tensor] | None = None, - semantics_map: dict[str, object] | None = None, - static_semantics_map: dict[str, object] | None = None): + def compile_trace( + self, + tensors: dict[str, "TracedTensor"], + backend=None, + backend_params={}, + static_tensors: dict[str, torch.Tensor] | None = None, + semantics_map: dict[str, object] | None = None, + static_semantics_map: dict[str, object] | None = None, + ): # Auto-collect buffer mutations before merging state outputs. # collect() may remove non-mutated buffers from _state_tensors and # bake them as constants, so it must run before get_state_outputs(). @@ -209,7 +221,9 @@ def compile_trace(self, tensors: dict[str, "TracedTensor"], backend=None, backen # Merge state outputs into the tensor dict before processing state_outputs = self.get_state_outputs() if state_outputs: - _get_logger().info(f"Adding {len(state_outputs)} state outputs: {list(state_outputs.keys())}") + _get_logger().info( + f"Adding {len(state_outputs)} state outputs: {list(state_outputs.keys())}" + ) tensors = {**tensors, **state_outputs} # Held for the post-pruning alias check below. Static outputs join @@ -219,11 +233,13 @@ def compile_trace(self, tensors: dict[str, "TracedTensor"], backend=None, backen unwrapped_tensors = [] for name, tensor in tensors.items(): - unwrapped = self.create_output(tensor, name, semantics=(semantics_map or {}).get(name)) + unwrapped = self.create_output( + tensor, name, semantics=(semantics_map or {}).get(name) + ) unwrapped_tensors.append(unwrapped) if name in self._state_tensors: - self._state_tensors[name]["output_desc"] = self._get_required_io_description( - name, self.outputs, "output" + self._state_tensors[name]["output_desc"] = ( + self._get_required_io_description(name, self.outputs, "output") ) # Static tensors go through the same create_output path with static=True @@ -240,8 +256,7 @@ def compile_trace(self, tensors: dict[str, "TracedTensor"], backend=None, backen self.build_graph_module(list(tensors.values())) self._validate_declared_input_aliases(traced_outputs) - self.m = fx.GraphModule( - self.tracer.root, self.graph) + self.m = fx.GraphModule(self.tracer.root, self.graph) # Sanitize the FX graph for compatibility with export backends. self._rewrite_aten_ops(self.m.graph) @@ -249,15 +264,20 @@ def compile_trace(self, tensors: dict[str, "TracedTensor"], backend=None, backen self._make_tensor_attrs_contiguous(self.m) self.m.recompile() + _get_logger().debug(f"Compiled graph module for {self.name}:\n{self.m.graph}") _get_logger().debug( - f"Compiled graph module for {self.name}:\n{self.m.graph}") - _get_logger().debug( - f"Graph module inputs: {[resolve_tensor_descriptions_to_names(input) for input in self.input_formats]}") + f"Graph module inputs: {[resolve_tensor_descriptions_to_names(input) for input in self.input_formats]}" + ) _get_logger().debug( - f"Graph module outputs: {[resolve_tensor_descriptions_to_names(output) for output in self.output_formats]}") + f"Graph module outputs: {[resolve_tensor_descriptions_to_names(output) for output in self.output_formats]}" + ) self.setup_backend(backend, backend_params) - return unwrapped_tensors[0] if len(unwrapped_tensors) == 1 else tuple(unwrapped_tensors) + return ( + unwrapped_tensors[0] + if len(unwrapped_tensors) == 1 + else tuple(unwrapped_tensors) + ) def state_feedback_pairs(self): """State tensors are an eager form of feedback. @@ -300,8 +320,8 @@ def _rewrite_method_descriptors(graph: fx.Graph): target_type_name = type(target).__name__ target_module = getattr(target, "__module__", "") is_tensor_method_target = ( - target_type_name == "method_descriptor" or - target_module.startswith("torch._tensor") + target_type_name == "method_descriptor" + or target_module.startswith("torch._tensor") ) if not is_tensor_method_target: continue @@ -314,7 +334,8 @@ def _rewrite_method_descriptors(graph: fx.Graph): if rewritten: _get_logger().debug( f"Rewrote {len(rewritten)} tensor-method call_function node(s) " - f"to call_method: {rewritten}") + f"to call_method: {rewritten}" + ) graph.lint() @staticmethod @@ -336,10 +357,10 @@ def _rewrite_aten_ops(graph: fx.Graph): # key = overloadpacket name, value = method name on the tensor. # ----------------------------------------------------------------- _TO_METHOD = { - 'view': 'view', - 'reshape': 'reshape', - 'contiguous': 'contiguous', - 'to': 'to', + "view": "view", + "reshape": "reshape", + "contiguous": "contiguous", + "to": "to", } # ----------------------------------------------------------------- @@ -352,25 +373,27 @@ def _upsample_bilinear2d_args(args, kwargs): # F.interpolate sig: (input, size=, mode=, align_corners=) inp, size, align_corners = args[0], args[1], args[2] return (inp,), { - 'size': size, 'mode': 'bilinear', - 'align_corners': align_corners, + "size": size, + "mode": "bilinear", + "align_corners": align_corners, } def _upsample_bicubic2d_args(args, kwargs): inp, size, align_corners = args[0], args[1], args[2] return (inp,), { - 'size': size, 'mode': 'bicubic', - 'align_corners': align_corners, + "size": size, + "mode": "bicubic", + "align_corners": align_corners, } def _upsample_nearest2d_args(args, kwargs): inp, size = args[0], args[1] - return (inp,), {'size': size, 'mode': 'nearest'} + return (inp,), {"size": size, "mode": "nearest"} _CUSTOM_FN = { - 'upsample_bilinear2d': (F.interpolate, _upsample_bilinear2d_args), - 'upsample_bicubic2d': (F.interpolate, _upsample_bicubic2d_args), - 'upsample_nearest2d': (F.interpolate, _upsample_nearest2d_args), + "upsample_bilinear2d": (F.interpolate, _upsample_bilinear2d_args), + "upsample_bicubic2d": (F.interpolate, _upsample_bicubic2d_args), + "upsample_nearest2d": (F.interpolate, _upsample_nearest2d_args), } # ----------------------------------------------------------------- @@ -379,7 +402,7 @@ def _upsample_nearest2d_args(args, kwargs): # When step==1 this is equivalent to torch.narrow. # Identity slices (start=0, end>=2^62) are replaced by the input. # ----------------------------------------------------------------- - _SLICE_OPS = {'slice', 'slice_copy'} + _SLICE_OPS = {"slice", "slice_copy"} rewritten = [] nodes_to_erase = [] @@ -408,8 +431,7 @@ def _upsample_nearest2d_args(args, kwargs): if op_name in _CUSTOM_FN: fn, arg_rewriter = _CUSTOM_FN[op_name] if arg_rewriter is not None: - node.args, node.kwargs = arg_rewriter( - node.args, node.kwargs) + node.args, node.kwargs = arg_rewriter(node.args, node.kwargs) node.target = fn rewritten.append(op_name) continue @@ -473,7 +495,8 @@ def _upsample_nearest2d_args(args, kwargs): else: _get_logger().warning( f"Could not find a high-level equivalent for aten op " - f"'{target._name}'. torch.jit.script may fail.") + f"'{target._name}'. torch.jit.script may fail." + ) # Erase identity-slice nodes (must be done after iteration) for node in nodes_to_erase: @@ -482,7 +505,8 @@ def _upsample_nearest2d_args(args, kwargs): if rewritten: _get_logger().debug( f"Rewrote {len(rewritten)} aten OpOverload node(s) to " - f"high-level equivalents: {rewritten}") + f"high-level equivalents: {rewritten}" + ) graph.lint() @staticmethod @@ -510,29 +534,31 @@ def _make_tensor_attrs_contiguous(graph_module: fx.GraphModule): for name, param in list(graph_module.named_parameters()): if not param.is_contiguous(): graph_module.register_parameter( - name, torch.nn.Parameter(param.contiguous(), - requires_grad=param.requires_grad)) + name, + torch.nn.Parameter( + param.contiguous(), requires_grad=param.requires_grad + ), + ) count += 1 if count: _get_logger().debug( - f"Made {count} non-contiguous tensor constant(s) contiguous") + f"Made {count} non-contiguous tensor constant(s) contiguous" + ) def _create_io_helper(self, data, name: str, to: str): if isinstance(data, collections.abc.Mapping): new_data = {} for key, value in data.items(): child_name = "_".join([name, key]) if name else key - new_data[key] = self._create_io_helper( - value, child_name, to) + new_data[key] = self._create_io_helper(value, child_name, to) return new_data elif isinstance(data, (list, tuple)): new_data = [] for idx, value in enumerate(data): child_name = "_".join([name, str(idx)]) if name else str(idx) - new_data.append(self._create_io_helper( - value, child_name, to)) + new_data.append(self._create_io_helper(value, child_name, to)) return type(data)(new_data) - + elif is_tracable_tensor_type(data): is_traced = isinstance(data, TracedData) is_active_traced = is_traced and data.is_tracing @@ -546,17 +572,21 @@ def _create_io_helper(self, data, name: str, to: str): f"detected data '{name}' is an active {data.__class__.__name__} from a different node '{data.context}'. \n" f"Mixing active contexts is not allowed. " f"Call output_tensors() on the source node first.", - error_type=Exception) + error_type=Exception, + ) - if to=="traced": + if to == "traced": # A root that is a placeholder means the value entered its node # as a declaration rather than being computed there. Declaring # promotes the caller's tensor in place, so one external tensor # reaching two nodes arrives here still carrying the first # node's placeholder, which is an ordinary shared input and not # a missing edge between the two. - entered_by_declaration = is_traced and getattr( - getattr(data.proxy, "node", None), "op", None) == "placeholder" + entered_by_declaration = ( + is_traced + and getattr(getattr(data.proxy, "node", None), "op", None) + == "placeholder" + ) if is_active_traced: # Same context: allow override with warning @@ -565,10 +595,13 @@ def _create_io_helper(self, data, name: str, to: str): f"from the same node. Creating fresh input placeholder " f"(previous trace will be discarded for this branch)." ) - elif (self.is_tracing and is_traced - and data.context_obj is not self - and data.output_port is None - and not entered_by_declaration): + elif ( + self.is_tracing + and is_traced + and data.context_obj is not self + and data.output_port is None + and not entered_by_declaration + ): # Came out of a previous node but was never registered as one # of that node's outputs, so there is no edge to connect to. # Only the first pass can say this, because a re-entry pass @@ -581,7 +614,8 @@ def _create_io_helper(self, data, name: str, to: str): f"node '{data.context}' but is not one of its registered outputs.\n" f"Add it to that node's output_tensors() to connect the two nodes, " f"or ignore this if '{name}' is meant to enter the graph from outside.\n" - f"Treating '{name}' as a dangling input.") + f"Treating '{name}' as a dangling input." + ) proxy = None if self.is_tracing: @@ -605,25 +639,28 @@ def _create_io_helper(self, data, name: str, to: str): self._layout_index.setdefault(key, traced) return traced - elif to=="tensor": + elif to == "tensor": if not is_traced: _get_logger().warning( f"Warning: when creating outputs for {self.name}, " - f"detected data {name} is not a TracedTensor") + f"detected data {name} is not a TracedTensor" + ) return data - - elif to=="static": + + elif to == "static": if is_traced: _get_logger().fatal( f"Cannot create static output from traced " f"{type(data).__name__} '{name}'. " "Static outputs must be raw tensors.", - error_type=Exception) + error_type=Exception, + ) # Create unique attribute name and store on root module attr_name = f"_static_{name}".replace(".", "_") self.tracer.root.register_buffer( - attr_name, to_export_torch_tensor(data).clone()) + attr_name, to_export_torch_tensor(data).clone() + ) # Create get_attr node that retrieves the stored constant node = self.graph.create_node("get_attr", attr_name, (), {}) proxy = Proxy(node, self.tracer) @@ -632,7 +669,8 @@ def _create_io_helper(self, data, name: str, to: str): if to == "traced": _get_logger().warning( f"Non-tracable input '{name}' (type={type(data).__name__}) in node '{self.name}' " - f"will be passed through as a constant.") + f"will be passed through as a constant." + ) return data def create_input(self, data, name: str, semantics=None) -> "TracedTensor": @@ -653,16 +691,14 @@ def create_input(self, data, name: str, semantics=None) -> "TracedTensor": getattr(data, "is_tracing", False) and getattr(data, "context_obj", None) is not self ) - adoptable = ( - getattr(data, "output_port", None) is None - and not foreign_active - ) + adoptable = getattr(data, "output_port", None) is None and not foreign_active declared = self.find_declared_alias(data) if adoptable else None if declared is not None: _get_logger().info( f"Input '{name}' for node '{self.name}' shares memory with declared " f"input '{declared.name}'. Both names describe one graph value, so " - f"'{declared.name}' will be the only port in the exported node interface.") + f"'{declared.name}' will be the only port in the exported node interface." + ) # Adoption skips add_input, which is where semantics are normally # attached, so remember the alias and preserve any semantics the # surviving port does not already have. @@ -701,19 +737,18 @@ def create_output(self, data, name: str, static: bool = False, semantics=None): self.publish_output_port(wrapped, existing.port) self.validate_output_and_update_sources(name, name, data) return wrapped - descriptions = self.add_output( - name, name, data, semantics=semantics) + descriptions = self.add_output(name, name, data, semantics=semantics) self.publish_output_ports(wrapped, name, descriptions) return wrapped else: unwrapped_data = self._create_io_helper(data, name, to="tensor") if existing is not None: self.publish_output_port(unwrapped_data, existing.port) - self.validate_output_and_update_sources( - name, name, unwrapped_data) + self.validate_output_and_update_sources(name, name, unwrapped_data) return unwrapped_data descriptions = self.add_output( - name, name, unwrapped_data, semantics=semantics) + name, name, unwrapped_data, semantics=semantics + ) self.publish_output_ports(unwrapped_data, name, descriptions) return unwrapped_data @@ -724,13 +759,15 @@ def _validate_static_tensor(self, tensor, name: str): "but expected a raw tracable data type.\n" "**Static outputs must be raw tensors, not derived from input tensors.**\n" "If this value depends on inputs, use it as a regular output tensor instead.", - error_type=Exception) + error_type=Exception, + ) if is_traced_type(tensor): _get_logger().fatal( f"Error: static output '{name}' is a traced " f"{type(tensor).__name__}. " "Static outputs should be constant tensors, not traced computations.", - error_type=Exception) + error_type=Exception, + ) def create_static_tensors(self, static_outputs): """Wrap raw tensors as static graph nodes (for register_buffer). @@ -739,11 +776,11 @@ def create_static_tensors(self, static_outputs): or register outputs — it only validates and wraps. Returns data in the same nested structure as the input payload. """ - flattened_static_outputs = flatten_io_structure(static_outputs, '') + flattened_static_outputs = flatten_io_structure(static_outputs, "") for tensor_name, tensor in flattened_static_outputs.items(): self._validate_static_tensor(tensor, tensor_name) - return self._create_io_helper(static_outputs, '', to="static") + return self._create_io_helper(static_outputs, "", to="static") def build_graph_module(self, outputs: list["TracedTensor"]) -> fx.GraphModule: """Convert the traced computation to a torch.fx.GraphModule. @@ -785,11 +822,13 @@ def mark_used(node): for node in reversed(nodes_to_remove): if node.op == "placeholder": input_description = LeappNode.get_io_description_by_name( - node.target, self.inputs) + node.target, self.inputs + ) if input_description is None: _get_logger().error( f"Error: when building the graph module for {self.name}, " - f"could not find input description for placeholder '{node.target}'") + f"could not find input description for placeholder '{node.target}'" + ) continue self.trimmed_inputs.add(input_description.name_str) @@ -797,14 +836,17 @@ def mark_used(node): self.graph.erase_node(node) if len(self.trimmed_inputs) > 0: - _get_logger().warning(f"Warning: when building the graph module for {self.name}, " - "detected the following inputs are not used in the computation or directly returned as output: " - f"{self.trimmed_inputs} \n" - "For clarity and efficiency, consider removing these inputs") + _get_logger().warning( + f"Warning: when building the graph module for {self.name}, " + "detected the following inputs are not used in the computation or directly returned as output: " + f"{self.trimmed_inputs} \n" + "For clarity and efficiency, consider removing these inputs" + ) # Remove trimmed inputs from self.inputs self.inputs = [ - inp for inp in self.inputs if inp.name_str not in self.trimmed_inputs] + inp for inp in self.inputs if inp.name_str not in self.trimmed_inputs + ] # After pruning, record which Warp-segment outputs survived (are actually # used) as a binary mask on the segment marker node's metadata. @@ -884,16 +926,21 @@ def _stamp_warp_used_output_masks(self) -> None: output_mask=mask, ) - op_node.update_arg(1, warp_operator.encode_runtime_metadata(runtime_metadata)) + op_node.update_arg( + 1, warp_operator.encode_runtime_metadata(runtime_metadata) + ) - def create_state_tensors(self, tensors: dict[str, torch.Tensor]) -> dict[str, TracedTensor]: + def create_state_tensors( + self, tensors: dict[str, torch.Tensor] + ) -> dict[str, TracedTensor]: """Create state tensors that are both inputs and outputs.""" result = {} for name, tensor in tensors.items(): if name in self._state_tensors: _get_logger().warning( f"State tensor '{name}' already registered for node '{self.name}'. " - "Returning existing TracedTensor.") + "Returning existing TracedTensor." + ) result[name] = self._state_tensors[name]["input"] continue @@ -903,7 +950,9 @@ def create_state_tensors(self, tensors: dict[str, torch.Tensor]) -> dict[str, Tr # Track as state tensor self._state_tensors[name] = { "input": traced_input, - "input_desc": self._get_required_io_description(name, self.inputs, "input"), + "input_desc": self._get_required_io_description( + name, self.inputs, "input" + ), "output": None, # Set via update_state_tensors() "output_desc": None, } @@ -919,25 +968,29 @@ def update_state_tensors(self, tensors: dict[str, TracedTensor]) -> None: _get_logger().fatal( f"Error: update_state called for '{name}' but it was not registered " f"as a state tensor. Call state_tensors() first.", - error_type=Exception) + error_type=Exception, + ) # Validate shape and dtype match the input state input_tensor = self._state_tensors[name]["input"] input_underlying = ( - input_tensor.tensor if isinstance(input_tensor, TracedData) - else input_tensor) - value_underlying = ( - value.tensor if isinstance(value, TracedData) else value) + input_tensor.tensor + if isinstance(input_tensor, TracedData) + else input_tensor + ) + value_underlying = value.tensor if isinstance(value, TracedData) else value if input_underlying.shape != value_underlying.shape: _get_logger().fatal( f"Error: update_state for '{name}' has shape {value_underlying.shape} " f"but expected {input_underlying.shape} (must match input state shape).", - error_type=Exception) + error_type=Exception, + ) if input_underlying.dtype != value_underlying.dtype: _get_logger().fatal( f"Error: update_state for '{name}' has dtype {value_underlying.dtype} " f"but expected {input_underlying.dtype} (must match input state dtype).", - error_type=Exception) + error_type=Exception, + ) self._state_tensors[name]["output"] = value @@ -948,8 +1001,7 @@ def get_state_outputs(self) -> dict[str, TracedTensor]: States declared without an update are treated as regular inputs. """ inactive_states = [ - name for name, info in self._state_tensors.items() - if info["output"] is None + name for name, info in self._state_tensors.items() if info["output"] is None ] if inactive_states: _get_logger().warning( @@ -967,7 +1019,9 @@ def get_state_outputs(self) -> dict[str, TracedTensor]: def _create_warp_bundle_proxy(self, wrp_archive: bytes, runner_name: str) -> Proxy: """Wire a pre-packed WRPB archive as a ``get_attr`` FX node.""" buffer_name = f"_{runner_name}_bundle" - bundle_tensor = torch.frombuffer(bytearray(wrp_archive), dtype=torch.uint8).clone() + bundle_tensor = torch.frombuffer( + bytearray(wrp_archive), dtype=torch.uint8 + ).clone() self.tracer.root.register_buffer(buffer_name, bundle_tensor, persistent=True) bundle_node = self.graph.create_node( "get_attr", @@ -1027,7 +1081,6 @@ def create_warp_proxy( return warp_runner, output_proxies - @property def state_names(self) -> list[str]: """Get list of registered state tensor names.""" diff --git a/leapp/leapp_graph/visualization_adapter.py b/leapp/leapp_graph/visualization_adapter.py index ec92f82a..d7bc67eb 100644 --- a/leapp/leapp_graph/visualization_adapter.py +++ b/leapp/leapp_graph/visualization_adapter.py @@ -20,7 +20,14 @@ from collections.abc import Mapping, Sequence from typing import Any -from leapp_visualization.model import VisualEdge, VisualGraph, VisualNode, VisualPort, VisualTerminal, visual_id +from leapp_visualization.model import ( + VisualEdge, + VisualGraph, + VisualNode, + VisualPort, + VisualTerminal, + visual_id, +) def build_visual_graph( @@ -30,7 +37,9 @@ def build_visual_graph( graph_inputs: Sequence[str], graph_outputs: Sequence[str], ) -> VisualGraph: - ordered_nodes = [_build_node(node) for _, node in sorted(nodes.items(), key=_node_sort_key)] + ordered_nodes = [ + _build_node(node) for _, node in sorted(nodes.items(), key=_node_sort_key) + ] node_name_map = {node.title: node for node in ordered_nodes} terminals: list[VisualTerminal] = [] @@ -43,16 +52,22 @@ def build_visual_graph( edges.extend(_build_internal_edges(connection, "feedback")) for sequence_index, terminal_edge in enumerate(graph_inputs): - terminal, edge = _build_graph_input_terminal_edge(terminal_edge, node_name_map, sequence_index) + terminal, edge = _build_graph_input_terminal_edge( + terminal_edge, node_name_map, sequence_index + ) terminals.append(terminal) edges.append(edge) for sequence_index, terminal_edge in enumerate(graph_outputs): - terminal, edge = _build_graph_output_terminal_edge(terminal_edge, node_name_map, sequence_index) + terminal, edge = _build_graph_output_terminal_edge( + terminal_edge, node_name_map, sequence_index + ) terminals.append(terminal) edges.append(edge) - return VisualGraph(nodes=tuple(ordered_nodes), terminals=tuple(terminals), edges=tuple(edges)) + return VisualGraph( + nodes=tuple(ordered_nodes), terminals=tuple(terminals), edges=tuple(edges) + ) def _build_internal_edges(connection: dict, kind: str) -> list[VisualEdge]: @@ -70,7 +85,15 @@ def _build_internal_edges(connection: dict, kind: str) -> list[VisualEdge]: target_id = visual_id("node", target_node.name) edges.append( VisualEdge( - id=visual_id("edge", kind, source_id, source_desc.name_str, target_id, target_desc.name_str, sequence_index), + id=visual_id( + "edge", + kind, + source_id, + source_desc.name_str, + target_id, + target_desc.name_str, + sequence_index, + ), kind=kind, source_id=source_id, target_id=target_id, @@ -82,7 +105,9 @@ def _build_internal_edges(connection: dict, kind: str) -> list[VisualEdge]: return edges -def _build_graph_input_terminal_edge(node_port: str, node_name_map: dict[str, VisualNode], sequence_index: int) -> tuple[VisualTerminal, VisualEdge]: +def _build_graph_input_terminal_edge( + node_port: str, node_name_map: dict[str, VisualNode], sequence_index: int +) -> tuple[VisualTerminal, VisualEdge]: node_name, port_name = node_port.split("/", 1) node = node_name_map[node_name] port_id = _find_port_id(node, "input", port_name) @@ -95,7 +120,15 @@ def _build_graph_input_terminal_edge(node_port: str, node_name_map: dict[str, Vi port_id=port_id, ) edge = VisualEdge( - id=visual_id("edge", "graph_input", terminal_id, port_name, node.id, port_name, sequence_index), + id=visual_id( + "edge", + "graph_input", + terminal_id, + port_name, + node.id, + port_name, + sequence_index, + ), kind="graph_input", source_id=terminal_id, target_id=node.id, @@ -106,7 +139,9 @@ def _build_graph_input_terminal_edge(node_port: str, node_name_map: dict[str, Vi return terminal, edge -def _build_graph_output_terminal_edge(node_port: str, node_name_map: dict[str, VisualNode], sequence_index: int) -> tuple[VisualTerminal, VisualEdge]: +def _build_graph_output_terminal_edge( + node_port: str, node_name_map: dict[str, VisualNode], sequence_index: int +) -> tuple[VisualTerminal, VisualEdge]: node_name, port_name = node_port.split("/", 1) node = node_name_map[node_name] port_id = _find_port_id(node, "output", port_name) @@ -119,7 +154,15 @@ def _build_graph_output_terminal_edge(node_port: str, node_name_map: dict[str, V port_id=port_id, ) edge = VisualEdge( - id=visual_id("edge", "graph_output", node.id, port_name, terminal_id, port_name, sequence_index), + id=visual_id( + "edge", + "graph_output", + node.id, + port_name, + terminal_id, + port_name, + sequence_index, + ), kind="graph_output", source_id=node.id, target_id=terminal_id, @@ -131,7 +174,11 @@ def _build_graph_output_terminal_edge(node_port: str, node_name_map: dict[str, V def _build_node(node: Any) -> VisualNode: - backend = node.get_backend() if hasattr(node, "get_backend") else getattr(node, "backend", None) + backend = ( + node.get_backend() + if hasattr(node, "get_backend") + else getattr(node, "backend", None) + ) return VisualNode( id=visual_id("node", node.name), title=node.name, @@ -179,7 +226,9 @@ def _find_port_id(node: Any, side: str, port_name: str) -> str: desc_name = getattr(desc, "name_str", getattr(desc, "name", None)) if desc_name == port_name: return visual_id("port", node_name, side, desc_name) - raise KeyError(f"Port '{port_name}' not found on node '{node_name}' for side '{side}'") + raise KeyError( + f"Port '{port_name}' not found on node '{node_name}' for side '{side}'" + ) def _node_sort_key(item: tuple[str, Any]) -> tuple[int, str]: diff --git a/leapp/leapp_graph/warp_op.py b/leapp/leapp_graph/warp_op.py index b0429a99..c2e2c512 100644 --- a/leapp/leapp_graph/warp_op.py +++ b/leapp/leapp_graph/warp_op.py @@ -40,6 +40,7 @@ wp = None WarpOp = None else: + def _save_warp_bundle(node_name: str, segment: WarpSegment) -> bytes: """Persist a captured Warp APIC graph and pack it as a WRPB archive. @@ -84,10 +85,7 @@ def _as_raw_warp_array(array): ) return archive - - def _update_output_ref_proxy( - node_ref: Any, ref, proxy - ) -> None: + def _update_output_ref_proxy(node_ref: Any, ref, proxy) -> None: """Point this segment output's value at the runner output that produced it. Updates the traced array's shared view so every Torch or NumPy alias of @@ -192,7 +190,6 @@ def _insert_warp_marker( segment.status = "closed" return segment - class WarpOp: @staticmethod def _verify_warp() -> None: @@ -256,8 +253,7 @@ def begin(self, call_stack, *, owner_token=None): def __exit__(self, exc_type, exc_value, traceback): if self._segment is None: _get_logger().fatal( - "WarpOp.__exit__ called before the Warp operation was " - "initialized.", + "WarpOp.__exit__ called before the Warp operation was initialized.", error_type=RuntimeError, ) return self._context_exit_impl(exc_type, exc_value, traceback) @@ -274,8 +270,10 @@ def terminate( if self._segment is None: return None return self._terminate_impl( - exc_type=exc_type, exc_value=exc_value, - traceback=traceback, close_call_stack=close_call_stack, + exc_type=exc_type, + exc_value=exc_value, + traceback=traceback, + close_call_stack=close_call_stack, ) def _context_exit_discovery(self, exc_type, exc_value, traceback): @@ -395,7 +393,8 @@ def _finalize_discovery(self, *, close_call_stack) -> WarpSegment: ) warp_call_events = tuple( - event for event in self._segment.events + event + for event in self._segment.events if event.get("kind") == "warp_call" ) self._segment.call_qualnames = tuple( @@ -469,9 +468,7 @@ def _terminate_capture( scope_result = False if self._scope is not None: with self._session.pause(): - scope_result = self._scope.__exit__( - exc_type, exc_value, traceback - ) + scope_result = self._scope.__exit__(exc_type, exc_value, traceback) self._validate_capture_close(close_call_stack) if exc_type is None and self._segment is not None: self._finalize_capture() diff --git a/leapp/utils/caller_identity.py b/leapp/utils/caller_identity.py index 80eb9921..8c429ce3 100644 --- a/leapp/utils/caller_identity.py +++ b/leapp/utils/caller_identity.py @@ -22,21 +22,34 @@ from dataclasses import dataclass _LEAPP_PKG_DIR = os.path.realpath( - os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -_STDLIB_DIRS = tuple(sorted({ - os.path.realpath(path) for path in ( - sysconfig.get_paths().get("stdlib"), - sysconfig.get_paths().get("platstdlib"), - ) if path -})) -_SITE_PACKAGE_DIRS = tuple(sorted({ - os.path.realpath(path) for path in ( - site.getusersitepackages(), - *site.getsitepackages(), - sysconfig.get_paths().get("purelib"), - sysconfig.get_paths().get("platlib"), - ) if path -})) + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +) +_STDLIB_DIRS = tuple( + sorted( + { + os.path.realpath(path) + for path in ( + sysconfig.get_paths().get("stdlib"), + sysconfig.get_paths().get("platstdlib"), + ) + if path + } + ) +) +_SITE_PACKAGE_DIRS = tuple( + sorted( + { + os.path.realpath(path) + for path in ( + site.getusersitepackages(), + *site.getsitepackages(), + sysconfig.get_paths().get("purelib"), + sysconfig.get_paths().get("platlib"), + ) + if path + } + ) +) @dataclass(frozen=True) @@ -101,8 +114,7 @@ def get_caller_stack_identity(): closest_first = list(reversed(frames)) anchor = closest_first[0] context = tuple( - frame for frame in closest_first[1:] - if _is_context_relevant_frame(frame) + frame for frame in closest_first[1:] if _is_context_relevant_frame(frame) )[:2] return CallerIdentity(anchor=anchor, context=context) @@ -117,14 +129,17 @@ def format_caller_identity(identity): if identity.context: lines.append(" context:") for frame in identity.context: - lines.append( - f" {frame.filename}:{frame.lineno} in {frame.function}") + lines.append(f" {frame.filename}:{frame.lineno} in {frame.function}") return "\n".join(lines) if isinstance(identity, CallerFrame): return f"{identity.filename}:{identity.lineno} in {identity.function}" if not isinstance(identity, tuple) or not identity: return str(identity) - if len(identity) == 2 and isinstance(identity[0], str) and isinstance(identity[1], int): + if ( + len(identity) == 2 + and isinstance(identity[0], str) + and isinstance(identity[1], int) + ): return f"{identity[0]}:{identity[1]}" lines = [] for frame in identity: diff --git a/leapp/utils/dtype.py b/leapp/utils/dtype.py index e6246582..c8fcd351 100644 --- a/leapp/utils/dtype.py +++ b/leapp/utils/dtype.py @@ -25,12 +25,23 @@ # Canonical common dtype-name vocabulary shared across backends. Backend codecs # map their dtype objects onto these names. -_VALID_DTYPE_NAMES = frozenset({ - "float16", "float32", "float64", "bfloat16", - "int8", "int16", "int32", "int64", - "uint8", "uint16", "uint32", "uint64", - "bool", -}) +_VALID_DTYPE_NAMES = frozenset( + { + "float16", + "float32", + "float64", + "bfloat16", + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", + "bool", + } +) def map_to_torch_dtype(string): @@ -53,6 +64,7 @@ def map_to_torch_dtype(string): # Backend dtype-codec registry # ============================================================================= + @dataclass(frozen=True) class DtypeCodec: # Backend label, e.g. "torch" / "numpy" / "warp". diff --git a/leapp/utils/enums.py b/leapp/utils/enums.py index 982c8027..00cb2d94 100644 --- a/leapp/utils/enums.py +++ b/leapp/utils/enums.py @@ -19,37 +19,37 @@ class InputKindEnum(Enum): - JOINT_POSITION= "state/joint/position" - JOINT_VELOCITY= "state/joint/velocity" - JOINT_EFFORT= "state/joint/effort" - BODY_POSE= "state/body/pose" - BODY_VEL= "state/body/velocity" - BODY_ACC= "state/body/acceleration" - BODY_LINEAR_ACCELERATION= "state/body/linear_acceleration" - BODY_LINEAR_VELOCITY= "state/body/linear_velocity" - BODY_ANGULAR_ACCELERATION= "state/body/angular_acceleration" - BODY_ANGULAR_VELOCITY= "state/body/angular_velocity" - BODY_ROTATION= "state/body/rotation" - BODY_POSITION= "state/body/position" - WRENCH= "state/wrench" + JOINT_POSITION = "state/joint/position" + JOINT_VELOCITY = "state/joint/velocity" + JOINT_EFFORT = "state/joint/effort" + BODY_POSE = "state/body/pose" + BODY_VEL = "state/body/velocity" + BODY_ACC = "state/body/acceleration" + BODY_LINEAR_ACCELERATION = "state/body/linear_acceleration" + BODY_LINEAR_VELOCITY = "state/body/linear_velocity" + BODY_ANGULAR_ACCELERATION = "state/body/angular_acceleration" + BODY_ANGULAR_VELOCITY = "state/body/angular_velocity" + BODY_ROTATION = "state/body/rotation" + BODY_POSITION = "state/body/position" + WRENCH = "state/wrench" VECTOR3D = "state/vector3d" - COMMAND_JOINT_POSITION= "command/joint/position" - COMMAND_JOINT_VELOCITY= "command/joint/velocity" - COMMAND_JOINT_TORQUES= "command/joint/torques" - COMMAND_BODY_ROTATION= "command/body/rotation" - COMMAND_BODY_VELOCITY= "command/body/velocity" - COMMAND_POSE= "command/body/pose" + COMMAND_JOINT_POSITION = "command/joint/position" + COMMAND_JOINT_VELOCITY = "command/joint/velocity" + COMMAND_JOINT_TORQUES = "command/joint/torques" + COMMAND_BODY_ROTATION = "command/body/rotation" + COMMAND_BODY_VELOCITY = "command/body/velocity" + COMMAND_POSE = "command/body/pose" -class OutputKindEnum(Enum): - KP="kp" - KD="kd" - JOINT_POSITION="target/joint/position" - JOINT_VELOCITY="target/joint/velocity" - JOINT_TORQUES="target/joint/torques" - JOINT_EFFORT="target/joint/effort" - BODY_POSITION="target/body/position" - BODY_LINEAR_ACCELERATION="target/body/linear_acceleration" - BODY_ORIENTATION="target/body/orientation" - BODY_LINEAR_VELOCITY="target/body/linear_velocity" - BODY_ANGULAR_ACCELERATION="target/body/angular_acceleration" +class OutputKindEnum(Enum): + KP = "kp" + KD = "kd" + JOINT_POSITION = "target/joint/position" + JOINT_VELOCITY = "target/joint/velocity" + JOINT_TORQUES = "target/joint/torques" + JOINT_EFFORT = "target/joint/effort" + BODY_POSITION = "target/body/position" + BODY_LINEAR_ACCELERATION = "target/body/linear_acceleration" + BODY_ORIENTATION = "target/body/orientation" + BODY_LINEAR_VELOCITY = "target/body/linear_velocity" + BODY_ANGULAR_ACCELERATION = "target/body/angular_acceleration" diff --git a/leapp/utils/logging.py b/leapp/utils/logging.py index c7684afa..8d507dd7 100644 --- a/leapp/utils/logging.py +++ b/leapp/utils/logging.py @@ -32,9 +32,9 @@ SECTION = 25 SYSTEM = 31 FATAL = logging.ERROR + 1 -logging.addLevelName(SECTION, 'SECTION') -logging.addLevelName(SYSTEM, 'SYSTEM') -logging.addLevelName(FATAL, 'FATAL') +logging.addLevelName(SECTION, "SECTION") +logging.addLevelName(SYSTEM, "SYSTEM") +logging.addLevelName(FATAL, "FATAL") class _ColoredFormatter(logging.Formatter): @@ -42,20 +42,20 @@ class _ColoredFormatter(logging.Formatter): # ANSI color codes (matching export_manager.py style) COLORS = { - 'DEBUG': '', # White (no color) - 'INFO': '', # White - 'SECTION': '\033[1;4;97m', # Bold white (for section headers) - 'SYSTEM': '', # White, but above WARNING for non-verbose console output - 'WARNING': '\033[1;33m', # Bold yellow (matching export_manager.py) - 'ERROR': '\033[91m', # Bright red - 'FATAL': '\033[91m', # Same as ERROR, but raises after logging - 'CRITICAL': '\033[1;31m', # Bold red - 'RESET': '\033[0m' # Reset + "DEBUG": "", # White (no color) + "INFO": "", # White + "SECTION": "\033[1;4;97m", # Bold white (for section headers) + "SYSTEM": "", # White, but above WARNING for non-verbose console output + "WARNING": "\033[1;33m", # Bold yellow (matching export_manager.py) + "ERROR": "\033[91m", # Bright red + "FATAL": "\033[91m", # Same as ERROR, but raises after logging + "CRITICAL": "\033[1;31m", # Bold red + "RESET": "\033[0m", # Reset } def format(self, record): # Add color to entire message - log_color = self.COLORS.get(record.levelname, self.COLORS['RESET']) + log_color = self.COLORS.get(record.levelname, self.COLORS["RESET"]) if log_color: record.levelname = f"{log_color}{record.levelname}{self.COLORS['RESET']}" record.msg = f"{log_color}{record.msg}{self.COLORS['RESET']}" @@ -66,7 +66,7 @@ class _LeappLogger: """Internal logger class for leapp.""" def __init__(self): - self.logger = logging.getLogger('leapp') + self.logger = logging.getLogger("leapp") self.logger.setLevel(logging.DEBUG) # Capture everything self.initialized = False self.console_handler = None @@ -81,14 +81,14 @@ def configure(self, savepath, verbose): # log file handler if os.path.exists(savepath): - filepath = os.path.join(savepath, 'log.txt') + filepath = os.path.join(savepath, "log.txt") self.log_path = filepath self.file_handler = logging.FileHandler( - filepath, mode='w', encoding='utf-8') + filepath, mode="w", encoding="utf-8" + ) self.file_handler.setLevel(logging.DEBUG) # Plain format for file (with timestamp, no colors) - file_formatter = logging.Formatter( - '[%(levelname)s]: %(message)s') + file_formatter = logging.Formatter("[%(levelname)s]: %(message)s") self.file_handler.setFormatter(file_formatter) self.logger.addHandler(self.file_handler) else: @@ -99,7 +99,8 @@ def configure(self, savepath, verbose): self.set_verbose(verbose) # Use simple format for console (no timestamp, just message) self.console_handler.setFormatter( - _ColoredFormatter('[%(levelname)s]: %(message)s')) + _ColoredFormatter("[%(levelname)s]: %(message)s") + ) self.logger.addHandler(self.console_handler) self.initialized = True @@ -109,8 +110,7 @@ def set_verbose(self, verbose): self.verbose = verbose if self.console_handler: # When verbose, show everything (DEBUG and above) - self.console_handler.setLevel( - logging.DEBUG if verbose else logging.WARNING) + self.console_handler.setLevel(logging.DEBUG if verbose else logging.WARNING) def is_verbose(self): return self.verbose @@ -147,7 +147,8 @@ def log_export_artifact_locations(self, save_path, graph_name, nodes): for node in nodes_to_report: if node.model_path: saved_models.append( - f"- {node.name}: {os.path.abspath(node.model_path)}") + f"- {node.name}: {os.path.abspath(node.model_path)}" + ) else: nodes_without_artifacts.append(node.name) @@ -165,7 +166,8 @@ def log_export_artifact_locations(self, save_path, graph_name, nodes): if nodes_without_artifacts: message_lines.append( - "Nodes without saved model artifacts: " + ", ".join(nodes_without_artifacts) + "Nodes without saved model artifacts: " + + ", ".join(nodes_without_artifacts) ) self.system("\n".join(message_lines)) diff --git a/leapp/utils/tensor_description.py b/leapp/utils/tensor_description.py index 29f42bb6..38d009b6 100644 --- a/leapp/utils/tensor_description.py +++ b/leapp/utils/tensor_description.py @@ -51,6 +51,7 @@ class CompactYamlList(list): formatted in a more compact way in the yaml file. """ + pass @@ -59,16 +60,23 @@ class CompactYamlDict(dict): This class will behave exactly like a dict. The only difference is that it will be formatted in a more compact way in the yaml file. """ + pass # Register custom YAML representers so these types are always serialized in flow style yaml.add_representer( CompactYamlList, - lambda dumper, data: dumper.represent_sequence('tag:yaml.org,2002:seq', data, flow_style=True)) + lambda dumper, data: dumper.represent_sequence( + "tag:yaml.org,2002:seq", data, flow_style=True + ), +) yaml.add_representer( CompactYamlDict, - lambda dumper, data: dumper.represent_mapping('tag:yaml.org,2002:map', data, flow_style=True)) + lambda dumper, data: dumper.represent_mapping( + "tag:yaml.org,2002:map", data, flow_style=True + ), +) TEMPORAL_AXIS_SENTINEL = "__temporal_axis__" @@ -129,22 +137,24 @@ class TensorSemantics: """ # Fields excluded from serialization (internal use only) - _INTERNAL_FIELDS = frozenset({'ref', 'name'}) + _INTERNAL_FIELDS = frozenset({"ref", "name"}) name: str = None ref: Any = None # Semantic fields kind: Optional[InputKindEnum | OutputKindEnum | str] = None - element_names: Optional[List] = None # deprecated + element_names: Optional[List] = None # deprecated temporal_period_ms: Optional[float] = field(default=None, init=False) extra: Optional[Dict[str, Any]] = None def __post_init__(self): - '''error checking, auto conditioning''' + """error checking, auto conditioning""" existing_period_ms = self.temporal_period_ms self.temporal_period_ms = None - if not is_tracable_tensor_type(self.ref): # this checks for both base types and traced types + if not is_tracable_tensor_type( + self.ref + ): # this checks for both base types and traced types _get_logger().fatal( f"TensorSemantics 'ref' must be a traceable tensor type " f"accepted types are: {TRACABLE_BASE_TYPES}" @@ -153,10 +163,11 @@ def __post_init__(self): ) if self.element_names is not None: self.element_names, detected_period_ms = self._normalize_element_names( - self.element_names, allow_temporal_sentinel=existing_period_ms is not None) + self.element_names, + allow_temporal_sentinel=existing_period_ms is not None, + ) self.temporal_period_ms = detected_period_ms or existing_period_ms - def to_dict(self) -> Dict[str, Any]: """Return all non-None semantic fields, excluding internal fields. @@ -165,7 +176,7 @@ def to_dict(self) -> Dict[str, Any]: """ result = {} for f in fields(self): - if f.name in self._INTERNAL_FIELDS or f.name == 'extra': + if f.name in self._INTERNAL_FIELDS or f.name == "extra": continue value = getattr(self, f.name) if value is not None: @@ -251,7 +262,8 @@ def _normalize_element_names(element_names, allow_temporal_sentinel=False): normalized.append(None) else: _get_logger().warning( - "element_names has mixed types, expected names or axis descriptors") + "element_names has mixed types, expected names or axis descriptors" + ) return element_names, temporal_period_ms if has_axis_descriptors: @@ -268,8 +280,9 @@ class TensorDescription: or in bulk via get_semantics()/set_semantics(). """ - def __init__(self, name: str, value: Any, - semantics: Optional[TensorSemantics] = None): + def __init__( + self, name: str, value: Any, semantics: Optional[TensorSemantics] = None + ): # Capture where the value came from before unwrapping loses that state. source_node = None port = None @@ -339,13 +352,13 @@ def dict(self, ignore_value=True) -> Dict[str, Any]: "name": self.name, "dtype": self.dtype, "shape": self.shape, - "type": self.type + "type": self.type, } if self.value is not None and not ignore_value: result["value"] = self.value # Merge in all non-None semantic fields for key, value in self.get_semantics().items(): - if hasattr(value, 'value'): # Enum → serialize to its .value string + if hasattr(value, "value"): # Enum → serialize to its .value string result[key] = value.value else: result[key] = value @@ -377,6 +390,7 @@ def name_str(self) -> str: """Return just the name as a string.""" return self.name + @dataclass class ParameterFormat: name: str @@ -384,14 +398,17 @@ class ParameterFormat: def get_format_string(self, data: Any) -> Tuple[bool, str]: - if isinstance(data, collections.abc.Sequence) and not isinstance(data, (str, bytes, torch.Tensor)): + if isinstance(data, collections.abc.Sequence) and not isinstance( + data, (str, bytes, torch.Tensor) + ): # Warn if not a native list if not isinstance(data, list): type_name = type(data).__name__ _get_logger().debug( f"Parameter '{self.name}' has list-like type '{type_name}' which will be " f"treated as 'list' in the generated function signature. This may cause " - f"issues if the exact type is required.") + f"issues if the exact type is required." + ) if len(data) == 0: return True, "list" # Empty list, can't determine child type @@ -413,7 +430,8 @@ def get_format_string(self, data: Any) -> Tuple[bool, str]: _get_logger().debug( f"Parameter '{self.name}' has dict-like type '{type_name}' which will be " f"treated as 'dict' in the generated function signature. This may cause " - f"issues if the exact type is required.") + f"issues if the exact type is required." + ) if len(data) == 0: return True, "dict" # Empty dict, can't determine child type @@ -425,8 +443,7 @@ def get_format_string(self, data: Any) -> Tuple[bool, str]: if len(value_types) > 1: return False, "INCONSISTENT DICT VALUE TYPES" - success, child = self.get_format_string( - list(data.values())[0]) + success, child = self.get_format_string(list(data.values())[0]) if not success: return False, "DICT CHILD FORMAT STRING NOT VALID" else: @@ -440,7 +457,7 @@ def get_format_string(self, data: Any) -> Tuple[bool, str]: @property def name_str(self) -> str: - """Return just the name as a string. + """Return just the name as a string. will do some modifications to me the vairable compliant with python""" return self.name.replace(".", "_") @@ -469,6 +486,7 @@ def tensor_expr_in_order(self) -> list[str]: """ Generate a list of tensor expressions in the order of the nested structure. """ + def _generate_expr(format_item) -> list[str]: if isinstance(format_item, ParameterFormat): return _generate_expr(format_item.formatting) @@ -480,6 +498,7 @@ def _generate_expr(format_item) -> list[str]: return [_generate_expr(value) for value in format_item.values()] else: return [] + return _generate_expr(self.formatting) @property @@ -491,9 +510,10 @@ def packed_tensor_expr(self) -> str: Handles nested structures (lists, dicts). Returns: - Assignment string like "inputA = [inputA_0, inputA_1]". + Assignment string like "inputA = [inputA_0, inputA_1]". Empty string if no packing needed. """ + def _generate_expr(format_item) -> str: if isinstance(format_item, ParameterFormat): return _generate_expr(format_item.formatting) @@ -503,8 +523,7 @@ def _generate_expr(format_item) -> str: elements = [_generate_expr(item) for item in format_item] return "[" + ", ".join(elements) + "]" elif isinstance(format_item, dict): - items = [ - f'"{k}": {_generate_expr(v)}' for k, v in format_item.items()] + items = [f'"{k}": {_generate_expr(v)}' for k, v in format_item.items()] return "{" + ", ".join(items) + "}" else: return "None" @@ -533,10 +552,10 @@ def unpacked_tensor_expr(self) -> str: Returns: Newline-separated assignment strings. Empty string if no unpacking needed. """ + def _generate_unpacking(format_item, accessor_path: str, assignments: list): if isinstance(format_item, ParameterFormat): - _generate_unpacking(format_item.formatting, - accessor_path, assignments) + _generate_unpacking(format_item.formatting, accessor_path, assignments) elif isinstance(format_item, TensorDescription): tensor_name = format_item.name_str if tensor_name != accessor_path: @@ -554,8 +573,10 @@ def _generate_unpacking(format_item, accessor_path: str, assignments: list): _generate_unpacking(self.formatting, self.name_str, assignments) return "\n".join(assignments) -def validate_connection_compatibility(source_name, source_shape, source_dtype, - target_name, target_shape, target_dtype): + +def validate_connection_compatibility( + source_name, source_shape, source_dtype, target_name, target_shape, target_dtype +): """Validate shape and dtype compatibility for a single pipeline edge.""" if list(source_shape) != list(target_shape): _get_logger().fatal( @@ -602,7 +623,7 @@ def verify_data_exact_match(source_data, target_data): return False source_data = to_export_torch_tensor(source_data) target_data = to_export_torch_tensor(target_data) - + # identity comparison if source_data.shape != target_data.shape: return False @@ -634,7 +655,9 @@ def verify_data_exact_match(source_data, target_data): def flatten_io_structure(data, name_str): flat_data = {} - if isinstance(data, collections.abc.Sequence) and not isinstance(data, (str, bytes, torch.Tensor)): + if isinstance(data, collections.abc.Sequence) and not isinstance( + data, (str, bytes, torch.Tensor) + ): for idx, item in enumerate(data): child_name = f"{name_str}_{idx}" if name_str else str(idx) flat_data.update(flatten_io_structure(item, child_name)) @@ -674,7 +697,11 @@ def unwrap_tensor_semantics(data): ) if not all(isinstance(item, TensorSemantics) for item in data): - bad_types = {type(item).__name__ for item in data if not isinstance(item, TensorSemantics)} + bad_types = { + type(item).__name__ + for item in data + if not isinstance(item, TensorSemantics) + } _get_logger().fatal( f"All items must be TensorSemantics when using semantic inputs. " f"Found non-TensorSemantics types: {bad_types}", @@ -703,18 +730,20 @@ def unwrap_tensor_semantics(data): def describe_io_helper(data, name_str): data_description = [] io_format = {} - if isinstance(data, collections.abc.Sequence) and not isinstance(data, (str, bytes, torch.Tensor)): + if isinstance(data, collections.abc.Sequence) and not isinstance( + data, (str, bytes, torch.Tensor) + ): if not isinstance(data, list): type_name = type(data).__name__ _get_logger().warning( f"Input/Output '{name_str}' has list-like type '{type_name}' which will be " - f"treated as 'list'. Ensure the runtime can handle this substitution.") + f"treated as 'list'. Ensure the runtime can handle this substitution." + ) io_format = [] for idx, item in enumerate(data): child_name = f"{name_str}_{idx}" if name_str else str(idx) - child_description, child_format = describe_io_helper( - item, child_name) + child_description, child_format = describe_io_helper(item, child_name) if child_format is not None: io_format.append(child_format) data_description.extend(child_description) @@ -724,13 +753,13 @@ def describe_io_helper(data, name_str): type_name = type(data).__name__ _get_logger().warning( f"Input/Output '{name_str}' has dict-like type '{type_name}' which will be " - f"treated as 'dict'. Ensure the runtime can handle this substitution.") + f"treated as 'dict'. Ensure the runtime can handle this substitution." + ) io_format = {} for k, v in data.items(): child_name = f"{name_str}_{k}" if name_str else k - child_description, child_format = describe_io_helper( - v, child_name) + child_description, child_format = describe_io_helper(v, child_name) if child_format is not None: io_format[k] = child_format data_description.extend(child_description) @@ -749,10 +778,8 @@ def describe_io_helper(data, name_str): def describe_io(name, raw_name, data): - data_description, io_format = describe_io_helper( - data, name) - parameter_description = ParameterFormat( - name=raw_name, formatting=io_format) + data_description, io_format = describe_io_helper(data, name) + parameter_description = ParameterFormat(name=raw_name, formatting=io_format) return data_description, parameter_description @@ -778,7 +805,10 @@ def _resolve_tensor_descriptions(io_format, extractor): elif isinstance(io_format, list): return [_resolve_tensor_descriptions(item, extractor) for item in io_format] elif isinstance(io_format, dict): - return {key: _resolve_tensor_descriptions(value, extractor) for key, value in io_format.items()} + return { + key: _resolve_tensor_descriptions(value, extractor) + for key, value in io_format.items() + } else: return io_format diff --git a/leapp/utils/tracing_lock.py b/leapp/utils/tracing_lock.py index 9f33a057..bbc476af 100644 --- a/leapp/utils/tracing_lock.py +++ b/leapp/utils/tracing_lock.py @@ -20,32 +20,33 @@ class TracingLock: """Singleton class to manage global function tracing lock. - + This acts as a global mutex to prevent re-entrant tracing across the entire process. When _is_active_function_tracing is True, sys.settrace is active for a node and no other tracing should be started. """ + _instance = None - + def __new__(cls): """Singleton implementation - only one instance allowed.""" if cls._instance is None: cls._instance = super(TracingLock, cls).__new__(cls) cls._instance._is_active_function_tracing = False return cls._instance - + @property def is_active(self) -> bool: """Check if function tracing is currently active. - + Returns: bool: True if sys.settrace is active for a node, False otherwise. """ return self._is_active_function_tracing - + def acquire(self): """Acquire the tracing lock. - + Raises: Exception: If tracing lock is already acquired. """ @@ -55,14 +56,14 @@ def acquire(self): error_type=Exception, ) self._is_active_function_tracing = True - + def release(self): """Release the tracing lock.""" self._is_active_function_tracing = False - + def reset(self): """Reset the tracing lock to False. - + This is useful for error recovery or reset scenarios. """ self._is_active_function_tracing = False diff --git a/leapp/utils/utils.py b/leapp/utils/utils.py index 39c7b346..f1f1fdf9 100644 --- a/leapp/utils/utils.py +++ b/leapp/utils/utils.py @@ -37,19 +37,19 @@ def find_with_block_end(filename, start_lineno): """Use AST to find the end line of the with block starting at start_lineno. - + This is deterministic and doesn't depend on runtime tracing behavior. - + Args: filename: Path to the source file start_lineno: Line number where the 'with' statement starts - + Returns: The end line number of the with block, or None if not found """ - with open(filename, 'r') as f: + with open(filename, "r") as f: source = f.read() - + tree = ast.parse(source) for node in ast.walk(tree): if isinstance(node, (ast.With, ast.AsyncWith)): @@ -74,16 +74,16 @@ def extract_source_from_line_range(executed_lines, context_name, is_function=Fal Returns: tuple: (source_code, message) where source_code is the extracted code or empty string """ - if executed_lines['filename'] is None: + if executed_lines["filename"] is None: return "", f"No filename provided for {context_name}" try: - filename = executed_lines['filename'] - min_line = executed_lines['min_line'] - max_line = executed_lines['max_line'] + filename = executed_lines["filename"] + min_line = executed_lines["min_line"] + max_line = executed_lines["max_line"] # Read the source file - with open(filename, 'r') as f: + with open(filename, "r") as f: source_lines = f.readlines() # Extract lines from min_line to max_line (convert to 0-based indexing) @@ -97,7 +97,7 @@ def extract_source_from_line_range(executed_lines, context_name, is_function=Fal ) # Determine the start of the executable body using AST (robust for multiline headers) - file_source = ''.join(source_lines) + file_source = "".join(source_lines) body_start_idx = None try: @@ -115,35 +115,39 @@ def extract_source_from_line_range(executed_lines, context_name, is_function=Fal # Find the node that contains our line range for n in ast.walk(tree): if isinstance(n, wanted): - node_start = getattr(n, 'lineno', None) - node_end = getattr(n, 'end_lineno', None) + node_start = getattr(n, "lineno", None) + node_end = getattr(n, "end_lineno", None) if node_start is None: continue if node_end is None: - if getattr(n, 'body', None): + if getattr(n, "body", None): last_child = n.body[-1] - node_end = getattr(last_child, 'end_lineno', getattr( - last_child, 'lineno', node_start)) + node_end = getattr( + last_child, + "end_lineno", + getattr(last_child, "lineno", node_start), + ) else: node_end = node_start # Check if this node overlaps with our line range # (handles decorators by checking if def line or any body line is in range) - if (min_line <= node_start <= max_line) or (node_start <= min_line <= node_end): + if (min_line <= node_start <= max_line) or ( + node_start <= min_line <= node_end + ): if node_start > best_start: best_candidate = n best_start = node_start - if best_candidate is not None and getattr(best_candidate, 'body', None): - body_start_line = getattr( - best_candidate.body[0], 'lineno', None) + if best_candidate is not None and getattr(best_candidate, "body", None): + body_start_line = getattr(best_candidate.body[0], "lineno", None) if body_start_line is not None: body_start_idx = body_start_line - 1 # Extend end_idx to capture multiline statements # Python's line tracer only fires once for multiline statements (e.g., dict comprehensions), # so max_line may not include the closing brace. Use AST's end_lineno instead. - ast_end_line = getattr(best_candidate, 'end_lineno', None) + ast_end_line = getattr(best_candidate, "end_lineno", None) if ast_end_line is not None and ast_end_line > end_idx: end_idx = ast_end_line except Exception: @@ -153,7 +157,7 @@ def extract_source_from_line_range(executed_lines, context_name, is_function=Fal # Fallback: scan to the end of the logical header by balancing symbols and strings open_paren = open_brack = open_brace = 0 in_string = False - string_char = '' + string_char = "" header_end_line_idx = start_idx i = start_idx while i < end_idx: @@ -164,23 +168,32 @@ def extract_source_from_line_range(executed_lines, context_name, is_function=Fal if not in_string and ch in ('"', "'"): in_string = True string_char = ch - elif in_string and ch == string_char and (j == 0 or line[j-1] != '\\'): + elif ( + in_string + and ch == string_char + and (j == 0 or line[j - 1] != "\\") + ): in_string = False - string_char = '' + string_char = "" elif not in_string: - if ch == '(': + if ch == "(": open_paren += 1 - elif ch == ')': - open_paren = max(0, open_paren-1) - elif ch == '[': + elif ch == ")": + open_paren = max(0, open_paren - 1) + elif ch == "[": open_brack += 1 - elif ch == ']': - open_brack = max(0, open_brack-1) - elif ch == '{': + elif ch == "]": + open_brack = max(0, open_brack - 1) + elif ch == "{": open_brace += 1 - elif ch == '}': - open_brace = max(0, open_brace-1) - elif ch == ':' and open_paren == 0 and open_brack == 0 and open_brace == 0: + elif ch == "}": + open_brace = max(0, open_brace - 1) + elif ( + ch == ":" + and open_paren == 0 + and open_brack == 0 + and open_brace == 0 + ): header_end_line_idx = i i = i + 1 break @@ -198,7 +211,7 @@ def extract_source_from_line_range(executed_lines, context_name, is_function=Fal extracted_lines = source_lines[start_idx:end_idx] if extracted_lines: - dedented_code = textwrap.dedent(''.join(extracted_lines)) + dedented_code = textwrap.dedent("".join(extracted_lines)) source_code = dedented_code.rstrip() else: source_code = "" @@ -236,7 +249,7 @@ def extract_return_names(func): statement_names.append(elt.id) else: # Complex expression, use generic name - statement_names.append(f"output{i+1}") + statement_names.append(f"output{i + 1}") elif isinstance(node.value, ast.Name): # Single return value: return var1 statement_names.append(node.value.id) @@ -255,8 +268,9 @@ def extract_return_names(func): if len(unique_counts) > 1: # Different number of return values - this is an error - count_details = {count: return_counts.count( - count) for count in unique_counts} + count_details = { + count: return_counts.count(count) for count in unique_counts + } _get_logger().fatal( f"Function {func.__name__} has inconsistent return statements with different numbers of values: " f"{count_details}. All return statements must return the same number of values.", @@ -268,9 +282,11 @@ def extract_return_names(func): if len(return_statements) > 1: # Multiple return statements - log warning and intelligently merge - _get_logger().warning(f"Warning: Function {func.__name__} has {len(return_statements)} return statements. " - f"This may cause unexpected behavior if the return statements return different " - f"data types, shapes, or dtypes.") + _get_logger().warning( + f"Warning: Function {func.__name__} has {len(return_statements)} return statements. " + f"This may cause unexpected behavior if the return statements return different " + f"data types, shapes, or dtypes." + ) # Intelligently merge return names by position, prioritizing named over generic final_names = [] @@ -280,9 +296,11 @@ def extract_return_names(func): # Separate named variables from generic ones named_vars = [ - name for name in names_at_position if not name.startswith("output")] + name for name in names_at_position if not name.startswith("output") + ] generic_vars = [ - name for name in names_at_position if name.startswith("output")] + name for name in names_at_position if name.startswith("output") + ] # Choose the best name for this position if named_vars: @@ -293,7 +311,7 @@ def extract_return_names(func): chosen_name = generic_vars[0] else: # Fallback - chosen_name = f"output{pos+1}" + chosen_name = f"output{pos + 1}" final_names.append(chosen_name) @@ -302,7 +320,8 @@ def extract_return_names(func): except Exception as e: # Fallback to generic names if parsing fails _get_logger().error( - f"Error extracting return names from {func.__name__}, reverting to defaults: {e}") + f"Error extracting return names from {func.__name__}, reverting to defaults: {e}" + ) return ["output1"] @@ -321,10 +340,10 @@ def safe_deepcopy(data): def frame_to_namespace(frame): """Convert a Python frame object to a namespace dictionary. - + Args: frame: A Python frame object from sys._getframe() or similar - + Returns: dict: A namespace combining globals and locals from the frame """ @@ -333,11 +352,11 @@ def frame_to_namespace(frame): def get_attribute_value_from_namespace(namespace, attr_name): """Look up a variable (possibly dotted like 'self.counter') from a namespace dict. - + Args: namespace: dict containing variables to look up attr_name: Variable name, possibly dotted (e.g., 'self.counter', 'my_var') - + Returns: tuple: (value, final_attr_name) where final_attr_name is the last part of dotted name """ @@ -349,7 +368,7 @@ def get_attribute_value_from_namespace(namespace, attr_name): f"Variable '{parts[0]}' not found in namespace", error_type=Exception, ) - + obj = namespace[parts[0]] for attr in parts[1:]: try: @@ -379,22 +398,24 @@ def get_attribute_value_from_namespace(namespace, attr_name): def mirror_all_traced_state(source, target): - ''' + """ Copy the traced boundary state of every tensor in source onto target. Torch and Warp targets are upgraded in place, so the caller's object is updated directly. A raw numpy target becomes a zero-copy view, which is a new object, so the rebuilt target is returned and numpy callers must use it. - ''' + """ if is_tracable_tensor_type(source) and is_tracable_tensor_type(target): if not is_traced_type(source): _get_logger().error( f"Error: cannot mirror traced state from a plain {type(source).__name__}.\n" "Only values registered as a node output carry the state needed " - "to connect nodes.") + "to connect nodes." + ) return target promoted = promote_in_place( - target, source.name, source.context_obj, source.proxy) + target, source.name, source.context_obj, source.proxy + ) promoted.output_port = source.output_port return promoted elif isinstance(source, collections.abc.Mapping): @@ -431,6 +452,7 @@ def get_relative_path(model_path, yaml_save_path): # If relative path calculation fails (e.g., different drives on Windows), keep absolute path return model_path + def _get_os_description(): freedesktop_os_release = getattr(platform, "freedesktop_os_release", None) if freedesktop_os_release is None: @@ -461,15 +483,16 @@ def _get_warp_version(): def get_system_info(): # lazy import to avoid circular import import leapp - metadata = {'system information': {}} - metadata['system information']['leapp version'] = leapp.__version__ - metadata['system information']['leapp config version'] = leapp.__config_version__ - metadata['system information']['torch version'] = str(torch.__version__) - metadata['system information']['python version'] = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" - metadata['system information']['cuda version'] = str(torch.version.cuda) - metadata['system information']['warp version'] = _get_warp_version() - metadata['system information']['os'] = _get_os_description() - - return metadata + metadata = {"system information": {}} + metadata["system information"]["leapp version"] = leapp.__version__ + metadata["system information"]["leapp config version"] = leapp.__config_version__ + metadata["system information"]["torch version"] = str(torch.__version__) + metadata["system information"]["python version"] = ( + f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" + ) + metadata["system information"]["cuda version"] = str(torch.version.cuda) + metadata["system information"]["warp version"] = _get_warp_version() + metadata["system information"]["os"] = _get_os_description() + return metadata diff --git a/leapp/warp_runtime.py b/leapp/warp_runtime.py index c144da09..c9769eb2 100644 --- a/leapp/warp_runtime.py +++ b/leapp/warp_runtime.py @@ -37,11 +37,7 @@ def warp_runtime_build_dir() -> Path: return cache_root / "leapp" / "warp-runtime" / "build" cache_home = os.environ.get("XDG_CACHE_HOME") - cache_root = ( - Path(cache_home).expanduser() - if cache_home - else Path.home() / ".cache" - ) + cache_root = Path(cache_home).expanduser() if cache_home else Path.home() / ".cache" return cache_root / "leapp" / "warp-runtime" / "build" @@ -54,10 +50,8 @@ def warp_runtime_artifact_paths( else warp_runtime_build_dir() ) return { - "onnx": output_dir - / _shared_library_name("leapp_wrp_onnx_custom_op"), - "pt2": output_dir - / _shared_library_name("leapp_wrp_torch_custom_op"), + "onnx": output_dir / _shared_library_name("leapp_wrp_onnx_custom_op"), + "pt2": output_dir / _shared_library_name("leapp_wrp_torch_custom_op"), } diff --git a/packages/leapp-visualization/src/leapp_visualization/geometry.py b/packages/leapp-visualization/src/leapp_visualization/geometry.py index be7da785..d4fb8ff6 100644 --- a/packages/leapp-visualization/src/leapp_visualization/geometry.py +++ b/packages/leapp-visualization/src/leapp_visualization/geometry.py @@ -115,13 +115,14 @@ def truncate_text(text: str, max_chars: int) -> str: return text[: max_chars - 1] + "..." -def resolve_geometry(graph: VisualGraph, layout: LayoutResult, graph_name: str) -> GraphGeometry: +def resolve_geometry( + graph: VisualGraph, layout: LayoutResult, graph_name: str +) -> GraphGeometry: visual_sizes = _measure_visual_sizes(graph) centers = _expand_layout_centers(layout.centers, visual_sizes) graph = _order_ports_for_geometry(graph, centers) node_geometries = { - node.id: _build_node_geometry(node, centers[node.id]) - for node in graph.nodes + node.id: _build_node_geometry(node, centers[node.id]) for node in graph.nodes } port_geometries: dict[str, PortGeometry] = {} @@ -129,7 +130,9 @@ def resolve_geometry(graph: VisualGraph, layout: LayoutResult, graph_name: str) port_geometries.update(_build_port_geometries(node, node_geometries[node.id])) terminal_geometries = { - terminal.id: _build_terminal_geometry(terminal, centers[terminal.id], port_geometries[terminal.port_id]) + terminal.id: _build_terminal_geometry( + terminal, centers[terminal.id], port_geometries[terminal.port_id] + ) for terminal in graph.terminals } @@ -138,8 +141,12 @@ def resolve_geometry(graph: VisualGraph, layout: LayoutResult, graph_name: str) + [geometry.rect for geometry in terminal_geometries.values()] ) - feedback_edge_ids = tuple(sorted(edge.id for edge in graph.edges if edge.kind == "feedback")) - feedback_lane_index = {edge_id: index for index, edge_id in enumerate(feedback_edge_ids)} + feedback_edge_ids = tuple( + sorted(edge.id for edge in graph.edges if edge.kind == "feedback") + ) + feedback_lane_index = { + edge_id: index for index, edge_id in enumerate(feedback_edge_ids) + } edge_geometries = { edge.id: _build_edge_geometry( @@ -199,7 +206,9 @@ def resolve_geometry(graph: VisualGraph, layout: LayoutResult, graph_name: str) id=geometry.id, kind=geometry.kind, label=geometry.label, - points=tuple(_shift_point(point, shift_x, shift_y) for point in geometry.points), + points=tuple( + _shift_point(point, shift_x, shift_y) for point in geometry.points + ), ) for edge_id, geometry in edge_geometries.items() } @@ -222,7 +231,9 @@ def resolve_geometry(graph: VisualGraph, layout: LayoutResult, graph_name: str) ) -def _order_ports_for_geometry(graph: VisualGraph, centers: dict[str, Point]) -> VisualGraph: +def _order_ports_for_geometry( + graph: VisualGraph, centers: dict[str, Point] +) -> VisualGraph: original_port_indices = _collect_original_port_indices(graph.nodes) input_order_keys: dict[str, list[tuple[float, float]]] = {} @@ -231,17 +242,23 @@ def _order_ports_for_geometry(graph: VisualGraph, centers: dict[str, Point]) -> continue if edge.source_port_id is None or edge.target_port_id is None: continue - source_key = _connected_port_order_key(edge.source_id, edge.source_port_id, centers, original_port_indices) + source_key = _connected_port_order_key( + edge.source_id, edge.source_port_id, centers, original_port_indices + ) input_order_keys.setdefault(edge.target_port_id, []).append(source_key) ordered_nodes = tuple( replace( node, - inputs=_reorder_connected_ports(node.inputs, input_order_keys, original_port_indices), + inputs=_reorder_connected_ports( + node.inputs, input_order_keys, original_port_indices + ), ) for node in graph.nodes ) - return VisualGraph(nodes=ordered_nodes, terminals=graph.terminals, edges=graph.edges) + return VisualGraph( + nodes=ordered_nodes, terminals=graph.terminals, edges=graph.edges + ) def _collect_original_port_indices(nodes: tuple[VisualNode, ...]) -> dict[str, int]: @@ -269,7 +286,9 @@ def _reorder_connected_ports( order_keys: dict[str, list[tuple[float, float]]], original_port_indices: dict[str, int], ) -> tuple[VisualPort, ...]: - connected_slots = [index for index, port in enumerate(ports) if port.id in order_keys] + connected_slots = [ + index for index, port in enumerate(ports) if port.id in order_keys + ] if len(connected_slots) < 2: return ports @@ -311,7 +330,9 @@ def _build_node_geometry(node: VisualNode, center: Point) -> NodeGeometry: ) -def _build_terminal_geometry(terminal: VisualTerminal, center: Point, connected_port: PortGeometry) -> TerminalGeometry: +def _build_terminal_geometry( + terminal: VisualTerminal, center: Point, connected_port: PortGeometry +) -> TerminalGeometry: width = _measure_terminal_width(terminal) rect = Rect( x=center.x - (width / 2.0), @@ -330,7 +351,9 @@ def _build_terminal_geometry(terminal: VisualTerminal, center: Point, connected_ ) -def _build_port_geometries(node: VisualNode, node_geometry: NodeGeometry) -> dict[str, PortGeometry]: +def _build_port_geometries( + node: VisualNode, node_geometry: NodeGeometry +) -> dict[str, PortGeometry]: port_geometries: dict[str, PortGeometry] = {} input_slot_width = _measure_side_slot_width(node.inputs) output_slot_width = _measure_side_slot_width(node.outputs) @@ -340,27 +363,43 @@ def _build_port_geometries(node: VisualNode, node_geometry: NodeGeometry) -> dic for index in range(row_count): row_y = first_row_y + (index * _PORT_ROW_HEIGHT) if index < len(node.inputs): - geometry = _build_port_geometry(node.inputs[index], node_geometry, row_y, input_slot_width) + geometry = _build_port_geometry( + node.inputs[index], node_geometry, row_y, input_slot_width + ) port_geometries[geometry.id] = geometry if index < len(node.outputs): - geometry = _build_port_geometry(node.outputs[index], node_geometry, row_y, output_slot_width) + geometry = _build_port_geometry( + node.outputs[index], node_geometry, row_y, output_slot_width + ) port_geometries[geometry.id] = geometry return port_geometries -def _build_port_geometry(port: VisualPort, node_geometry: NodeGeometry, row_y: float, slot_width: float) -> PortGeometry: +def _build_port_geometry( + port: VisualPort, node_geometry: NodeGeometry, row_y: float, slot_width: float +) -> PortGeometry: name = truncate_text(port.name, _PORT_NAME_MAX_CHARS) detail = truncate_text(_port_detail_text(port), _PORT_DETAIL_MAX_CHARS) - full_label = "\n".join(part for part in (port.name, _port_detail_text(port), port.kind) if part) - rect_x = node_geometry.rect.x if port.side == "input" else node_geometry.rect.x + node_geometry.rect.width - slot_width + full_label = "\n".join( + part for part in (port.name, _port_detail_text(port), port.kind) if part + ) + rect_x = ( + node_geometry.rect.x + if port.side == "input" + else node_geometry.rect.x + node_geometry.rect.width - slot_width + ) rect = Rect( x=rect_x, y=row_y, width=slot_width, height=_PORT_ROW_HEIGHT, ) - anchor_x = node_geometry.rect.x if port.side == "input" else node_geometry.rect.x + node_geometry.rect.width + anchor_x = ( + node_geometry.rect.x + if port.side == "input" + else node_geometry.rect.x + node_geometry.rect.width + ) anchor = Point(anchor_x, rect.y + (rect.height / 2.0)) return PortGeometry( id=port.id, @@ -436,7 +475,9 @@ def _measure_visual_sizes(graph: VisualGraph) -> dict[str, _Size]: return sizes -def _expand_layout_centers(centers: dict[str, Point], sizes: dict[str, _Size]) -> dict[str, Point]: +def _expand_layout_centers( + centers: dict[str, Point], sizes: dict[str, _Size] +) -> dict[str, Point]: if not centers: return {} @@ -447,8 +488,14 @@ def _expand_layout_centers(centers: dict[str, Point], sizes: dict[str, _Size]) - layer_x: dict[float, float] = {} previous_right: float | None = None for layer_key in sorted(layer_ids): - half_width = max(sizes[visual_id].width / 2.0 for visual_id in layer_ids[layer_key]) - center_x = half_width if previous_right is None else previous_right + _LAYER_GAP + half_width + half_width = max( + sizes[visual_id].width / 2.0 for visual_id in layer_ids[layer_key] + ) + center_x = ( + half_width + if previous_right is None + else previous_right + _LAYER_GAP + half_width + ) layer_x[layer_key] = center_x previous_right = center_x + half_width @@ -458,7 +505,10 @@ def _expand_layout_centers(centers: dict[str, Point], sizes: dict[str, _Size]) - for visual_id in sorted(visual_ids, key=lambda item: (centers[item].y, item)): half_height = sizes[visual_id].height / 2.0 center_y = centers[visual_id].y - if previous_bottom is not None and center_y - half_height < previous_bottom + _LAYER_ITEM_GAP: + if ( + previous_bottom is not None + and center_y - half_height < previous_bottom + _LAYER_ITEM_GAP + ): center_y = previous_bottom + _LAYER_ITEM_GAP + half_height adjusted[visual_id] = Point(layer_x[layer_key], center_y) previous_bottom = center_y + half_height @@ -469,7 +519,9 @@ def _expand_layout_centers(centers: dict[str, Point], sizes: dict[str, _Size]) - def _measure_node_size(node: VisualNode) -> _Size: input_slot_width = _measure_side_slot_width(node.inputs) output_slot_width = _measure_side_slot_width(node.outputs) - port_widths = [width for width in (input_slot_width, output_slot_width) if width > 0.0] + port_widths = [ + width for width in (input_slot_width, output_slot_width) if width > 0.0 + ] if len(port_widths) == 2: port_width = input_slot_width + _PORT_COLUMN_GAP + output_slot_width elif port_widths: @@ -479,7 +531,9 @@ def _measure_node_size(node: VisualNode) -> _Size: width = max(_MIN_NODE_WIDTH, _measure_node_header_width(node), port_width) row_count = max(len(node.inputs), len(node.outputs)) - height = _HEADER_HEIGHT + _NODE_PADDING + (row_count * _PORT_ROW_HEIGHT) + _NODE_PADDING + height = ( + _HEADER_HEIGHT + _NODE_PADDING + (row_count * _PORT_ROW_HEIGHT) + _NODE_PADDING + ) return _Size(width, height) @@ -507,7 +561,8 @@ def _measure_side_slot_width(ports: tuple[VisualPort, ...]) -> float: def _measure_port_content_width(port: VisualPort) -> float: widths = [ len(truncate_text(port.name, _PORT_NAME_MAX_CHARS)) * _PORT_PRIMARY_CHAR_WIDTH, - len(truncate_text(_port_detail_text(port), _PORT_DETAIL_MAX_CHARS)) * _PORT_DETAIL_CHAR_WIDTH, + len(truncate_text(_port_detail_text(port), _PORT_DETAIL_MAX_CHARS)) + * _PORT_DETAIL_CHAR_WIDTH, ] visible_kind = _visible_port_kind(port) if visible_kind: diff --git a/packages/leapp-visualization/src/leapp_visualization/layout.py b/packages/leapp-visualization/src/leapp_visualization/layout.py index fd4ea18a..45919090 100644 --- a/packages/leapp-visualization/src/leapp_visualization/layout.py +++ b/packages/leapp-visualization/src/leapp_visualization/layout.py @@ -46,7 +46,9 @@ def compute_layered_layout(graph: VisualGraph) -> LayoutResult: if not visual_ids: return LayoutResult(centers={}, forward_edge_ids=()) - layout_edges = tuple(edge for edge in graph.edges if edge.kind in _LAYOUT_EDGE_KINDS) + layout_edges = tuple( + edge for edge in graph.edges if edge.kind in _LAYOUT_EDGE_KINDS + ) if not layout_edges: return LayoutResult( centers=_deterministic_grid(visual_ids), @@ -55,11 +57,14 @@ def compute_layered_layout(graph: VisualGraph) -> LayoutResult: id_to_index = {visual_id: index for index, visual_id in enumerate(visual_ids)} edge_pairs = sorted( - (id_to_index[edge.source_id], id_to_index[edge.target_id]) for edge in layout_edges + (id_to_index[edge.source_id], id_to_index[edge.target_id]) + for edge in layout_edges ) raw_positions = _compute_sugiyama_positions(edge_pairs) centers = _normalize_points(_decode_positions(visual_ids, raw_positions)) - missing_ids = tuple(visual_id for visual_id in visual_ids if visual_id not in centers) + missing_ids = tuple( + visual_id for visual_id in visual_ids if visual_id not in centers + ) if missing_ids: centers.update(_place_missing_ids(centers, missing_ids)) @@ -69,7 +74,9 @@ def compute_layered_layout(graph: VisualGraph) -> LayoutResult: ) -def _compute_sugiyama_positions(edge_pairs: list[tuple[int, int]]) -> dict[int, tuple[float, float]]: +def _compute_sugiyama_positions( + edge_pairs: list[tuple[int, int]], +) -> dict[int, tuple[float, float]]: layouts = from_edges( edge_pairs, vertex_spacing=96, @@ -103,7 +110,9 @@ def _normalize_points(points: dict[str, Point]) -> dict[str, Point]: } -def _place_missing_ids(points: dict[str, Point], missing_ids: tuple[str, ...]) -> dict[str, Point]: +def _place_missing_ids( + points: dict[str, Point], missing_ids: tuple[str, ...] +) -> dict[str, Point]: next_x = max(point.x for point in points.values()) + _GRID_X_STEP return { visual_id: Point(x=next_x, y=row * _GRID_Y_STEP) diff --git a/packages/leapp-visualization/src/leapp_visualization/model.py b/packages/leapp-visualization/src/leapp_visualization/model.py index 1ec11598..d2fa7cc1 100644 --- a/packages/leapp-visualization/src/leapp_visualization/model.py +++ b/packages/leapp-visualization/src/leapp_visualization/model.py @@ -72,7 +72,9 @@ class VisualGraph: edges: tuple[VisualEdge, ...] def visual_ids(self) -> tuple[str, ...]: - return tuple(node.id for node in self.nodes) + tuple(terminal.id for terminal in self.terminals) + return tuple(node.id for node in self.nodes) + tuple( + terminal.id for terminal in self.terminals + ) def visual_id(*parts: object) -> str: diff --git a/packages/leapp-visualization/src/leapp_visualization/png_renderer.py b/packages/leapp-visualization/src/leapp_visualization/png_renderer.py index eca5a637..1a167a1c 100644 --- a/packages/leapp-visualization/src/leapp_visualization/png_renderer.py +++ b/packages/leapp-visualization/src/leapp_visualization/png_renderer.py @@ -69,7 +69,9 @@ def write_png(path: str, geometry: GraphGeometry) -> None: _draw_edge(draw, edge) for terminal in geometry.terminals.values(): - _draw_terminal(draw, terminal.rect, terminal.anchor, terminal.title, terminal.kind, fonts) + _draw_terminal( + draw, terminal.rect, terminal.anchor, terminal.title, terminal.kind, fonts + ) for node in geometry.nodes.values(): _draw_node_card(draw, node.rect, node.header_rect) @@ -84,30 +86,34 @@ def write_png(path: str, geometry: GraphGeometry) -> None: anchor="ls", ) - for port in sorted(geometry.ports.values(), key=lambda port: (port.rect.y, port.rect.x, port.id)): + for port in sorted( + geometry.ports.values(), key=lambda port: (port.rect.y, port.rect.x, port.id) + ): _draw_port(draw, port, fonts) final_image = image.convert("RGB").resize((width, height), Image.Resampling.LANCZOS) final_image.save(path, format="PNG") -def sample_cubic(points: tuple[Point, Point, Point, Point], samples: int = 24) -> list[Point]: +def sample_cubic( + points: tuple[Point, Point, Point, Point], samples: int = 24 +) -> list[Point]: start, control1, control2, end = points sampled: list[Point] = [] for index in range(samples + 1): t = index / samples one_minus_t = 1.0 - t x = ( - (one_minus_t ** 3) * start.x - + 3.0 * (one_minus_t ** 2) * t * control1.x - + 3.0 * one_minus_t * (t ** 2) * control2.x - + (t ** 3) * end.x + (one_minus_t**3) * start.x + + 3.0 * (one_minus_t**2) * t * control1.x + + 3.0 * one_minus_t * (t**2) * control2.x + + (t**3) * end.x ) y = ( - (one_minus_t ** 3) * start.y - + 3.0 * (one_minus_t ** 2) * t * control1.y - + 3.0 * one_minus_t * (t ** 2) * control2.y - + (t ** 3) * end.y + (one_minus_t**3) * start.y + + 3.0 * (one_minus_t**2) * t * control1.y + + 3.0 * one_minus_t * (t**2) * control2.y + + (t**3) * end.y ) sampled.append(Point(x, y)) return sampled @@ -149,12 +155,12 @@ def sample_rounded_polyline( one_minus_t = 1.0 - t sampled.append( Point( - (one_minus_t ** 2) * before.x + (one_minus_t**2) * before.x + 2.0 * one_minus_t * t * corner.x - + (t ** 2) * after.x, - (one_minus_t ** 2) * before.y + + (t**2) * after.x, + (one_minus_t**2) * before.y + 2.0 * one_minus_t * t * corner.y - + (t ** 2) * after.y, + + (t**2) * after.y, ) ) sampled.append(points[-1]) @@ -168,11 +174,21 @@ def _draw_edge(draw: ImageDraw.ImageDraw, edge: EdgeGeometry) -> None: else: route = sample_cubic(edge.points) # type: ignore[arg-type] - draw.line(_scaled_points(route), fill=color, width=_scaled_width(_EDGE_WIDTH), joint="curve") + draw.line( + _scaled_points(route), + fill=color, + width=_scaled_width(_EDGE_WIDTH), + joint="curve", + ) _draw_arrowhead(draw, route[-2], route[-1], color) -def _draw_arrowhead(draw: ImageDraw.ImageDraw, start: Point, end: Point, color: tuple[int, int, int, int]) -> None: +def _draw_arrowhead( + draw: ImageDraw.ImageDraw, + start: Point, + end: Point, + color: tuple[int, int, int, int], +) -> None: dx = end.x - start.x dy = end.y - start.y length = math.hypot(dx, dy) @@ -186,8 +202,16 @@ def _draw_arrowhead(draw: ImageDraw.ImageDraw, start: Point, end: Point, color: tip = _scaled_point(end) base_x = end.x - (unit_x * _ARROW_LENGTH) base_y = end.y - (unit_y * _ARROW_LENGTH) - left = _scaled_point(Point(base_x + (perp_x * _ARROW_HALF_WIDTH), base_y + (perp_y * _ARROW_HALF_WIDTH))) - right = _scaled_point(Point(base_x - (perp_x * _ARROW_HALF_WIDTH), base_y - (perp_y * _ARROW_HALF_WIDTH))) + left = _scaled_point( + Point( + base_x + (perp_x * _ARROW_HALF_WIDTH), base_y + (perp_y * _ARROW_HALF_WIDTH) + ) + ) + right = _scaled_point( + Point( + base_x - (perp_x * _ARROW_HALF_WIDTH), base_y - (perp_y * _ARROW_HALF_WIDTH) + ) + ) draw.polygon([left, tip, right], fill=color) @@ -209,7 +233,11 @@ def _draw_terminal( radius=_TERMINAL_RADIUS, ) anchor_on_right = anchor.x >= rect.x + rect.width - text_x = rect.x + _TERMINAL_PADDING if anchor_on_right else rect.x + rect.width - _TERMINAL_PADDING + text_x = ( + rect.x + _TERMINAL_PADDING + if anchor_on_right + else rect.x + rect.width - _TERMINAL_PADDING + ) _draw_text( draw, Point(text_x, rect.y + rect.height / 2.0 + 5.0), @@ -221,9 +249,25 @@ def _draw_terminal( def _draw_node_card(draw: ImageDraw.ImageDraw, rect: Rect, header_rect: Rect) -> None: - _rounded_rect(draw, rect, fill=_color("node"), outline=_color("node_border"), width=_NODE_BORDER_WIDTH, radius=_NODE_RADIUS) - _rounded_rect(draw, header_rect, fill=_color("header"), outline=None, width=0, radius=_NODE_RADIUS) - bottom_rect = Rect(header_rect.x, header_rect.y + header_rect.height - 8.0, header_rect.width, 8.0) + _rounded_rect( + draw, + rect, + fill=_color("node"), + outline=_color("node_border"), + width=_NODE_BORDER_WIDTH, + radius=_NODE_RADIUS, + ) + _rounded_rect( + draw, + header_rect, + fill=_color("header"), + outline=None, + width=0, + radius=_NODE_RADIUS, + ) + bottom_rect = Rect( + header_rect.x, header_rect.y + header_rect.height - 8.0, header_rect.width, 8.0 + ) _scaled_rect(draw, bottom_rect, fill=_color("header")) @@ -255,8 +299,14 @@ def _draw_port( fonts: dict[str, ImageFont.ImageFont | ImageFont.FreeTypeFont], ) -> None: accent = _accent_color(port.kind) - line_x = port.rect.x if port.side == "input" else port.rect.x + port.rect.width - 1.5 - text_x = port.rect.x + _NODE_PADDING if port.side == "input" else port.rect.x + port.rect.width - _NODE_PADDING + line_x = ( + port.rect.x if port.side == "input" else port.rect.x + port.rect.width - 1.5 + ) + text_x = ( + port.rect.x + _NODE_PADDING + if port.side == "input" + else port.rect.x + port.rect.width - _NODE_PADDING + ) anchor = "ls" if port.side == "input" else "rs" draw.ellipse(_ellipse_bounds(port.anchor, _PORT_DOT_RADIUS), fill=accent) @@ -270,14 +320,39 @@ def _draw_port( fill=accent, width=_scaled_width(_PORT_ACCENT_WIDTH), ) - _draw_text(draw, Point(text_x, port.rect.y + 18.0), port.name, fill=_color("text"), font=fonts["label_bold"], anchor=anchor) - _draw_text(draw, Point(text_x, port.rect.y + 34.0), port.detail, fill=_color("secondary_text"), font=fonts["label"], anchor=anchor) + _draw_text( + draw, + Point(text_x, port.rect.y + 18.0), + port.name, + fill=_color("text"), + font=fonts["label_bold"], + anchor=anchor, + ) + _draw_text( + draw, + Point(text_x, port.rect.y + 34.0), + port.detail, + fill=_color("secondary_text"), + font=fonts["label"], + anchor=anchor, + ) visible_kind = _visible_port_kind(port) if visible_kind: - kind_x = port.rect.x + _NODE_PADDING if port.side == "input" else port.rect.x + port.rect.width - _NODE_PADDING + kind_x = ( + port.rect.x + _NODE_PADDING + if port.side == "input" + else port.rect.x + port.rect.width - _NODE_PADDING + ) kind_anchor = "ls" if port.side == "input" else "rs" - _draw_text(draw, Point(kind_x, port.rect.y + 50.0), visible_kind, fill=accent, font=fonts["label"], anchor=kind_anchor) + _draw_text( + draw, + Point(kind_x, port.rect.y + 50.0), + visible_kind, + fill=accent, + font=fonts["label"], + anchor=kind_anchor, + ) def _edge_color_name(kind: str) -> str: @@ -352,12 +427,16 @@ def _rounded_rect( ) -def _scaled_rect(draw: ImageDraw.ImageDraw, rect: Rect, *, fill: tuple[int, int, int, int]) -> None: +def _scaled_rect( + draw: ImageDraw.ImageDraw, rect: Rect, *, fill: tuple[int, int, int, int] +) -> None: draw.rectangle(_scaled_box(rect), fill=fill) def _ellipse_bounds(center: Point, radius: float) -> tuple[float, float, float, float]: - return _scaled_box(Rect(center.x - radius, center.y - radius, radius * 2.0, radius * 2.0)) + return _scaled_box( + Rect(center.x - radius, center.y - radius, radius * 2.0, radius * 2.0) + ) def _scaled_box(rect: Rect) -> tuple[int, int, int, int]: @@ -395,7 +474,9 @@ def _accent_color(name: str | None) -> tuple[int, int, int, int]: return _hex_to_rgba(COLORS.get(name, COLORS["node_border"])) -def _blend_with_background(color: tuple[int, int, int, int], alpha: int) -> tuple[int, int, int, int]: +def _blend_with_background( + color: tuple[int, int, int, int], alpha: int +) -> tuple[int, int, int, int]: """Flatten a translucent tint onto the canvas background. ImageDraw overwrites pixels instead of compositing them, so a low-alpha fill diff --git a/packages/leapp-visualization/tests/test_graph_visualization_builder.py b/packages/leapp-visualization/tests/test_graph_visualization_builder.py index d4b5a7fb..9b8fd415 100644 --- a/packages/leapp-visualization/tests/test_graph_visualization_builder.py +++ b/packages/leapp-visualization/tests/test_graph_visualization_builder.py @@ -43,8 +43,16 @@ class FakeNode: def test_build_visual_graph_preserves_ports_semantics_and_external_io(): policy = FakeNode( name="policy", - inputs=[FakeTensorDescription("obs", "float32", (1, 12), semantics={"kind": "state"})], - outputs=[FakeTensorDescription("action", "float32", (1, 4), semantics={"kind": "command"})], + inputs=[ + FakeTensorDescription( + "obs", "float32", (1, 12), semantics={"kind": "state"} + ) + ], + outputs=[ + FakeTensorDescription( + "action", "float32", (1, 4), semantics={"kind": "command"} + ) + ], backend="onnx-dynamo", node_index=0, ) @@ -106,8 +114,18 @@ def test_build_visual_graph_marks_feedback_edges_without_adding_them_to_forward_ graph = build_visual_graph( nodes={"first": first, "second": second}, - connections=[{"source": {"node": first, "idx": 0}, "targets": [{"node": second, "idx": 0}]}], - feedback_connections=[{"source": {"node": second, "idx": 0}, "targets": [{"node": first, "idx": 0}]}], + connections=[ + { + "source": {"node": first, "idx": 0}, + "targets": [{"node": second, "idx": 0}], + } + ], + feedback_connections=[ + { + "source": {"node": second, "idx": 0}, + "targets": [{"node": first, "idx": 0}], + } + ], graph_inputs=[], graph_outputs=[], ) diff --git a/packages/leapp-visualization/tests/test_graph_visualization_compatibility.py b/packages/leapp-visualization/tests/test_graph_visualization_compatibility.py index 7cf5d7a8..99400623 100644 --- a/packages/leapp-visualization/tests/test_graph_visualization_compatibility.py +++ b/packages/leapp-visualization/tests/test_graph_visualization_compatibility.py @@ -32,10 +32,8 @@ def test_visualization_dependency_is_conditional_on_python_311(): root_pyproject = (ROOT / "pyproject.toml").read_text(encoding="utf-8") assert 'requires-python = ">=3.10"' in root_pyproject - assert ( - '"fast-sugiyama>=0.5.3; python_version >= \'3.11\'"' in root_pyproject - ) - assert '"Pillow>=10.0.0; python_version >= \'3.11\'"' in root_pyproject + assert "\"fast-sugiyama>=0.5.3; python_version >= '3.11'\"" in root_pyproject + assert "\"Pillow>=10.0.0; python_version >= '3.11'\"" in root_pyproject def test_visualization_modules_are_shipped_with_leapp(): diff --git a/packages/leapp-visualization/tests/test_graph_visualization_geometry.py b/packages/leapp-visualization/tests/test_graph_visualization_geometry.py index 2c8356be..474bf599 100644 --- a/packages/leapp-visualization/tests/test_graph_visualization_geometry.py +++ b/packages/leapp-visualization/tests/test_graph_visualization_geometry.py @@ -18,7 +18,13 @@ from leapp_visualization import geometry as geometry_module from leapp_visualization.geometry import resolve_geometry from leapp_visualization.layout import LayoutResult, Point -from leapp_visualization.model import VisualEdge, VisualGraph, VisualNode, VisualPort, VisualTerminal +from leapp_visualization.model import ( + VisualEdge, + VisualGraph, + VisualNode, + VisualPort, + VisualTerminal, +) def test_geometry_places_input_and_output_ports_on_node_edges(): @@ -26,8 +32,28 @@ def test_geometry_places_input_and_output_ports_on_node_edges(): id="node:policy", title="policy", backend="onnx-dynamo", - inputs=(VisualPort("port:policy:input:obs", "node:policy", "input", "obs", ("1", "12"), "float32", "state"),), - outputs=(VisualPort("port:policy:output:action", "node:policy", "output", "action", ("1", "4"), "float32", "command"),), + inputs=( + VisualPort( + "port:policy:input:obs", + "node:policy", + "input", + "obs", + ("1", "12"), + "float32", + "state", + ), + ), + outputs=( + VisualPort( + "port:policy:output:action", + "node:policy", + "output", + "action", + ("1", "4"), + "float32", + "command", + ), + ), ) graph = VisualGraph(nodes=(node,), terminals=(), edges=()) layout = LayoutResult(centers={"node:policy": Point(0.0, 0.0)}, forward_edge_ids=()) @@ -51,12 +77,44 @@ def test_geometry_pairs_input_and_output_ports_by_row(): title="policy", backend="onnx", inputs=( - VisualPort("port:policy:input:obs", "node:policy", "input", "obs", ("1", "12"), "float32", "state/joint/position"), - VisualPort("port:policy:input:vel", "node:policy", "input", "vel", ("1", "12"), "float32", "state/joint/velocity"), + VisualPort( + "port:policy:input:obs", + "node:policy", + "input", + "obs", + ("1", "12"), + "float32", + "state/joint/position", + ), + VisualPort( + "port:policy:input:vel", + "node:policy", + "input", + "vel", + ("1", "12"), + "float32", + "state/joint/velocity", + ), ), outputs=( - VisualPort("port:policy:output:action", "node:policy", "output", "action", ("1", "12"), "float32", "target/joint/position"), - VisualPort("port:policy:output:effort", "node:policy", "output", "effort", ("1", "12"), "float32", "target/joint/effort"), + VisualPort( + "port:policy:output:action", + "node:policy", + "output", + "action", + ("1", "12"), + "float32", + "target/joint/position", + ), + VisualPort( + "port:policy:output:effort", + "node:policy", + "output", + "effort", + ("1", "12"), + "float32", + "target/joint/effort", + ), ), ) graph = VisualGraph(nodes=(node,), terminals=(), edges=()) @@ -64,8 +122,14 @@ def test_geometry_pairs_input_and_output_ports_by_row(): geometry = resolve_geometry(graph, layout, graph_name="demo") - assert geometry.ports["port:policy:input:obs"].anchor.y == geometry.ports["port:policy:output:action"].anchor.y - assert geometry.ports["port:policy:input:vel"].anchor.y == geometry.ports["port:policy:output:effort"].anchor.y + assert ( + geometry.ports["port:policy:input:obs"].anchor.y + == geometry.ports["port:policy:output:action"].anchor.y + ) + assert ( + geometry.ports["port:policy:input:vel"].anchor.y + == geometry.ports["port:policy:output:effort"].anchor.y + ) def test_geometry_keeps_two_sided_port_columns_from_overlapping(): @@ -101,8 +165,12 @@ def test_geometry_keeps_two_sided_port_columns_from_overlapping(): geometry = resolve_geometry(graph, layout, graph_name="demo") - input_rect = geometry.ports["port:policy:input:backbone_outputs_attention_mask"].rect - output_rect = geometry.ports["port:policy:output:converted_outputs_backbone_attention_mask"].rect + input_rect = geometry.ports[ + "port:policy:input:backbone_outputs_attention_mask" + ].rect + output_rect = geometry.ports[ + "port:policy:output:converted_outputs_backbone_attention_mask" + ].rect assert input_rect.x + input_rect.width < output_rect.x @@ -162,18 +230,66 @@ def test_geometry_keeps_generous_canvas_margin_around_content(): id="node:policy", title="policy", backend=None, - inputs=(VisualPort("port:policy:input:obs", "node:policy", "input", "obs", ("1", "12"), "float32", None),), - outputs=(VisualPort("port:policy:output:action", "node:policy", "output", "action", ("1", "12"), "float32", None),), + inputs=( + VisualPort( + "port:policy:input:obs", + "node:policy", + "input", + "obs", + ("1", "12"), + "float32", + None, + ), + ), + outputs=( + VisualPort( + "port:policy:output:action", + "node:policy", + "output", + "action", + ("1", "12"), + "float32", + None, + ), + ), ) graph = VisualGraph( nodes=(node,), terminals=( - VisualTerminal("terminal:input:policy:obs", "graph_input", "obs", "node:policy", "port:policy:input:obs"), - VisualTerminal("terminal:output:policy:action", "graph_output", "action", "node:policy", "port:policy:output:action"), + VisualTerminal( + "terminal:input:policy:obs", + "graph_input", + "obs", + "node:policy", + "port:policy:input:obs", + ), + VisualTerminal( + "terminal:output:policy:action", + "graph_output", + "action", + "node:policy", + "port:policy:output:action", + ), ), edges=( - VisualEdge("edge:input", "graph_input", "terminal:input:policy:obs", "node:policy", None, "port:policy:input:obs", "obs"), - VisualEdge("edge:output", "graph_output", "node:policy", "terminal:output:policy:action", "port:policy:output:action", None, "action"), + VisualEdge( + "edge:input", + "graph_input", + "terminal:input:policy:obs", + "node:policy", + None, + "port:policy:input:obs", + "obs", + ), + VisualEdge( + "edge:output", + "graph_output", + "node:policy", + "terminal:output:policy:action", + "port:policy:output:action", + None, + "action", + ), ), ) layout = LayoutResult( @@ -201,23 +317,93 @@ def test_geometry_aligns_terminal_pills_to_connected_port_rows(): title="preprocess_state", backend="onnx", inputs=( - VisualPort("port:preprocess_state:input:left_leg", "node:preprocess_state", "input", "left_leg", ("1", "6"), "float32", None), - VisualPort("port:preprocess_state:input:right_leg", "node:preprocess_state", "input", "right_leg", ("1", "6"), "float32", None), + VisualPort( + "port:preprocess_state:input:left_leg", + "node:preprocess_state", + "input", + "left_leg", + ("1", "6"), + "float32", + None, + ), + VisualPort( + "port:preprocess_state:input:right_leg", + "node:preprocess_state", + "input", + "right_leg", + ("1", "6"), + "float32", + None, + ), ), outputs=( - VisualPort("port:preprocess_state:output:reference_0_left_leg", "node:preprocess_state", "output", "reference_0_left_leg", ("1", "6"), "float32", None), - VisualPort("port:preprocess_state:output:reference_0_right_leg", "node:preprocess_state", "output", "reference_0_right_leg", ("1", "6"), "float32", None), - VisualPort("port:preprocess_state:output:reference_0_waist", "node:preprocess_state", "output", "reference_0_waist", ("1", "3"), "float32", None), + VisualPort( + "port:preprocess_state:output:reference_0_left_leg", + "node:preprocess_state", + "output", + "reference_0_left_leg", + ("1", "6"), + "float32", + None, + ), + VisualPort( + "port:preprocess_state:output:reference_0_right_leg", + "node:preprocess_state", + "output", + "reference_0_right_leg", + ("1", "6"), + "float32", + None, + ), + VisualPort( + "port:preprocess_state:output:reference_0_waist", + "node:preprocess_state", + "output", + "reference_0_waist", + ("1", "3"), + "float32", + None, + ), ), ) graph = VisualGraph( nodes=(node,), terminals=( - VisualTerminal("terminal:input:preprocess_state:left_leg", "graph_input", "left_leg", "node:preprocess_state", "port:preprocess_state:input:left_leg"), - VisualTerminal("terminal:input:preprocess_state:right_leg", "graph_input", "right_leg", "node:preprocess_state", "port:preprocess_state:input:right_leg"), - VisualTerminal("terminal:output:preprocess_state:reference_0_left_leg", "graph_output", "reference_0_left_leg", "node:preprocess_state", "port:preprocess_state:output:reference_0_left_leg"), - VisualTerminal("terminal:output:preprocess_state:reference_0_right_leg", "graph_output", "reference_0_right_leg", "node:preprocess_state", "port:preprocess_state:output:reference_0_right_leg"), - VisualTerminal("terminal:output:preprocess_state:reference_0_waist", "graph_output", "reference_0_waist", "node:preprocess_state", "port:preprocess_state:output:reference_0_waist"), + VisualTerminal( + "terminal:input:preprocess_state:left_leg", + "graph_input", + "left_leg", + "node:preprocess_state", + "port:preprocess_state:input:left_leg", + ), + VisualTerminal( + "terminal:input:preprocess_state:right_leg", + "graph_input", + "right_leg", + "node:preprocess_state", + "port:preprocess_state:input:right_leg", + ), + VisualTerminal( + "terminal:output:preprocess_state:reference_0_left_leg", + "graph_output", + "reference_0_left_leg", + "node:preprocess_state", + "port:preprocess_state:output:reference_0_left_leg", + ), + VisualTerminal( + "terminal:output:preprocess_state:reference_0_right_leg", + "graph_output", + "reference_0_right_leg", + "node:preprocess_state", + "port:preprocess_state:output:reference_0_right_leg", + ), + VisualTerminal( + "terminal:output:preprocess_state:reference_0_waist", + "graph_output", + "reference_0_waist", + "node:preprocess_state", + "port:preprocess_state:output:reference_0_waist", + ), ), edges=(), ) @@ -226,8 +412,12 @@ def test_geometry_aligns_terminal_pills_to_connected_port_rows(): "terminal:input:preprocess_state:left_leg": Point(0.0, 240.0), "terminal:input:preprocess_state:right_leg": Point(0.0, 80.0), "node:preprocess_state": Point(260.0, 160.0), - "terminal:output:preprocess_state:reference_0_left_leg": Point(520.0, 300.0), - "terminal:output:preprocess_state:reference_0_right_leg": Point(520.0, 80.0), + "terminal:output:preprocess_state:reference_0_left_leg": Point( + 520.0, 300.0 + ), + "terminal:output:preprocess_state:reference_0_right_leg": Point( + 520.0, 80.0 + ), "terminal:output:preprocess_state:reference_0_waist": Point(520.0, 180.0), }, forward_edge_ids=(), @@ -236,7 +426,10 @@ def test_geometry_aligns_terminal_pills_to_connected_port_rows(): geometry = resolve_geometry(graph, layout, graph_name="demo") for terminal in graph.terminals: - assert geometry.terminals[terminal.id].anchor.y == geometry.ports[terminal.port_id].anchor.y + assert ( + geometry.terminals[terminal.id].anchor.y + == geometry.ports[terminal.port_id].anchor.y + ) def test_geometry_reorders_target_inputs_by_connected_source_output_order(): @@ -246,11 +439,51 @@ def test_geometry_reorders_target_inputs_by_connected_source_output_order(): backend="onnx", inputs=(), outputs=( - VisualPort("port:preprocess_video:output:pixel_values", "node:preprocess_video", "output", "pixel_values", ("352", "1536"), "float32", None), - VisualPort("port:preprocess_video:output:input_ids", "node:preprocess_video", "output", "input_ids", ("1", "101"), "int64", None), - VisualPort("port:preprocess_video:output:attention_mask", "node:preprocess_video", "output", "attention_mask", ("1", "101"), "int64", None), - VisualPort("port:preprocess_video:output:image_grid_thw", "node:preprocess_video", "output", "image_grid_thw", ("1", "3"), "int64", None), - VisualPort("port:preprocess_video:output:embodiment_id", "node:preprocess_video", "output", "embodiment_id", ("1",), "int64", None), + VisualPort( + "port:preprocess_video:output:pixel_values", + "node:preprocess_video", + "output", + "pixel_values", + ("352", "1536"), + "float32", + None, + ), + VisualPort( + "port:preprocess_video:output:input_ids", + "node:preprocess_video", + "output", + "input_ids", + ("1", "101"), + "int64", + None, + ), + VisualPort( + "port:preprocess_video:output:attention_mask", + "node:preprocess_video", + "output", + "attention_mask", + ("1", "101"), + "int64", + None, + ), + VisualPort( + "port:preprocess_video:output:image_grid_thw", + "node:preprocess_video", + "output", + "image_grid_thw", + ("1", "3"), + "int64", + None, + ), + VisualPort( + "port:preprocess_video:output:embodiment_id", + "node:preprocess_video", + "output", + "embodiment_id", + ("1",), + "int64", + None, + ), ), ) target = VisualNode( @@ -258,11 +491,51 @@ def test_geometry_reorders_target_inputs_by_connected_source_output_order(): title="action_head", backend="onnx", inputs=( - VisualPort("port:action_head:input:action_inputs_embodiment_id", "node:action_head", "input", "action_inputs_embodiment_id", ("1",), "int64", None), - VisualPort("port:action_head:input:action_inputs_image_grid_thw", "node:action_head", "input", "action_inputs_image_grid_thw", ("1", "3"), "int64", None), - VisualPort("port:action_head:input:action_inputs_attention_mask", "node:action_head", "input", "action_inputs_attention_mask", ("1", "101"), "int64", None), - VisualPort("port:action_head:input:action_inputs_input_ids", "node:action_head", "input", "action_inputs_input_ids", ("1", "101"), "int64", None), - VisualPort("port:action_head:input:action_inputs_pixel_values", "node:action_head", "input", "action_inputs_pixel_values", ("352", "1536"), "float32", None), + VisualPort( + "port:action_head:input:action_inputs_embodiment_id", + "node:action_head", + "input", + "action_inputs_embodiment_id", + ("1",), + "int64", + None, + ), + VisualPort( + "port:action_head:input:action_inputs_image_grid_thw", + "node:action_head", + "input", + "action_inputs_image_grid_thw", + ("1", "3"), + "int64", + None, + ), + VisualPort( + "port:action_head:input:action_inputs_attention_mask", + "node:action_head", + "input", + "action_inputs_attention_mask", + ("1", "101"), + "int64", + None, + ), + VisualPort( + "port:action_head:input:action_inputs_input_ids", + "node:action_head", + "input", + "action_inputs_input_ids", + ("1", "101"), + "int64", + None, + ), + VisualPort( + "port:action_head:input:action_inputs_pixel_values", + "node:action_head", + "input", + "action_inputs_pixel_values", + ("352", "1536"), + "float32", + None, + ), ), outputs=(), ) @@ -270,11 +543,51 @@ def test_geometry_reorders_target_inputs_by_connected_source_output_order(): nodes=(source, target), terminals=(), edges=( - VisualEdge("edge:pixel", "forward", "node:preprocess_video", "node:action_head", "port:preprocess_video:output:pixel_values", "port:action_head:input:action_inputs_pixel_values", "pixel_values"), - VisualEdge("edge:input_ids", "forward", "node:preprocess_video", "node:action_head", "port:preprocess_video:output:input_ids", "port:action_head:input:action_inputs_input_ids", "input_ids"), - VisualEdge("edge:attention_mask", "forward", "node:preprocess_video", "node:action_head", "port:preprocess_video:output:attention_mask", "port:action_head:input:action_inputs_attention_mask", "attention_mask"), - VisualEdge("edge:image_grid_thw", "forward", "node:preprocess_video", "node:action_head", "port:preprocess_video:output:image_grid_thw", "port:action_head:input:action_inputs_image_grid_thw", "image_grid_thw"), - VisualEdge("edge:embodiment_id", "forward", "node:preprocess_video", "node:action_head", "port:preprocess_video:output:embodiment_id", "port:action_head:input:action_inputs_embodiment_id", "embodiment_id"), + VisualEdge( + "edge:pixel", + "forward", + "node:preprocess_video", + "node:action_head", + "port:preprocess_video:output:pixel_values", + "port:action_head:input:action_inputs_pixel_values", + "pixel_values", + ), + VisualEdge( + "edge:input_ids", + "forward", + "node:preprocess_video", + "node:action_head", + "port:preprocess_video:output:input_ids", + "port:action_head:input:action_inputs_input_ids", + "input_ids", + ), + VisualEdge( + "edge:attention_mask", + "forward", + "node:preprocess_video", + "node:action_head", + "port:preprocess_video:output:attention_mask", + "port:action_head:input:action_inputs_attention_mask", + "attention_mask", + ), + VisualEdge( + "edge:image_grid_thw", + "forward", + "node:preprocess_video", + "node:action_head", + "port:preprocess_video:output:image_grid_thw", + "port:action_head:input:action_inputs_image_grid_thw", + "image_grid_thw", + ), + VisualEdge( + "edge:embodiment_id", + "forward", + "node:preprocess_video", + "node:action_head", + "port:preprocess_video:output:embodiment_id", + "port:action_head:input:action_inputs_embodiment_id", + "embodiment_id", + ), ), ) layout = LayoutResult( @@ -305,9 +618,33 @@ def test_geometry_preserves_source_outputs_and_reorders_target_inputs_to_match() backend="onnx", inputs=(), outputs=( - VisualPort("port:producer:output:c", "node:producer", "output", "c", ("1",), "float32", None), - VisualPort("port:producer:output:b", "node:producer", "output", "b", ("1",), "float32", None), - VisualPort("port:producer:output:a", "node:producer", "output", "a", ("1",), "float32", None), + VisualPort( + "port:producer:output:c", + "node:producer", + "output", + "c", + ("1",), + "float32", + None, + ), + VisualPort( + "port:producer:output:b", + "node:producer", + "output", + "b", + ("1",), + "float32", + None, + ), + VisualPort( + "port:producer:output:a", + "node:producer", + "output", + "a", + ("1",), + "float32", + None, + ), ), ) target = VisualNode( @@ -315,9 +652,33 @@ def test_geometry_preserves_source_outputs_and_reorders_target_inputs_to_match() title="consumer", backend="onnx", inputs=( - VisualPort("port:consumer:input:a", "node:consumer", "input", "a", ("1",), "float32", None), - VisualPort("port:consumer:input:b", "node:consumer", "input", "b", ("1",), "float32", None), - VisualPort("port:consumer:input:c", "node:consumer", "input", "c", ("1",), "float32", None), + VisualPort( + "port:consumer:input:a", + "node:consumer", + "input", + "a", + ("1",), + "float32", + None, + ), + VisualPort( + "port:consumer:input:b", + "node:consumer", + "input", + "b", + ("1",), + "float32", + None, + ), + VisualPort( + "port:consumer:input:c", + "node:consumer", + "input", + "c", + ("1",), + "float32", + None, + ), ), outputs=(), ) @@ -325,9 +686,33 @@ def test_geometry_preserves_source_outputs_and_reorders_target_inputs_to_match() nodes=(source, target), terminals=(), edges=( - VisualEdge("edge:a", "forward", "node:producer", "node:consumer", "port:producer:output:a", "port:consumer:input:a", "a"), - VisualEdge("edge:b", "forward", "node:producer", "node:consumer", "port:producer:output:b", "port:consumer:input:b", "b"), - VisualEdge("edge:c", "forward", "node:producer", "node:consumer", "port:producer:output:c", "port:consumer:input:c", "c"), + VisualEdge( + "edge:a", + "forward", + "node:producer", + "node:consumer", + "port:producer:output:a", + "port:consumer:input:a", + "a", + ), + VisualEdge( + "edge:b", + "forward", + "node:producer", + "node:consumer", + "port:producer:output:b", + "port:consumer:input:b", + "b", + ), + VisualEdge( + "edge:c", + "forward", + "node:producer", + "node:consumer", + "port:producer:output:c", + "port:consumer:input:c", + "c", + ), ), ) layout = LayoutResult( @@ -364,9 +749,33 @@ def test_geometry_reorders_ports_iteratively_for_fanout_targets(): backend="onnx", inputs=(), outputs=( - VisualPort("port:preprocess_video:output:pixel_values", "node:preprocess_video", "output", "pixel_values", ("1",), "float32", None), - VisualPort("port:preprocess_video:output:input_ids", "node:preprocess_video", "output", "input_ids", ("1",), "int64", None), - VisualPort("port:preprocess_video:output:attention_mask", "node:preprocess_video", "output", "attention_mask", ("1",), "int64", None), + VisualPort( + "port:preprocess_video:output:pixel_values", + "node:preprocess_video", + "output", + "pixel_values", + ("1",), + "float32", + None, + ), + VisualPort( + "port:preprocess_video:output:input_ids", + "node:preprocess_video", + "output", + "input_ids", + ("1",), + "int64", + None, + ), + VisualPort( + "port:preprocess_video:output:attention_mask", + "node:preprocess_video", + "output", + "attention_mask", + ("1",), + "int64", + None, + ), ), ) first_target = VisualNode( @@ -374,9 +783,33 @@ def test_geometry_reorders_ports_iteratively_for_fanout_targets(): title="backbone", backend="onnx", inputs=( - VisualPort("port:backbone:input:attention_mask", "node:backbone", "input", "attention_mask", ("1",), "int64", None), - VisualPort("port:backbone:input:input_ids", "node:backbone", "input", "input_ids", ("1",), "int64", None), - VisualPort("port:backbone:input:pixel_values", "node:backbone", "input", "pixel_values", ("1",), "float32", None), + VisualPort( + "port:backbone:input:attention_mask", + "node:backbone", + "input", + "attention_mask", + ("1",), + "int64", + None, + ), + VisualPort( + "port:backbone:input:input_ids", + "node:backbone", + "input", + "input_ids", + ("1",), + "int64", + None, + ), + VisualPort( + "port:backbone:input:pixel_values", + "node:backbone", + "input", + "pixel_values", + ("1",), + "float32", + None, + ), ), outputs=(), ) @@ -385,9 +818,33 @@ def test_geometry_reorders_ports_iteratively_for_fanout_targets(): title="action_head", backend="onnx", inputs=( - VisualPort("port:action_head:input:action_inputs_attention_mask", "node:action_head", "input", "action_inputs_attention_mask", ("1",), "int64", None), - VisualPort("port:action_head:input:action_inputs_input_ids", "node:action_head", "input", "action_inputs_input_ids", ("1",), "int64", None), - VisualPort("port:action_head:input:action_inputs_pixel_values", "node:action_head", "input", "action_inputs_pixel_values", ("1",), "float32", None), + VisualPort( + "port:action_head:input:action_inputs_attention_mask", + "node:action_head", + "input", + "action_inputs_attention_mask", + ("1",), + "int64", + None, + ), + VisualPort( + "port:action_head:input:action_inputs_input_ids", + "node:action_head", + "input", + "action_inputs_input_ids", + ("1",), + "int64", + None, + ), + VisualPort( + "port:action_head:input:action_inputs_pixel_values", + "node:action_head", + "input", + "action_inputs_pixel_values", + ("1",), + "float32", + None, + ), ), outputs=(), ) @@ -395,12 +852,60 @@ def test_geometry_reorders_ports_iteratively_for_fanout_targets(): nodes=(source, first_target, second_target), terminals=(), edges=( - VisualEdge("edge:backbone_pixel", "forward", "node:preprocess_video", "node:backbone", "port:preprocess_video:output:pixel_values", "port:backbone:input:pixel_values", "pixel_values"), - VisualEdge("edge:backbone_input_ids", "forward", "node:preprocess_video", "node:backbone", "port:preprocess_video:output:input_ids", "port:backbone:input:input_ids", "input_ids"), - VisualEdge("edge:backbone_attention", "forward", "node:preprocess_video", "node:backbone", "port:preprocess_video:output:attention_mask", "port:backbone:input:attention_mask", "attention_mask"), - VisualEdge("edge:action_pixel", "forward", "node:preprocess_video", "node:action_head", "port:preprocess_video:output:pixel_values", "port:action_head:input:action_inputs_pixel_values", "pixel_values"), - VisualEdge("edge:action_input_ids", "forward", "node:preprocess_video", "node:action_head", "port:preprocess_video:output:input_ids", "port:action_head:input:action_inputs_input_ids", "input_ids"), - VisualEdge("edge:action_attention", "forward", "node:preprocess_video", "node:action_head", "port:preprocess_video:output:attention_mask", "port:action_head:input:action_inputs_attention_mask", "attention_mask"), + VisualEdge( + "edge:backbone_pixel", + "forward", + "node:preprocess_video", + "node:backbone", + "port:preprocess_video:output:pixel_values", + "port:backbone:input:pixel_values", + "pixel_values", + ), + VisualEdge( + "edge:backbone_input_ids", + "forward", + "node:preprocess_video", + "node:backbone", + "port:preprocess_video:output:input_ids", + "port:backbone:input:input_ids", + "input_ids", + ), + VisualEdge( + "edge:backbone_attention", + "forward", + "node:preprocess_video", + "node:backbone", + "port:preprocess_video:output:attention_mask", + "port:backbone:input:attention_mask", + "attention_mask", + ), + VisualEdge( + "edge:action_pixel", + "forward", + "node:preprocess_video", + "node:action_head", + "port:preprocess_video:output:pixel_values", + "port:action_head:input:action_inputs_pixel_values", + "pixel_values", + ), + VisualEdge( + "edge:action_input_ids", + "forward", + "node:preprocess_video", + "node:action_head", + "port:preprocess_video:output:input_ids", + "port:action_head:input:action_inputs_input_ids", + "input_ids", + ), + VisualEdge( + "edge:action_attention", + "forward", + "node:preprocess_video", + "node:action_head", + "port:preprocess_video:output:attention_mask", + "port:action_head:input:action_inputs_attention_mask", + "attention_mask", + ), ), ) layout = LayoutResult( @@ -428,25 +933,72 @@ def test_geometry_routes_forward_and_feedback_edges_differently(): id="node:a", title="a", backend=None, - inputs=(VisualPort("port:a:input:state", "node:a", "input", "state", ("2",), "float32", None),), - outputs=(VisualPort("port:a:output:y", "node:a", "output", "y", ("2",), "float32", None),), + inputs=( + VisualPort( + "port:a:input:state", + "node:a", + "input", + "state", + ("2",), + "float32", + None, + ), + ), + outputs=( + VisualPort( + "port:a:output:y", "node:a", "output", "y", ("2",), "float32", None + ), + ), ) b = VisualNode( id="node:b", title="b", backend=None, - inputs=(VisualPort("port:b:input:y", "node:b", "input", "y", ("2",), "float32", None),), - outputs=(VisualPort("port:b:output:state", "node:b", "output", "state", ("2",), "float32", None),), + inputs=( + VisualPort( + "port:b:input:y", "node:b", "input", "y", ("2",), "float32", None + ), + ), + outputs=( + VisualPort( + "port:b:output:state", + "node:b", + "output", + "state", + ("2",), + "float32", + None, + ), + ), ) graph = VisualGraph( nodes=(a, b), terminals=(), edges=( - VisualEdge("edge:forward", "forward", "node:a", "node:b", "port:a:output:y", "port:b:input:y", "y"), - VisualEdge("edge:feedback", "feedback", "node:b", "node:a", "port:b:output:state", "port:a:input:state", "state"), + VisualEdge( + "edge:forward", + "forward", + "node:a", + "node:b", + "port:a:output:y", + "port:b:input:y", + "y", + ), + VisualEdge( + "edge:feedback", + "feedback", + "node:b", + "node:a", + "port:b:output:state", + "port:a:input:state", + "state", + ), ), ) - layout = LayoutResult(centers={"node:a": Point(0.0, 0.0), "node:b": Point(300.0, 0.0)}, forward_edge_ids=("edge:forward",)) + layout = LayoutResult( + centers={"node:a": Point(0.0, 0.0), "node:b": Point(300.0, 0.0)}, + forward_edge_ids=("edge:forward",), + ) geometry = resolve_geometry(graph, layout, graph_name="demo") @@ -487,7 +1039,9 @@ def test_geometry_budgets_output_row_width_for_kind_labels_on_their_own_row(): geometry = resolve_geometry(graph, layout, graph_name="demo") - visible_kind = geometry_module.truncate_text(long_kind, geometry_module._PORT_NAME_MAX_CHARS) + visible_kind = geometry_module.truncate_text( + long_kind, geometry_module._PORT_NAME_MAX_CHARS + ) output_rect = geometry.ports["port:policy:output:action"].rect expected_kind_width = len(visible_kind) * geometry_module._PORT_PRIMARY_CHAR_WIDTH diff --git a/packages/leapp-visualization/tests/test_graph_visualization_imports.py b/packages/leapp-visualization/tests/test_graph_visualization_imports.py index eba84a54..d95847c6 100644 --- a/packages/leapp-visualization/tests/test_graph_visualization_imports.py +++ b/packages/leapp-visualization/tests/test_graph_visualization_imports.py @@ -15,6 +15,7 @@ # limitations under the License. # + def test_visualization_package_exports_public_entrypoint(): from leapp_visualization import render_graph diff --git a/packages/leapp-visualization/tests/test_graph_visualization_integration.py b/packages/leapp-visualization/tests/test_graph_visualization_integration.py index d36f5a9d..0215069e 100644 --- a/packages/leapp-visualization/tests/test_graph_visualization_integration.py +++ b/packages/leapp-visualization/tests/test_graph_visualization_integration.py @@ -61,7 +61,11 @@ def test_leapp_graph_visualize_writes_png(tmp_path: Path): node = FakeNode( name="policy", inputs=[FakeTensorDescription("obs", "float32", (1, 12))], - outputs=[FakeTensorDescription("action", "float32", (1, 4), semantics={"kind": "command"})], + outputs=[ + FakeTensorDescription( + "action", "float32", (1, 4), semantics={"kind": "command"} + ) + ], backend="jit-script", ) diff --git a/packages/leapp-visualization/tests/test_graph_visualization_layout.py b/packages/leapp-visualization/tests/test_graph_visualization_layout.py index 640a908d..41764838 100644 --- a/packages/leapp-visualization/tests/test_graph_visualization_layout.py +++ b/packages/leapp-visualization/tests/test_graph_visualization_layout.py @@ -16,7 +16,13 @@ # from leapp_visualization.layout import compute_layered_layout -from leapp_visualization.model import VisualEdge, VisualGraph, VisualNode, VisualPort, VisualTerminal +from leapp_visualization.model import ( + VisualEdge, + VisualGraph, + VisualNode, + VisualPort, + VisualTerminal, +) def _node(name): @@ -27,13 +33,41 @@ def test_layered_layout_places_pipeline_left_to_right(): graph = VisualGraph( nodes=(_node("a"), _node("b")), terminals=( - VisualTerminal("terminal:input:a:x", "graph_input", "x", "node:a", "port:a:input:x"), - VisualTerminal("terminal:output:b:y", "graph_output", "y", "node:b", "port:b:output:y"), + VisualTerminal( + "terminal:input:a:x", "graph_input", "x", "node:a", "port:a:input:x" + ), + VisualTerminal( + "terminal:output:b:y", "graph_output", "y", "node:b", "port:b:output:y" + ), ), edges=( - VisualEdge("e0", "graph_input", "terminal:input:a:x", "node:a", None, "port:a:input:x", "x"), - VisualEdge("e1", "forward", "node:a", "node:b", "port:a:output:y", "port:b:input:y", "y"), - VisualEdge("e2", "graph_output", "node:b", "terminal:output:b:y", "port:b:output:y", None, "y"), + VisualEdge( + "e0", + "graph_input", + "terminal:input:a:x", + "node:a", + None, + "port:a:input:x", + "x", + ), + VisualEdge( + "e1", + "forward", + "node:a", + "node:b", + "port:a:output:y", + "port:b:input:y", + "y", + ), + VisualEdge( + "e2", + "graph_output", + "node:b", + "terminal:output:b:y", + "port:b:output:y", + None, + "y", + ), ), ) @@ -99,7 +133,11 @@ def fake_from_edges(edge_pairs, **kwargs): graph = VisualGraph( nodes=(_node("b"), _node("a"), _node("c")), - terminals=(VisualTerminal("terminal:input:c:z", "graph_input", "z", "node:c", "port:c:input:z"),), + terminals=( + VisualTerminal( + "terminal:input:c:z", "graph_input", "z", "node:c", "port:c:input:z" + ), + ), edges=( VisualEdge("e1", "forward", "node:b", "node:c"), VisualEdge("e0", "forward", "node:a", "node:c"), diff --git a/packages/leapp-visualization/tests/test_graph_visualization_png.py b/packages/leapp-visualization/tests/test_graph_visualization_png.py index 5320780f..bdf8822a 100644 --- a/packages/leapp-visualization/tests/test_graph_visualization_png.py +++ b/packages/leapp-visualization/tests/test_graph_visualization_png.py @@ -92,19 +92,35 @@ def _geometry() -> GraphGeometry: "edge:input", "graph_input", "obs", - (Point(128.0, 140.0), Point(134.0, 140.0), Point(134.0, 153.0), Point(140.0, 153.0)), + ( + Point(128.0, 140.0), + Point(134.0, 140.0), + Point(134.0, 153.0), + Point(140.0, 153.0), + ), ), "edge:feedback": EdgeGeometry( "edge:feedback", "feedback", "state", - (Point(380.0, 153.0), Point(410.0, 153.0), Point(410.0, 92.0), Point(110.0, 92.0), Point(110.0, 153.0)), + ( + Point(380.0, 153.0), + Point(410.0, 153.0), + Point(410.0, 92.0), + Point(110.0, 92.0), + Point(110.0, 153.0), + ), ), "edge:output": EdgeGeometry( "edge:output", "graph_output", "action", - (Point(380.0, 176.0), Point(382.0, 176.0), Point(382.0, 176.0), Point(384.0, 176.0)), + ( + Point(380.0, 176.0), + Point(382.0, 176.0), + Point(382.0, 176.0), + Point(384.0, 176.0), + ), ), }, ) @@ -197,7 +213,9 @@ def fake_arrowhead(draw, start, end, color): ] -def test_write_png_renders_on_scaled_canvas_then_downsamples(monkeypatch, tmp_path: Path): +def test_write_png_renders_on_scaled_canvas_then_downsamples( + monkeypatch, tmp_path: Path +): calls = {"new": None, "convert": None, "resize": None, "save": None} class FinalImageStub: @@ -216,14 +234,27 @@ def resize(self, size, resample): monkeypatch.setattr( png_renderer.Image, "new", - lambda mode, size, color: calls.__setitem__("new", (mode, size, color)) or ImageStub(), + lambda mode, size, color: ( + calls.__setitem__("new", (mode, size, color)) or ImageStub() + ), ) monkeypatch.setattr(png_renderer.ImageDraw, "Draw", lambda image: object()) - monkeypatch.setattr(png_renderer, "_load_fonts", lambda: {"label": object(), "label_bold": object(), "label_bold_small": object(), "title": object()}) + monkeypatch.setattr( + png_renderer, + "_load_fonts", + lambda: { + "label": object(), + "label_bold": object(), + "label_bold_small": object(), + "title": object(), + }, + ) monkeypatch.setattr(png_renderer, "_draw_edge", lambda draw, edge: None) monkeypatch.setattr(png_renderer, "_draw_terminal", lambda *args, **kwargs: None) monkeypatch.setattr(png_renderer, "_draw_node_card", lambda *args, **kwargs: None) - monkeypatch.setattr(png_renderer, "_draw_backend_badge", lambda *args, **kwargs: None) + monkeypatch.setattr( + png_renderer, "_draw_backend_badge", lambda *args, **kwargs: None + ) monkeypatch.setattr(png_renderer, "_draw_text", lambda *args, **kwargs: None) monkeypatch.setattr(png_renderer, "_draw_port", lambda *args, **kwargs: None) @@ -237,17 +268,40 @@ def resize(self, size, resample): assert calls["save"] == (str(path), "PNG") -def test_write_png_passes_terminal_and_port_geometry_to_drawers(monkeypatch, tmp_path: Path): +def test_write_png_passes_terminal_and_port_geometry_to_drawers( + monkeypatch, tmp_path: Path +): terminal_calls = [] port_calls = [] - monkeypatch.setattr(png_renderer, "_load_fonts", lambda: {"label": object(), "label_bold": object(), "label_bold_small": object(), "title": object()}) + monkeypatch.setattr( + png_renderer, + "_load_fonts", + lambda: { + "label": object(), + "label_bold": object(), + "label_bold_small": object(), + "title": object(), + }, + ) monkeypatch.setattr(png_renderer, "_draw_edge", lambda draw, edge: None) - monkeypatch.setattr(png_renderer, "_draw_terminal", lambda draw, rect, anchor, title, kind, fonts: terminal_calls.append((rect, anchor, title, kind, fonts))) + monkeypatch.setattr( + png_renderer, + "_draw_terminal", + lambda draw, rect, anchor, title, kind, fonts: terminal_calls.append( + (rect, anchor, title, kind, fonts) + ), + ) monkeypatch.setattr(png_renderer, "_draw_node_card", lambda *args, **kwargs: None) - monkeypatch.setattr(png_renderer, "_draw_backend_badge", lambda *args, **kwargs: None) + monkeypatch.setattr( + png_renderer, "_draw_backend_badge", lambda *args, **kwargs: None + ) monkeypatch.setattr(png_renderer, "_draw_text", lambda *args, **kwargs: None) - monkeypatch.setattr(png_renderer, "_draw_port", lambda draw, port, fonts: port_calls.append((port, fonts))) + monkeypatch.setattr( + png_renderer, + "_draw_port", + lambda draw, port, fonts: port_calls.append((port, fonts)), + ) path = tmp_path / "graph.png" geometry = _geometry() diff --git a/pyproject.toml b/pyproject.toml index 4307f0ab..9ca99c61 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,7 +64,7 @@ dependencies = [ [project.optional-dependencies] dev = [ "pytest", - "black", + "ruff==0.17.1", "flake8", "mypy", ] @@ -84,6 +84,16 @@ warp-cu13 = [ "onnxruntime-gpu[cuda,cudnn]>=1.27.0; python_version >= '3.11'", ] +[tool.ruff] +required-version = "==0.17.1" +target-version = "py310" +line-length = 88 +include = ["*.py", "*.pyi", "*.ipynb", "scripts/leapp-build-warp-runtime"] +force-exclude = true + +[tool.ruff.format] +exclude = ["examples/**"] + [tool.uv] conflicts = [ [ diff --git a/scripts/leapp-build-warp-runtime b/scripts/leapp-build-warp-runtime index e2a891c2..03159a6b 100755 --- a/scripts/leapp-build-warp-runtime +++ b/scripts/leapp-build-warp-runtime @@ -97,20 +97,16 @@ def _validate_build_resources() -> tuple[Path, Path | None]: apic_header = warp_dir / "native" / "apic.h" if not warp_library.is_file(): raise RuntimeError( - "Warp is installed, but its native library was not found: " - f"{warp_library}" + f"Warp is installed, but its native library was not found: {warp_library}" ) if not apic_header.is_file(): raise RuntimeError( - "Warp is installed, but its APIC header was not found: " - f"{apic_header}" + f"Warp is installed, but its APIC header was not found: {apic_header}" ) warp_import_library = None if sys.platform == "win32": - configured_import_library = os.environ.get( - "LEAPP_WARP_IMPORT_LIBRARY" - ) + configured_import_library = os.environ.get("LEAPP_WARP_IMPORT_LIBRARY") warp_import_library = ( Path(configured_import_library).expanduser() if configured_import_library @@ -136,13 +132,9 @@ def _validate_build_resources() -> tuple[Path, Path | None]: ) torch = _require_module("torch", "pip install torch") - cmake_prefix = getattr( - getattr(torch, "utils", None), "cmake_prefix_path", None - ) + cmake_prefix = getattr(getattr(torch, "utils", None), "cmake_prefix_path", None) if not cmake_prefix: - raise RuntimeError( - "PyTorch is installed, but its CMake path is unavailable." - ) + raise RuntimeError("PyTorch is installed, but its CMake path is unavailable.") torch_cmake_dir = Path(cmake_prefix) / "Torch" if not (torch_cmake_dir / "TorchConfig.cmake").is_file(): raise RuntimeError( @@ -183,9 +175,7 @@ def build_warp_runtime( "-DLEAPP_WARP_BUILD_TORCH=ON", ] if warp_import_library is not None: - configure_command.append( - f"-DWARP_IMPORT_LIBRARY={warp_import_library}" - ) + configure_command.append(f"-DWARP_IMPORT_LIBRARY={warp_import_library}") subprocess.run(configure_command, check=True) subprocess.run( ["cmake", "--build", str(output_dir), "--config", "Release", "-j"], @@ -195,8 +185,7 @@ def build_warp_runtime( missing = [str(path) for path in artifacts.values() if not path.is_file()] if missing: raise RuntimeError( - "Warp runtime build completed without producing: " - + ", ".join(missing) + "Warp runtime build completed without producing: " + ", ".join(missing) ) return artifacts diff --git a/tests/functional_tests/base.py b/tests/functional_tests/base.py index ae9a6e54..11e896b8 100644 --- a/tests/functional_tests/base.py +++ b/tests/functional_tests/base.py @@ -20,7 +20,9 @@ import torch from safetensors.torch import load_file from leapp.export_manager import ExportManager -from leapp.leapp import _MANAGER as annotate # non-protected access to annotate singleton +from leapp.leapp import ( + _MANAGER as annotate, +) # non-protected access to annotate singleton import yaml @@ -35,15 +37,20 @@ def tearDown(self): ExportManager.set_interpret_graph(False) annotate.reset_nodes() annotate.set_dry_run_and_non_traced(False, []) - + def verify_all_models_exist(self, *model_names): for model_name in model_names: model_exists = False - model_exists |= os.path.exists(os.path.join(self.TEST_GRAPH_NAME, f"{model_name}.pt")) - model_exists |= os.path.exists(os.path.join(self.TEST_GRAPH_NAME, f"{model_name}.onnx")) - model_exists |= os.path.exists(os.path.join(self.TEST_GRAPH_NAME, f"{model_name}.pt2")) - self.assertTrue(model_exists, - f"Model {model_name} does not exist") + model_exists |= os.path.exists( + os.path.join(self.TEST_GRAPH_NAME, f"{model_name}.pt") + ) + model_exists |= os.path.exists( + os.path.join(self.TEST_GRAPH_NAME, f"{model_name}.onnx") + ) + model_exists |= os.path.exists( + os.path.join(self.TEST_GRAPH_NAME, f"{model_name}.pt2") + ) + self.assertTrue(model_exists, f"Model {model_name} does not exist") def verify_node_io(self, node, *, inputs=None, outputs=None): if inputs is not None: @@ -59,29 +66,56 @@ def verify_node_io(self, node, *, inputs=None, outputs=None): f"Expected {outputs} node output(s), got {len(node.outputs)}", ) - def verify_num_connections(self, leapp_annotation, nodes=None, inputs=None, outputs=None, - internal_connections=None, feedback_connections=None): + def verify_num_connections( + self, + leapp_annotation, + nodes=None, + inputs=None, + outputs=None, + internal_connections=None, + feedback_connections=None, + ): if nodes is not None: - self.assertEqual(nodes, len(leapp_annotation.detected_nodes), - "Number of nodes do not match") + self.assertEqual( + nodes, + len(leapp_annotation.detected_nodes), + "Number of nodes do not match", + ) if inputs is not None: total_inputs = sum( - [len(graph_inputs) for graph_inputs in leapp_annotation.detected_pipeline['inputs'].values()]) - self.assertEqual(inputs, total_inputs, - "Number of inputs do not match") + [ + len(graph_inputs) + for graph_inputs in leapp_annotation.detected_pipeline[ + "inputs" + ].values() + ] + ) + self.assertEqual(inputs, total_inputs, "Number of inputs do not match") if outputs is not None: total_outputs = sum( - [len(graph_outputs) for graph_outputs in leapp_annotation.detected_pipeline['outputs'].values()]) - self.assertEqual(outputs, total_outputs, - "Number of outputs do not match") + [ + len(graph_outputs) + for graph_outputs in leapp_annotation.detected_pipeline[ + "outputs" + ].values() + ] + ) + self.assertEqual(outputs, total_outputs, "Number of outputs do not match") if internal_connections is not None: num_connections = 0 - for _, targets in leapp_annotation.detected_pipeline['data_flow'].items(): + for _, targets in leapp_annotation.detected_pipeline["data_flow"].items(): num_connections += len(targets) - self.assertEqual(internal_connections, num_connections, "Number of internal connections do not match") + self.assertEqual( + internal_connections, + num_connections, + "Number of internal connections do not match", + ) if feedback_connections is not None: - self.assertEqual(feedback_connections, len( - leapp_annotation.detected_pipeline['feedback_flow']), "Number of feedback connections do not match") + self.assertEqual( + feedback_connections, + len(leapp_annotation.detected_pipeline["feedback_flow"]), + "Number of feedback connections do not match", + ) def inspect_torchscript_model(self, model_name, model_path=None): """ @@ -101,119 +135,114 @@ def inspect_torchscript_model(self, model_name, model_path=None): if model_path is None: model_path = self.TEST_GRAPH_NAME model = torch.jit.load(os.path.join(model_path, f"{model_name}.pt")) - result = { - 'inputs': [], - 'outputs': [], - 'graph': None, - 'code': None - } + result = {"inputs": [], "outputs": [], "graph": None, "code": None} try: # Get the computation graph graph = model.graph - result['graph'] = graph + result["graph"] = graph # Extract input information inputs = list(graph.inputs()) for i, inp in enumerate(inputs): input_info = { - 'index': i, - 'debug_name': inp.debugName(), - 'type': str(inp.type()), + "index": i, + "debug_name": inp.debugName(), + "type": str(inp.type()), } # Try to get shape information if it's a tensor - if hasattr(inp.type(), 'sizes'): + if hasattr(inp.type(), "sizes"): try: - input_info['shape'] = list(inp.type().sizes()) + input_info["shape"] = list(inp.type().sizes()) except Exception: - input_info['shape'] = None + input_info["shape"] = None else: - input_info['shape'] = None + input_info["shape"] = None - result['inputs'].append(input_info) + result["inputs"].append(input_info) # Extract output information outputs = list(graph.outputs()) - + # Check if the output is a single tuple - if so, unwrap it if len(outputs) == 1: output_type = outputs[0].type() # Check if it's a tuple type by looking at the string representation type_str = str(output_type) - if type_str.startswith('Tuple[') or type_str.startswith('('): + if type_str.startswith("Tuple[") or type_str.startswith("("): # It's a tuple - unwrap it by getting the element types - if hasattr(output_type, 'elements'): + if hasattr(output_type, "elements"): # This is a TupleType, get the individual elements tuple_elements = output_type.elements() for i, elem_type in enumerate(tuple_elements): output_info = { - 'index': i, - 'debug_name': f'tuple_element_{i}', - 'type': str(elem_type), + "index": i, + "debug_name": f"tuple_element_{i}", + "type": str(elem_type), } - + # Try to get shape information if it's a tensor - if hasattr(elem_type, 'sizes'): + if hasattr(elem_type, "sizes"): try: - output_info['shape'] = list(elem_type.sizes()) + output_info["shape"] = list(elem_type.sizes()) except Exception: - output_info['shape'] = None + output_info["shape"] = None else: - output_info['shape'] = None - - result['outputs'].append(output_info) + output_info["shape"] = None + + result["outputs"].append(output_info) else: # Fallback: treat as single output output_info = { - 'index': 0, - 'debug_name': outputs[0].debugName(), - 'type': type_str, - 'shape': None + "index": 0, + "debug_name": outputs[0].debugName(), + "type": type_str, + "shape": None, } - result['outputs'].append(output_info) + result["outputs"].append(output_info) else: # Not a tuple, process normally for i, out in enumerate(outputs): output_info = { - 'index': i, - 'debug_name': out.debugName(), - 'type': str(out.type()), + "index": i, + "debug_name": out.debugName(), + "type": str(out.type()), } # Try to get shape information if it's a tensor - if hasattr(out.type(), 'sizes'): + if hasattr(out.type(), "sizes"): try: - output_info['shape'] = list(out.type().sizes()) + output_info["shape"] = list(out.type().sizes()) except Exception: - output_info['shape'] = None + output_info["shape"] = None else: - output_info['shape'] = None + output_info["shape"] = None - result['outputs'].append(output_info) + result["outputs"].append(output_info) else: # Multiple outputs, process normally for i, out in enumerate(outputs): output_info = { - 'index': i, - 'debug_name': out.debugName(), - 'type': str(out.type()), + "index": i, + "debug_name": out.debugName(), + "type": str(out.type()), } # Try to get shape information if it's a tensor - if hasattr(out.type(), 'sizes'): + if hasattr(out.type(), "sizes"): try: - output_info['shape'] = list(out.type().sizes()) + output_info["shape"] = list(out.type().sizes()) except Exception: - output_info['shape'] = None + output_info["shape"] = None else: - output_info['shape'] = None + output_info["shape"] = None - result['outputs'].append(output_info) + result["outputs"].append(output_info) # Get the code representation if available - if hasattr(model, 'code'): - result['code'] = model.code + if hasattr(model, "code"): + result["code"] = model.code except Exception as e: print(f"Error inspecting TorchScript model: {e}") @@ -231,44 +260,45 @@ def print_torchscript_model_info(self, model_name, model_path=None, verbose=True """ info = self.inspect_torchscript_model(model_name, model_path) - print("\n" + "="*60) + print("\n" + "=" * 60) print("TorchScript Model Information") - print("="*60) + print("=" * 60) print("\n[Inputs]") - if not info['inputs']: + if not info["inputs"]: print(" No inputs found") else: - for inp in info['inputs']: + for inp in info["inputs"]: print(f" Input {inp['index']}:") print(f" Name: {inp['debug_name']}") print(f" Type: {inp['type']}") - if inp['shape']: + if inp["shape"]: print(f" Shape: {inp['shape']}") print("\n[Outputs]") - if not info['outputs']: + if not info["outputs"]: print(" No outputs found") else: - for out in info['outputs']: + for out in info["outputs"]: print(f" Output {out['index']}:") print(f" Name: {out['debug_name']}") print(f" Type: {out['type']}") - if out['shape']: + if out["shape"]: print(f" Shape: {out['shape']}") - if verbose and info['graph']: + if verbose and info["graph"]: print("\n[Computation Graph]") - print(info['graph']) + print(info["graph"]) - if info['code']: + if info["code"]: print("\n[Model Code]") - print(info['code']) + print(info["code"]) - print("\n" + "="*60) + print("\n" + "=" * 60) def verify_data_exact_match(self, source_data, target_data): from leapp.leapp_graph.datatypes.traced_data import TracedData + if isinstance(source_data, TracedData): source_data = source_data.data if isinstance(target_data, TracedData): @@ -311,7 +341,7 @@ def _flatten_to_tensors(self, data): Uses the same traversal order as describe_io_helper. """ tensors = [] - + def _flatten(item): if isinstance(item, torch.Tensor): tensors.append(item) @@ -321,24 +351,25 @@ def _flatten(item): elif isinstance(item, (list, tuple)): for elem in item: _flatten(elem) - + _flatten(data) return tensors def _assert_artifacts(self, save_dir, graph_name): yaml_path = os.path.join(save_dir, f"{graph_name}.yaml") model_path = os.path.join(save_dir, "identity.pt") - self.assertTrue(os.path.exists(yaml_path), - f"YAML not found at {yaml_path}") - self.assertTrue(os.path.exists(model_path), - f"Model not found at {model_path}") + self.assertTrue(os.path.exists(yaml_path), f"YAML not found at {yaml_path}") + self.assertTrue(os.path.exists(model_path), f"Model not found at {model_path}") # The exported model_path entry in YAML must be relative to the YAML dir. with open(yaml_path, "r") as f: config = yaml.safe_load(f) identity_params = config["models"]["identity"]["parameters"] - self.assertEqual(identity_params["model_path"], "identity.pt", - f"model_path should be relative to YAML directory, " - f"got {identity_params['model_path']!r}") + self.assertEqual( + identity_params["model_path"], + "identity.pt", + f"model_path should be relative to YAML directory, " + f"got {identity_params['model_path']!r}", + ) def verify_feedback_initial_values(self, expected_values, graph_name=None): """Verify that the feedback initial values safetensors file exists and contains expected values. @@ -352,21 +383,32 @@ def verify_feedback_initial_values(self, expected_values, graph_name=None): """ if graph_name is None: graph_name = self.TEST_GRAPH_NAME - safetensors_path = os.path.join(graph_name, f"{graph_name}_initial_values.safetensors") - self.assertTrue(os.path.exists(safetensors_path), - f"Feedback initial values safetensors file not found at {safetensors_path}") + safetensors_path = os.path.join( + graph_name, f"{graph_name}_initial_values.safetensors" + ) + self.assertTrue( + os.path.exists(safetensors_path), + f"Feedback initial values safetensors file not found at {safetensors_path}", + ) loaded = load_file(safetensors_path) for key, expected in expected_values.items(): - self.assertIn(key, loaded, f"Key '{key}' not found in feedback initial values") + self.assertIn( + key, loaded, f"Key '{key}' not found in feedback initial values" + ) if isinstance(expected, tuple): # Shape-only check - self.assertEqual(loaded[key].shape, torch.Size(expected), - f"Shape mismatch for '{key}': got {loaded[key].shape}, expected {expected}") + self.assertEqual( + loaded[key].shape, + torch.Size(expected), + f"Shape mismatch for '{key}': got {loaded[key].shape}, expected {expected}", + ) else: # Exact value check - self.assertTrue(torch.equal(loaded[key], expected), - f"Value mismatch for '{key}': got {loaded[key]}, expected {expected}") + self.assertTrue( + torch.equal(loaded[key], expected), + f"Value mismatch for '{key}': got {loaded[key]}, expected {expected}", + ) def verify_safetensors_matches_feedback(self, leapp_annotation, graph_name=None): """Verify safetensors keys exactly match the feedback targets in the pipeline. @@ -378,10 +420,10 @@ def verify_safetensors_matches_feedback(self, leapp_annotation, graph_name=None) if graph_name is None: graph_name = self.TEST_GRAPH_NAME safetensors_path = os.path.join( - graph_name, f"{graph_name}_initial_values.safetensors") + graph_name, f"{graph_name}_initial_values.safetensors" + ) - feedback_flow = leapp_annotation.detected_pipeline.get( - 'feedback_flow', {}) + feedback_flow = leapp_annotation.detected_pipeline.get("feedback_flow", {}) # Collect all feedback target keys (node_name/input_name) expected_keys = set() @@ -392,24 +434,29 @@ def verify_safetensors_matches_feedback(self, leapp_annotation, graph_name=None) if not expected_keys: self.assertFalse( os.path.exists(safetensors_path), - "No feedback connections exist, but safetensors file was created") + "No feedback connections exist, but safetensors file was created", + ) return self.assertTrue( os.path.exists(safetensors_path), - f"Feedback connections exist but safetensors file not found at {safetensors_path}") + f"Feedback connections exist but safetensors file not found at {safetensors_path}", + ) loaded = load_file(safetensors_path) actual_keys = set(loaded.keys()) self.assertEqual( - expected_keys, actual_keys, + expected_keys, + actual_keys, f"Safetensors keys do not match feedback targets.\n" f" Expected: {sorted(expected_keys)}\n" - f" Actual: {sorted(actual_keys)}") + f" Actual: {sorted(actual_keys)}", + ) - def verify_inference_manager(self, source_inputs, source_outputs, - graph_name=None, rtol=1e-3, atol=1e-5): + def verify_inference_manager( + self, source_inputs, source_outputs, graph_name=None, rtol=1e-3, atol=1e-5 + ): """Load the exported graph via InferenceManager, run it, and compare outputs. Args: @@ -424,23 +471,28 @@ def verify_inference_manager(self, source_inputs, source_outputs, if graph_name is None: graph_name = self.TEST_GRAPH_NAME yaml_path = os.path.join(graph_name, f"{graph_name}.yaml") - self.assertTrue(os.path.exists(yaml_path), - f"YAML not found at {yaml_path}") + self.assertTrue(os.path.exists(yaml_path), f"YAML not found at {yaml_path}") manager = InferenceManager(yaml_path) # Verify expected input/output keys exist for key in source_inputs: - self.assertIn(key, manager.inputs, - f"Input key '{key}' not found in InferenceManager inputs: {manager.inputs}") + self.assertIn( + key, + manager.inputs, + f"Input key '{key}' not found in InferenceManager inputs: {manager.inputs}", + ) for key in source_outputs: - self.assertIn(key, manager.outputs, - f"Output key '{key}' not found in InferenceManager outputs: {manager.outputs}") + self.assertIn( + key, + manager.outputs, + f"Output key '{key}' not found in InferenceManager outputs: {manager.outputs}", + ) # Move inputs to the device the InferenceManager expects device_inputs = {} for key, tensor in source_inputs.items(): - node_name = key.split('/')[0] + node_name = key.split("/")[0] target_device = manager.nodes[node_name].device device_inputs[key] = tensor.to(target_device) @@ -449,8 +501,11 @@ def verify_inference_manager(self, source_inputs, source_outputs, # Compare each expected output for key, expected in source_outputs.items(): - self.assertIn(key, exported_outputs, - f"Output key '{key}' missing from InferenceManager results") + self.assertIn( + key, + exported_outputs, + f"Output key '{key}' missing from InferenceManager results", + ) actual = exported_outputs[key] if expected.device != actual.device: actual = actual.to(expected.device) @@ -459,13 +514,16 @@ def verify_inference_manager(self, source_inputs, source_outputs, f"Output mismatch for '{key}':\n" f" Expected: {expected}\n" f" Actual: {actual}\n" - f" Max diff: {(expected - actual).abs().max().item():.6e}") + f" Max diff: {(expected - actual).abs().max().item():.6e}", + ) - def verify_single_torchscript_model_expected_value(self, inputs, expected_outputs, model_name, model_path=None): + def verify_single_torchscript_model_expected_value( + self, inputs, expected_outputs, model_name, model_path=None + ): if model_path is None: model_path = self.TEST_GRAPH_NAME model = torch.jit.load(os.path.join(model_path, f"{model_name}.pt")) - + # Flatten all inputs since models now expect flat tensor arguments flat_inputs = [] for inp in inputs: @@ -479,14 +537,20 @@ def verify_single_torchscript_model_expected_value(self, inputs, expected_output outputs = (outputs,) for output, expected_output in zip(outputs, flat_expected_outputs): - self.assertTrue(self.verify_data_exact_match(output, expected_output), "An output value does not match expected value: " - f"got {output} but expected {expected_output}") + self.assertTrue( + self.verify_data_exact_match(output, expected_output), + "An output value does not match expected value: " + f"got {output} but expected {expected_output}", + ) - def verify_single_exported_program_model_expected_value(self, inputs, expected_outputs, model_name, model_path=None): + def verify_single_exported_program_model_expected_value( + self, inputs, expected_outputs, model_name, model_path=None + ): if model_path is None: model_path = self.TEST_GRAPH_NAME exported_program = torch.export.load( - os.path.join(model_path, f"{model_name}.pt2")) + os.path.join(model_path, f"{model_name}.pt2") + ) model = exported_program.module() flat_inputs = [] @@ -501,5 +565,8 @@ def verify_single_exported_program_model_expected_value(self, inputs, expected_o outputs = (outputs,) for output, expected_output in zip(outputs, flat_expected_outputs): - self.assertTrue(self.verify_data_exact_match(output, expected_output), "An output value does not match expected value: " - f"got {output} but expected {expected_output}") + self.assertTrue( + self.verify_data_exact_match(output, expected_output), + "An output value does not match expected value: " + f"got {output} but expected {expected_output}", + ) diff --git a/tests/functional_tests/test_annotate.py b/tests/functional_tests/test_annotate.py index 9783a712..21e24ef3 100644 --- a/tests/functional_tests/test_annotate.py +++ b/tests/functional_tests/test_annotate.py @@ -27,13 +27,14 @@ class TestAnnotateMethod(LEAPPFunctionalTestBase): def test_annotate_method(self): """tests the basic situation of using the annotate.method decorator""" + @annotate.method(export_with="jit") def funcA(inputA: torch.Tensor): return inputA @annotate.method(export_with="jit") def funcC(inputB: torch.Tensor): - return inputB+5.0 + return inputB + 5.0 leapp.start(name=self.TEST_GRAPH_NAME) for i in range(5): @@ -47,15 +48,16 @@ def funcC(inputB: torch.Tensor): self.assertEqual(len(annotate.nodes[funcC.__name__].inputs), 1) self.assertEqual(len(annotate.nodes[funcA.__name__].outputs), 1) self.assertEqual(len(annotate.nodes[funcC.__name__].outputs), 1) - self.assertEqual( - annotate.nodes[funcA.__name__].inputs[0].name, "inputA") - self.assertEqual( - annotate.nodes[funcC.__name__].inputs[0].name, "inputB") + self.assertEqual(annotate.nodes[funcA.__name__].inputs[0].name, "inputA") + self.assertEqual(annotate.nodes[funcC.__name__].inputs[0].name, "inputB") def test_annotate_method_with_kwargs_and_default_value(self): """tests the situation where the function has and default values""" + @annotate.method(export_with="jit") - def funcA(inputA: torch.Tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)): + def funcA( + inputA: torch.Tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32), + ): return inputA leapp.start(name=self.TEST_GRAPH_NAME) @@ -68,16 +70,20 @@ def funcA(inputA: torch.Tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float expected_output = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) # check if using default self.verify_single_torchscript_model_expected_value( - [], [expected_output], funcA.__name__) + [], [expected_output], funcA.__name__ + ) # Or get structured data model_info = self.inspect_torchscript_model(funcA.__name__) - self.assertEqual(len(model_info['inputs']), 1) - self.assertEqual(len(model_info['outputs']), 1) + self.assertEqual(len(model_info["inputs"]), 1) + self.assertEqual(len(model_info["outputs"]), 1) def test_annotate_method_ignoring_default_values(self): """tests the situation where we pass in a value overriding the default""" + @annotate.method(export_with="jit") - def funcA(inputA: torch.Tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)): + def funcA( + inputA: torch.Tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32), + ): return inputA leapp.start(name=self.TEST_GRAPH_NAME) @@ -87,16 +93,21 @@ def funcA(inputA: torch.Tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float self.assertEqual(len(annotate.nodes), 1) self.assertEqual(len(annotate.nodes[funcA.__name__].inputs), 1) model_info = self.inspect_torchscript_model(funcA.__name__) - self.assertEqual(len(model_info['inputs']), 2) - self.assertEqual(len(model_info['outputs']), 1) + self.assertEqual(len(model_info["inputs"]), 2) + self.assertEqual(len(model_info["outputs"]), 1) def test_annotate_method_ignoring_middle_kwargs(self): """tests the situation where the user provides kwargs out of order""" default_tensor = torch.tensor([0]) @annotate.method(export_with="jit") - def funcA(input1=default_tensor, input2=default_tensor, input3=default_tensor, - input4=default_tensor, input5=default_tensor): + def funcA( + input1=default_tensor, + input2=default_tensor, + input3=default_tensor, + input4=default_tensor, + input5=default_tensor, + ): output = torch.cat([input1, input2, input3, input4, input5], dim=0) return output @@ -107,14 +118,16 @@ def funcA(input1=default_tensor, input2=default_tensor, input3=default_tensor, self.assertEqual(len(annotate.nodes), 1) self.assertEqual(len(annotate.nodes[funcA.__name__].inputs), 2) model_info = self.inspect_torchscript_model(funcA.__name__) - self.assertEqual(len(model_info['inputs']), 3) - self.assertEqual(len(model_info['outputs']), 1) + self.assertEqual(len(model_info["inputs"]), 3) + self.assertEqual(len(model_info["outputs"]), 1) self.verify_single_torchscript_model_expected_value( - [torch.tensor([1]), torch.tensor([1])], [outputA], funcA.__name__) + [torch.tensor([1]), torch.tensor([1])], [outputA], funcA.__name__ + ) def test_annotate_method_var_positional_with_mixed_types(self): """tests *args with mixed types (int, tensor, int) preserving positions""" + @annotate.method(export_with="jit") def funcA(*inputs): scale_before, tensor, scale_after = inputs @@ -126,14 +139,19 @@ def funcA(*inputs): leapp.compile_graph(visualize=False) self.assertEqual(len(annotate.nodes), 1) expected = torch.tensor([3, 4, 5]) * 2 + 10 - self.assertTrue(torch.equal(outputA.data if hasattr(outputA, 'data') else outputA, expected), - f"got {outputA} but expected {expected}") + self.assertTrue( + torch.equal( + outputA.data if hasattr(outputA, "data") else outputA, expected + ), + f"got {outputA} but expected {expected}", + ) def test_annotate_method_with_multiple_unnamed_returns(self): """tests the situation where the function has multiple unnamed returns""" + @annotate.method(export_with="jit") def funcA(input1: torch.Tensor): - return input1+1, input1+2, input1+3 + return input1 + 1, input1 + 2, input1 + 3 leapp.start(name=self.TEST_GRAPH_NAME) outputA, outputB, outputC = funcA(torch.tensor([1])) @@ -145,7 +163,8 @@ def funcA(input1: torch.Tensor): self.assertEqual(len(annotate.nodes[funcA.__name__].outputs), 3) self.verify_single_torchscript_model_expected_value( - [torch.tensor([1])], [outputA, outputB, outputC], funcA.__name__) + [torch.tensor([1])], [outputA, outputB, outputC], funcA.__name__ + ) def test_annotate_method_with_torch_no_grad_wrapper(self): """Tests that annotate.method works with functions wrapped by multiple decorators. @@ -172,22 +191,22 @@ def wrapped_func(inputA: torch.Tensor): # Verify node was created correctly self.assertEqual(len(annotate.nodes), 1) - self.assertIn('wrapped_func', annotate.nodes) - self.assertEqual(len(annotate.nodes['wrapped_func'].inputs), 1) - self.assertEqual(len(annotate.nodes['wrapped_func'].outputs), 1) - self.assertEqual( - annotate.nodes['wrapped_func'].inputs[0].name, "inputA") + self.assertIn("wrapped_func", annotate.nodes) + self.assertEqual(len(annotate.nodes["wrapped_func"].inputs), 1) + self.assertEqual(len(annotate.nodes["wrapped_func"].outputs), 1) + self.assertEqual(annotate.nodes["wrapped_func"].inputs[0].name, "inputA") # Compile the graph leapp.compile_graph(visualize=False) # Verify model was generated - self.verify_all_models_exist('wrapped_func') + self.verify_all_models_exist("wrapped_func") # Verify the model produces correct results expected_output = input_tensor * 2 + 1 self.verify_single_torchscript_model_expected_value( - [input_tensor], [expected_output], 'wrapped_func') + [input_tensor], [expected_output], "wrapped_func" + ) def test_annotate_method_with_torch_no_grad_wrapper_on_class_method(self): """Tests that annotate.method works with class methods wrapped by @torch.no_grad()""" @@ -204,10 +223,9 @@ def forward(self, x: torch.Tensor): model = MockModel() # Apply annotate.method to the bound method (simulating runtime decoration) - model.forward = annotate.method( - node_name='model_forward', - export_with="jit" - )(model.forward) + model.forward = annotate.method(node_name="model_forward", export_with="jit")( + model.forward + ) leapp.start(name=self.TEST_GRAPH_NAME) input_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) @@ -216,24 +234,26 @@ def forward(self, x: torch.Tensor): # Verify node was created correctly self.assertEqual(len(annotate.nodes), 1) - self.assertIn('model_forward', annotate.nodes) - self.assertEqual(len(annotate.nodes['model_forward'].inputs), 1) - self.assertEqual(len(annotate.nodes['model_forward'].outputs), 1) - self.assertEqual(annotate.nodes['model_forward'].inputs[0].name, "x") + self.assertIn("model_forward", annotate.nodes) + self.assertEqual(len(annotate.nodes["model_forward"].inputs), 1) + self.assertEqual(len(annotate.nodes["model_forward"].outputs), 1) + self.assertEqual(annotate.nodes["model_forward"].inputs[0].name, "x") # Compile the graph leapp.compile_graph(visualize=False) # Verify model was generated - self.verify_all_models_exist('model_forward') + self.verify_all_models_exist("model_forward") # Verify the model produces correct results expected_output = input_tensor * 3.0 self.verify_single_torchscript_model_expected_value( - [input_tensor], [expected_output], 'model_forward') + [input_tensor], [expected_output], "model_forward" + ) def test_annotate_method_forwards_backend_params(self): """backend_params passed to method() should reach the export backend.""" + @annotate.method( export_with="onnx-torchscript", backend_params={"prescript": True, "report": True}, @@ -254,7 +274,6 @@ def funcA(inputA: torch.Tensor): class TestAnnotateTensor(LEAPPFunctionalTestBase): - def _assert_nested_traced_structure(self, source, returned): """Validate nested container shape is preserved and tensor leaves are TracedTensor.""" if isinstance(source, torch.Tensor): @@ -283,64 +302,73 @@ def _assert_nested_traced_structure(self, source, returned): self._assert_nested_traced_structure(source[key], returned[key]) return - self.fail(f"Unsupported nested leaf type in test payload: {type(source).__name__}") + self.fail( + f"Unsupported nested leaf type in test payload: {type(source).__name__}" + ) def test_annotate_tensor_single(self): leapp.start(name=self.TEST_GRAPH_NAME, verbose=True) tensor1 = annotate.input_tensors( - 'func1', {'input1': torch.tensor([1.0, 2.0, 3.0])}) + "func1", {"input1": torch.tensor([1.0, 2.0, 3.0])} + ) tensor1 = tensor1 + 1.0 tensor1 = tensor1 - 2.0 tensor1 = tensor1 * 3.0 tensor1 = tensor1 / 4.0 - tensor1 = tensor1.matmul(torch.tensor( - [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])) + tensor1 = tensor1.matmul( + torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) + ) returned_output = annotate.output_tensors( - 'func1', {'output1': tensor1}, export_with="jit") + "func1", {"output1": tensor1}, export_with="jit" + ) self.assertIs(returned_output, tensor1) self.assertIsInstance(returned_output, TracedTensor) leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0) + annotate, nodes=1, inputs=1, outputs=1, internal_connections=0 + ) self.verify_single_torchscript_model_expected_value( - [torch.tensor([1.0, 2.0, 3.0])], [tensor1.tensor], 'func1') + [torch.tensor([1.0, 2.0, 3.0])], [tensor1.tensor], "func1" + ) def test_annotate_tensor_with_static_outputs(self): """Test that static outputs (constant tensors not derived from inputs) are returned by the model.""" leapp.start(name=self.TEST_GRAPH_NAME, verbose=True) - + # Input tensor input_tensor = torch.tensor([1.0, 2.0, 3.0]) - traced_input = annotate.input_tensors('static_test', {'input': input_tensor}) - + traced_input = annotate.input_tensors("static_test", {"input": input_tensor}) + # Static output - a constant tensor not derived from input static_tensor = torch.tensor([4.0, 5.0, 6.0]) - + # Computed output derived from input computed_output = traced_input + 1.0 - + # Output both computed and static tensors annotate.output_tensors( - 'static_test', - {'computed': computed_output}, - static_outputs={'static': static_tensor}, - export_with="jit" + "static_test", + {"computed": computed_output}, + static_outputs={"static": static_tensor}, + export_with="jit", ) - + leapp.stop() leapp.compile_graph(visualize=False) - + # Verify graph structure: 1 input, 2 outputs (computed + static) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=2, internal_connections=0) - + annotate, nodes=1, inputs=1, outputs=2, internal_connections=0 + ) + # Verify model outputs correct values # Expected: computed = input + 1 = [2, 3, 4], static = [4, 5, 6] expected_computed = input_tensor + 1.0 self.verify_single_torchscript_model_expected_value( - [input_tensor], [expected_computed, static_tensor], 'static_test') + [input_tensor], [expected_computed, static_tensor], "static_test" + ) def test_annotate_tensor_with_static_outputs_multiple_runs_validation(self): """Static outputs should validate correctly across cached re-entry examples.""" @@ -351,15 +379,17 @@ def test_annotate_tensor_with_static_outputs_multiple_runs_validation(self): for i in range(3): input_tensor = torch.tensor( - [1.0 + i, 2.0 + i, 3.0 + i], dtype=torch.float32) + [1.0 + i, 2.0 + i, 3.0 + i], dtype=torch.float32 + ) traced_input = annotate.input_tensors( - 'static_test_cached', {'input': input_tensor}) + "static_test_cached", {"input": input_tensor} + ) computed_output = traced_input + 1.0 annotate.output_tensors( - 'static_test_cached', - {'computed': computed_output}, - static_outputs={'static': static_tensor}, - export_with="jit" + "static_test_cached", + {"computed": computed_output}, + static_outputs={"static": static_tensor}, + export_with="jit", ) last_input = input_tensor @@ -367,53 +397,59 @@ def test_annotate_tensor_with_static_outputs_multiple_runs_validation(self): leapp.compile_graph(visualize=False, validate=True) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=2, internal_connections=0) + annotate, nodes=1, inputs=1, outputs=2, internal_connections=0 + ) self.verify_single_torchscript_model_expected_value( - [last_input], [last_input + 1.0, static_tensor], 'static_test_cached') + [last_input], [last_input + 1.0, static_tensor], "static_test_cached" + ) def test_annotate_register_buffer_with_inplace_assignment(self): - """Test that register_buffer allows a tensor to participate in tracing with in-place assignment. - """ + """Test that register_buffer allows a tensor to participate in tracing with in-place assignment.""" + class MockModule: def __init__(self): self.values = torch.tensor([1.0, 2.0, 3.0]) - + def run(self, traced_input): # Register the buffer - makes self.values a TracedTensor - self.values = annotate.register_buffer('buffer_test', {'values': self.values}) + self.values = annotate.register_buffer( + "buffer_test", {"values": self.values} + ) # In-place assignment - now traced because self.values is a TracedTensor self.values[:] = traced_input # Operations on the buffer are traced return self.values * 100.0 - + module = MockModule() - + leapp.start(name=self.TEST_GRAPH_NAME, verbose=True) - + # Create traced input input_tensor = torch.tensor([4.0, 5.0, 6.0]) - traced_input = annotate.input_tensors('buffer_test', {'input': input_tensor}) - + traced_input = annotate.input_tensors("buffer_test", {"input": input_tensor}) + # Run the module - this registers the buffer and performs traced operations result = module.run(traced_input) - + # Output the result - annotate.output_tensors('buffer_test', {'result': result}, export_with="jit") - + annotate.output_tensors("buffer_test", {"result": result}, export_with="jit") + leapp.stop() leapp.compile_graph(visualize=False) - + # Verify graph structure: 1 input, 1 output self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0) - + annotate, nodes=1, inputs=1, outputs=1, internal_connections=0 + ) + # Verify model produces correct results # The buffer receives the input via in-place assignment, then multiplied by 100 # Expected: input * 100 = [400, 500, 600] input_tensor = torch.tensor([1.0, 2.0, 1.0]) expected_output = input_tensor * 100.0 self.verify_single_torchscript_model_expected_value( - [input_tensor], [expected_output], 'buffer_test') + [input_tensor], [expected_output], "buffer_test" + ) def test_api_returns_preserve_complex_nested_format(self): """All return-producing APIs preserve complex dict/list/tuple payload format.""" @@ -471,7 +507,10 @@ def test_api_returns_preserve_complex_nested_format(self): self._assert_nested_traced_structure(buffer_payload, traced_buffer_payload) annotate.output_tensors( "complex_returns_buffer", - {"buffer_out": traced_buffer_payload["buffer_a"][0] + traced_buffer_payload["buffer_b"]["inner"]}, + { + "buffer_out": traced_buffer_payload["buffer_a"][0] + + traced_buffer_payload["buffer_b"]["inner"] + }, export_with="jit", ) @@ -498,97 +537,107 @@ def test_api_returns_preserve_complex_nested_format(self): def test_annotate_tensor_with_dict_io(self): leapp.start(name=self.TEST_GRAPH_NAME, verbose=True) - input_dict = {'input1': torch.tensor([1.0, 2.0, 3.0])} - input_dict = annotate.input_tensors( - 'func1', {'input_dict': input_dict}) - input_dict['input1'] = input_dict['input1'].matmul(torch.tensor( - [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])) - annotate.output_tensors('func1', input_dict, export_with="jit") + input_dict = {"input1": torch.tensor([1.0, 2.0, 3.0])} + input_dict = annotate.input_tensors("func1", {"input_dict": input_dict}) + input_dict["input1"] = input_dict["input1"].matmul( + torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) + ) + annotate.output_tensors("func1", input_dict, export_with="jit") leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0) + annotate, nodes=1, inputs=1, outputs=1, internal_connections=0 + ) def test_annotate_tensor_with_list_input_dict_output(self): """Test with list input and dict output.""" leapp.start(name=self.TEST_GRAPH_NAME, verbose=True) - tensor_list = [torch.tensor([1.0, 2.0, 3.0]), - torch.tensor([4.0, 5.0, 6.0])] - tensor_list = annotate.input_tensors( - 'func1', {'tensor_list': tensor_list}) + tensor_list = [torch.tensor([1.0, 2.0, 3.0]), torch.tensor([4.0, 5.0, 6.0])] + tensor_list = annotate.input_tensors("func1", {"tensor_list": tensor_list}) # Process tensors processed_0 = tensor_list[0] * 2.0 processed_1 = tensor_list[1] * 3.0 # Output as dict - output_dict = {'result_a': processed_0, 'result_b': processed_1} - annotate.output_tensors('func1', output_dict, export_with="jit") + output_dict = {"result_a": processed_0, "result_b": processed_1} + annotate.output_tensors("func1", output_dict, export_with="jit") leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=2, outputs=2, internal_connections=0) + annotate, nodes=1, inputs=2, outputs=2, internal_connections=0 + ) def test_annotate_tensor_with_list_and_single_input_concat_output(self): """Test with 2 inputs: one list of 2 tensors, one regular tensor. Output is all tensors concatenated together.""" leapp.start(name=self.TEST_GRAPH_NAME, verbose=True) # First input: a list of 2 tensors - tensor_list = [torch.tensor([1.0, 2.0, 3.0]), - torch.tensor([4.0, 5.0, 6.0])] - tensor_list = annotate.input_tensors( - 'func1', {'tensor_list': tensor_list}) + tensor_list = [torch.tensor([1.0, 2.0, 3.0]), torch.tensor([4.0, 5.0, 6.0])] + tensor_list = annotate.input_tensors("func1", {"tensor_list": tensor_list}) # Second input: a single tensor single_tensor = torch.tensor([7.0, 8.0, 9.0]) single_tensor = annotate.input_tensors( - 'func1', {'single_tensor': single_tensor}) + "func1", {"single_tensor": single_tensor} + ) # Concatenate all tensors together - concatenated = torch.cat( - [tensor_list[0], tensor_list[1], single_tensor], dim=0) + concatenated = torch.cat([tensor_list[0], tensor_list[1], single_tensor], dim=0) # Output as single tensor annotate.output_tensors( - 'func1', {'concatenated': concatenated}, export_with="jit") + "func1", {"concatenated": concatenated}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=3, outputs=1, internal_connections=0) + annotate, nodes=1, inputs=3, outputs=1, internal_connections=0 + ) def test_annotate_tensor_sequential_context(self): leapp.start(name=self.TEST_GRAPH_NAME, verbose=True) tensor1 = annotate.input_tensors( - 'func1', {'input1': torch.tensor([1.0, 2.0, 3.0])}) + "func1", {"input1": torch.tensor([1.0, 2.0, 3.0])} + ) tensor2 = tensor1 + 1.0 - annotate.output_tensors('func1', {'output1': tensor2}, - export_with="jit") - tensor2 = annotate.input_tensors('func2', {'input2': tensor2}) + annotate.output_tensors("func1", {"output1": tensor2}, export_with="jit") + tensor2 = annotate.input_tensors("func2", {"input2": tensor2}) tensor3 = tensor2 + 2.0 - annotate.output_tensors('func2', {'output2': tensor3}, - export_with="jit") - tensor3 = annotate.input_tensors('func3', {'input3': tensor3}) + annotate.output_tensors("func2", {"output2": tensor3}, export_with="jit") + tensor3 = annotate.input_tensors("func3", {"input3": tensor3}) tensor4 = tensor3 + 3.0 - annotate.output_tensors('func3', {'output3': tensor4}, - export_with="jit") + annotate.output_tensors("func3", {"output3": tensor4}, export_with="jit") leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=3, inputs=1, outputs=1, internal_connections=2) + annotate, nodes=3, inputs=1, outputs=1, internal_connections=2 + ) def test_annotate_multiple_parallel_contexts(self): leapp.start(name=self.TEST_GRAPH_NAME) for i in range(10): tensor1 = annotate.input_tensors( - f'func{i}', {'input1': torch.tensor([1.0, 2.0, 3.0])}) + f"func{i}", {"input1": torch.tensor([1.0, 2.0, 3.0])} + ) tensor1 += i - annotate.output_tensors( - f'func{i}', {'output1': tensor1}, export_with="jit") + annotate.output_tensors(f"func{i}", {"output1": tensor1}, export_with="jit") leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=10, inputs=10, outputs=10, internal_connections=0) + annotate, nodes=10, inputs=10, outputs=10, internal_connections=0 + ) self.verify_all_models_exist( - 'func0', 'func1', 'func2', 'func3', 'func4', 'func5', 'func6', 'func7', 'func8', 'func9') + "func0", + "func1", + "func2", + "func3", + "func4", + "func5", + "func6", + "func7", + "func8", + "func9", + ) def test_annotate_multiple_parallel_inputs_as_one_context(self): @@ -596,85 +645,101 @@ def test_annotate_multiple_parallel_inputs_as_one_context(self): output_tensors = [] for i in range(10): tensor1 = annotate.input_tensors( - 'func_combined', {f'input{i}': torch.tensor([1.0, 2.0, 3.0])}) + "func_combined", {f"input{i}": torch.tensor([1.0, 2.0, 3.0])} + ) tensor1 += i output_tensors.append(tensor1) annotate.output_tensors( - 'func_combined', {'outputs': output_tensors}, export_with="jit") + "func_combined", {"outputs": output_tensors}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=10, outputs=10, internal_connections=0) - self.verify_all_models_exist('func_combined') + annotate, nodes=1, inputs=10, outputs=10, internal_connections=0 + ) + self.verify_all_models_exist("func_combined") def test_annotate_node_with_trimmed_inputs(self): - '''test the situation where the node has inputs that are not used in the computation''' + """test the situation where the node has inputs that are not used in the computation""" leapp.start(name=self.TEST_GRAPH_NAME) for i in range(10): tensor1 = annotate.input_tensors( - 'func1', {'input1': torch.tensor([1.0, 2.0, 3.0])}) + "func1", {"input1": torch.tensor([1.0, 2.0, 3.0])} + ) tensor2 = annotate.input_tensors( - 'func1', {'input2': torch.tensor([4.0, 5.0, 6.0])}) + "func1", {"input2": torch.tensor([4.0, 5.0, 6.0])} + ) tensor3 = tensor1 + 20 - annotate.output_tensors( - 'func1', {'output1': tensor3}, export_with="jit") + annotate.output_tensors("func1", {"output1": tensor3}, export_with="jit") leapp.stop() leapp.compile_graph(visualize=False) - result = self.inspect_torchscript_model('func1') + result = self.inspect_torchscript_model("func1") self.assertTrue( - 'input2' not in result['inputs'], "input2 should be trimmed from the model") - self.assertTrue(len(result['inputs']) == 2, - "Unexpected number of inputs in the model") - self.assertTrue(len(result['outputs']) == 1, - "Unexpected number of outputs in the model") + "input2" not in result["inputs"], "input2 should be trimmed from the model" + ) + self.assertTrue( + len(result["inputs"]) == 2, "Unexpected number of inputs in the model" + ) + self.assertTrue( + len(result["outputs"]) == 1, "Unexpected number of outputs in the model" + ) - self.verify_all_models_exist('func1') + self.verify_all_models_exist("func1") self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0) + annotate, nodes=1, inputs=1, outputs=1, internal_connections=0 + ) def test_annotate_node_with_many_trimmed_inputs(self): leapp.start(name=self.TEST_GRAPH_NAME) for i in range(5): - inputs = [torch.tensor( - [float(i), float(i+1), float(i+2)], dtype=torch.float32) for i in range(20)] - inputs = annotate.input_tensors('func1', {'inputs': inputs}) + inputs = [ + torch.tensor( + [float(i), float(i + 1), float(i + 2)], dtype=torch.float32 + ) + for i in range(20) + ] + inputs = annotate.input_tensors("func1", {"inputs": inputs}) # Start with first traced input, then accumulate (keeps output in traced graph) output = torch.zeros(3, dtype=torch.float32) for i in range(0, 15): output = output + inputs[i] - annotate.output_tensors( - 'func1', {'output': output}, export_with="jit") + annotate.output_tensors("func1", {"output": output}, export_with="jit") leapp.stop() leapp.compile_graph(visualize=False) - result = self.inspect_torchscript_model('func1') - self.assertEqual(len(result['inputs']), 16, - "Unexpected number of inputs in the model") - self.assertEqual(len(result['outputs']), 1, - "Unexpected number of outputs in the model") + result = self.inspect_torchscript_model("func1") + self.assertEqual( + len(result["inputs"]), 16, "Unexpected number of inputs in the model" + ) + self.assertEqual( + len(result["outputs"]), 1, "Unexpected number of outputs in the model" + ) - self.verify_all_models_exist('func1') + self.verify_all_models_exist("func1") # 15 inputs used (inputs 0-14), 5 trimmed (inputs 15-19) self.verify_num_connections( - annotate, nodes=1, inputs=15, outputs=1, internal_connections=0) + annotate, nodes=1, inputs=15, outputs=1, internal_connections=0 + ) - inputs = [torch.tensor( - [float(i), float(i+1), float(i+2)], dtype=torch.float32) for i in range(15)] - self.verify_single_torchscript_model_expected_value( - [inputs], [output], 'func1') + inputs = [ + torch.tensor([float(i), float(i + 1), float(i + 2)], dtype=torch.float32) + for i in range(15) + ] + self.verify_single_torchscript_model_expected_value([inputs], [output], "func1") def test_annotate_multiple_runs(self): leapp.start(name=self.TEST_GRAPH_NAME) for i in range(10): tensor1 = annotate.input_tensors( - 'func1', {'input1': torch.tensor([1.0, 2.0, 3.0])}) + "func1", {"input1": torch.tensor([1.0, 2.0, 3.0])} + ) if i == 0: self.assertTrue(type(tensor1) is TracedTensor) if i > 0: @@ -682,15 +747,15 @@ def test_annotate_multiple_runs(self): self.assertFalse(tensor1.is_tracing) self.assertEqual(tensor1.context, "func1") tensor1 += 0 - annotate.output_tensors( - 'func1', {'output1': tensor1}, export_with="jit") + annotate.output_tensors("func1", {"output1": tensor1}, export_with="jit") leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0) - self.verify_all_models_exist('func1') + annotate, nodes=1, inputs=1, outputs=1, internal_connections=0 + ) + self.verify_all_models_exist("func1") def test_annotate_traced_tensor_override_same_node(self): """Test that passing an active TracedTensor from the same node into input_tensors @@ -704,62 +769,67 @@ def test_annotate_traced_tensor_override_same_node(self): # First input - creates a TracedTensor input1 = annotate.input_tensors( - 'physics_node', {'velocity': torch.tensor([1.0, 2.0, 3.0])}) - self.assertTrue(isinstance(input1, TracedTensor), - "First input should be a TracedTensor") + "physics_node", {"velocity": torch.tensor([1.0, 2.0, 3.0])} + ) + self.assertTrue( + isinstance(input1, TracedTensor), "First input should be a TracedTensor" + ) # Some processing on input1 processed1 = input1 * 2.0 # Now pass the processed TracedTensor back into input_tensors for the SAME node # This should create a fresh input placeholder (override behavior) with a warning - input2 = annotate.input_tensors( - 'physics_node', {'acceleration': processed1}) - self.assertTrue(isinstance(input2, TracedTensor), - "Second input should also be a TracedTensor") + input2 = annotate.input_tensors("physics_node", {"acceleration": processed1}) + self.assertTrue( + isinstance(input2, TracedTensor), + "Second input should also be a TracedTensor", + ) # More processing on input2 final_output = input2 + 10.0 - annotate.output_tensors('physics_node', {'force': final_output}, - export_with="jit") + annotate.output_tensors( + "physics_node", {"force": final_output}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) # Verify we have 2 inputs (velocity and acceleration) - self.assertEqual(len(annotate.nodes['physics_node'].inputs), 1) - input_names = [ - inp.name for inp in annotate.nodes['physics_node'].inputs] - self.assertIn('acceleration', input_names) + self.assertEqual(len(annotate.nodes["physics_node"].inputs), 1) + input_names = [inp.name for inp in annotate.nodes["physics_node"].inputs] + self.assertIn("acceleration", input_names) # Verify 1 output - self.assertEqual(len(annotate.nodes['physics_node'].outputs), 1) + self.assertEqual(len(annotate.nodes["physics_node"].outputs), 1) # Verify model exists and connections self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0) - self.verify_all_models_exist('physics_node') + annotate, nodes=1, inputs=1, outputs=1, internal_connections=0 + ) + self.verify_all_models_exist("physics_node") # Verify the torchscript model works correctly # The model should take 2 inputs: velocity and acceleration # Since input1 (velocity) is overridden by input2 (acceleration), the computation # only uses input2 (acceleration) -> acceleration + 10.0 - model_info = self.inspect_torchscript_model('physics_node') + model_info = self.inspect_torchscript_model("physics_node") # Note: velocity input will be trimmed since it's not used in the final output # (it was overridden by acceleration) - self.assertEqual(len(model_info['outputs']), 1) + self.assertEqual(len(model_info["outputs"]), 1) # Test model execution - only acceleration is used in the graph # Expected: acceleration + 10.0 test_acceleration = torch.tensor([5.0, 6.0, 7.0]) expected_output = test_acceleration + 10.0 self.verify_single_torchscript_model_expected_value( - [test_acceleration], [expected_output], 'physics_node') + [test_acceleration], [expected_output], "physics_node" + ) def test_annotate_traced_tensors_with_feedback(self): - '''two traced tensor nodes. - the first one has 2 inputs 1 output which feeds into the second one. - the output of the second one going into one of the inputs of the first one''' + """two traced tensor nodes. + the first one has 2 inputs 1 output which feeds into the second one. + the output of the second one going into one of the inputs of the first one""" leapp.start(name=self.TEST_GRAPH_NAME) @@ -768,19 +838,20 @@ def test_annotate_traced_tensors_with_feedback(self): for i in range(5): # func1: 2 inputs (input1 and feedback from func2), 1 output - input1, feedback_tensor = annotate.input_tensors('func1', { - 'input1': torch.tensor([1.0, 2.0, 3.0]), - 'feedback_in': feedback_tensor - }) + input1, feedback_tensor = annotate.input_tensors( + "func1", + { + "input1": torch.tensor([1.0, 2.0, 3.0]), + "feedback_in": feedback_tensor, + }, + ) output1 = input1 + feedback_tensor - annotate.output_tensors( - 'func1', {'output1': output1}, export_with="jit") + annotate.output_tensors("func1", {"output1": output1}, export_with="jit") # func2: takes output from func1, produces output that feeds back to func1 - input2 = annotate.input_tensors('func2', {'input2': output1}) + input2 = annotate.input_tensors("func2", {"input2": output1}) output2 = input2 * 2.0 - annotate.output_tensors( - 'func2', {'output2': output2}, export_with="jit") + annotate.output_tensors("func2", {"output2": output2}, export_with="jit") # Update feedback tensor for next iteration feedback_tensor = output2 @@ -792,13 +863,18 @@ def test_annotate_traced_tensors_with_feedback(self): # 1 internal connection (func1 -> func2) # 1 feedback connection (func2 -> func1) self.verify_num_connections( - annotate, nodes=2, inputs=1, outputs=0, - internal_connections=1, feedback_connections=1) - self.verify_all_models_exist('func1', 'func2') + annotate, + nodes=2, + inputs=1, + outputs=0, + internal_connections=1, + feedback_connections=1, + ) + self.verify_all_models_exist("func1", "func2") def test_annotate_traced_tensors_three_node_chain_with_feedback(self): - '''Three node chain where the last node feeds back to the first. - func1 -> func2 -> func3 -> (feedback to func1)''' + """Three node chain where the last node feeds back to the first. + func1 -> func2 -> func3 -> (feedback to func1)""" leapp.start(name=self.TEST_GRAPH_NAME) @@ -806,25 +882,22 @@ def test_annotate_traced_tensors_three_node_chain_with_feedback(self): for i in range(10): # func1: external input + feedback from func3 - input1, fb = annotate.input_tensors('func1', { - 'input1': torch.tensor([1.0, 2.0, 3.0]), - 'feedback': feedback_tensor - }) + input1, fb = annotate.input_tensors( + "func1", + {"input1": torch.tensor([1.0, 2.0, 3.0]), "feedback": feedback_tensor}, + ) out1 = input1 + fb - annotate.output_tensors( - 'func1', {'out1': out1}, export_with="jit") + annotate.output_tensors("func1", {"out1": out1}, export_with="jit") # func2: middle of the chain - in2 = annotate.input_tensors('func2', {'in2': out1}) + in2 = annotate.input_tensors("func2", {"in2": out1}) out2 = in2 * 2.0 - annotate.output_tensors( - 'func2', {'out2': out2}, export_with="jit") + annotate.output_tensors("func2", {"out2": out2}, export_with="jit") # func3: end of chain, output feeds back to func1 - in3 = annotate.input_tensors('func3', {'in3': out2}) + in3 = annotate.input_tensors("func3", {"in3": out2}) out3 = in3 - 1.0 - annotate.output_tensors( - 'func3', {'out3': out3}, export_with="jit") + annotate.output_tensors("func3", {"out3": out3}, export_with="jit") feedback_tensor = out3 @@ -835,50 +908,61 @@ def test_annotate_traced_tensors_three_node_chain_with_feedback(self): # 2 internal connections (func1->func2, func2->func3) # 1 feedback connection (func3->func1) self.verify_num_connections( - annotate, nodes=3, inputs=1, outputs=0, - internal_connections=2, feedback_connections=1) - self.verify_all_models_exist('func1', 'func2', 'func3') + annotate, + nodes=3, + inputs=1, + outputs=0, + internal_connections=2, + feedback_connections=1, + ) + self.verify_all_models_exist("func1", "func2", "func3") def test_annotate_traced_tensors_with_complex_nested_io(self): - '''Single node with complex nested dict/list IO run multiple times''' + """Single node with complex nested dict/list IO run multiple times""" leapp.start(name=self.TEST_GRAPH_NAME) for i in range(10): # Complex nested input structure - tensor_list, single_tensor = annotate.input_tensors('complex_node', { - 'tensor_list': [ - torch.tensor([1.0, 2.0, 3.0]), - torch.tensor([4.0, 5.0, 6.0]) - ], - 'single_tensor': torch.tensor([7.0, 8.0, 9.0]) - }) + tensor_list, single_tensor = annotate.input_tensors( + "complex_node", + { + "tensor_list": [ + torch.tensor([1.0, 2.0, 3.0]), + torch.tensor([4.0, 5.0, 6.0]), + ], + "single_tensor": torch.tensor([7.0, 8.0, 9.0]), + }, + ) # Process the nested inputs list_sum = tensor_list[0] + tensor_list[1] combined = list_sum + single_tensor # Complex nested output structure - annotate.output_tensors('complex_node', { - 'sum': list_sum, - 'combined': combined}, export_with="jit") + annotate.output_tensors( + "complex_node", + {"sum": list_sum, "combined": combined}, + export_with="jit", + ) leapp.stop() leapp.compile_graph(visualize=False) # 1 node, 3 flattened inputs, 2 flattened outputs self.verify_num_connections( - annotate, nodes=1, inputs=3, outputs=2, internal_connections=0) - self.verify_all_models_exist('complex_node') + annotate, nodes=1, inputs=3, outputs=2, internal_connections=0 + ) + self.verify_all_models_exist("complex_node") def test_annotate_traced_tensors_diamond_with_feedback(self): - '''Diamond pattern: func1 splits into func2a and func2b, + """Diamond pattern: func1 splits into func2a and func2b, which merge at func3. func3 feeds back to func1. -> func2a -> func1 func3 -> (feedback to func1) -> func2b -> - ''' + """ leapp.start(name=self.TEST_GRAPH_NAME) @@ -886,34 +970,27 @@ def test_annotate_traced_tensors_diamond_with_feedback(self): for i in range(10): # func1: source node with feedback - in1, fb = annotate.input_tensors('func1', { - 'input': torch.tensor([1.0, 2.0, 3.0]), - 'feedback': feedback_tensor - }) + in1, fb = annotate.input_tensors( + "func1", + {"input": torch.tensor([1.0, 2.0, 3.0]), "feedback": feedback_tensor}, + ) out1 = in1 + fb - annotate.output_tensors( - 'func1', {'out': out1}, export_with="jit") + annotate.output_tensors("func1", {"out": out1}, export_with="jit") # func2a: first parallel branch - in2a = annotate.input_tensors('func2a', {'input2a': out1}) + in2a = annotate.input_tensors("func2a", {"input2a": out1}) out2a = in2a * 2.0 - annotate.output_tensors( - 'func2a', {'out': out2a}, export_with="jit") + annotate.output_tensors("func2a", {"out": out2a}, export_with="jit") # func2b: second parallel branch - in2b = annotate.input_tensors('func2b', {'input2b': out1}) + in2b = annotate.input_tensors("func2b", {"input2b": out1}) out2b = in2b * 3.0 - annotate.output_tensors( - 'func2b', {'out': out2b}, export_with="jit") + annotate.output_tensors("func2b", {"out": out2b}, export_with="jit") # func3: merge node, output feeds back - in3a, in3b = annotate.input_tensors('func3', { - 'in_a': out2a, - 'in_b': out2b - }) + in3a, in3b = annotate.input_tensors("func3", {"in_a": out2a, "in_b": out2b}) out3 = in3a + in3b - annotate.output_tensors( - 'func3', {'out': out3}, export_with="jit") + annotate.output_tensors("func3", {"out": out3}, export_with="jit") feedback_tensor = out3 @@ -931,15 +1008,20 @@ def test_annotate_traced_tensors_diamond_with_feedback(self): # Actually func2a->func3 and func2b->func3 are 2 separate connections # 1 feedback connection (func3->func1) self.verify_num_connections( - annotate, nodes=4, inputs=1, outputs=0, - internal_connections=4, feedback_connections=1) - self.verify_all_models_exist('func1', 'func2a', 'func2b', 'func3') + annotate, + nodes=4, + inputs=1, + outputs=0, + internal_connections=4, + feedback_connections=1, + ) + self.verify_all_models_exist("func1", "func2a", "func2b", "func3") def test_annotate_traced_tensors_two_independent_feedback_loops(self): - '''Two independent feedback loops running in parallel. + """Two independent feedback loops running in parallel. Loop A: funcA1 -> funcA2 -> (feedback to funcA1) Loop B: funcB1 -> funcB2 -> (feedback to funcB1) - ''' + """ leapp.start(name=self.TEST_GRAPH_NAME) @@ -948,33 +1030,29 @@ def test_annotate_traced_tensors_two_independent_feedback_loops(self): for i in range(10): # Loop A - inA1, fbA = annotate.input_tensors('funcA1', { - 'input': torch.tensor([1.0, 2.0, 3.0]), - 'feedback': feedback_a - }) + inA1, fbA = annotate.input_tensors( + "funcA1", + {"input": torch.tensor([1.0, 2.0, 3.0]), "feedback": feedback_a}, + ) outA1 = inA1 + fbA - annotate.output_tensors( - 'funcA1', {'out': outA1}, export_with="jit") + annotate.output_tensors("funcA1", {"out": outA1}, export_with="jit") - inA2 = annotate.input_tensors('funcA2', {'inputA2': outA1}) + inA2 = annotate.input_tensors("funcA2", {"inputA2": outA1}) outA2 = inA2 * 2.0 - annotate.output_tensors( - 'funcA2', {'out': outA2}, export_with="jit") + annotate.output_tensors("funcA2", {"out": outA2}, export_with="jit") feedback_a = outA2 # Loop B (independent) - inB1, fbB = annotate.input_tensors('funcB1', { - 'input': torch.tensor([4.0, 5.0, 6.0]), - 'feedback': feedback_b - }) + inB1, fbB = annotate.input_tensors( + "funcB1", + {"input": torch.tensor([4.0, 5.0, 6.0]), "feedback": feedback_b}, + ) outB1 = inB1 - fbB - annotate.output_tensors( - 'funcB1', {'out': outB1}, export_with="jit") + annotate.output_tensors("funcB1", {"out": outB1}, export_with="jit") - inB2 = annotate.input_tensors('funcB2', {'inputB2': outB1}) + inB2 = annotate.input_tensors("funcB2", {"inputB2": outB1}) outB2 = inB2 / 2.0 - annotate.output_tensors( - 'funcB2', {'out': outB2}, export_with="jit") + annotate.output_tensors("funcB2", {"out": outB2}, export_with="jit") feedback_b = outB2 leapp.stop() @@ -984,102 +1062,116 @@ def test_annotate_traced_tensors_two_independent_feedback_loops(self): # 2 internal connections (funcA1->funcA2, funcB1->funcB2) # 2 feedback connections (funcA2->funcA1, funcB2->funcB1) self.verify_num_connections( - annotate, nodes=4, inputs=2, outputs=0, - internal_connections=2, feedback_connections=2) - self.verify_all_models_exist('funcA1', 'funcA2', 'funcB1', 'funcB2') + annotate, + nodes=4, + inputs=2, + outputs=0, + internal_connections=2, + feedback_connections=2, + ) + self.verify_all_models_exist("funcA1", "funcA2", "funcB1", "funcB2") def test_annotate_traced_tensors_nested_dict_multiple_runs(self): - '''Test nested dict input structure with multiple iterations. - This tests that validation correctly handles nested structures on reentry.''' + """Test nested dict input structure with multiple iterations. + This tests that validation correctly handles nested structures on reentry.""" leapp.start(name=self.TEST_GRAPH_NAME) for i in range(5): # Nested dict input - {'group': {'a': tensor, 'b': tensor}} - inputs = annotate.input_tensors('nested_node', { - 'group': { - 'x': torch.tensor([1.0, 2.0, 3.0]), - 'y': torch.tensor([4.0, 5.0, 6.0]) - } - }) + inputs = annotate.input_tensors( + "nested_node", + { + "group": { + "x": torch.tensor([1.0, 2.0, 3.0]), + "y": torch.tensor([4.0, 5.0, 6.0]), + } + }, + ) # Access nested structure - result = inputs['x'] + inputs['y'] + result = inputs["x"] + inputs["y"] annotate.output_tensors( - 'nested_node', {'result': result}, export_with="jit") + "nested_node", {"result": result}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False, validate=True) # 1 node, 2 flattened inputs (group_x, group_y), 1 output self.verify_num_connections( - annotate, nodes=1, inputs=2, outputs=1, internal_connections=0) - self.verify_all_models_exist('nested_node') + annotate, nodes=1, inputs=2, outputs=1, internal_connections=0 + ) + self.verify_all_models_exist("nested_node") def test_annotate_traced_tensors_nested_list_multiple_runs(self): - '''Test nested list input structure with multiple iterations.''' + """Test nested list input structure with multiple iterations.""" leapp.start(name=self.TEST_GRAPH_NAME) for i in range(10): # List input that gets flattened - tensor_list = annotate.input_tensors('list_node', { - 'tensors': [ - torch.tensor([1.0, 2.0, 3.0]), - torch.tensor([4.0, 5.0, 6.0]), - torch.tensor([7.0, 8.0, 9.0]) - ] - }) + tensor_list = annotate.input_tensors( + "list_node", + { + "tensors": [ + torch.tensor([1.0, 2.0, 3.0]), + torch.tensor([4.0, 5.0, 6.0]), + torch.tensor([7.0, 8.0, 9.0]), + ] + }, + ) # Access list elements result = tensor_list[0] + tensor_list[1] + tensor_list[2] - annotate.output_tensors( - 'list_node', {'result': result}, export_with="jit") + annotate.output_tensors("list_node", {"result": result}, export_with="jit") leapp.stop() leapp.compile_graph(visualize=False) # 1 node, 3 flattened inputs, 1 output self.verify_num_connections( - annotate, nodes=1, inputs=3, outputs=1, internal_connections=0) - self.verify_all_models_exist('list_node') + annotate, nodes=1, inputs=3, outputs=1, internal_connections=0 + ) + self.verify_all_models_exist("list_node") def test_annotate_traced_tensors_mixed_nested_structure_multiple_runs(self): - '''Test mixed nested structure (dict containing list) with multiple iterations.''' + """Test mixed nested structure (dict containing list) with multiple iterations.""" leapp.start(name=self.TEST_GRAPH_NAME) for i in range(10): # Mixed nested: dict with both list and scalar tensor - inputs = annotate.input_tensors('mixed_node', { - 'batch': [ - torch.tensor([1.0, 2.0]), - torch.tensor([3.0, 4.0]) - ], - 'scale': torch.tensor([2.0]) - }) + inputs = annotate.input_tensors( + "mixed_node", + { + "batch": [torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0])], + "scale": torch.tensor([2.0]), + }, + ) # Access mixed structure - inputs should be (list, tensor) tuple batch_list, scale = inputs result = (batch_list[0] + batch_list[1]) * scale - annotate.output_tensors( - 'mixed_node', {'result': result}, export_with="jit") + annotate.output_tensors("mixed_node", {"result": result}, export_with="jit") leapp.stop() leapp.compile_graph(visualize=False) # 1 node, 3 flattened inputs (batch_0, batch_1, scale), 1 output self.verify_num_connections( - annotate, nodes=1, inputs=3, outputs=1, internal_connections=0) - self.verify_all_models_exist('mixed_node') + annotate, nodes=1, inputs=3, outputs=1, internal_connections=0 + ) + self.verify_all_models_exist("mixed_node") class TestAnnotateMixed(LEAPPFunctionalTestBase): def test_mixing_annotated_tensors_and_method_nodes(self): """Test: traced_tensors → method""" + def run_function(tensor): return tensor + 1.0 @@ -1090,20 +1182,24 @@ def funcA(inputA: torch.Tensor): input_tensor = torch.tensor([1.0, 2.0, 3.0]) leapp.start(name=self.TEST_GRAPH_NAME) input_tensor = annotate.input_tensors( - 'run_function', {'input_tensor': input_tensor}) + "run_function", {"input_tensor": input_tensor} + ) output_tensor = run_function(input_tensor) annotate.output_tensors( - 'run_function', {'output_tensor': output_tensor}, export_with="jit") + "run_function", {"output_tensor": output_tensor}, export_with="jit" + ) outputA = funcA(output_tensor) leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=2, inputs=1, outputs=1, internal_connections=1) - self.verify_all_models_exist('run_function', 'funcA') + annotate, nodes=2, inputs=1, outputs=1, internal_connections=1 + ) + self.verify_all_models_exist("run_function", "funcA") def test_method_then_traced_tensors(self): """Test: method → traced_tensors""" + @annotate.method(export_with="jit") def funcA(inputA: torch.Tensor): return inputA * 2.0 @@ -1116,19 +1212,23 @@ def run_function(tensor): output_funcA = funcA(input_tensor) # Now use traced tensors for the next step traced_input = annotate.input_tensors( - 'run_function', {'traced_input': output_funcA}) + "run_function", {"traced_input": output_funcA} + ) output_tensor = run_function(traced_input) annotate.output_tensors( - 'run_function', {'output_tensor': output_tensor}, export_with="jit") + "run_function", {"output_tensor": output_tensor}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=2, inputs=1, outputs=1, internal_connections=1) - self.verify_all_models_exist('funcA', 'run_function') + annotate, nodes=2, inputs=1, outputs=1, internal_connections=1 + ) + self.verify_all_models_exist("funcA", "run_function") def test_traced_tensors_sandwich_method(self): """Test: traced_tensors → method → traced_tensors""" + def preprocess(tensor): return tensor * 2.0 @@ -1142,26 +1242,28 @@ def postprocess(tensor): input_tensor = torch.tensor([1.0, 2.0, 3.0]) leapp.start(name=self.TEST_GRAPH_NAME) # First traced tensor node - traced_input = annotate.input_tensors( - 'preprocess', {'input': input_tensor}) + traced_input = annotate.input_tensors("preprocess", {"input": input_tensor}) preprocessed = preprocess(traced_input) annotate.output_tensors( - 'preprocess', {'preprocessed': preprocessed}, export_with="jit") + "preprocess", {"preprocessed": preprocessed}, export_with="jit" + ) # Method in the middle method_output = funcA(preprocessed) # Second traced tensor node traced_method_output = annotate.input_tensors( - 'postprocess', {'method_out': method_output}) + "postprocess", {"method_out": method_output} + ) postprocessed = postprocess(traced_method_output) annotate.output_tensors( - 'postprocess', {'postprocessed': postprocessed}, export_with="jit") + "postprocess", {"postprocessed": postprocessed}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_num_connections( - annotate, nodes=3, inputs=1, outputs=1, internal_connections=2) - self.verify_all_models_exist('preprocess', 'funcA', 'postprocess') + annotate, nodes=3, inputs=1, outputs=1, internal_connections=2 + ) + self.verify_all_models_exist("preprocess", "funcA", "postprocess") # def test_block_then_traced_tensors(self): # """Test: block → traced_tensors""" @@ -1240,6 +1342,7 @@ def postprocess(tensor): def test_two_parallel_traced_tensors_merge_to_method(self): """Test: two parallel traced_tensor nodes feeding into one method""" + def process_a(tensor): return tensor * 2.0 @@ -1254,25 +1357,24 @@ def combine(inputA: torch.Tensor, inputB: torch.Tensor): input_b = torch.tensor([4.0, 5.0, 6.0]) leapp.start(name=self.TEST_GRAPH_NAME) # First traced tensor path - traced_a = annotate.input_tensors('process_a', {'input_a': input_a}) + traced_a = annotate.input_tensors("process_a", {"input_a": input_a}) # First traced tensor path - traced_b = annotate.input_tensors('process_b', {'input_b': input_b}) + traced_b = annotate.input_tensors("process_b", {"input_b": input_b}) output_a = process_a(traced_a) output_b = process_b(traced_b) - annotate.output_tensors('process_a', {'output_a': output_a}, - export_with="jit") - annotate.output_tensors('process_b', {'output_b': output_b}, - export_with="jit") + annotate.output_tensors("process_a", {"output_a": output_a}, export_with="jit") + annotate.output_tensors("process_b", {"output_b": output_b}, export_with="jit") # Combine with method combined = combine(output_a, output_b) leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=3, inputs=2, outputs=1, internal_connections=2) - self.verify_all_models_exist('process_a', 'process_b', 'combine') + annotate, nodes=3, inputs=2, outputs=1, internal_connections=2 + ) + self.verify_all_models_exist("process_a", "process_b", "combine") # def test_block_with_multiline_dict_comprehension(self): # """Test that block context correctly handles multiline dict comprehensions. @@ -1307,5 +1409,5 @@ def combine(inputA: torch.Tensor, inputB: torch.Tensor): # self.verify_all_models_exist('multiline_dict') -if __name__ == '__main__': +if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/functional_tests/test_config_generation.py b/tests/functional_tests/test_config_generation.py index 96344edf..20d8cd16 100644 --- a/tests/functional_tests/test_config_generation.py +++ b/tests/functional_tests/test_config_generation.py @@ -34,20 +34,20 @@ def _load_yaml(self): """Load the generated YAML config file.""" yaml_path = os.path.join(self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}.yaml") self.assertTrue(os.path.exists(yaml_path), f"YAML file not found: {yaml_path}") - with open(yaml_path, 'r') as f: + with open(yaml_path, "r") as f: return yaml.safe_load(f) def _get_node_io_from_yaml(self, config, node_name): """Get inputs and outputs dicts for a node from the YAML config.""" - self.assertIn('models', config) - self.assertIn(node_name, config['models']) - node = config['models'][node_name] - return node.get('inputs', []), node.get('outputs', []) + self.assertIn("models", config) + self.assertIn(node_name, config["models"]) + node = config["models"][node_name] + return node.get("inputs", []), node.get("outputs", []) def _find_io_by_name(self, io_list, name): """Find an input/output entry by name in a YAML io list.""" for entry in io_list: - if entry['name'] == name: + if entry["name"] == name: return entry return None @@ -61,10 +61,17 @@ def test_input_td_list_with_kind(self): joint_vel = torch.randn(1, 12) leapp.start(name=self.TEST_GRAPH_NAME) - traced_pos, traced_vel = annotate.input_tensors("policy", [ - TensorSemantics(name="joint_pos", ref=joint_pos, kind=InputKindEnum.JOINT_POSITION), - TensorSemantics(name="joint_vel", ref=joint_vel, kind=InputKindEnum.JOINT_VELOCITY), - ]) + traced_pos, traced_vel = annotate.input_tensors( + "policy", + [ + TensorSemantics( + name="joint_pos", ref=joint_pos, kind=InputKindEnum.JOINT_POSITION + ), + TensorSemantics( + name="joint_vel", ref=joint_vel, kind=InputKindEnum.JOINT_VELOCITY + ), + ], + ) output = traced_pos + traced_vel annotate.output_tensors("policy", {"command": output}) @@ -80,10 +87,10 @@ def test_input_td_list_with_kind(self): self.assertIsNotNone(pos_entry) self.assertIsNotNone(vel_entry) - self.assertEqual(pos_entry['kind'], "state/joint/position") - self.assertEqual(vel_entry['kind'], "state/joint/velocity") + self.assertEqual(pos_entry["kind"], "state/joint/position") + self.assertEqual(vel_entry["kind"], "state/joint/velocity") # Output without semantics should have no kind - self.assertNotIn('kind', cmd_entry) + self.assertNotIn("kind", cmd_entry) def test_input_single_td(self): """Test passing a single TensorSemantics directly (not in a list).""" @@ -105,8 +112,8 @@ def test_input_single_td(self): pos_entry = self._find_io_by_name(inputs, "pos") self.assertIsNotNone(pos_entry) - self.assertEqual(pos_entry['kind'], "state/joint/position") - self.assertEqual(pos_entry['shape'], [1, 4]) + self.assertEqual(pos_entry["kind"], "state/joint/position") + self.assertEqual(pos_entry["shape"], [1, 4]) def test_input_td_with_element_names(self): """Test that element_names from input TensorSemantics appears in YAML.""" @@ -130,7 +137,7 @@ def test_input_td_with_element_names(self): pos_entry = self._find_io_by_name(inputs, "position") self.assertIsNotNone(pos_entry) # element_names should be normalized to [["x", "y", "z"]] - self.assertEqual(pos_entry['element_names'], [["x", "y", "z"]]) + self.assertEqual(pos_entry["element_names"], [["x", "y", "z"]]) def test_input_td_with_kind_and_element_names(self): """Test that both kind and element_names appear together in YAML.""" @@ -140,9 +147,12 @@ def test_input_td_with_kind_and_element_names(self): leapp.start(name=self.TEST_GRAPH_NAME) traced = annotate.input_tensors( "full_meta_node", - TensorSemantics(name="joint_pos", ref=tensor, - kind=InputKindEnum.JOINT_POSITION, - element_names=joint_names), + TensorSemantics( + name="joint_pos", + ref=tensor, + kind=InputKindEnum.JOINT_POSITION, + element_names=joint_names, + ), ) output = traced * 0.5 @@ -155,10 +165,10 @@ def test_input_td_with_kind_and_element_names(self): entry = self._find_io_by_name(inputs, "joint_pos") self.assertIsNotNone(entry) - self.assertEqual(entry['kind'], "state/joint/position") - self.assertEqual(entry['element_names'], [joint_names]) - self.assertEqual(entry['dtype'], "float32") - self.assertEqual(entry['shape'], [1, 6]) + self.assertEqual(entry["kind"], "state/joint/position") + self.assertEqual(entry["element_names"], [joint_names]) + self.assertEqual(entry["dtype"], "float32") + self.assertEqual(entry["shape"], [1, 6]) def test_input_and_output_td_without_source(self): """Test that TensorSemantics YAML no longer includes a source field.""" @@ -174,13 +184,16 @@ def test_input_and_output_td_without_source(self): ), ) - annotate.output_tensors("source_node", [ - TensorSemantics( - name="filtered_imu", - ref=traced * 0.5, - kind=OutputKindEnum.BODY_ANGULAR_ACCELERATION, - ), - ]) + annotate.output_tensors( + "source_node", + [ + TensorSemantics( + name="filtered_imu", + ref=traced * 0.5, + kind=OutputKindEnum.BODY_ANGULAR_ACCELERATION, + ), + ], + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -196,7 +209,7 @@ def test_input_and_output_td_without_source(self): self.assertNotIn("source", output_entry) self.assertEqual(input_entry["kind"], "state/body/angular_velocity") self.assertEqual(output_entry["kind"], "target/body/angular_acceleration") - + def test_input_td_with_string_kind(self): """Test that a string kind appears in YAML.""" tensor = torch.randn(1, 6) @@ -205,9 +218,12 @@ def test_input_td_with_string_kind(self): leapp.start(name=self.TEST_GRAPH_NAME) traced = annotate.input_tensors( "full_meta_node", - TensorSemantics(name="joint_pos", ref=tensor, - kind="my/custom/kind", - element_names=joint_names), + TensorSemantics( + name="joint_pos", + ref=tensor, + kind="my/custom/kind", + element_names=joint_names, + ), ) output = traced * 0.5 @@ -220,10 +236,10 @@ def test_input_td_with_string_kind(self): entry = self._find_io_by_name(inputs, "joint_pos") self.assertIsNotNone(entry) - self.assertEqual(entry['kind'], "my/custom/kind") - self.assertEqual(entry['element_names'], [joint_names]) - self.assertEqual(entry['dtype'], "float32") - self.assertEqual(entry['shape'], [1, 6]) + self.assertEqual(entry["kind"], "my/custom/kind") + self.assertEqual(entry["element_names"], [joint_names]) + self.assertEqual(entry["dtype"], "float32") + self.assertEqual(entry["shape"], [1, 6]) # ========================================================================= # Output TensorSemantics tests @@ -238,9 +254,14 @@ def test_output_td_with_kind(self): command = traced * 2.0 - annotate.output_tensors("out_kind_node", [ - TensorSemantics(name="command", ref=command, kind=OutputKindEnum.JOINT_TORQUES), - ]) + annotate.output_tensors( + "out_kind_node", + [ + TensorSemantics( + name="command", ref=command, kind=OutputKindEnum.JOINT_TORQUES + ), + ], + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -249,7 +270,7 @@ def test_output_td_with_kind(self): cmd_entry = self._find_io_by_name(outputs, "command") self.assertIsNotNone(cmd_entry) - self.assertEqual(cmd_entry['kind'], "target/joint/torques") + self.assertEqual(cmd_entry["kind"], "target/joint/torques") def test_output_td_with_element_names(self): """Test that element_names from output TensorSemantics appears in YAML.""" @@ -260,9 +281,12 @@ def test_output_td_with_element_names(self): result = traced + 1.0 - annotate.output_tensors("out_elem_node", [ - TensorSemantics(name="rgb", ref=result, element_names=["r", "g", "b"]), - ]) + annotate.output_tensors( + "out_elem_node", + [ + TensorSemantics(name="rgb", ref=result, element_names=["r", "g", "b"]), + ], + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -271,7 +295,7 @@ def test_output_td_with_element_names(self): rgb_entry = self._find_io_by_name(outputs, "rgb") self.assertIsNotNone(rgb_entry) - self.assertEqual(rgb_entry['element_names'], [["r", "g", "b"]]) + self.assertEqual(rgb_entry["element_names"], [["r", "g", "b"]]) def test_static_output_td_with_kind(self): """Test that kind metadata from static output TensorSemantics appears in YAML.""" @@ -298,7 +322,7 @@ def test_static_output_td_with_kind(self): static_entry = self._find_io_by_name(outputs, "command_bias") self.assertIsNotNone(static_entry) - self.assertEqual(static_entry['kind'], "target/joint/torques") + self.assertEqual(static_entry["kind"], "target/joint/torques") # ========================================================================= # Both inputs and outputs @@ -310,25 +334,39 @@ def test_both_input_and_output_td(self): vel = torch.randn(1, 6) leapp.start(name=self.TEST_GRAPH_NAME) - traced_pos, traced_vel = annotate.input_tensors("both_node", [ - TensorSemantics(name="pos", ref=pos, kind=InputKindEnum.JOINT_POSITION), - TensorSemantics(name="vel", ref=vel, kind=InputKindEnum.JOINT_VELOCITY), - ]) + traced_pos, traced_vel = annotate.input_tensors( + "both_node", + [ + TensorSemantics(name="pos", ref=pos, kind=InputKindEnum.JOINT_POSITION), + TensorSemantics(name="vel", ref=vel, kind=InputKindEnum.JOINT_VELOCITY), + ], + ) command = traced_pos + traced_vel - annotate.output_tensors("both_node", [ - TensorSemantics(name="torques", ref=command, kind=OutputKindEnum.JOINT_TORQUES), - ]) + annotate.output_tensors( + "both_node", + [ + TensorSemantics( + name="torques", ref=command, kind=OutputKindEnum.JOINT_TORQUES + ), + ], + ) leapp.stop() leapp.compile_graph(visualize=False) config = self._load_yaml() inputs, outputs = self._get_node_io_from_yaml(config, "both_node") - self.assertEqual(self._find_io_by_name(inputs, "pos")['kind'], "state/joint/position") - self.assertEqual(self._find_io_by_name(inputs, "vel")['kind'], "state/joint/velocity") - self.assertEqual(self._find_io_by_name(outputs, "torques")['kind'], "target/joint/torques") + self.assertEqual( + self._find_io_by_name(inputs, "pos")["kind"], "state/joint/position" + ) + self.assertEqual( + self._find_io_by_name(inputs, "vel")["kind"], "state/joint/velocity" + ) + self.assertEqual( + self._find_io_by_name(outputs, "torques")["kind"], "target/joint/torques" + ) def test_input_td_with_extra_fields_flattened_into_yaml(self): """Test that TensorSemantics.extra fields are emitted as top-level YAML keys.""" @@ -435,14 +473,17 @@ def test_temporal_period_marker_appears_in_output_yaml(self): leapp.start(name=self.TEST_GRAPH_NAME) traced = annotate.input_tensors("chunk_node", {"input": tensor}) - annotate.output_tensors("chunk_node", [ - TensorSemantics( - name="actions", - ref=traced + 1.0, - element_names=[TemporalAxis(period_ms=100), names], - ), - TensorSemantics(name="plain_output", ref=traced - 1.0), - ]) + annotate.output_tensors( + "chunk_node", + [ + TensorSemantics( + name="actions", + ref=traced + 1.0, + element_names=[TemporalAxis(period_ms=100), names], + ), + TensorSemantics(name="plain_output", ref=traced - 1.0), + ], + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -479,12 +520,12 @@ def test_no_metadata_no_extra_yaml_fields(self): input_entry = self._find_io_by_name(inputs, "input") output_entry = self._find_io_by_name(outputs, "output") - self.assertNotIn('kind', input_entry) - self.assertNotIn('source', input_entry) - self.assertNotIn('element_names', input_entry) - self.assertNotIn('kind', output_entry) - self.assertNotIn('source', output_entry) - self.assertNotIn('element_names', output_entry) + self.assertNotIn("kind", input_entry) + self.assertNotIn("source", input_entry) + self.assertNotIn("element_names", input_entry) + self.assertNotIn("kind", output_entry) + self.assertNotIn("source", output_entry) + self.assertNotIn("element_names", output_entry) def test_tensor_semantics_rejects_temporal_period_ms_argument(self): """Test temporal period is only set through TemporalAxis.""" @@ -509,10 +550,15 @@ def test_mixed_td_and_raw_raises(self): leapp.start(name=self.TEST_GRAPH_NAME) with self.assertRaises(TypeError): - annotate.input_tensors("fail_node", [ - t1, - TensorSemantics(name="td", ref=t2, kind=InputKindEnum.JOINT_POSITION), - ]) + annotate.input_tensors( + "fail_node", + [ + t1, + TensorSemantics( + name="td", ref=t2, kind=InputKindEnum.JOINT_POSITION + ), + ], + ) leapp.stop() def test_duplicate_td_names_raises(self): @@ -522,10 +568,13 @@ def test_duplicate_td_names_raises(self): leapp.start(name=self.TEST_GRAPH_NAME) with self.assertRaises(Exception): - annotate.input_tensors("dup_node", [ - TensorSemantics(name="same_name", ref=t1), - TensorSemantics(name="same_name", ref=t2), - ]) + annotate.input_tensors( + "dup_node", + [ + TensorSemantics(name="same_name", ref=t1), + TensorSemantics(name="same_name", ref=t2), + ], + ) leapp.stop() # ========================================================================= @@ -538,10 +587,13 @@ def test_td_inputs_graph_structure(self): vel = torch.randn(1, 4) leapp.start(name=self.TEST_GRAPH_NAME) - traced_pos, traced_vel = annotate.input_tensors("struct_node", [ - TensorSemantics(name="pos", ref=pos, kind=InputKindEnum.JOINT_POSITION), - TensorSemantics(name="vel", ref=vel, kind=InputKindEnum.JOINT_VELOCITY), - ]) + traced_pos, traced_vel = annotate.input_tensors( + "struct_node", + [ + TensorSemantics(name="pos", ref=pos, kind=InputKindEnum.JOINT_POSITION), + TensorSemantics(name="vel", ref=vel, kind=InputKindEnum.JOINT_VELOCITY), + ], + ) output = traced_pos + traced_vel annotate.output_tensors("struct_node", {"command": output}) @@ -549,7 +601,8 @@ def test_td_inputs_graph_structure(self): leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=2, outputs=1, internal_connections=0) + annotate, nodes=1, inputs=2, outputs=1, internal_connections=0 + ) # ========================================================================= # Reentry with TensorSemantics @@ -562,10 +615,16 @@ def test_semantic_input_reentry_does_not_crash(self): leapp.start(name=self.TEST_GRAPH_NAME) for _ in range(2): - traced = annotate.input_tensors("policy", [ - TensorSemantics(name="joint_pos", ref=joint_pos, - kind=InputKindEnum.JOINT_POSITION), - ]) + traced = annotate.input_tensors( + "policy", + [ + TensorSemantics( + name="joint_pos", + ref=joint_pos, + kind=InputKindEnum.JOINT_POSITION, + ), + ], + ) out = traced * 2.0 annotate.output_tensors("policy", {"cmd": out}, export_with="jit") @@ -576,7 +635,7 @@ def test_semantic_input_reentry_does_not_crash(self): inputs, _ = self._get_node_io_from_yaml(config, "policy") entry = self._find_io_by_name(inputs, "joint_pos") self.assertIsNotNone(entry) - self.assertEqual(entry['kind'], "state/joint/position") + self.assertEqual(entry["kind"], "state/joint/position") def test_semantic_output_reentry_does_not_crash(self): """Reentry with TensorSemantics outputs must not KeyError on semantic-only keys.""" @@ -587,10 +646,15 @@ def test_semantic_output_reentry_does_not_crash(self): for _ in range(2): traced = annotate.input_tensors("out_reentry", {"x": tensor}) result = traced + 1.0 - annotate.output_tensors("out_reentry", [ - TensorSemantics(name="torques", ref=result, - kind=OutputKindEnum.JOINT_TORQUES), - ], export_with="jit") + annotate.output_tensors( + "out_reentry", + [ + TensorSemantics( + name="torques", ref=result, kind=OutputKindEnum.JOINT_TORQUES + ), + ], + export_with="jit", + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -599,8 +663,8 @@ def test_semantic_output_reentry_does_not_crash(self): _, outputs = self._get_node_io_from_yaml(config, "out_reentry") entry = self._find_io_by_name(outputs, "torques") self.assertIsNotNone(entry) - self.assertEqual(entry['kind'], "target/joint/torques") + self.assertEqual(entry["kind"], "target/joint/torques") -if __name__ == '__main__': +if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/functional_tests/test_connection_case.py b/tests/functional_tests/test_connection_case.py index 72ab0f06..2facae27 100644 --- a/tests/functional_tests/test_connection_case.py +++ b/tests/functional_tests/test_connection_case.py @@ -74,22 +74,22 @@ def pipeline_views(self): class TestConnectionCase(ConnectivityTestBase): - def test_multiple_runs_of_same_graph(self): """tests the situation where the same graph is run multiple times""" + @annotate.method(export_with="jit") def funcA(inputA: torch.Tensor): return inputA @annotate.method(export_with="jit") def funcC(inputB: torch.Tensor): - return inputB+5.0 + return inputB + 5.0 leapp.start(name=self.TEST_GRAPH_NAME) for _ in range(10): outputA = funcA(torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)) outputA = annotate.input_tensors("blockA", {"outputA": outputA}) - outputB = outputA*2. + outputB = outputA * 2.0 annotate.output_tensors("blockA", {"outputB": outputB}, export_with="jit") funcC(outputB) leapp.stop() @@ -98,6 +98,7 @@ def funcC(inputB: torch.Tensor): def test_feedback_connections(self): """tests the situation where the graph has feedback connections""" + @annotate.method() def funcA(inputA: torch.Tensor, loop_back: torch.Tensor): return inputA + loop_back @@ -127,7 +128,8 @@ def funcD(inputD: torch.Tensor): leapp.stop() leapp.compile_graph(visualize=False) self.verify_connectivity( - nodes=4, internal_connections=3, feedback_connections=1) + nodes=4, internal_connections=3, feedback_connections=1 + ) def test_interleaved_traced_nodes_keep_forward_execution_order(self): """Interleaved traced nodes should classify completed dependencies as forward flow.""" @@ -181,6 +183,7 @@ def test_interleaved_traced_nodes_keep_forward_execution_order(self): def test_mirror_leapp_tags_with_inplace_assignment(self): """mirror_leapp_tags with a single in-place buffer between nodes.""" + @annotate.method(export_with="jit") def funcA(inputA: torch.Tensor): return inputA * 2.0 @@ -206,6 +209,7 @@ def funcC(inputC: torch.Tensor): def test_mirror_leapp_tags_multiple_buffers(self): """Test mirror_leapp_tags with multiple buffers between nodes""" + @annotate.method(export_with="jit") def funcA(inputA: torch.Tensor): return inputA, inputA + 1.0 @@ -221,19 +225,19 @@ def funcC(inputC1: torch.Tensor, inputC2: torch.Tensor): leapp.start(name=self.TEST_GRAPH_NAME) out_A1, out_A2 = funcA(torch.tensor([1.0, 2.0, 3.0])) - + # Create two separate buffers BETWEEN nodes buffer1 = torch.zeros_like(out_A1) buffer2 = torch.zeros_like(out_A2) - + # In-place copy buffer1[:] = out_A1 buffer2[:] = out_A2 - + # Transfer traced state for both annotate.mirror_leapp_tags(out_A1, buffer1) annotate.mirror_leapp_tags(out_A2, buffer2) - + out_B1, out_B2 = funcB(buffer1, buffer2) funcC(out_B1, out_B2) leapp.stop() @@ -244,6 +248,7 @@ def funcC(inputC1: torch.Tensor, inputC2: torch.Tensor): def test_mirror_leapp_tags_preserves_output_port_through_chain(self): """Test that mirror_leapp_tags preserves output ports through a long chain of buffer operations""" + @annotate.method(export_with="jit") def source_node(inputA: torch.Tensor): return inputA * 2.0 @@ -262,19 +267,19 @@ def sink_node(inputD: torch.Tensor): leapp.start(name=self.TEST_GRAPH_NAME) out1 = source_node(torch.tensor([1.0, 2.0, 3.0])) - + # Buffer between nodes 1 and 2 buffer1 = torch.empty_like(out1) buffer1[:] = out1 annotate.mirror_leapp_tags(out1, buffer1) - + out2 = process_node1(buffer1) - + # Buffer between nodes 2 and 3 buffer2 = torch.empty_like(out2) buffer2[:] = out2 annotate.mirror_leapp_tags(out2, buffer2) - + out3 = process_node2(buffer2) sink_node(out3) leapp.stop() @@ -297,12 +302,10 @@ def test_one_external_tensor_declared_on_two_nodes(self): leapp.start(name=self.TEST_GRAPH_NAME) first = annotate.input_tensors("node_a", {"obs": shared}) - annotate.output_tensors( - "node_a", {"out_a": first * 2.0}, export_with="jit") + annotate.output_tensors("node_a", {"out_a": first * 2.0}, export_with="jit") with self.assertNoLogs("leapp", level="ERROR"): second = annotate.input_tensors("node_b", {"obs": shared}) - annotate.output_tensors( - "node_b", {"out_b": second * 3.0}, export_with="jit") + annotate.output_tensors("node_b", {"out_b": second * 3.0}, export_with="jit") leapp.stop() leapp.compile_graph(visualize=False) @@ -313,10 +316,10 @@ def test_mirror_leapp_tags_noop_outside_tracing(self): source = torch.tensor([1.0, 2.0, 3.0]) target = torch.zeros(3) target[:] = source - + # Should not raise an error when called outside tracing annotate.mirror_leapp_tags(source, target) - + # Verify data is still correct self.assertTrue(torch.allclose(source, target)) @@ -383,9 +386,11 @@ def test_numpy_allocating_copy_keeps_output_port(self): def test_numpy_traced_buffer_assignment_adopts_new_source(self): """A traced numpy buffer takes over the identity it was overwritten with.""" out_a = self._finished_output( - np.array([1.0, 2.0, 3.0], dtype=np.float32), "node_a", "out_a") + np.array([1.0, 2.0, 3.0], dtype=np.float32), "node_a", "out_a" + ) out_b = self._finished_output( - np.array([4.0, 5.0, 6.0], dtype=np.float32), "node_b", "out_b") + np.array([4.0, 5.0, 6.0], dtype=np.float32), "node_b", "out_b" + ) buffer = np.copy(out_a) self.assertEqual("out_a", buffer.output_port) @@ -411,11 +416,13 @@ def test_numpy_raw_buffer_assignment_stays_raw(self): def test_cross_backend_copy_keeps_output_port(self): """A finished output handed to the other backend without changing values.""" np_output = self._finished_output( - np.array([1.0, 2.0, 3.0], dtype=np.float32), "np_node", "y") + np.array([1.0, 2.0, 3.0], dtype=np.float32), "np_node", "y" + ) self.assertEqual("y", torch.as_tensor(np.asarray(np_output)).output_port) torch_output = self._finished_output( - torch.tensor([1.0, 2.0, 3.0]), "torch_node", "z") + torch.tensor([1.0, 2.0, 3.0]), "torch_node", "z" + ) self.assertEqual("z", torch_output.numpy().output_port) def test_transformed_copy_reports_no_output_port(self): @@ -441,6 +448,7 @@ class TestAutomaticCopyConnectivity(ConnectivityTestBase): def test_buffer_chain_connects_without_mirroring(self): """Preallocated buffers between four nodes carry the edges themselves.""" + @annotate.method(export_with="jit") def source_node(value: torch.Tensor): return value * 2.0 @@ -475,6 +483,7 @@ def sink_node(value: torch.Tensor): def test_clone_fans_out_to_multiple_consumers(self): """One output cloned separately for two consumers keeps both edges.""" + @annotate.method(export_with="jit") def source_node(value: torch.Tensor): return value * 2.0 @@ -493,11 +502,11 @@ def consumer_b(value: torch.Tensor): consumer_b(out.detach()) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_connectivity( - nodes=3, internal_connections=2, outputs=2) + self.verify_connectivity(nodes=3, internal_connections=2, outputs=2) def test_clone_carries_a_feedback_edge(self): """A cloned output closes a feedback loop across iterations.""" + @annotate.method() def funcA(inputA: torch.Tensor, loop_back: torch.Tensor): return inputA + loop_back @@ -516,10 +525,12 @@ def funcB(inputB: torch.Tensor): leapp.stop() leapp.compile_graph(visualize=False) self.verify_connectivity( - nodes=2, internal_connections=1, feedback_connections=1, outputs=0) + nodes=2, internal_connections=1, feedback_connections=1, outputs=0 + ) def test_repeated_reentry_through_buffers(self): """Re-tracing the same buffered graph does not accumulate edges.""" + @annotate.method(export_with="jit") def source_node(value: torch.Tensor): return value * 2.0 @@ -540,6 +551,7 @@ def sink_node(value: torch.Tensor): def test_transformed_copy_does_not_create_a_connection(self): """A value-changing copy leaves the consumer dangling, not falsely wired.""" + @annotate.method(export_with="jit") def source_node(value: torch.Tensor): return value * 2.0 @@ -553,8 +565,7 @@ def sink_node(value: torch.Tensor): sink_node(out + 1.0) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_connectivity( - nodes=2, internal_connections=0, inputs=2, outputs=2) + self.verify_connectivity(nodes=2, internal_connections=0, inputs=2, outputs=2) class TestNumpyConnectionCase(ConnectivityTestBase): @@ -568,12 +579,14 @@ class TestNumpyConnectionCase(ConnectivityTestBase): def _run_scale_node(self, node_name, input_name, output_name, src, scale): traced = annotate.input_tensors(node_name, {input_name: src}) return annotate.output_tensors( - node_name, {output_name: traced * scale}, export_with=None) + node_name, {output_name: traced * scale}, export_with=None + ) def _run_add_node(self, node_name, input_name, output_name, src, value): traced = annotate.input_tensors(node_name, {input_name: src}) return annotate.output_tensors( - node_name, {output_name: traced + value}, export_with=None) + node_name, {output_name: traced + value}, export_with=None + ) def test_numpy_nodes_chain_via_output_ports(self): """Two numpy nodes connected by a published ndarray output -> input edge.""" @@ -600,8 +613,12 @@ def test_multiple_runs_of_same_graph(self): leapp.start(name=self.TEST_GRAPH_NAME) for _ in range(10): arr = self._run_scale_node( - "numpy_a", "arr", "out", - np.array([1.0, 2.0, 3.0], dtype=np.float32), 2.0) + "numpy_a", + "arr", + "out", + np.array([1.0, 2.0, 3.0], dtype=np.float32), + 2.0, + ) arr = self._run_add_node("numpy_b", "arr", "out", arr, 1.0) arr = self._run_scale_node("numpy_c", "arr", "out", arr, 3.0) leapp.stop() @@ -615,16 +632,20 @@ def test_feedback_connections(self): external = np.array([1.0, 2.0, 3.0], dtype=np.float32) for _ in range(2): in_a, loop = annotate.input_tensors( - "node_a", {"inputA": external, "loop_back": loop_back}) + "node_a", {"inputA": external, "loop_back": loop_back} + ) out_a = annotate.output_tensors( - "node_a", {"out_a": in_a + loop}, export_with=None) + "node_a", {"out_a": in_a + loop}, export_with=None + ) out_b = self._run_scale_node("node_b", "in_b", "out_b", out_a, 1.0) in_c, loop = annotate.input_tensors( - "node_c", {"inputC": out_b, "loop_back": loop_back}) + "node_c", {"inputC": out_b, "loop_back": loop_back} + ) out_c = annotate.output_tensors( - "node_c", {"out_c": in_c + loop}, export_with=None) + "node_c", {"out_c": in_c + loop}, export_with=None + ) self._run_scale_node("node_d", "in_d", "out_d", out_c, 1.0) loop_back = out_c @@ -632,15 +653,16 @@ def test_feedback_connections(self): leapp.stop() leapp.compile_graph(visualize=False) self.verify_connectivity( - nodes=4, internal_connections=3, feedback_connections=1) + nodes=4, internal_connections=3, feedback_connections=1 + ) def test_mirror_leapp_tags_with_inplace_assignment(self): """mirror_leapp_tags with a single in-place buffer between numpy nodes.""" leapp.start(name=self.TEST_GRAPH_NAME) out_a = self._run_scale_node( - "node_a", "in_a", "out_a", - np.array([1.0, 2.0, 3.0], dtype=np.float32), 2.0) + "node_a", "in_a", "out_a", np.array([1.0, 2.0, 3.0], dtype=np.float32), 2.0 + ) buffer = np.zeros_like(np.asarray(out_a)) buffer[:] = out_a @@ -658,8 +680,8 @@ def test_mirror_leapp_tags_preserves_output_port_through_chain(self): leapp.start(name=self.TEST_GRAPH_NAME) out1 = self._run_scale_node( - "source", "in_a", "out_a", - np.array([1.0, 2.0, 3.0], dtype=np.float32), 2.0) + "source", "in_a", "out_a", np.array([1.0, 2.0, 3.0], dtype=np.float32), 2.0 + ) buffer1 = np.empty_like(np.asarray(out1)) buffer1[:] = out1 @@ -677,6 +699,7 @@ def test_mirror_leapp_tags_preserves_output_port_through_chain(self): leapp.compile_graph(visualize=False) self.verify_connectivity(nodes=4, internal_connections=3) + class TestWarpConnectionCase(WarpTestCase, ConnectivityTestBase): """Warp equivalents of the torch connection/traced-state tests in TestConnectionCase.""" @@ -802,7 +825,10 @@ def test_multiple_runs_of_same_graph(self): leapp.start(name=self.TEST_GRAPH_NAME) for _ in range(10): arr = self._run_copy_node( - "warp_a", "arr", "out", wp.array([1.0, 2.0, 3.0], dtype=wp.float32, device=self.DEVICE) + "warp_a", + "arr", + "out", + wp.array([1.0, 2.0, 3.0], dtype=wp.float32, device=self.DEVICE), ) arr = self._run_copy_node("warp_b", "arr", "out", arr) arr = self._run_copy_node("warp_c", "arr", "out", arr) @@ -821,7 +847,9 @@ def test_feedback_connections(self): ) with annotate.warp_op("node_a"): out_a = self._warp_add_arrays(in_a, loop) - out_a = annotate.output_tensors("node_a", {"out_a": out_a}, export_with="onnx") + out_a = annotate.output_tensors( + "node_a", {"out_a": out_a}, export_with="onnx" + ) out_b = self._run_copy_node("node_b", "in_b", "out_b", out_a) @@ -830,7 +858,9 @@ def test_feedback_connections(self): ) with annotate.warp_op("node_c"): out_c = self._warp_add_arrays(in_c, loop) - out_c = annotate.output_tensors("node_c", {"out_c": out_c}, export_with="onnx") + out_c = annotate.output_tensors( + "node_c", {"out_c": out_c}, export_with="onnx" + ) self._run_copy_node("node_d", "in_d", "out_d", out_c) loop_back = out_c @@ -838,7 +868,8 @@ def test_feedback_connections(self): leapp.stop() leapp.compile_graph(visualize=False) self.verify_connectivity( - nodes=4, internal_connections=3, feedback_connections=1) + nodes=4, internal_connections=3, feedback_connections=1 + ) def test_interleaved_traced_nodes_keep_forward_execution_order(self): """Completed warp node output consumed by a later node is forward flow.""" @@ -924,15 +955,16 @@ def test_warp_output_port_persistence_through_operations(self): leapp.stop() leapp.compile_graph(visualize=False) - self.verify_connectivity( - nodes=5, outputs=2, internal_connections=4) + self.verify_connectivity(nodes=5, outputs=2, internal_connections=4) def test_mirror_leapp_tags_with_inplace_assignment(self): """mirror_leapp_tags with a single in-place wp.copy between warp nodes.""" leapp.start(name=self.TEST_GRAPH_NAME) out_a = self._run_copy_node( - "node_a", "in_a", "out_a", + "node_a", + "in_a", + "out_a", wp.array([1.0, 2.0, 3.0], dtype=wp.float32, device=self.DEVICE), ) buffer = wp.empty_like(out_a) @@ -947,5 +979,5 @@ def test_mirror_leapp_tags_with_inplace_assignment(self): self.verify_connectivity(nodes=3, internal_connections=2) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/functional_tests/test_export_backend.py b/tests/functional_tests/test_export_backend.py index 2ee1d385..e68b5283 100644 --- a/tests/functional_tests/test_export_backend.py +++ b/tests/functional_tests/test_export_backend.py @@ -28,7 +28,6 @@ import pytest - class TestOnnxProviderSelection(unittest.TestCase): def test_cuda_provider_is_always_preferred(self): program = SimplifiedONNXProgram("fake.onnx") @@ -49,10 +48,13 @@ def test_cuda_provider_is_always_preferred(self): def test_warns_when_cuda_provider_is_unavailable(self): program = SimplifiedONNXProgram("fake.onnx") - with mock.patch( - "leapp.backends.export_backend.ort.get_available_providers", - return_value=["CPUExecutionProvider"], - ), mock.patch("leapp.backends.export_backend._get_logger") as get_logger: + with ( + mock.patch( + "leapp.backends.export_backend.ort.get_available_providers", + return_value=["CPUExecutionProvider"], + ), + mock.patch("leapp.backends.export_backend._get_logger") as get_logger, + ): providers = program._get_providers() self.assertEqual(providers, ["CPUExecutionProvider"]) @@ -86,7 +88,10 @@ def test_cuda_iobinding_helper_binds_real_session_calls(self): self.assertEqual(input_kwargs["name"], "input") self.assertEqual(input_kwargs["device_type"], "cuda") self.assertEqual(input_kwargs["device_id"], 0) - self.assertEqual(input_kwargs["element_type"], torch.tensor([], dtype=torch.float32).numpy().dtype.type) + self.assertEqual( + input_kwargs["element_type"], + torch.tensor([], dtype=torch.float32).numpy().dtype.type, + ) self.assertEqual(input_kwargs["shape"], (2, 3)) self.assertIsInstance(input_kwargs["buffer_ptr"], int) self.assertGreater(input_kwargs["buffer_ptr"], 0) @@ -95,7 +100,10 @@ def test_cuda_iobinding_helper_binds_real_session_calls(self): self.assertEqual(output_kwargs["name"], "output") self.assertEqual(output_kwargs["device_type"], "cuda") self.assertEqual(output_kwargs["device_id"], 0) - self.assertEqual(output_kwargs["element_type"], torch.tensor([], dtype=torch.float32).numpy().dtype.type) + self.assertEqual( + output_kwargs["element_type"], + torch.tensor([], dtype=torch.float32).numpy().dtype.type, + ) self.assertEqual(output_kwargs["shape"], (2, 3)) self.assertIsInstance(output_kwargs["buffer_ptr"], int) self.assertGreater(output_kwargs["buffer_ptr"], 0) @@ -108,19 +116,24 @@ def test_cuda_iobinding_helper_binds_real_session_calls(self): def test_onnx_program_uses_iobinding_when_cuda_provider_is_active(self): program = SimplifiedONNXProgram("fake.onnx") fake_session = mock.Mock() - fake_session.get_providers.return_value = ["CUDAExecutionProvider", "CPUExecutionProvider"] + fake_session.get_providers.return_value = [ + "CUDAExecutionProvider", + "CPUExecutionProvider", + ] program._session = fake_session program._input_names = ["input"] program._output_metas = [] program._active_provider = "CUDAExecutionProvider" - with mock.patch.object( - program, "_can_use_cuda_iobinding", return_value=True - ) as can_use, mock.patch.object( - program, "_run_with_cuda_iobinding", return_value=("fast-path",) - ) as run_cuda, mock.patch.object( - program, "_run_with_standard_inference" - ) as run_standard: + with ( + mock.patch.object( + program, "_can_use_cuda_iobinding", return_value=True + ) as can_use, + mock.patch.object( + program, "_run_with_cuda_iobinding", return_value=("fast-path",) + ) as run_cuda, + mock.patch.object(program, "_run_with_standard_inference") as run_standard, + ): result = program(torch.tensor([1.0], dtype=torch.float32)) can_use.assert_called_once() @@ -137,13 +150,15 @@ def test_onnx_program_falls_back_without_cuda_iobinding(self): program._output_metas = [] program._active_provider = "CPUExecutionProvider" - with mock.patch.object( - program, "_can_use_cuda_iobinding", return_value=False - ) as can_use, mock.patch.object( - program, "_run_with_cuda_iobinding" - ) as run_cuda, mock.patch.object( - program, "_run_with_standard_inference", return_value=("cpu-path",) - ) as run_standard: + with ( + mock.patch.object( + program, "_can_use_cuda_iobinding", return_value=False + ) as can_use, + mock.patch.object(program, "_run_with_cuda_iobinding") as run_cuda, + mock.patch.object( + program, "_run_with_standard_inference", return_value=("cpu-path",) + ) as run_standard, + ): result = program(torch.tensor([1.0], dtype=torch.float32)) can_use.assert_called_once() @@ -164,7 +179,7 @@ class TestOnnxBackend(LEAPPFunctionalTestBase): def test_onnx_backend(self): @annotate.method(export_with="onnx") def funcA(inputA: torch.Tensor): - return inputA*2.0 + return inputA * 2.0 input_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) @@ -197,12 +212,14 @@ def funcA(inputA: torch.Tensor): leapp.stop() leapp.compile_graph(visualize=False) - @pytest.mark.filterwarnings("ignore:You are using the legacy TorchScript-based ONNX export") + @pytest.mark.filterwarnings( + "ignore:You are using the legacy TorchScript-based ONNX export" + ) @pytest.mark.filterwarnings("ignore:The feature will be removed") def test_onnx_backend_script(self): @annotate.method(export_with="onnx-torchscript") def funcA(inputA: torch.Tensor): - return inputA*2.0 + return inputA * 2.0 input_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) @@ -211,7 +228,9 @@ def funcA(inputA: torch.Tensor): leapp.stop() leapp.compile_graph(visualize=False) - @pytest.mark.filterwarnings("ignore:You are using the legacy TorchScript-based ONNX export") + @pytest.mark.filterwarnings( + "ignore:You are using the legacy TorchScript-based ONNX export" + ) @pytest.mark.filterwarnings("ignore:The feature will be removed") def test_onnx_backend_name_overlap_is_renamed(self): """Overlapping input/output names should be disambiguated for ONNX export.""" @@ -243,7 +262,9 @@ def test_onnx_backend_name_overlap_is_renamed(self): leapp.stop() leapp.compile_graph(visualize=False) - yaml_path = os.path.join(self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}.yaml") + yaml_path = os.path.join( + self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}.yaml" + ) with open(yaml_path) as f: exported = yaml.safe_load(f) @@ -265,7 +286,8 @@ def test_onnx_backend_name_overlap_is_renamed(self): model = onnx.load(model_path) initializer_names = {init.name for init in model.graph.initializer} onnx_input_names = [ - value.name for value in model.graph.input + value.name + for value in model.graph.input if value.name not in initializer_names ] onnx_output_names = [value.name for value in model.graph.output] @@ -278,10 +300,7 @@ def test_onnx_connected_node_names_match_artifact_names(self): source_input = annotate.input_tensors( "producer", - { - "source_input": torch.tensor( - [1.0, 2.0, 3.0], dtype=torch.float32) - }, + {"source_input": torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)}, ) producer_output = source_input * 2.0 annotate.output_tensors( @@ -303,26 +322,19 @@ def test_onnx_connected_node_names_match_artifact_names(self): leapp.stop() leapp.compile_graph(visualize=False, validate=False) - yaml_path = os.path.join( - self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}.yaml") + yaml_path = os.path.join(self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}.yaml") with open(yaml_path) as f: exported = yaml.safe_load(f) producer_desc = exported["models"]["producer"] - producer_yaml_outputs = [ - desc["name"] for desc in producer_desc["outputs"] - ] + producer_yaml_outputs = [desc["name"] for desc in producer_desc["outputs"]] self.assertEqual( producer_yaml_outputs, ["producer_original_output"], ) self.assertEqual( exported["pipeline"]["data_flow"], - { - "producer/producer_original_output": [ - "consumer/consumer_input_name" - ] - }, + {"producer/producer_original_output": ["consumer/consumer_input_name"]}, ) model_path = os.path.join(self.TEST_GRAPH_NAME, "producer.onnx") @@ -330,7 +342,9 @@ def test_onnx_connected_node_names_match_artifact_names(self): onnx_output_names = [value.name for value in model.graph.output] self.assertEqual(onnx_output_names, producer_yaml_outputs) - @pytest.mark.filterwarnings("ignore:You are using the legacy TorchScript-based ONNX export") + @pytest.mark.filterwarnings( + "ignore:You are using the legacy TorchScript-based ONNX export" + ) @pytest.mark.filterwarnings("ignore:The feature will be removed") def test_onnx_backend_prescript_with_traced_model(self): """Test ONNX export with a torch.jit.trace model, prescript=True, and conditional flow.""" @@ -353,7 +367,7 @@ def forward(self, x): @annotate.method( export_with="onnx-torchscript", backend_params={"prescript": True}, - environment_constants=["traced_model"] + environment_constants=["traced_model"], ) def funcA(inputA: torch.Tensor, threshold: torch.Tensor): # Apply the traced model @@ -370,7 +384,9 @@ def funcA(inputA: torch.Tensor, threshold: torch.Tensor): leapp.stop() leapp.compile_graph(visualize=False) - @pytest.mark.filterwarnings("ignore:You are using the legacy TorchScript-based ONNX export") + @pytest.mark.filterwarnings( + "ignore:You are using the legacy TorchScript-based ONNX export" + ) @pytest.mark.filterwarnings("ignore:The feature will be removed") def test_onnx_with_multiple_traced_modules(self): """Test ONNX export with multiple torch.jit.trace models as environment constants.""" @@ -395,17 +411,17 @@ def forward(self, x): encoder = EncoderModel() decoder = DecoderModel() - + encoder_input = torch.randn(4, dtype=torch.float32) decoder_input = torch.randn(8, dtype=torch.float32) - + traced_encoder = torch.jit.trace(encoder, encoder_input) traced_decoder = torch.jit.trace(decoder, decoder_input) @annotate.method( export_with="onnx-torchscript", backend_params={"prescript": True}, - environment_constants=["traced_encoder", "traced_decoder"] + environment_constants=["traced_encoder", "traced_decoder"], ) def encode_decode(inputA: torch.Tensor): encoded = traced_encoder(inputA) @@ -418,9 +434,11 @@ def encode_decode(inputA: torch.Tensor): encode_decode(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('encode_decode') + self.verify_all_models_exist("encode_decode") - @pytest.mark.filterwarnings("ignore:You are using the legacy TorchScript-based ONNX export") + @pytest.mark.filterwarnings( + "ignore:You are using the legacy TorchScript-based ONNX export" + ) @pytest.mark.filterwarnings("ignore:The feature will be removed") def test_onnx_with_traced_conv_model(self): """Test ONNX export with a traced convolutional model.""" @@ -443,7 +461,7 @@ def forward(self, x): conv_model = ConvModel().eval() example_input = torch.randn(1, 1, 8, 8, dtype=torch.float32) traced_conv = torch.jit.trace(conv_model, example_input) - + # Check if LLVM backend is available try: traced_conv(example_input) @@ -455,7 +473,7 @@ def forward(self, x): @annotate.method( export_with="onnx-torchscript", backend_params={"prescript": True}, - environment_constants=["traced_conv"] + environment_constants=["traced_conv"], ) def process_image(image: torch.Tensor): features = traced_conv(image) @@ -468,9 +486,11 @@ def process_image(image: torch.Tensor): process_image(input_image) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('process_image') + self.verify_all_models_exist("process_image") - @pytest.mark.filterwarnings("ignore:You are using the legacy TorchScript-based ONNX export") + @pytest.mark.filterwarnings( + "ignore:You are using the legacy TorchScript-based ONNX export" + ) @pytest.mark.filterwarnings("ignore:The feature will be removed") def test_onnx_with_complex_tensor_operations(self): """Test ONNX export with complex tensor operations (einsum, scatter, gather).""" @@ -481,7 +501,7 @@ def __init__(self, dim=8): self.query = torch.nn.Linear(dim, dim) self.key = torch.nn.Linear(dim, dim) self.value = torch.nn.Linear(dim, dim) - self.scale = dim ** -0.5 + self.scale = dim**-0.5 def forward(self, x): q = self.query(x) @@ -499,7 +519,7 @@ def forward(self, x): @annotate.method( export_with="onnx-torchscript", backend_params={"prescript": True}, - environment_constants=["traced_attn"] + environment_constants=["traced_attn"], ) def apply_attention(x: torch.Tensor): attended = traced_attn(x) @@ -514,13 +534,15 @@ def apply_attention(x: torch.Tensor): apply_attention(input_tensor) leapp.stop() leapp.compile_graph(visualize=False, atol=1e-3) - self.verify_all_models_exist('apply_attention') + self.verify_all_models_exist("apply_attention") - @pytest.mark.filterwarnings("ignore:You are using the legacy TorchScript-based ONNX export") + @pytest.mark.filterwarnings( + "ignore:You are using the legacy TorchScript-based ONNX export" + ) @pytest.mark.filterwarnings("ignore:The feature will be removed") def test_onnx_with_traced_tensor_node_chain(self): """Test ONNX export with chained traced tensor nodes. - + Note: Uses detach() on intermediate tensor to avoid deepcopy issues with non-leaf tensors during ONNX export. """ @@ -550,7 +572,7 @@ def forward(self, x): @annotate.method( export_with="onnx-torchscript", backend_params={"prescript": True}, - environment_constants=["traced_a"] + environment_constants=["traced_a"], ) def step_a(inputA: torch.Tensor): return traced_a(inputA) @@ -558,7 +580,7 @@ def step_a(inputA: torch.Tensor): @annotate.method( export_with="onnx-torchscript", backend_params={"prescript": True}, - environment_constants=["traced_b"] + environment_constants=["traced_b"], ) def step_b(inputB: torch.Tensor): return traced_b(inputB) @@ -571,9 +593,11 @@ def step_b(inputB: torch.Tensor): step_b(intermediate.detach()) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('step_a', 'step_b') + self.verify_all_models_exist("step_a", "step_b") - @pytest.mark.filterwarnings("ignore:You are using the legacy TorchScript-based ONNX export") + @pytest.mark.filterwarnings( + "ignore:You are using the legacy TorchScript-based ONNX export" + ) @pytest.mark.filterwarnings("ignore:The feature will be removed") def test_onnx_with_traced_model_and_tensor_broadcasting(self): """Test ONNX export with traced model combined with tensor broadcasting operations.""" @@ -593,15 +617,13 @@ def forward(self, x): @annotate.method( export_with="onnx-torchscript", backend_params={"prescript": True}, - environment_constants=["traced_scale"] + environment_constants=["traced_scale"], ) def scale_and_broadcast(inputA: torch.Tensor, mask: torch.Tensor): scaled = traced_scale(inputA) # Broadcasting: mask is (2,) and scaled is (2, 4) masked_output = torch.where( - mask.unsqueeze(-1) > 0, - scaled, - torch.zeros_like(scaled) + mask.unsqueeze(-1) > 0, scaled, torch.zeros_like(scaled) ) return masked_output @@ -612,9 +634,11 @@ def scale_and_broadcast(inputA: torch.Tensor, mask: torch.Tensor): scale_and_broadcast(input_tensor, mask) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('scale_and_broadcast') + self.verify_all_models_exist("scale_and_broadcast") - @pytest.mark.filterwarnings("ignore:You are using the legacy TorchScript-based ONNX export") + @pytest.mark.filterwarnings( + "ignore:You are using the legacy TorchScript-based ONNX export" + ) @pytest.mark.filterwarnings("ignore:The feature will be removed") def test_onnx_with_traced_model_multiple_outputs(self): """Test ONNX export with traced model that has multiple outputs.""" @@ -639,7 +663,7 @@ def forward(self, x): @annotate.method( export_with="onnx-torchscript", backend_params={"prescript": True}, - environment_constants=["traced_multi"] + environment_constants=["traced_multi"], ) def sample_gaussian(inputA: torch.Tensor, noise: torch.Tensor): mean, var = traced_multi(inputA) @@ -654,7 +678,7 @@ def sample_gaussian(inputA: torch.Tensor, noise: torch.Tensor): sample_gaussian(input_tensor, noise) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('sample_gaussian') + self.verify_all_models_exist("sample_gaussian") def test_onnx_dynamo_with_complex_math(self): """Test ONNX dynamo export with complex mathematical operations.""" @@ -675,7 +699,7 @@ def complex_math(x: torch.Tensor, y: torch.Tensor): complex_math(x, y) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('complex_math') + self.verify_all_models_exist("complex_math") def test_onnx_dynamo_with_reduction_operations(self): """Test ONNX dynamo export with reduction operations.""" @@ -685,11 +709,11 @@ def reduction_ops(x: torch.Tensor): mean = x.mean(dim=-1, keepdim=True) std = x.std(dim=-1, keepdim=True) normalized = (x - mean) / (std + 1e-5) - + # More reductions - keep shapes compatible max_val = normalized.max(dim=1).values # shape (4,) - sum_val = normalized.sum(dim=1) # shape (4,) - + sum_val = normalized.sum(dim=1) # shape (4,) + return max_val + sum_val x = torch.randn(4, 8, dtype=torch.float32) @@ -698,11 +722,11 @@ def reduction_ops(x: torch.Tensor): reduction_ops(x) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('reduction_ops') + self.verify_all_models_exist("reduction_ops") def test_onnx_dict_inputs_to_list_outputs(self): """Test ONNX export with dict-like inputs and list-like outputs. - + This tests that the ONNX exporter correctly handles: - Flattening dict inputs into individual tensor inputs - Packing multiple outputs into a list @@ -710,27 +734,27 @@ def test_onnx_dict_inputs_to_list_outputs(self): @annotate.method(export_with="onnx") def dict_to_list(inputs: dict): - a = inputs['a'] - b = inputs['b'] + a = inputs["a"] + b = inputs["b"] # Return as a list of outputs sum_result = a + b diff_result = a - b return [sum_result, diff_result] input_dict = { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), } leapp.start(name=self.TEST_GRAPH_NAME) dict_to_list(input_dict) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('dict_to_list') + self.verify_all_models_exist("dict_to_list") def test_onnx_list_inputs_to_dict_outputs(self): """Test ONNX export with list-like inputs and dict-like outputs. - + This tests that the ONNX exporter correctly handles: - Flattening list inputs into individual tensor inputs - Unpacking dict outputs into individual tensors @@ -742,8 +766,8 @@ def list_to_dict(inputs: list): b = inputs[1] # Return as a dict of outputs return { - 'sum': a + b, - 'product': a * b, + "sum": a + b, + "product": a * b, } input_list = [ @@ -755,7 +779,7 @@ def list_to_dict(inputs: list): list_to_dict(input_list) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('list_to_dict') + self.verify_all_models_exist("list_to_dict") def test_onnx_nested_dict_input(self): """Test ONNX export with deeply nested dict input structure.""" @@ -763,26 +787,26 @@ def test_onnx_nested_dict_input(self): @annotate.method(export_with="onnx") def process_nested_dict(nested: dict): # Access deeply nested values - a = nested['level1']['a'] - b = nested['level1']['b'] - c = nested['level2']['c'] + a = nested["level1"]["a"] + b = nested["level1"]["b"] + c = nested["level2"]["c"] return a + b + c nested_input = { - 'level1': { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "level1": { + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), + }, + "level2": { + "c": torch.randn(3, 4, dtype=torch.float32), }, - 'level2': { - 'c': torch.randn(3, 4, dtype=torch.float32), - } } leapp.start(name=self.TEST_GRAPH_NAME) process_nested_dict(nested_input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('process_nested_dict') + self.verify_all_models_exist("process_nested_dict") def test_onnx_nested_list_input(self): """Test ONNX export with nested list input structure.""" @@ -796,7 +820,10 @@ def process_nested_list(nested: list): return a * b + c nested_input = [ - [torch.randn(3, 4, dtype=torch.float32), torch.randn(3, 4, dtype=torch.float32)], + [ + torch.randn(3, 4, dtype=torch.float32), + torch.randn(3, 4, dtype=torch.float32), + ], [torch.randn(3, 4, dtype=torch.float32)], ] @@ -804,22 +831,22 @@ def process_nested_list(nested: list): process_nested_list(nested_input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('process_nested_list') + self.verify_all_models_exist("process_nested_list") def test_onnx_mixed_dict_list_input(self): """Test ONNX export with mixed dict and list inputs.""" @annotate.method(export_with="onnx") def process_mixed_inputs(dict_input: dict, list_input: list): - a = dict_input['a'] - b = dict_input['b'] + a = dict_input["a"] + b = dict_input["b"] c = list_input[0] d = list_input[1] return a + b, c * d dict_input = { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), } list_input = [ torch.randn(3, 4, dtype=torch.float32), @@ -830,7 +857,7 @@ def process_mixed_inputs(dict_input: dict, list_input: list): process_mixed_inputs(dict_input, list_input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('process_mixed_inputs') + self.verify_all_models_exist("process_mixed_inputs") def test_onnx_nested_list_output(self): """Test ONNX export with nested list output structure.""" @@ -846,7 +873,7 @@ def create_nested_list_output(x: torch.Tensor): create_nested_list_output(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('create_nested_list_output') + self.verify_all_models_exist("create_nested_list_output") def test_onnx_dict_of_lists_output(self): """Test ONNX export with dict containing list values as output.""" @@ -854,8 +881,8 @@ def test_onnx_dict_of_lists_output(self): @annotate.method(export_with="onnx") def create_dict_of_lists(x: torch.Tensor): return { - 'group_a': [x[0:1], x[1:2]], - 'group_b': [x[2:3], x[3:4]], + "group_a": [x[0:1], x[1:2]], + "group_b": [x[2:3], x[3:4]], } input_tensor = torch.randn(4, 4, dtype=torch.float32) @@ -864,7 +891,7 @@ def create_dict_of_lists(x: torch.Tensor): create_dict_of_lists(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('create_dict_of_lists') + self.verify_all_models_exist("create_dict_of_lists") def test_onnx_list_of_dicts_output(self): """Test ONNX export with list containing dict values as output.""" @@ -872,8 +899,8 @@ def test_onnx_list_of_dicts_output(self): @annotate.method(export_with="onnx") def create_list_of_dicts(x: torch.Tensor): return [ - {'a': x[0:1], 'b': x[1:2]}, - {'a': x[2:3], 'b': x[3:4]}, + {"a": x[0:1], "b": x[1:2]}, + {"a": x[2:3], "b": x[3:4]}, ] input_tensor = torch.randn(4, 4, dtype=torch.float32) @@ -882,7 +909,7 @@ def create_list_of_dicts(x: torch.Tensor): create_list_of_dicts(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('create_list_of_dicts') + self.verify_all_models_exist("create_list_of_dicts") def test_onnx_mixed_return_types(self): """Test ONNX export with mixed return types: tensor, list, and dict.""" @@ -891,7 +918,7 @@ def test_onnx_mixed_return_types(self): def mixed_outputs(x: torch.Tensor): single = x[0:1] list_out = [x[1:2], x[2:3]] - dict_out = {'a': x[3:4], 'b': x[4:5]} + dict_out = {"a": x[3:4], "b": x[4:5]} return single, list_out, dict_out input_tensor = torch.randn(5, 4, dtype=torch.float32) @@ -900,7 +927,7 @@ def mixed_outputs(x: torch.Tensor): mixed_outputs(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('mixed_outputs') + self.verify_all_models_exist("mixed_outputs") def test_onnx_large_nested_structure(self): """Test ONNX export with large nested dict/list structure (stress test).""" @@ -910,7 +937,7 @@ def process_large_structure(data: dict): # Process a large nested structure results = [] for i in range(4): - group = data[f'group_{i}'] + group = data[f"group_{i}"] for j in range(2): results.append(group[j] * 2.0) # Return sum of all processed tensors @@ -918,16 +945,15 @@ def process_large_structure(data: dict): # Create large nested structure: 4 groups, each with 2 tensors large_input = { - f'group_{i}': [ - torch.randn(3, 4, dtype=torch.float32) for _ in range(2) - ] for i in range(4) + f"group_{i}": [torch.randn(3, 4, dtype=torch.float32) for _ in range(2)] + for i in range(4) } leapp.start(name=self.TEST_GRAPH_NAME) process_large_structure(large_input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('process_large_structure') + self.verify_all_models_exist("process_large_structure") def test_onnx_deeply_nested_dict_input(self): """Test ONNX export with deeply nested dict input (3+ levels).""" @@ -935,16 +961,16 @@ def test_onnx_deeply_nested_dict_input(self): @annotate.method(export_with="onnx") def process_deep_dict(nested: dict): # Access 3-level deep nested values - a = nested['l1']['l2']['l3']['a'] - b = nested['l1']['l2']['l3']['b'] + a = nested["l1"]["l2"]["l3"]["a"] + b = nested["l1"]["l2"]["l3"]["b"] return a + b deep_input = { - 'l1': { - 'l2': { - 'l3': { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "l1": { + "l2": { + "l3": { + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), } } } @@ -954,16 +980,24 @@ def process_deep_dict(nested: dict): process_deep_dict(deep_input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('process_deep_dict') + self.verify_all_models_exist("process_deep_dict") def test_onnx_many_tensor_inputs(self): """Test ONNX export with many individual tensor inputs (stress test).""" @annotate.method(export_with="onnx") - def sum_many_tensors(t0: torch.Tensor, t1: torch.Tensor, t2: torch.Tensor, - t3: torch.Tensor, t4: torch.Tensor, t5: torch.Tensor, - t6: torch.Tensor, t7: torch.Tensor, t8: torch.Tensor, - t9: torch.Tensor): + def sum_many_tensors( + t0: torch.Tensor, + t1: torch.Tensor, + t2: torch.Tensor, + t3: torch.Tensor, + t4: torch.Tensor, + t5: torch.Tensor, + t6: torch.Tensor, + t7: torch.Tensor, + t8: torch.Tensor, + t9: torch.Tensor, + ): return t0 + t1 + t2 + t3 + t4 + t5 + t6 + t7 + t8 + t9 tensors = [torch.randn(3, 4, dtype=torch.float32) for _ in range(10)] @@ -972,25 +1006,27 @@ def sum_many_tensors(t0: torch.Tensor, t1: torch.Tensor, t2: torch.Tensor, sum_many_tensors(*tensors) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('sum_many_tensors') + self.verify_all_models_exist("sum_many_tensors") def test_onnx_many_dict_tensor_inputs(self): """Test ONNX export with dict containing many tensors.""" @annotate.method(export_with="onnx") def sum_dict_tensors(data: dict): - total = data['t0'] + total = data["t0"] for i in range(1, 10): - total = total + data[f't{i}'] + total = total + data[f"t{i}"] return total - input_dict = {f't{i}': torch.randn(3, 4, dtype=torch.float32) for i in range(10)} + input_dict = { + f"t{i}": torch.randn(3, 4, dtype=torch.float32) for i in range(10) + } leapp.start(name=self.TEST_GRAPH_NAME) sum_dict_tensors(input_dict) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('sum_dict_tensors') + self.verify_all_models_exist("sum_dict_tensors") def test_onnx_split_tensor_to_many_outputs(self): """Test ONNX export that splits one tensor into many outputs.""" @@ -998,7 +1034,7 @@ def test_onnx_split_tensor_to_many_outputs(self): @annotate.method(export_with="onnx") def split_to_many(x: torch.Tensor): # Split tensor into 8 parts - return [x[i:i+1] for i in range(8)] + return [x[i : i + 1] for i in range(8)] input_tensor = torch.randn(8, 4, dtype=torch.float32) @@ -1006,26 +1042,26 @@ def split_to_many(x: torch.Tensor): split_to_many(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('split_to_many') + self.verify_all_models_exist("split_to_many") def test_onnx_bidirectional_dict_list_io(self): """Test ONNX export with dict and list inputs producing list and dict outputs. - + Mirrors test_export_dict_and_list_bidirectional_io from test_export_situation.py """ @annotate.method(export_with="onnx") def bidirectional_io(dict_input: dict, list_input: list): # Dict input -> list output - list_out = [v+1 for v in dict_input.values()] + list_out = [v + 1 for v in dict_input.values()] # List input -> dict output - dict_out = {f'item_{i}': v+1 for i, v in enumerate(list_input)} + dict_out = {f"item_{i}": v + 1 for i, v in enumerate(list_input)} return list_out, dict_out dict_input = { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), - 'c': torch.randn(3, 4, dtype=torch.float32), + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), + "c": torch.randn(3, 4, dtype=torch.float32), } list_input = [ torch.randn(3, 4, dtype=torch.float32), @@ -1036,28 +1072,28 @@ def bidirectional_io(dict_input: dict, list_input: list): bidirectional_io(dict_input, list_input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('bidirectional_io') + self.verify_all_models_exist("bidirectional_io") def test_onnx_complex_nested_dict_with_extra_input(self): """Test ONNX export with nested dict input plus additional tensor input. - + Similar to test_export_nnModule_with_large_nested_dict_io from test_export_situation.py """ @annotate.method(export_with="onnx") def nested_with_extra(nested_input: dict, extra_tensor: torch.Tensor): # Access nested values - nested = nested_input[0]['nested'] - a = nested['a'] - b = nested['b'] + nested = nested_input[0]["nested"] + a = nested["a"] + b = nested["b"] # Combine with extra tensor return [a + extra_tensor, b * extra_tensor] nested = [ { - 'nested': { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "nested": { + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), } } ] @@ -1067,31 +1103,35 @@ def nested_with_extra(nested_input: dict, extra_tensor: torch.Tensor): nested_with_extra(nested, extra) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('nested_with_extra') + self.verify_all_models_exist("nested_with_extra") - @pytest.mark.filterwarnings("ignore:You are using the legacy TorchScript-based ONNX export") + @pytest.mark.filterwarnings( + "ignore:You are using the legacy TorchScript-based ONNX export" + ) @pytest.mark.filterwarnings("ignore:The feature will be removed") def test_onnx_torchscript_dict_list_io(self): """Test ONNX TorchScript export with dict/list I/O (non-dynamo).""" @annotate.method(export_with="onnx-torchscript") def torchscript_io(inputs: dict): - a = inputs['a'] - b = inputs['b'] + a = inputs["a"] + b = inputs["b"] return [a + b, a - b] input_dict = { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), } leapp.start(name=self.TEST_GRAPH_NAME) torchscript_io(input_dict) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('torchscript_io') + self.verify_all_models_exist("torchscript_io") - @pytest.mark.filterwarnings("ignore:You are using the legacy TorchScript-based ONNX export") + @pytest.mark.filterwarnings( + "ignore:You are using the legacy TorchScript-based ONNX export" + ) @pytest.mark.filterwarnings("ignore:The feature will be removed") def test_onnx_torchscript_nested_inputs(self): """Test ONNX TorchScript export with nested list/dict inputs (non-dynamo).""" @@ -1105,7 +1145,10 @@ def torchscript_nested(data: list): return a + b + c nested_input = [ - [torch.randn(3, 4, dtype=torch.float32), torch.randn(3, 4, dtype=torch.float32)], + [ + torch.randn(3, 4, dtype=torch.float32), + torch.randn(3, 4, dtype=torch.float32), + ], [torch.randn(3, 4, dtype=torch.float32)], ] @@ -1113,7 +1156,7 @@ def torchscript_nested(data: list): torchscript_nested(nested_input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('torchscript_nested') + self.verify_all_models_exist("torchscript_nested") def test_onnx_with_nn_module_and_dict_input(self): """Test ONNX export with nn.Module method that takes dict input.""" @@ -1125,21 +1168,21 @@ def __init__(self): @annotate.method(export_with="onnx") def process(self, inputs: dict): - a = inputs['a'] - b = inputs['b'] + a = inputs["a"] + b = inputs["b"] return self.linear(a + b) module = ProcessingModule() input_dict = { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), } leapp.start(name=self.TEST_GRAPH_NAME) module.process(input_dict) leapp.stop() leapp.compile_graph(visualize=False, atol=1e-3) - self.verify_all_models_exist('process') + self.verify_all_models_exist("process") def test_onnx_with_nn_module_and_list_output(self): """Test ONNX export with nn.Module method that returns list output.""" @@ -1161,8 +1204,9 @@ def split_process(self, x: torch.Tensor): module.split_process(input_tensor) leapp.stop() leapp.compile_graph(visualize=False, atol=1e-3) - self.verify_all_models_exist('split_process') - + self.verify_all_models_exist("split_process") + + class TestTorchBackend(LEAPPFunctionalTestBase): """ Unit tests to see if export situation is properly handled @@ -1175,32 +1219,34 @@ class TestTorchBackend(LEAPPFunctionalTestBase): def test_torch_trace_backend(self): @annotate.method(export_with="jit-trace") def funcA(inputA: torch.Tensor): - return inputA*2.0 + return inputA * 2.0 input_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) - expected_output = input_tensor*2.0 + expected_output = input_tensor * 2.0 leapp.start(name=self.TEST_GRAPH_NAME) funcA(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) self.verify_single_torchscript_model_expected_value( - [input_tensor], [expected_output], funcA.__name__) + [input_tensor], [expected_output], funcA.__name__ + ) def test_torch_script_backend(self): @annotate.method(export_with="jit") def funcA(inputA: torch.Tensor): - return inputA*2 + return inputA * 2 input_tensor = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32) - expected_output = input_tensor*2.0 + expected_output = input_tensor * 2.0 leapp.start(name=self.TEST_GRAPH_NAME) funcA(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) self.verify_single_torchscript_model_expected_value( - [input_tensor], [expected_output], funcA.__name__) + [input_tensor], [expected_output], funcA.__name__ + ) def test_exported_program_backend(self): @annotate.method(export_with="exported-program") @@ -1215,7 +1261,8 @@ def funcA(inputA: torch.Tensor): leapp.stop() leapp.compile_graph(visualize=False, validate=True) self.verify_single_exported_program_model_expected_value( - [input_tensor], [expected_output], funcA.__name__) + [input_tensor], [expected_output], funcA.__name__ + ) def test_pt2_alias_backend(self): @annotate.method(export_with="pt2") @@ -1230,12 +1277,16 @@ def funcA(inputA: torch.Tensor): leapp.stop() leapp.compile_graph(visualize=False, validate=True) self.verify_single_exported_program_model_expected_value( - [input_tensor], [expected_output], funcA.__name__) + [input_tensor], [expected_output], funcA.__name__ + ) - with open(os.path.join(self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}.yaml")) as f: + with open( + os.path.join(self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}.yaml") + ) as f: yaml_data = yaml.safe_load(f) self.assertEqual( - yaml_data["models"][funcA.__name__]["parameters"]["backend"], "pt2") + yaml_data["models"][funcA.__name__]["parameters"]["backend"], "pt2" + ) # ----------------------------------------------------------------- # Complex math operations @@ -1259,7 +1310,7 @@ def complex_math(x: torch.Tensor, y: torch.Tensor): complex_math(x, y) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('complex_math') + self.verify_all_models_exist("complex_math") def test_torch_script_complex_math(self): """Test jit-script export with complex mathematical operations.""" @@ -1279,7 +1330,7 @@ def complex_math(x: torch.Tensor, y: torch.Tensor): complex_math(x, y) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('complex_math') + self.verify_all_models_exist("complex_math") # ----------------------------------------------------------------- # Reduction operations @@ -1303,7 +1354,7 @@ def reduction_ops(x: torch.Tensor): reduction_ops(x) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('reduction_ops') + self.verify_all_models_exist("reduction_ops") def test_torch_trace_tensor_method_norm(self): """Test jit-trace export with tensor.norm() method calls.""" @@ -1320,7 +1371,8 @@ def tensor_method_norm(x: torch.Tensor): leapp.stop() leapp.compile_graph(visualize=False) self.verify_single_torchscript_model_expected_value( - [x], [expected], 'tensor_method_norm') + [x], [expected], "tensor_method_norm" + ) def test_torch_script_reduction_operations(self): """Test jit-script export with reduction operations.""" @@ -1340,7 +1392,7 @@ def reduction_ops(x: torch.Tensor): reduction_ops(x) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('reduction_ops') + self.verify_all_models_exist("reduction_ops") # ----------------------------------------------------------------- # Dict / list I/O @@ -1351,22 +1403,22 @@ def test_torch_trace_dict_inputs_to_list_outputs(self): @annotate.method(export_with="jit-trace") def dict_to_list(inputs: dict): - a = inputs['a'] - b = inputs['b'] + a = inputs["a"] + b = inputs["b"] sum_result = a + b diff_result = a - b return [sum_result, diff_result] input_dict = { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), } leapp.start(name=self.TEST_GRAPH_NAME) dict_to_list(input_dict) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('dict_to_list') + self.verify_all_models_exist("dict_to_list") def test_torch_trace_list_inputs_to_dict_outputs(self): """Test jit-trace export with list inputs and dict outputs.""" @@ -1376,8 +1428,8 @@ def list_to_dict(inputs: list): a = inputs[0] b = inputs[1] return { - 'sum': a + b, - 'product': a * b, + "sum": a + b, + "product": a * b, } input_list = [ @@ -1389,33 +1441,33 @@ def list_to_dict(inputs: list): list_to_dict(input_list) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('list_to_dict') + self.verify_all_models_exist("list_to_dict") def test_torch_trace_nested_dict_input(self): """Test jit-trace export with nested dict input.""" @annotate.method(export_with="jit-trace") def process_nested_dict(nested: dict): - a = nested['level1']['a'] - b = nested['level1']['b'] - c = nested['level2']['c'] + a = nested["level1"]["a"] + b = nested["level1"]["b"] + c = nested["level2"]["c"] return a + b + c nested_input = { - 'level1': { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "level1": { + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), + }, + "level2": { + "c": torch.randn(3, 4, dtype=torch.float32), }, - 'level2': { - 'c': torch.randn(3, 4, dtype=torch.float32), - } } leapp.start(name=self.TEST_GRAPH_NAME) process_nested_dict(nested_input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('process_nested_dict') + self.verify_all_models_exist("process_nested_dict") def test_torch_trace_nested_list_input(self): """Test jit-trace export with nested list input.""" @@ -1428,7 +1480,10 @@ def process_nested_list(nested: list): return a * b + c nested_input = [ - [torch.randn(3, 4, dtype=torch.float32), torch.randn(3, 4, dtype=torch.float32)], + [ + torch.randn(3, 4, dtype=torch.float32), + torch.randn(3, 4, dtype=torch.float32), + ], [torch.randn(3, 4, dtype=torch.float32)], ] @@ -1436,22 +1491,22 @@ def process_nested_list(nested: list): process_nested_list(nested_input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('process_nested_list') + self.verify_all_models_exist("process_nested_list") def test_torch_trace_mixed_dict_list_input(self): """Test jit-trace export with mixed dict and list inputs.""" @annotate.method(export_with="jit-trace") def process_mixed_inputs(dict_input: dict, list_input: list): - a = dict_input['a'] - b = dict_input['b'] + a = dict_input["a"] + b = dict_input["b"] c = list_input[0] d = list_input[1] return a + b, c * d dict_input = { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), } list_input = [ torch.randn(3, 4, dtype=torch.float32), @@ -1462,7 +1517,7 @@ def process_mixed_inputs(dict_input: dict, list_input: list): process_mixed_inputs(dict_input, list_input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('process_mixed_inputs') + self.verify_all_models_exist("process_mixed_inputs") def test_torch_trace_nested_list_output(self): """Test jit-trace export with nested list output structure.""" @@ -1477,7 +1532,7 @@ def create_nested_list_output(x: torch.Tensor): create_nested_list_output(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('create_nested_list_output') + self.verify_all_models_exist("create_nested_list_output") def test_torch_trace_dict_of_lists_output(self): """Test jit-trace export with dict containing list values as output.""" @@ -1485,8 +1540,8 @@ def test_torch_trace_dict_of_lists_output(self): @annotate.method(export_with="jit-trace") def create_dict_of_lists(x: torch.Tensor): return { - 'group_a': [x[0:1], x[1:2]], - 'group_b': [x[2:3], x[3:4]], + "group_a": [x[0:1], x[1:2]], + "group_b": [x[2:3], x[3:4]], } input_tensor = torch.randn(4, 4, dtype=torch.float32) @@ -1495,7 +1550,7 @@ def create_dict_of_lists(x: torch.Tensor): create_dict_of_lists(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('create_dict_of_lists') + self.verify_all_models_exist("create_dict_of_lists") def test_torch_trace_list_of_dicts_output(self): """Test jit-trace export with list containing dict values as output.""" @@ -1503,8 +1558,8 @@ def test_torch_trace_list_of_dicts_output(self): @annotate.method(export_with="jit-trace") def create_list_of_dicts(x: torch.Tensor): return [ - {'a': x[0:1], 'b': x[1:2]}, - {'a': x[2:3], 'b': x[3:4]}, + {"a": x[0:1], "b": x[1:2]}, + {"a": x[2:3], "b": x[3:4]}, ] input_tensor = torch.randn(4, 4, dtype=torch.float32) @@ -1513,7 +1568,7 @@ def create_list_of_dicts(x: torch.Tensor): create_list_of_dicts(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('create_list_of_dicts') + self.verify_all_models_exist("create_list_of_dicts") def test_torch_trace_mixed_return_types(self): """Test jit-trace export with mixed return types: tensor, list, and dict.""" @@ -1522,7 +1577,7 @@ def test_torch_trace_mixed_return_types(self): def mixed_outputs(x: torch.Tensor): single = x[0:1] list_out = [x[1:2], x[2:3]] - dict_out = {'a': x[3:4], 'b': x[4:5]} + dict_out = {"a": x[3:4], "b": x[4:5]} return single, list_out, dict_out input_tensor = torch.randn(5, 4, dtype=torch.float32) @@ -1531,7 +1586,7 @@ def mixed_outputs(x: torch.Tensor): mixed_outputs(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('mixed_outputs') + self.verify_all_models_exist("mixed_outputs") def test_torch_trace_bidirectional_dict_list_io(self): """Test jit-trace export with dict and list inputs producing list and dict outputs.""" @@ -1539,13 +1594,13 @@ def test_torch_trace_bidirectional_dict_list_io(self): @annotate.method(export_with="jit-trace") def bidirectional_io(dict_input: dict, list_input: list): list_out = [v + 1 for v in dict_input.values()] - dict_out = {f'item_{i}': v + 1 for i, v in enumerate(list_input)} + dict_out = {f"item_{i}": v + 1 for i, v in enumerate(list_input)} return list_out, dict_out dict_input = { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), - 'c': torch.randn(3, 4, dtype=torch.float32), + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), + "c": torch.randn(3, 4, dtype=torch.float32), } list_input = [ torch.randn(3, 4, dtype=torch.float32), @@ -1556,7 +1611,7 @@ def bidirectional_io(dict_input: dict, list_input: list): bidirectional_io(dict_input, list_input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('bidirectional_io') + self.verify_all_models_exist("bidirectional_io") # ----------------------------------------------------------------- # Deeply nested and stress tests @@ -1567,16 +1622,16 @@ def test_torch_trace_deeply_nested_dict_input(self): @annotate.method(export_with="jit-trace") def process_deep_dict(nested: dict): - a = nested['l1']['l2']['l3']['a'] - b = nested['l1']['l2']['l3']['b'] + a = nested["l1"]["l2"]["l3"]["a"] + b = nested["l1"]["l2"]["l3"]["b"] return a + b deep_input = { - 'l1': { - 'l2': { - 'l3': { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "l1": { + "l2": { + "l3": { + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), } } } @@ -1586,16 +1641,24 @@ def process_deep_dict(nested: dict): process_deep_dict(deep_input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('process_deep_dict') + self.verify_all_models_exist("process_deep_dict") def test_torch_trace_many_tensor_inputs(self): """Test jit-trace export with many individual tensor inputs.""" @annotate.method(export_with="jit-trace") - def sum_many_tensors(t0: torch.Tensor, t1: torch.Tensor, t2: torch.Tensor, - t3: torch.Tensor, t4: torch.Tensor, t5: torch.Tensor, - t6: torch.Tensor, t7: torch.Tensor, t8: torch.Tensor, - t9: torch.Tensor): + def sum_many_tensors( + t0: torch.Tensor, + t1: torch.Tensor, + t2: torch.Tensor, + t3: torch.Tensor, + t4: torch.Tensor, + t5: torch.Tensor, + t6: torch.Tensor, + t7: torch.Tensor, + t8: torch.Tensor, + t9: torch.Tensor, + ): return t0 + t1 + t2 + t3 + t4 + t5 + t6 + t7 + t8 + t9 tensors = [torch.randn(3, 4, dtype=torch.float32) for _ in range(10)] @@ -1604,32 +1667,34 @@ def sum_many_tensors(t0: torch.Tensor, t1: torch.Tensor, t2: torch.Tensor, sum_many_tensors(*tensors) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('sum_many_tensors') + self.verify_all_models_exist("sum_many_tensors") def test_torch_trace_many_dict_tensor_inputs(self): """Test jit-trace export with dict containing many tensors.""" @annotate.method(export_with="jit-trace") def sum_dict_tensors(data: dict): - total = data['t0'] + total = data["t0"] for i in range(1, 10): - total = total + data[f't{i}'] + total = total + data[f"t{i}"] return total - input_dict = {f't{i}': torch.randn(3, 4, dtype=torch.float32) for i in range(10)} + input_dict = { + f"t{i}": torch.randn(3, 4, dtype=torch.float32) for i in range(10) + } leapp.start(name=self.TEST_GRAPH_NAME) sum_dict_tensors(input_dict) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('sum_dict_tensors') + self.verify_all_models_exist("sum_dict_tensors") def test_torch_trace_split_tensor_to_many_outputs(self): """Test jit-trace export that splits one tensor into many outputs.""" @annotate.method(export_with="jit-trace") def split_to_many(x: torch.Tensor): - return [x[i:i+1] for i in range(8)] + return [x[i : i + 1] for i in range(8)] input_tensor = torch.randn(8, 4, dtype=torch.float32) @@ -1637,7 +1702,7 @@ def split_to_many(x: torch.Tensor): split_to_many(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('split_to_many') + self.verify_all_models_exist("split_to_many") def test_torch_trace_large_nested_structure(self): """Test jit-trace export with large nested dict/list structure.""" @@ -1646,38 +1711,37 @@ def test_torch_trace_large_nested_structure(self): def process_large_structure(data: dict): results = [] for i in range(4): - group = data[f'group_{i}'] + group = data[f"group_{i}"] for j in range(2): results.append(group[j] * 2.0) return torch.stack(results).sum(dim=0) large_input = { - f'group_{i}': [ - torch.randn(3, 4, dtype=torch.float32) for _ in range(2) - ] for i in range(4) + f"group_{i}": [torch.randn(3, 4, dtype=torch.float32) for _ in range(2)] + for i in range(4) } leapp.start(name=self.TEST_GRAPH_NAME) process_large_structure(large_input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('process_large_structure') + self.verify_all_models_exist("process_large_structure") def test_torch_trace_complex_nested_dict_with_extra_input(self): """Test jit-trace export with nested dict input plus additional tensor input.""" @annotate.method(export_with="jit-trace") def nested_with_extra(nested_input: dict, extra_tensor: torch.Tensor): - nested = nested_input[0]['nested'] - a = nested['a'] - b = nested['b'] + nested = nested_input[0]["nested"] + a = nested["a"] + b = nested["b"] return [a + extra_tensor, b * extra_tensor] nested = [ { - 'nested': { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "nested": { + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), } } ] @@ -1687,7 +1751,7 @@ def nested_with_extra(nested_input: dict, extra_tensor: torch.Tensor): nested_with_extra(nested, extra) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('nested_with_extra') + self.verify_all_models_exist("nested_with_extra") # ----------------------------------------------------------------- # Chained nodes @@ -1711,7 +1775,7 @@ def step_b(inputB: torch.Tensor): step_b(intermediate.detach()) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('step_a', 'step_b') + self.verify_all_models_exist("step_a", "step_b") def test_torch_script_chained_nodes(self): """Test jit-script export with two chained scripted nodes.""" @@ -1731,7 +1795,7 @@ def step_b(inputB: torch.Tensor): step_b(intermediate.detach()) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('step_a', 'step_b') + self.verify_all_models_exist("step_a", "step_b") # ----------------------------------------------------------------- # nn.Module tests @@ -1747,21 +1811,21 @@ def __init__(self): @annotate.method(export_with="jit-trace") def process(self, inputs: dict): - a = inputs['a'] - b = inputs['b'] + a = inputs["a"] + b = inputs["b"] return self.linear(a + b) module = ProcessingModule() input_dict = { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), } leapp.start(name=self.TEST_GRAPH_NAME) module.process(input_dict) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('process') + self.verify_all_models_exist("process") def test_torch_trace_nn_module_with_list_output(self): """Test jit-trace export with nn.Module method that returns list output.""" @@ -1783,7 +1847,7 @@ def split_process(self, x: torch.Tensor): module.split_process(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('split_process') + self.verify_all_models_exist("split_process") def test_torch_script_nn_module_with_dict_input(self): """Test jit-script export with nn.Module method that takes dict input.""" @@ -1795,21 +1859,21 @@ def __init__(self): @annotate.method(export_with="jit") def process(self, inputs: dict): - a = inputs['a'] - b = inputs['b'] + a = inputs["a"] + b = inputs["b"] return self.linear(a + b) module = ProcessingModule() input_dict = { - 'a': torch.randn(3, 4, dtype=torch.float32), - 'b': torch.randn(3, 4, dtype=torch.float32), + "a": torch.randn(3, 4, dtype=torch.float32), + "b": torch.randn(3, 4, dtype=torch.float32), } leapp.start(name=self.TEST_GRAPH_NAME) module.process(input_dict) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('process') + self.verify_all_models_exist("process") def test_torch_script_nn_module_with_list_output(self): """Test jit-script export with nn.Module method that returns list output.""" @@ -1831,7 +1895,7 @@ def split_process(self, x: torch.Tensor): module.split_process(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('split_process') + self.verify_all_models_exist("split_process") # ----------------------------------------------------------------- # Multiple outputs / broadcasting @@ -1850,7 +1914,7 @@ def multi_output(x: torch.Tensor): multi_output(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('multi_output') + self.verify_all_models_exist("multi_output") def test_torch_script_multiple_outputs(self): """Test jit-script export with multiple tensor outputs (tuple return).""" @@ -1865,7 +1929,7 @@ def multi_output(x: torch.Tensor): multi_output(input_tensor) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('multi_output') + self.verify_all_models_exist("multi_output") def test_torch_trace_tensor_broadcasting(self): """Test jit-trace export with tensor broadcasting operations.""" @@ -1874,9 +1938,7 @@ def test_torch_trace_tensor_broadcasting(self): def broadcast_ops(inputA: torch.Tensor, mask: torch.Tensor): scaled = inputA * 2.0 masked_output = torch.where( - mask.unsqueeze(-1) > 0, - scaled, - torch.zeros_like(scaled) + mask.unsqueeze(-1) > 0, scaled, torch.zeros_like(scaled) ) return masked_output @@ -1887,7 +1949,7 @@ def broadcast_ops(inputA: torch.Tensor, mask: torch.Tensor): broadcast_ops(input_tensor, mask) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('broadcast_ops') + self.verify_all_models_exist("broadcast_ops") def test_torch_script_tensor_broadcasting(self): """Test jit-script export with tensor broadcasting operations.""" @@ -1896,9 +1958,7 @@ def test_torch_script_tensor_broadcasting(self): def broadcast_ops(inputA: torch.Tensor, mask: torch.Tensor): scaled = inputA * 2.0 masked_output = torch.where( - mask.unsqueeze(-1) > 0, - scaled, - torch.zeros_like(scaled) + mask.unsqueeze(-1) > 0, scaled, torch.zeros_like(scaled) ) return masked_output @@ -1909,7 +1969,7 @@ def broadcast_ops(inputA: torch.Tensor, mask: torch.Tensor): broadcast_ops(input_tensor, mask) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_all_models_exist('broadcast_ops') + self.verify_all_models_exist("broadcast_ops") # ----------------------------------------------------------------- # Value verification tests @@ -1931,7 +1991,8 @@ def compute(x: torch.Tensor, y: torch.Tensor): leapp.stop() leapp.compile_graph(visualize=False) self.verify_single_torchscript_model_expected_value( - [x, y], [expected], 'compute') + [x, y], [expected], "compute" + ) def test_torch_script_value_verification_complex(self): """Test jit-script export and verify output values for a non-trivial function.""" @@ -1949,7 +2010,8 @@ def compute(x: torch.Tensor, y: torch.Tensor): leapp.stop() leapp.compile_graph(visualize=False) self.verify_single_torchscript_model_expected_value( - [x, y], [expected], 'compute') + [x, y], [expected], "compute" + ) def test_torch_trace_value_verification_multi_output(self): """Test jit-trace export and verify values for a multi-output function.""" @@ -1969,13 +2031,14 @@ def split_and_scale(x: torch.Tensor): leapp.stop() leapp.compile_graph(visualize=False) self.verify_single_torchscript_model_expected_value( - [x], [expected_half1, expected_half2], 'split_and_scale') + [x], [expected_half1, expected_half2], "split_and_scale" + ) # ============================================================================ # PROPOSED TESTS - POTENTIAL EDGE CASES THAT MIGHT FAIL # ============================================================================ -# +# # The following tests are proposed but NOT implemented. They represent edge # cases that might reveal issues in the ONNX export pipeline. # @@ -2049,12 +2112,12 @@ def split_and_scale(x: torch.Tensor): # serialize correctly to ONNX format # # 15. test_onnx_with_torch_scatter_ops -# - Test ONNX export with scatter/gather operations like scatter_add, +# - Test ONNX export with scatter/gather operations like scatter_add, # index_select, etc. # - May fail because: Some scatter operations have limited ONNX support # or require specific opset versions # ============================================================================ -if __name__ == '__main__': +if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/functional_tests/test_export_situation.py b/tests/functional_tests/test_export_situation.py index c2e4ef7c..f08cf471 100644 --- a/tests/functional_tests/test_export_situation.py +++ b/tests/functional_tests/test_export_situation.py @@ -45,17 +45,24 @@ def module(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: out = torch.softmax(out, dim=-1) out = out * 2.0 return out - - x, y = annotate.input_tensors('scripted_computation', { - 'x': a, - 'y': b, - }) + + x, y = annotate.input_tensors( + "scripted_computation", + { + "x": a, + "y": b, + }, + ) result = module(x, y) - annotate.output_tensors('scripted_computation', { - 'result': result, - }, export_with="jit-script") + annotate.output_tensors( + "scripted_computation", + { + "result": result, + }, + export_with="jit-script", + ) leapp.stop() leapp.compile_graph() @@ -63,15 +70,20 @@ def module(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: # Use allclose instead of exact match because torch.jit.freeze # (optimize_numerics=True) can cause tiny floating-point differences import os - model = torch.jit.load(os.path.join(self.TEST_GRAPH_NAME, "scripted_computation.pt")) + + model = torch.jit.load( + os.path.join(self.TEST_GRAPH_NAME, "scripted_computation.pt") + ) output = model(a, b) - self.assertTrue(torch.allclose(output, result, atol=1e-6), - f"Output mismatch: max diff = {(output - result).abs().max().item():.2e}") + self.assertTrue( + torch.allclose(output, result, atol=1e-6), + f"Output mismatch: max diff = {(output - result).abs().max().item():.2e}", + ) def test_export_nnModule_function(self): linear = torch.nn.Linear(3, 3) - @annotate.method(export_with="jit", environment_constants=['linear']) + @annotate.method(export_with="jit", environment_constants=["linear"]) def funcA(inputA: torch.Tensor): output = linear(inputA) return output @@ -103,72 +115,97 @@ def funcA(inputA: dict): list_conversion = list(inputA.values()) return list_conversion - input = {'a': torch.tensor([1]), 'b': torch.tensor( - [2]), 'c': torch.tensor([3]), 'd': torch.tensor([4])} + input = { + "a": torch.tensor([1]), + "b": torch.tensor([2]), + "c": torch.tensor([3]), + "d": torch.tensor([4]), + } leapp.start(name=self.TEST_GRAPH_NAME) expected_output = funcA(input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_num_connections(annotate, nodes=1, inputs=4, outputs=4, - internal_connections=0) + self.verify_num_connections( + annotate, nodes=1, inputs=4, outputs=4, internal_connections=0 + ) self.verify_single_torchscript_model_expected_value( - [input], [expected_output], funcA.__name__) + [input], [expected_output], funcA.__name__ + ) model_info = self.inspect_torchscript_model(funcA.__name__) # self + 4 flat tensors (inputA_a, inputA_b, inputA_c, inputA_d) - self.assertEqual(len(model_info['inputs']), 5) - self.assertEqual(len(model_info['outputs']), 4) # list_conversion + self.assertEqual(len(model_info["inputs"]), 5) + self.assertEqual(len(model_info["outputs"]), 4) # list_conversion def test_export_nnModule_with_large_nested_dict_io(self): @annotate.method(export_with="jit") def funcA(inputA: dict, inputB): - underlying_values = inputA[0]['nested'] + underlying_values = inputA[0]["nested"] list_conversion = list(underlying_values.values()) return_values = list_conversion + [inputB] - return {'return_values': [return_values]} - - input = [[{'nested': {'a': torch.tensor([1]), 'b': torch.tensor( - [2]), 'c': torch.tensor([3]), 'd': torch.tensor([4])}}], torch.tensor([5])] + return {"return_values": [return_values]} + + input = [ + [ + { + "nested": { + "a": torch.tensor([1]), + "b": torch.tensor([2]), + "c": torch.tensor([3]), + "d": torch.tensor([4]), + } + } + ], + torch.tensor([5]), + ] leapp.start(name=self.TEST_GRAPH_NAME) expected_output = funcA(*input) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_num_connections(annotate, nodes=1, inputs=5, outputs=5, - internal_connections=0) + self.verify_num_connections( + annotate, nodes=1, inputs=5, outputs=5, internal_connections=0 + ) self.verify_single_torchscript_model_expected_value( - input, [expected_output], funcA.__name__) + input, [expected_output], funcA.__name__ + ) model_info = self.inspect_torchscript_model(funcA.__name__) # self + 5 flat tensors (inputA contains 4 nested tensors + inputB is 1 tensor) - self.assertEqual(len(model_info['inputs']), 6) - self.assertEqual(len(model_info['outputs']), 5) # return_values + self.assertEqual(len(model_info["inputs"]), 6) + self.assertEqual(len(model_info["outputs"]), 5) # return_values def test_export_class_method_that_relies_on_dynamic_variable(self): - class moduleA(): + class moduleA: def __init__(self): self.idx = 0 self.stride = 3 - @annotate.method(export_with="jit", environment_constants=['self.idx', 'self.stride']) + @annotate.method( + export_with="jit", environment_constants=["self.idx", "self.stride"] + ) def get_subset(self, inputA: torch.Tensor): - retval = inputA[self.idx:self.idx+self.stride] + retval = inputA[self.idx : self.idx + self.stride] self.idx += self.stride return retval + moduleA = moduleA() leapp.start(name=self.TEST_GRAPH_NAME) subset = moduleA.get_subset(torch.tensor([1, 2, 3])) leapp.stop() leapp.compile_graph(visualize=False) - self.verify_num_connections(annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0) + self.verify_num_connections( + annotate, nodes=1, inputs=1, outputs=1, internal_connections=0 + ) self.verify_single_torchscript_model_expected_value( - [torch.tensor([1, 2, 3])], [subset], moduleA.get_subset.__name__) + [torch.tensor([1, 2, 3])], [subset], moduleA.get_subset.__name__ + ) def test_export_dict_and_list_bidirectional_io(self): """Test function that takes dict and list inputs, returns list and dict outputs""" + @annotate.method(export_with="jit") def test_complex_io(input: dict, input_2: list): dictionary_output = {} @@ -177,13 +214,17 @@ def test_complex_io(input: dict, input_2: list): return [value for value in input.values()], dictionary_output input_dict = { - 'a': torch.tensor([1]), - 'b': torch.tensor([2]), - 'c': torch.tensor([3]), - 'd': torch.tensor([4]), + "a": torch.tensor([1]), + "b": torch.tensor([2]), + "c": torch.tensor([3]), + "d": torch.tensor([4]), } - input_list = [torch.tensor([5]), torch.tensor( - [6]), torch.tensor([7]), torch.tensor([8])] + input_list = [ + torch.tensor([5]), + torch.tensor([6]), + torch.tensor([7]), + torch.tensor([8]), + ] leapp.start(name=self.TEST_GRAPH_NAME) expected_output = test_complex_io(input_dict, input_list) @@ -191,34 +232,36 @@ def test_complex_io(input: dict, input_2: list): leapp.compile_graph(visualize=False) # Verify graph statistics: 1 node, 8 dangling inputs, 8 dangling outputs, 0 internal connections - self.verify_num_connections(annotate, nodes=1, inputs=8, outputs=8, - internal_connections=0) + self.verify_num_connections( + annotate, nodes=1, inputs=8, outputs=8, internal_connections=0 + ) # Verify the model produces the correct output self.verify_single_torchscript_model_expected_value( - [input_dict, input_list], [expected_output], test_complex_io.__name__) + [input_dict, input_list], [expected_output], test_complex_io.__name__ + ) # Verify model structure: forward should have 8 flat tensor inputs + self model_info = self.inspect_torchscript_model(test_complex_io.__name__) # self + 4 dict tensors + 4 list tensors - self.assertEqual(len(model_info['inputs']), 9) + self.assertEqual(len(model_info["inputs"]), 9) # list output + dict output - self.assertEqual(len(model_info['outputs']), 8) + self.assertEqual(len(model_info["outputs"]), 8) # Verify output structure output_list, output_dict = expected_output self.assertEqual(len(output_list), 4) self.assertEqual(len(output_dict), 4) self.assertTrue(all(isinstance(v, torch.Tensor) for v in output_list)) - self.assertTrue(all(isinstance(v, torch.Tensor) - for v in output_dict.values())) + self.assertTrue(all(isinstance(v, torch.Tensor) for v in output_dict.values())) def test_export_split_tensor_to_list(self): """Test function that takes a single tensor and splits it into a list of tensors""" + @annotate.method(export_with="jit", verbose=True) def split_tensor(input_tensor: torch.Tensor): # Split the tensor into individual elements - return [input_tensor[i:i+1] for i in range(len(input_tensor))] + return [input_tensor[i : i + 1] for i in range(len(input_tensor))] # Create input tensor with shape (20,) with values 0, 10, 20, ..., 190 input_tensor = torch.arange(0, 200, 10) @@ -229,25 +272,25 @@ def split_tensor(input_tensor: torch.Tensor): leapp.compile_graph(visualize=False) # Verify graph statistics: 1 node, 1 dangling input, 20 dangling outputs, 0 internal connections - self.verify_num_connections(annotate, nodes=1, inputs=1, outputs=20, - internal_connections=0) + self.verify_num_connections( + annotate, nodes=1, inputs=1, outputs=20, internal_connections=0 + ) # Verify the model produces the correct output self.verify_single_torchscript_model_expected_value( - [input_tensor], [expected_output], split_tensor.__name__) + [input_tensor], [expected_output], split_tensor.__name__ + ) # Verify model structure: forward should have 1 tensor input + self model_info = self.inspect_torchscript_model(split_tensor.__name__) - self.assertEqual(len(model_info['inputs']), 2) # self + 1 input tensor + self.assertEqual(len(model_info["inputs"]), 2) # self + 1 input tensor # 20 individual tensor outputs - self.assertEqual(len(model_info['outputs']), 20) + self.assertEqual(len(model_info["outputs"]), 20) # Verify output structure: should be a list of 20 tensors, each with shape (1,) self.assertEqual(len(expected_output), 20) - self.assertTrue(all(isinstance(t, torch.Tensor) - for t in expected_output)) - self.assertTrue( - all(t.shape == torch.Size([1]) for t in expected_output)) + self.assertTrue(all(isinstance(t, torch.Tensor) for t in expected_output)) + self.assertTrue(all(t.shape == torch.Size([1]) for t in expected_output)) # Verify output values for i, tensor in enumerate(expected_output): @@ -255,10 +298,14 @@ def split_tensor(input_tensor: torch.Tensor): def test_export_nested_list_of_lists_output(self): """Test function that returns a nested list of lists of tensors""" + @annotate.method(export_with="jit") def create_nested_lists(input_tensor: torch.Tensor): # Split into groups of 2 - return [[input_tensor[i:i+1], input_tensor[i+1:i+2]] for i in range(0, 4, 2)] + return [ + [input_tensor[i : i + 1], input_tensor[i + 1 : i + 2]] + for i in range(0, 4, 2) + ] input_tensor = torch.arange(4) @@ -268,28 +315,30 @@ def create_nested_lists(input_tensor: torch.Tensor): leapp.compile_graph(visualize=False) # Should have 1 input, 4 outputs (flattened from nested structure) - self.verify_num_connections(annotate, nodes=1, inputs=1, outputs=4, - internal_connections=0) + self.verify_num_connections( + annotate, nodes=1, inputs=1, outputs=4, internal_connections=0 + ) self.verify_single_torchscript_model_expected_value( - [input_tensor], [expected_output], create_nested_lists.__name__) + [input_tensor], [expected_output], create_nested_lists.__name__ + ) # Verify structure: 2 lists, each containing 2 tensors self.assertEqual(len(expected_output), 2) self.assertTrue(all(len(sublist) == 2 for sublist in expected_output)) - model_info = self.inspect_torchscript_model( - create_nested_lists.__name__) - self.assertEqual(len(model_info['inputs']), 2) # self + 1 input - self.assertEqual(len(model_info['outputs']), 4) # 4 flattened outputs + model_info = self.inspect_torchscript_model(create_nested_lists.__name__) + self.assertEqual(len(model_info["inputs"]), 2) # self + 1 input + self.assertEqual(len(model_info["outputs"]), 4) # 4 flattened outputs def test_export_dict_with_list_values_output(self): """Test function that returns a dict where values are lists of tensors""" + @annotate.method(export_with="jit") def create_dict_of_lists(input_list: list): return { - 'first_half': [input_list[0], input_list[1]], - 'second_half': [input_list[2], input_list[3]] + "first_half": [input_list[0], input_list[1]], + "second_half": [input_list[2], input_list[3]], } input_list = [torch.tensor([i]) for i in range(4)] @@ -300,28 +349,30 @@ def create_dict_of_lists(input_list: list): leapp.compile_graph(visualize=False) # 4 inputs, 4 outputs (dict values flattened) - self.verify_num_connections(annotate, nodes=1, inputs=4, outputs=4, - internal_connections=0) + self.verify_num_connections( + annotate, nodes=1, inputs=4, outputs=4, internal_connections=0 + ) self.verify_single_torchscript_model_expected_value( - [input_list], [expected_output], create_dict_of_lists.__name__) + [input_list], [expected_output], create_dict_of_lists.__name__ + ) # Verify structure - self.assertEqual(len(expected_output['first_half']), 2) - self.assertEqual(len(expected_output['second_half']), 2) + self.assertEqual(len(expected_output["first_half"]), 2) + self.assertEqual(len(expected_output["second_half"]), 2) - model_info = self.inspect_torchscript_model( - create_dict_of_lists.__name__) - self.assertEqual(len(model_info['inputs']), 5) # self + 4 inputs - self.assertEqual(len(model_info['outputs']), 4) # 4 flattened outputs + model_info = self.inspect_torchscript_model(create_dict_of_lists.__name__) + self.assertEqual(len(model_info["inputs"]), 5) # self + 4 inputs + self.assertEqual(len(model_info["outputs"]), 4) # 4 flattened outputs def test_export_mixed_return_types(self): """Test function that returns a single tensor, a list, and a dict all together""" + @annotate.method(export_with="jit") def mixed_outputs(input_tensor: torch.Tensor): single_tensor = input_tensor[0:1] # Still a tensor, shape (1,) list_out = [input_tensor[1:2], input_tensor[2:3]] - dict_out = {'a': input_tensor[3:4], 'b': input_tensor[4:5]} + dict_out = {"a": input_tensor[3:4], "b": input_tensor[4:5]} return single_tensor, list_out, dict_out input_tensor = torch.arange(5) @@ -332,11 +383,13 @@ def mixed_outputs(input_tensor: torch.Tensor): leapp.compile_graph(visualize=False) # 1 input, 5 outputs (1 tensor + 2 list tensors + 2 dict tensors) - self.verify_num_connections(annotate, nodes=1, inputs=1, outputs=5, - internal_connections=0) + self.verify_num_connections( + annotate, nodes=1, inputs=1, outputs=5, internal_connections=0 + ) self.verify_single_torchscript_model_expected_value( - [input_tensor], [expected_output], mixed_outputs.__name__) + [input_tensor], [expected_output], mixed_outputs.__name__ + ) # Verify structure: all outputs are tensors single_tensor, list_out, dict_out = expected_output @@ -344,20 +397,20 @@ def mixed_outputs(input_tensor: torch.Tensor): self.assertEqual(len(list_out), 2) self.assertTrue(all(isinstance(t, torch.Tensor) for t in list_out)) self.assertEqual(len(dict_out), 2) - self.assertTrue(all(isinstance(t, torch.Tensor) - for t in dict_out.values())) + self.assertTrue(all(isinstance(t, torch.Tensor) for t in dict_out.values())) model_info = self.inspect_torchscript_model(mixed_outputs.__name__) - self.assertEqual(len(model_info['inputs']), 2) # self + 1 input - self.assertEqual(len(model_info['outputs']), 5) # 5 flattened outputs + self.assertEqual(len(model_info["inputs"]), 2) # self + 1 input + self.assertEqual(len(model_info["outputs"]), 5) # 5 flattened outputs def test_export_list_of_dicts_output(self): """Test function that returns a list of dictionaries""" + @annotate.method(export_with="jit") def create_list_of_dicts(input_tensor: torch.Tensor): return [ - {'x': input_tensor[0:1], 'y': input_tensor[1:2]}, - {'x': input_tensor[2:3], 'y': input_tensor[3:4]} + {"x": input_tensor[0:1], "y": input_tensor[1:2]}, + {"x": input_tensor[2:3], "y": input_tensor[3:4]}, ] input_tensor = torch.arange(4) @@ -368,37 +421,39 @@ def create_list_of_dicts(input_tensor: torch.Tensor): leapp.compile_graph(visualize=False) # 1 input, 4 outputs (flattened from list of dicts) - self.verify_num_connections(annotate, nodes=1, inputs=1, outputs=4, - internal_connections=0) + self.verify_num_connections( + annotate, nodes=1, inputs=1, outputs=4, internal_connections=0 + ) self.verify_single_torchscript_model_expected_value( - [input_tensor], [expected_output], create_list_of_dicts.__name__) + [input_tensor], [expected_output], create_list_of_dicts.__name__ + ) # Verify structure: list of 2 dicts, each with 2 keys self.assertEqual(len(expected_output), 2) self.assertTrue(all(len(d) == 2 for d in expected_output)) - self.assertTrue(all('x' in d and 'y' in d for d in expected_output)) + self.assertTrue(all("x" in d and "y" in d for d in expected_output)) - model_info = self.inspect_torchscript_model( - create_list_of_dicts.__name__) - self.assertEqual(len(model_info['inputs']), 2) # self + 1 input - self.assertEqual(len(model_info['outputs']), 4) # 4 flattened outputs + model_info = self.inspect_torchscript_model(create_list_of_dicts.__name__) + self.assertEqual(len(model_info["inputs"]), 2) # self + 1 input + self.assertEqual(len(model_info["outputs"]), 4) # 4 flattened outputs def test_export_deeply_nested_dict_to_flat_list(self): """Test function with deeply nested dict input and flat list output""" + @annotate.method(export_with="jit") def flatten_nested_structure(nested: dict): # Extract from deeply nested structure: nested['level1']['level2']['data'] - level2 = nested['level1']['level2'] - return [level2['data'][key] for key in sorted(level2['data'].keys())] + level2 = nested["level1"]["level2"] + return [level2["data"][key] for key in sorted(level2["data"].keys())] nested_input = { - 'level1': { - 'level2': { - 'data': { - 'a': torch.tensor([1]), - 'b': torch.tensor([2]), - 'c': torch.tensor([3]) + "level1": { + "level2": { + "data": { + "a": torch.tensor([1]), + "b": torch.tensor([2]), + "c": torch.tensor([3]), } } } @@ -410,22 +465,23 @@ def flatten_nested_structure(nested: dict): leapp.compile_graph(visualize=False) # 3 inputs (flattened from nested dict), 3 outputs - self.verify_num_connections(annotate, nodes=1, inputs=3, outputs=3, - internal_connections=0) + self.verify_num_connections( + annotate, nodes=1, inputs=3, outputs=3, internal_connections=0 + ) self.verify_single_torchscript_model_expected_value( - [nested_input], [expected_output], flatten_nested_structure.__name__) + [nested_input], [expected_output], flatten_nested_structure.__name__ + ) # Verify output is a list of 3 tensors self.assertEqual(len(expected_output), 3) - self.assertTrue(all(isinstance(t, torch.Tensor) - for t in expected_output)) + self.assertTrue(all(isinstance(t, torch.Tensor) for t in expected_output)) - model_info = self.inspect_torchscript_model( - flatten_nested_structure.__name__) + model_info = self.inspect_torchscript_model(flatten_nested_structure.__name__) # self + 3 nested inputs - self.assertEqual(len(model_info['inputs']), 4) - self.assertEqual(len(model_info['outputs']), 3) # 3 outputs + self.assertEqual(len(model_info["inputs"]), 4) + self.assertEqual(len(model_info["outputs"]), 3) # 3 outputs + -if __name__ == '__main__': +if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/functional_tests/test_inference.py b/tests/functional_tests/test_inference.py index 0f64e32f..d92ca999 100644 --- a/tests/functional_tests/test_inference.py +++ b/tests/functional_tests/test_inference.py @@ -26,11 +26,14 @@ import yaml import shutil import hashlib + + class TestInferenceManagerRobustness(LEAPPFunctionalTestBase): """Tests for InferenceManager resilience to YAML variations.""" def _export_simple_model(self): """Export a minimal jit model and return the yaml path.""" + @annotate.method(export_with="jit") def simple_model(x: torch.Tensor): return x * 2.0 @@ -58,7 +61,7 @@ def simple_model(x: torch.Tensor): # Corrupt the model file by appending garbage bytes model_path = os.path.join(self.TEST_GRAPH_NAME, "simple_model.pt") with open(model_path, "ab") as f: - f.write(b"\x00\xFF\x00\xFF") + f.write(b"\x00\xff\x00\xff") yaml_path = os.path.join(self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}.yaml") with self.assertRaises(ValueError) as ctx: @@ -134,8 +137,8 @@ def test_inference_manager_missing_feedback_flow(self): yaml_path = self._export_simple_model() with open(yaml_path) as f: data = yaml.safe_load(f) - data['pipeline'].pop('feedback_flow', None) - with open(yaml_path, 'w') as f: + data["pipeline"].pop("feedback_flow", None) + with open(yaml_path, "w") as f: yaml.dump(data, f) manager = InferenceManager(yaml_path) @@ -147,14 +150,13 @@ def test_inference_manager_missing_data_flow(self): yaml_path = self._export_simple_model() with open(yaml_path) as f: data = yaml.safe_load(f) - data['pipeline'].pop('data_flow', None) - with open(yaml_path, 'w') as f: + data["pipeline"].pop("data_flow", None) + with open(yaml_path, "w") as f: yaml.dump(data, f) manager = InferenceManager(yaml_path) self.assertIsNotNone(manager) - # ----------------------------------------------------------------- # TG3: Feedback state loading # ----------------------------------------------------------------- @@ -169,7 +171,9 @@ def test_feedback_initial_values_loaded_into_value_dict(self): counter = annotate.state_tensors("policy", {"counter": initial_counter}) new_counter = counter + obs_traced annotate.update_state("policy", {"counter": new_counter}) - annotate.output_tensors("policy", {"action": obs_traced * 2.0}, export_with="jit") + annotate.output_tensors( + "policy", {"action": obs_traced * 2.0}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -179,7 +183,8 @@ def test_feedback_initial_values_loaded_into_value_dict(self): loaded = manager.value_dict["policy"]["counter"] self.assertTrue( torch.allclose(loaded.cpu(), initial_counter), - f"Expected initial_counter {initial_counter}, got {loaded}") + f"Expected initial_counter {initial_counter}, got {loaded}", + ) def test_reset_zeros_runtime_buffers_then_repopulates_feedback(self): """reset() zeroes runtime caches and then restores feedback initial values.""" @@ -191,28 +196,38 @@ def test_reset_zeros_runtime_buffers_then_repopulates_feedback(self): counter = annotate.state_tensors("policy", {"counter": initial_counter}) new_counter = counter + obs_traced annotate.update_state("policy", {"counter": new_counter}) - annotate.output_tensors("policy", {"action": obs_traced * 2.0}, export_with="jit") + annotate.output_tensors( + "policy", {"action": obs_traced * 2.0}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) yaml_path = os.path.join(self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}.yaml") manager = InferenceManager(yaml_path) - manager.run_policy({"policy/observation": obs.to(manager.nodes["policy"].device)}) + manager.run_policy( + {"policy/observation": obs.to(manager.nodes["policy"].device)} + ) manager.reset() - self.assertTrue(torch.allclose( - manager.value_dict["policy"]["observation"].cpu(), - torch.zeros_like(obs), - )) - self.assertTrue(torch.allclose( - manager.value_dict["policy"]["counter"].cpu(), - initial_counter, - )) - self.assertTrue(torch.allclose( - manager.value_dict["==out=="]["policy/action"].cpu(), - torch.zeros_like(obs), - )) + self.assertTrue( + torch.allclose( + manager.value_dict["policy"]["observation"].cpu(), + torch.zeros_like(obs), + ) + ) + self.assertTrue( + torch.allclose( + manager.value_dict["policy"]["counter"].cpu(), + initial_counter, + ) + ) + self.assertTrue( + torch.allclose( + manager.value_dict["==out=="]["policy/action"].cpu(), + torch.zeros_like(obs), + ) + ) def test_feedback_malformed_safetensors_key_raises(self): """A safetensors key with more than one slash raises an error on load.""" @@ -226,13 +241,16 @@ def test_feedback_malformed_safetensors_key_raises(self): obs_traced = annotate.input_tensors("policy", {"observation": obs}) counter = annotate.state_tensors("policy", {"counter": initial_counter}) annotate.update_state("policy", {"counter": counter + obs_traced}) - annotate.output_tensors("policy", {"action": obs_traced * 2.0}, export_with="jit") + annotate.output_tensors( + "policy", {"action": obs_traced * 2.0}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) # Overwrite the safetensors file with a malformed key (extra slash) safetensors_path = os.path.join( - self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}_initial_values.safetensors") + self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}_initial_values.safetensors" + ) save_file({"policy/counter/extra": initial_counter}, safetensors_path) yaml_path = os.path.join(self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}.yaml") @@ -258,7 +276,8 @@ def test_fanout_clone_independence(self): # Copy node_a's model as node_b (same architecture, different name in YAML) shutil.copy( os.path.join(self.TEST_GRAPH_NAME, "node_a.pt"), - os.path.join(self.TEST_GRAPH_NAME, "node_b.pt")) + os.path.join(self.TEST_GRAPH_NAME, "node_b.pt"), + ) yaml_path = os.path.join(self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}.yaml") with open(yaml_path) as f: @@ -268,17 +287,25 @@ def test_fanout_clone_independence(self): model_bytes = open(os.path.join(self.TEST_GRAPH_NAME, "node_b.pt"), "rb").read() sha256 = hashlib.sha256(model_bytes).hexdigest() md5 = hashlib.md5(model_bytes).hexdigest() - data['models']['node_b'] = { - 'inputs': [{'name': 'x', 'dtype': 'float32', 'shape': [3], 'type': 'tensor'}], - 'outputs': [{'name': 'out', 'dtype': 'float32', 'shape': [3], 'type': 'tensor'}], - 'parameters': {'model_path': 'node_b.pt', 'backend': 'jit', - 'md5sum': md5, 'sha256sum': sha256}, + data["models"]["node_b"] = { + "inputs": [ + {"name": "x", "dtype": "float32", "shape": [3], "type": "tensor"} + ], + "outputs": [ + {"name": "out", "dtype": "float32", "shape": [3], "type": "tensor"} + ], + "parameters": { + "model_path": "node_b.pt", + "backend": "jit", + "md5sum": md5, + "sha256sum": sha256, + }, } # node_a/out fans out: to node_b/x (data_flow) AND stays as pipeline output - data['pipeline']['data_flow'] = {'node_a/out': ['node_b/x']} - data['pipeline']['outputs'] = {'node_a': ['out'], 'node_b': ['out']} - data['pipeline']['inputs'] = {'node_a': ['x']} - with open(yaml_path, 'w') as f: + data["pipeline"]["data_flow"] = {"node_a/out": ["node_b/x"]} + data["pipeline"]["outputs"] = {"node_a": ["out"], "node_b": ["out"]} + data["pipeline"]["inputs"] = {"node_a": ["x"]} + with open(yaml_path, "w") as f: yaml.dump(data, f) manager = InferenceManager(yaml_path) @@ -326,7 +353,9 @@ def test_split_output_node_runs_without_keyerror(self): result = manager.run_policy({"node_a/x": trace_input}) self.assertTrue(torch.allclose(result["node_a/final_a"], trace_input + 10.0)) - self.assertTrue(torch.allclose(result["node_b/final_b"], trace_input * 2.0 - 1.0)) + self.assertTrue( + torch.allclose(result["node_b/final_b"], trace_input * 2.0 - 1.0) + ) def test_unroutable_output_raises_on_load(self): """Malformed YAML with an output removed from routing should fail at load time.""" @@ -362,5 +391,5 @@ def test_unroutable_output_raises_on_load(self): self.assertIn("node_a", str(ctx.exception)) -if __name__ == '__main__': - unittest.main(verbosity=2) \ No newline at end of file +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/functional_tests/test_leapp_tools.py b/tests/functional_tests/test_leapp_tools.py index f8169abd..49a22377 100644 --- a/tests/functional_tests/test_leapp_tools.py +++ b/tests/functional_tests/test_leapp_tools.py @@ -55,8 +55,10 @@ def _track_cwd_dir(self, top_level_dir): def test_simple_name_preserves_legacy_layout(self): leapp.start(self.TEST_GRAPH_NAME) self.assertEqual(annotate.get_graph_name(), self.TEST_GRAPH_NAME) - self.assertEqual(os.path.normpath(annotate.get_save_path()), - os.path.normpath(self.TEST_GRAPH_NAME)) + self.assertEqual( + os.path.normpath(annotate.get_save_path()), + os.path.normpath(self.TEST_GRAPH_NAME), + ) self._build_trivial_graph() self._assert_artifacts(self.TEST_GRAPH_NAME, self.TEST_GRAPH_NAME) @@ -65,8 +67,9 @@ def test_relative_path_uses_basename_as_graph_name(self): leapp.start("foo/bar") self.assertEqual(annotate.get_graph_name(), "bar") - self.assertEqual(os.path.normpath(annotate.get_save_path()), - os.path.join("foo", "bar")) + self.assertEqual( + os.path.normpath(annotate.get_save_path()), os.path.join("foo", "bar") + ) self._build_trivial_graph() self._assert_artifacts(os.path.join("foo", "bar"), "bar") @@ -75,15 +78,19 @@ def test_absolute_path_overrides_save_path(self): leapp.start(absolute_target) self.assertEqual(annotate.get_graph_name(), "exported_gr00t") - self.assertEqual(os.path.normpath(annotate.get_save_path()), - os.path.normpath(absolute_target)) + self.assertEqual( + os.path.normpath(annotate.get_save_path()), + os.path.normpath(absolute_target), + ) self._build_trivial_graph() self._assert_artifacts(absolute_target, "exported_gr00t") # Critically, the YAML must NOT land one dir above the models, which # was the original bug. bad_path = os.path.join(self._tmp_root, "exported_gr00t.yaml") - self.assertFalse(os.path.exists(bad_path), - f"YAML must not be emitted to sibling-dir path {bad_path}") + self.assertFalse( + os.path.exists(bad_path), + f"YAML must not be emitted to sibling-dir path {bad_path}", + ) def test_trailing_slash_path_is_normalized(self): absolute_target = os.path.join(self._tmp_root, "exported_gr00t") @@ -91,8 +98,10 @@ def test_trailing_slash_path_is_normalized(self): leapp.start(name_with_slash) self.assertEqual(annotate.get_graph_name(), "exported_gr00t") - self.assertEqual(os.path.normpath(annotate.get_save_path()), - os.path.normpath(absolute_target)) + self.assertEqual( + os.path.normpath(annotate.get_save_path()), + os.path.normpath(absolute_target), + ) self._build_trivial_graph() self._assert_artifacts(absolute_target, "exported_gr00t") @@ -101,8 +110,9 @@ def test_save_path_combined_with_path_like_name(self): leapp.start("nested/bar", save_path=save_path) expected_dir = os.path.join(save_path, "nested", "bar") self.assertEqual(annotate.get_graph_name(), "bar") - self.assertEqual(os.path.normpath(annotate.get_save_path()), - os.path.normpath(expected_dir)) + self.assertEqual( + os.path.normpath(annotate.get_save_path()), os.path.normpath(expected_dir) + ) self._build_trivial_graph() self._assert_artifacts(expected_dir, "bar") @@ -130,7 +140,8 @@ def identity(inputA: torch.Tensor): leapp.stop() expected_model_path = os.path.abspath( - os.path.join(self.TEST_GRAPH_NAME, "identity.pt")) + os.path.join(self.TEST_GRAPH_NAME, "identity.pt") + ) with self.assertLogs("leapp", level=25) as logs: leapp.compile_graph(visualize=False, validate=False) diff --git a/tests/functional_tests/test_legacy_method.py b/tests/functional_tests/test_legacy_method.py index c51148d9..87373cee 100644 --- a/tests/functional_tests/test_legacy_method.py +++ b/tests/functional_tests/test_legacy_method.py @@ -33,6 +33,7 @@ class TestLegacyMethod(LEAPPFunctionalTestBase): def test_decorator_basic(self): """Wrapping a function with @annotate._method() decorator.""" + @annotate._method(export_with="jit") def double(x: torch.Tensor): return x * 2.0 @@ -55,6 +56,7 @@ def double(x: torch.Tensor): def test_direct_call(self): """Using annotate._method()(func) directly instead of as a decorator.""" + def triple(x: torch.Tensor): return x * 3.0 @@ -78,6 +80,7 @@ def triple(x: torch.Tensor): def test_chain_functions_pipeline(self): """Chain two _method() functions and verify pipeline connections.""" + @annotate._method(export_with="jit") def step_a(x: torch.Tensor): return x + 1.0 @@ -95,14 +98,14 @@ def step_b(y: torch.Tensor): self.assertEqual(len(annotate.nodes), 2) self.verify_all_models_exist("step_a", "step_b") - self.verify_num_connections(annotate, nodes=2, - internal_connections=1) + self.verify_num_connections(annotate, nodes=2, internal_connections=1) # ── Test 3: Feedback into a function ───────────────────────────────── def test_feedback_into_function(self): """A function whose output is fed back as input on the next iteration should produce a feedback connection.""" + @annotate._method(export_with="jit") def accumulate(state: torch.Tensor): return state + 1.0 @@ -123,8 +126,8 @@ def accumulate(state: torch.Tensor): def test_global_constant(self): """A function that references a global constant can still export when the constant is declared via environment_constants.""" - @annotate._method(export_with="jit", - environment_constants=["SCALE_FACTOR"]) + + @annotate._method(export_with="jit", environment_constants=["SCALE_FACTOR"]) def scale(x: torch.Tensor): return x * SCALE_FACTOR @@ -147,6 +150,7 @@ def scale(x: torch.Tensor): def test_nested_annotated_functions_error(self): """Calling one _method()-annotated function from inside another should fail because the tracing lock is already acquired.""" + @annotate._method(export_with="jit") def inner(x: torch.Tensor): return x * 2.0 @@ -169,14 +173,16 @@ def outer(x: torch.Tensor): def test_traced_data_as_input_error(self): """Passing an active TracedTensor (from input_tensors) as an argument to a _method()-decorated function should raise an error.""" + @annotate._method(export_with="jit") def func(x: torch.Tensor): return x + 5.0 leapp.start(name=self.TEST_GRAPH_NAME) try: - traced = annotate.input_tensors("other_node", - {"val": torch.tensor([1.0, 2.0, 3.0])}) + traced = annotate.input_tensors( + "other_node", {"val": torch.tensor([1.0, 2.0, 3.0])} + ) traced = traced + 100.0 func(traced) annotate.output_tensors("other_node", {"out": traced}, export_with="jit") @@ -192,10 +198,10 @@ def test_traced_operation_while_annotated_error(self): """Creating traced tensors (via input_tensors) while inside a _method()-decorated function should fail because the tracing lock prevents mixing contexts.""" + @annotate._method(export_with="jit") def func(x: torch.Tensor): - annotate.input_tensors("rogue_node", - {"y": torch.randn(3)}) + annotate.input_tensors("rogue_node", {"y": torch.randn(3)}) return x * 2.0 leapp.start(name=self.TEST_GRAPH_NAME) @@ -207,19 +213,19 @@ def func(x: torch.Tensor): leapp.stop() self.assertIn("Mixing active contexts is not allowed", str(e)) - # ── Test 8: Dict input → list output ──────────────────────────────── def test_dict_input_list_output(self): """Dict input flattened to 3 tensors, returned as a list of 3.""" + @annotate._method(export_with="jit") def fuse(sensors: dict): return list(sensors.values()) input_dict = { - 'lidar': torch.tensor([1.0, 2.0]), - 'camera': torch.tensor([3.0, 4.0]), - 'imu': torch.tensor([5.0, 6.0]), + "lidar": torch.tensor([1.0, 2.0]), + "camera": torch.tensor([3.0, 4.0]), + "imu": torch.tensor([5.0, 6.0]), } leapp.start(name=self.TEST_GRAPH_NAME) @@ -232,12 +238,14 @@ def fuse(sensors: dict): self.assertEqual(len(node.outputs), 3) self.verify_all_models_exist("fuse") self.verify_single_torchscript_model_expected_value( - [input_dict], [expected], "fuse") + [input_dict], [expected], "fuse" + ) # ── Test 9: List input → single output ─────────────────────────────── def test_list_input_single_output(self): """List input flattened to 3 tensors, reduced to a single output.""" + @annotate._method(export_with="jit") def sum_all(items: list): return items[0] + items[1] + items[2] @@ -258,20 +266,22 @@ def sum_all(items: list): self.assertEqual(len(node.outputs), 1) self.verify_all_models_exist("sum_all") self.verify_single_torchscript_model_expected_value( - [input_list], [expected], "sum_all") + [input_list], [expected], "sum_all" + ) # ── Test 10: Dict + list inputs → dict output ──────────────────────── def test_dict_and_list_inputs_dict_output(self): """Mixed dict and list args producing a dict return value. 4 flat input tensors, 2 flat output tensors.""" + @annotate._method(export_with="jit") def combine(state: dict, commands: list): - pos = state['pos'] + commands[0] - vel = state['vel'] + commands[1] - return {'pos': pos, 'vel': vel} + pos = state["pos"] + commands[0] + vel = state["vel"] + commands[1] + return {"pos": pos, "vel": vel} - state = {'pos': torch.tensor([1.0, 2.0]), 'vel': torch.tensor([0.1, 0.2])} + state = {"pos": torch.tensor([1.0, 2.0]), "vel": torch.tensor([0.1, 0.2])} commands = [torch.tensor([0.5, 0.5]), torch.tensor([0.01, 0.02])] leapp.start(name=self.TEST_GRAPH_NAME) @@ -284,22 +294,24 @@ def combine(state: dict, commands: list): self.assertEqual(len(node.outputs), 2) self.verify_all_models_exist("combine") self.verify_single_torchscript_model_expected_value( - [state, commands], [expected], "combine") + [state, commands], [expected], "combine" + ) # ── Test 11: Nested dict input → single output ─────────────────────── def test_nested_dict_input(self): """Nested dict-of-dicts flattened to 3 leaf tensors, single output.""" + @annotate._method(export_with="jit") def process(data: dict): - a = data['group_a']['x'] - b = data['group_a']['y'] - c = data['group_b']['z'] + a = data["group_a"]["x"] + b = data["group_a"]["y"] + c = data["group_b"]["z"] return a + b + c nested = { - 'group_a': {'x': torch.tensor([1.0]), 'y': torch.tensor([2.0])}, - 'group_b': {'z': torch.tensor([3.0])}, + "group_a": {"x": torch.tensor([1.0]), "y": torch.tensor([2.0])}, + "group_b": {"z": torch.tensor([3.0])}, } leapp.start(name=self.TEST_GRAPH_NAME) @@ -312,22 +324,24 @@ def process(data: dict): self.assertEqual(len(node.outputs), 1) self.verify_all_models_exist("process") self.verify_single_torchscript_model_expected_value( - [nested], [expected], "process") + [nested], [expected], "process" + ) # ── Test 12: Dict input → dict output (bidirectional) ──────────────── def test_dict_input_dict_output(self): """Dict in, dict out — both sides flatten to the same leaf count.""" + @annotate._method(export_with="jit") def transform(sensors: dict): return { - 'scaled_lidar': sensors['lidar'] * 2.0, - 'scaled_camera': sensors['camera'] * 3.0, + "scaled_lidar": sensors["lidar"] * 2.0, + "scaled_camera": sensors["camera"] * 3.0, } input_dict = { - 'lidar': torch.tensor([1.0, 2.0]), - 'camera': torch.tensor([3.0, 4.0]), + "lidar": torch.tensor([1.0, 2.0]), + "camera": torch.tensor([3.0, 4.0]), } leapp.start(name=self.TEST_GRAPH_NAME) @@ -340,12 +354,14 @@ def transform(sensors: dict): self.assertEqual(len(node.outputs), 2) self.verify_all_models_exist("transform") self.verify_single_torchscript_model_expected_value( - [input_dict], [expected], "transform") + [input_dict], [expected], "transform" + ) # ── Test 13: List input → list output ──────────────────────────────── def test_list_input_list_output(self): """List in, list out — element-wise transform preserving count.""" + @annotate._method(export_with="jit") def per_element(items: list): return [items[0] + 10.0, items[1] + 20.0, items[2] + 30.0] @@ -366,20 +382,22 @@ def per_element(items: list): self.assertEqual(len(node.outputs), 3) self.verify_all_models_exist("per_element") self.verify_single_torchscript_model_expected_value( - [input_list], [expected], "per_element") + [input_list], [expected], "per_element" + ) # ── Test 14: Dict + list inputs → list + dict outputs ──────────────── def test_bidirectional_complex_io(self): """Dict and list inputs producing both list and dict outputs. Verifies all 8 leaf tensors (4 in, 4 out) are handled.""" + @annotate._method(export_with="jit") def cross(state: dict, cmds: list): - list_out = [state['a'] + cmds[0], state['b'] + cmds[1]] - dict_out = {'sum_a': state['a'] + cmds[0], 'sum_b': state['b'] + cmds[1]} + list_out = [state["a"] + cmds[0], state["b"] + cmds[1]] + dict_out = {"sum_a": state["a"] + cmds[0], "sum_b": state["b"] + cmds[1]} return list_out, dict_out - state = {'a': torch.tensor([1.0]), 'b': torch.tensor([2.0])} + state = {"a": torch.tensor([1.0]), "b": torch.tensor([2.0])} cmds = [torch.tensor([10.0]), torch.tensor([20.0])] leapp.start(name=self.TEST_GRAPH_NAME) @@ -392,7 +410,8 @@ def cross(state: dict, cmds: list): self.assertEqual(len(node.outputs), 4) self.verify_all_models_exist("cross") self.verify_single_torchscript_model_expected_value( - [state, cmds], [expected], "cross") + [state, cmds], [expected], "cross" + ) if __name__ == "__main__": diff --git a/tests/functional_tests/test_logging.py b/tests/functional_tests/test_logging.py index 7b126631..e5f5cf56 100644 --- a/tests/functional_tests/test_logging.py +++ b/tests/functional_tests/test_logging.py @@ -42,7 +42,9 @@ def test_mirror_leapp_tags_logs_mismatch_with_boundary_context(self): leapp.start(name=self.TEST_GRAPH_NAME) try: with self.assertLogs("leapp", level="ERROR") as logs: - with self.assertRaisesRegex(Exception, "source and target do not match"): + with self.assertRaisesRegex( + Exception, "source and target do not match" + ): annotate.mirror_leapp_tags( torch.tensor([1.0]), torch.tensor([2.0]), diff --git a/tests/functional_tests/test_passthrough.py b/tests/functional_tests/test_passthrough.py index 262c6518..be632b58 100644 --- a/tests/functional_tests/test_passthrough.py +++ b/tests/functional_tests/test_passthrough.py @@ -13,7 +13,6 @@ class TestPassthrough(LEAPPFunctionalTestBase): - def test_method_passthrough(self): seen_traced_inputs = [] @@ -27,9 +26,13 @@ def policy(x: torch.Tensor): y = policy(x) self.assertEqual(len(seen_traced_inputs), 1) - self.assertFalse(seen_traced_inputs[0], "Input should remain a raw tensor outside tracing") + self.assertFalse( + seen_traced_inputs[0], "Input should remain a raw tensor outside tracing" + ) self.assertNotIsInstance(y, TracedTensor) - self.assertEqual(len(annotate.nodes), 0, "No nodes should be created outside start/stop") + self.assertEqual( + len(annotate.nodes), 0, "No nodes should be created outside start/stop" + ) def test_input_output_tensor_passthrough(self): a = torch.tensor([1.0, 2.0, 3.0]) @@ -42,11 +45,15 @@ def test_input_output_tensor_passthrough(self): self.assertNotIsInstance(traced_b, TracedTensor) out = traced_a + traced_b - returned_out = annotate.output_tensors("node_a", {"out": out}, export_with="jit") + returned_out = annotate.output_tensors( + "node_a", {"out": out}, export_with="jit" + ) self.assertIs(returned_out, out) self.assertNotIsInstance(out, TracedTensor) - self.assertEqual(len(annotate.nodes), 0, "No nodes should be created outside start/stop") + self.assertEqual( + len(annotate.nodes), 0, "No nodes should be created outside start/stop" + ) def test_state_tensor_passthrough(self): obs = torch.tensor([1.0, 2.0, 3.0]) @@ -61,7 +68,9 @@ def test_state_tensor_passthrough(self): self.assertIs(returned_state, new_state) self.assertNotIsInstance(new_state, TracedTensor) - self.assertEqual(len(annotate.nodes), 0, "No nodes should be created outside start/stop") + self.assertEqual( + len(annotate.nodes), 0, "No nodes should be created outside start/stop" + ) def test_module_passthrough(self): class TinyStatefulModel(torch.nn.Module): @@ -84,14 +93,16 @@ def forward(self, x): for name, buf in current_buffers.items(): self.assertNotIsInstance(buf, TracedTensor) self.assertEqual( - id(buf), original_buffer_ids[name], - f"Buffer '{name}' should not be replaced outside tracing" + id(buf), + original_buffer_ids[name], + f"Buffer '{name}' should not be replaced outside tracing", ) - self.assertEqual(len(annotate.nodes), 0, "No nodes should be created outside start/stop") + self.assertEqual( + len(annotate.nodes), 0, "No nodes should be created outside start/stop" + ) class TestPersistentCarrier(LEAPPFunctionalTestBase): - def test_numpy_output_crosses_node_boundary(self): leapp.start(name=self.TEST_GRAPH_NAME) @@ -107,9 +118,7 @@ def test_numpy_output_crosses_node_boundary(self): self.assertEqual(source_output.output_port, "y") self.assertIs(source_output.context_obj, annotate.nodes["numpy_source"]) - consumer_input = annotate.input_tensors( - "numpy_consumer", {"y": source_output} - ) + consumer_input = annotate.input_tensors("numpy_consumer", {"y": source_output}) self.assertIsInstance(consumer_input, TracedNpArray) self.assertIsNot(consumer_input, source_output) self.assertTrue(consumer_input.is_tracing) @@ -117,14 +126,11 @@ def test_numpy_output_crosses_node_boundary(self): self.assertEqual(source_output.context, "numpy_source") self.assertTrue(np.shares_memory(consumer_input.data, source_output.data)) - annotate.output_tensors( - "numpy_consumer", {"z": consumer_input + 1.0} - ) + annotate.output_tensors("numpy_consumer", {"z": consumer_input + 1.0}) leapp.stop() class TestDryrun(LEAPPFunctionalTestBase): - def test_complex_nested_input_and_state_tensors_dryrun_carriers(self): """dry_run should promote nested inputs, while nested states are rejected.""" leapp.start(name=self.TEST_GRAPH_NAME, dry_run=True) @@ -145,7 +151,9 @@ def _assert_nested_passthrough(original, returned): original_tensors = self._flatten_to_tensors(original) returned_tensors = self._flatten_to_tensors(returned) self.assertEqual(len(original_tensors), len(returned_tensors)) - for original_tensor, returned_tensor in zip(original_tensors, returned_tensors): + for original_tensor, returned_tensor in zip( + original_tensors, returned_tensors + ): # Declaring promotes the caller's tensor in place, so the # returned carrier is that same object at every nesting depth. self.assertIs(original_tensor, returned_tensor) @@ -155,12 +163,16 @@ def _assert_nested_passthrough(original, returned): returned_inputs = [] for idx, payload in enumerate(nested_cases): - returned_input = annotate.input_tensors("complex_node", {f"in_{idx}": payload}) + returned_input = annotate.input_tensors( + "complex_node", {f"in_{idx}": payload} + ) _assert_nested_passthrough(payload, returned_input) returned_inputs.append(returned_input) with self.assertRaises(TypeError) as exc: annotate.state_tensors("complex_node", {f"state_{idx}": payload}) - self.assertIn("does not support nested state structures", str(exc.exception)) + self.assertIn( + "does not support nested state structures", str(exc.exception) + ) annotate.output_tensors( "complex_node", @@ -170,8 +182,12 @@ def _assert_nested_passthrough(original, returned): leapp.stop() leapp.compile_graph(visualize=False) - self.assertFalse(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "complex_node.pt"))) - self.assertFalse(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "complex_node.onnx"))) + self.assertFalse( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "complex_node.pt")) + ) + self.assertFalse( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "complex_node.onnx")) + ) def test_method_dryrun_carrier_and_yaml(self): """dry_run should use inactive traced carriers and still produce graph YAML.""" @@ -193,8 +209,11 @@ def policy(x: torch.Tensor): self.assertTrue(seen_traced_inputs[0]) self.assertIsInstance(y, TracedTensor) self.assertFalse(y.is_tracing) - self.assertEqual(y.output_port, "y", - "Dryrun should still publish output ports for connections") + self.assertEqual( + y.output_port, + "y", + "Dryrun should still publish output ports for connections", + ) self.assertTrue(hasattr(annotate, "detected_pipeline")) self.assertTrue(hasattr(annotate, "detected_nodes")) @@ -207,9 +226,17 @@ def policy(x: torch.Tensor): torch.allclose(y, x * 2.0 + 1.0), "Dryrun method output should preserve functional behavior", ) - self.assertTrue(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}.yaml"))) - self.assertFalse(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "policy.pt"))) - self.assertFalse(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "policy.onnx"))) + self.assertTrue( + os.path.exists( + os.path.join(self.TEST_GRAPH_NAME, f"{self.TEST_GRAPH_NAME}.yaml") + ) + ) + self.assertFalse( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "policy.pt")) + ) + self.assertFalse( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "policy.onnx")) + ) def test_input_output_tensor_dryrun_passthrough(self): """dry_run at start should still build connectivity but skip model export.""" @@ -224,32 +251,47 @@ def test_input_output_tensor_dryrun_passthrough(self): self.assertTrue(traced_b.is_tracing) out = traced_a + traced_b - returned_out = annotate.output_tensors("node_a", {"out": out}, export_with="jit") + returned_out = annotate.output_tensors( + "node_a", {"out": out}, export_with="jit" + ) self.assertIs(returned_out, out) downstream_in = annotate.input_tensors("node_b", {"out": out}) self.assertIsInstance(downstream_in, TracedTensor) self.assertTrue(downstream_in.is_tracing) - annotate.output_tensors("node_b", {"final": downstream_in * 3.0}, export_with="jit") + annotate.output_tensors( + "node_b", {"final": downstream_in * 3.0}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) - self.assertEqual(out.output_port, "out", - "Dryrun should publish node output ports") + self.assertEqual( + out.output_port, "out", "Dryrun should publish node output ports" + ) self.assertIn("data_flow", annotate.detected_pipeline) self.assertEqual(1, len(annotate.detected_pipeline["data_flow"])) self.assertIn("node_a/out", annotate.detected_pipeline["data_flow"]) - self.assertEqual(["node_b/out"], list(annotate.detected_pipeline["data_flow"]["node_a/out"])) - self.assertFalse(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "node_a.pt"))) - self.assertFalse(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "node_b.pt"))) + self.assertEqual( + ["node_b/out"], list(annotate.detected_pipeline["data_flow"]["node_a/out"]) + ) + self.assertFalse( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "node_a.pt")) + ) + self.assertFalse( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "node_b.pt")) + ) def test_state_tensor_dryrun_passthrough(self): """dry_run at start should preserve feedback detection and skip export.""" leapp.start(name=self.TEST_GRAPH_NAME, dry_run=True) - obs = annotate.input_tensors("node_state", {"obs": torch.tensor([1.0, 2.0, 3.0])}) - state = annotate.state_tensors("node_state", {"state": torch.tensor([0.1, 0.2, 0.3])}) + obs = annotate.input_tensors( + "node_state", {"obs": torch.tensor([1.0, 2.0, 3.0])} + ) + state = annotate.state_tensors( + "node_state", {"state": torch.tensor([0.1, 0.2, 0.3])} + ) self.assertIsInstance(obs, TracedTensor) self.assertIsInstance(state, TracedTensor) self.assertTrue(obs.is_tracing) @@ -258,7 +300,9 @@ def test_state_tensor_dryrun_passthrough(self): new_state = state + obs returned_state = annotate.update_state("node_state", {"state": new_state}) self.assertIs(returned_state, new_state) - annotate.output_tensors("node_state", {"action": obs + state}, export_with="jit") + annotate.output_tensors( + "node_state", {"action": obs + state}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -269,10 +313,13 @@ def test_state_tensor_dryrun_passthrough(self): len(annotate.detected_pipeline["feedback_flow"]) >= 1, "Dryrun state path should still populate feedback_flow", ) - self.assertFalse(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "node_state.pt"))) + self.assertFalse( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "node_state.pt")) + ) def test_module_dryrun_passthrough(self): """dry_run at start allows module tracking but should restore buffers and skip export.""" + class TinyStatefulModel(torch.nn.Module): def __init__(self): super().__init__() @@ -297,7 +344,9 @@ def forward(self, x): self.assertEqual(set(original_buffers.keys()), set(current_buffers.keys())) for name, buf in current_buffers.items(): self.assertNotIsInstance(buf, TracedTensor) - self.assertFalse(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "policy.pt"))) + self.assertFalse( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "policy.pt")) + ) def test_non_traced_node_is_selective_and_preserves_connectivity(self): """non_traced should skip export only for selected nodes while preserving graph connectivity.""" @@ -310,8 +359,9 @@ def test_non_traced_node_is_selective_and_preserves_connectivity(self): raw_y = raw_x * 2.0 annotate.output_tensors("raw_node", {"y": raw_y}, export_with="jit") - self.assertEqual(raw_y.output_port, "y", - "non_traced outputs should still publish their port") + self.assertEqual( + raw_y.output_port, "y", "non_traced outputs should still publish their port" + ) traced_y = annotate.input_tensors("traced_node", {"y": raw_y}) self.assertIsInstance(traced_y, TracedTensor) @@ -326,9 +376,15 @@ def test_non_traced_node_is_selective_and_preserves_connectivity(self): self.assertTrue(results["traced_node"]) self.assertIn("raw_node", annotate.detected_nodes) self.assertIn("traced_node", annotate.detected_nodes) - self.assertTrue(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "traced_node.pt"))) - self.assertFalse(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "raw_node.pt"))) - self.assertFalse(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "raw_node.onnx"))) + self.assertTrue( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "traced_node.pt")) + ) + self.assertFalse( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "raw_node.pt")) + ) + self.assertFalse( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "raw_node.onnx")) + ) self.assertIn("raw_node/y", annotate.detected_pipeline["data_flow"]) self.assertEqual( ["traced_node/y"], @@ -344,7 +400,9 @@ def test_non_traced_node_validation_is_skipped_under_validate_true(self): annotate.output_tensors("raw_node", {"y": raw_y}, export_with="jit") traced_y = annotate.input_tensors("traced_node", {"y": raw_y}) - annotate.output_tensors("traced_node", {"z": traced_y.square()}, export_with="jit") + annotate.output_tensors( + "traced_node", {"z": traced_y.square()}, export_with="jit" + ) leapp.stop() results = leapp.compile_graph(visualize=False, validate=True) @@ -352,15 +410,23 @@ def test_non_traced_node_validation_is_skipped_under_validate_true(self): self.assertEqual({"raw_node", "traced_node"}, set(results.keys())) self.assertTrue(results["raw_node"]) self.assertTrue(results["traced_node"]) - self.assertFalse(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "raw_node.pt"))) - self.assertTrue(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "traced_node.pt"))) + self.assertFalse( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "raw_node.pt")) + ) + self.assertTrue( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "traced_node.pt")) + ) def test_non_traced_state_node_still_populates_feedback_flow(self): """non_traced state nodes should still register feedback connections without exporting a model.""" leapp.start(name=self.TEST_GRAPH_NAME, non_traced=["state_node"]) - obs = annotate.input_tensors("state_node", {"obs": torch.tensor([1.0, 2.0, 3.0])}) - state = annotate.state_tensors("state_node", {"state": torch.tensor([0.1, 0.2, 0.3])}) + obs = annotate.input_tensors( + "state_node", {"obs": torch.tensor([1.0, 2.0, 3.0])} + ) + state = annotate.state_tensors( + "state_node", {"state": torch.tensor([0.1, 0.2, 0.3])} + ) self.assertIsInstance(obs, TracedTensor) self.assertIsInstance(state, TracedTensor) self.assertTrue(obs.is_tracing) @@ -368,7 +434,9 @@ def test_non_traced_state_node_still_populates_feedback_flow(self): new_state = state + obs annotate.update_state("state_node", {"state": new_state}) - annotate.output_tensors("state_node", {"action": obs - state}, export_with="jit") + annotate.output_tensors( + "state_node", {"action": obs - state}, export_with="jit" + ) leapp.stop() results = leapp.compile_graph(visualize=False, validate=True) @@ -376,5 +444,9 @@ def test_non_traced_state_node_still_populates_feedback_flow(self): self.assertTrue(results["state_node"]) self.assertIn("feedback_flow", annotate.detected_pipeline) self.assertGreaterEqual(len(annotate.detected_pipeline["feedback_flow"]), 1) - self.assertFalse(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "state_node.pt"))) - self.assertFalse(os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "state_node.onnx"))) + self.assertFalse( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "state_node.pt")) + ) + self.assertFalse( + os.path.exists(os.path.join(self.TEST_GRAPH_NAME, "state_node.onnx")) + ) diff --git a/tests/functional_tests/test_shared_memory_arrays.py b/tests/functional_tests/test_shared_memory_arrays.py index 0f8dfae6..c571ae86 100644 --- a/tests/functional_tests/test_shared_memory_arrays.py +++ b/tests/functional_tests/test_shared_memory_arrays.py @@ -61,15 +61,21 @@ def _verify_alias_mutation_propagates(self, make_alias): # exported graph has to reproduce. self.assertTrue( torch.equal(action.tensor, expected), - f"eager value diverged: got {action.tensor}, expected {expected}") + f"eager value diverged: got {action.tensor}, expected {expected}", + ) annotate.output_tensors("policy", {"action": action}, export_with="jit") leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0, - feedback_connections=0) + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=0, + ) self.verify_inference_manager( source_inputs={"policy/obs": obs_value}, source_outputs={"policy/action": expected}, @@ -85,13 +91,13 @@ def test_full_slice_alias_shares_provenance(self): def test_noop_conversion_alias_shares_provenance(self): """A ``to`` that neither casts nor moves allocates nothing.""" - self._verify_alias_mutation_propagates( - lambda tensor: tensor.to(torch.float32)) + self._verify_alias_mutation_propagates(lambda tensor: tensor.to(torch.float32)) def test_chained_aliases_share_one_provenance(self): """An alias of an alias still names the one original value.""" self._verify_alias_mutation_propagates( - lambda tensor: tensor.detach()[:].to(torch.float32)) + lambda tensor: tensor.detach()[:].to(torch.float32) + ) def test_inplace_dispatch_mutation_reaches_the_alias(self): """An in-place operation reaching dispatch rebinds the shared root. @@ -112,12 +118,12 @@ def test_inplace_dispatch_mutation_reaches_the_alias(self): alias = obs.detach() torch.relu_(obs) action = alias * 2.0 - targets = [str(node.target) - for node in annotate.nodes["policy"].graph.nodes] + targets = [str(node.target) for node in annotate.nodes["policy"].graph.nodes] self.assertTrue( torch.equal(action.tensor, expected), - f"eager value diverged: got {action.tensor}, expected {expected}") + f"eager value diverged: got {action.tensor}, expected {expected}", + ) annotate.output_tensors("policy", {"action": action}, export_with="jit") leapp.stop() @@ -126,7 +132,8 @@ def test_inplace_dispatch_mutation_reaches_the_alias(self): # The mutation is a real operation, so unlike an alias it is recorded. self.assertTrue( any("relu" in target for target in targets), - f"the in-place mutation was not recorded: {targets}") + f"the in-place mutation was not recorded: {targets}", + ) self.verify_inference_manager( source_inputs={"policy/obs": obs_value}, source_outputs={"policy/action": expected}, @@ -150,7 +157,8 @@ def test_clone_is_a_separate_value(self): self.assertTrue( torch.equal(action.tensor, expected), - f"eager value diverged: got {action.tensor}, expected {expected}") + f"eager value diverged: got {action.tensor}, expected {expected}", + ) annotate.output_tensors("policy", {"action": action}, export_with="jit") leapp.stop() @@ -174,20 +182,24 @@ def test_aliasing_records_no_operation(self): obs = annotate.input_tensors("policy", {"obs": obs_value.clone()}) action = obs.detach()[:].to(torch.float32) * 2.0 - targets = [str(node.target) - for node in annotate.nodes["policy"].graph.nodes] + targets = [str(node.target) for node in annotate.nodes["policy"].graph.nodes] annotate.output_tensors("policy", {"action": action}, export_with="jit") leapp.stop() leapp.compile_graph(visualize=False) self.assertEqual( - 2, len(targets), - f"expected only a placeholder and a multiply, got {targets}") + 2, + len(targets), + f"expected only a placeholder and a multiply, got {targets}", + ) self.assertFalse( - any("detach" in target or "getitem" in target or target == "to" - for target in targets), - f"an aliasing operation was recorded: {targets}") + any( + "detach" in target or "getitem" in target or target == "to" + for target in targets + ), + f"an aliasing operation was recorded: {targets}", + ) self.verify_inference_manager( source_inputs={"policy/obs": obs_value}, source_outputs={"policy/action": obs_value * 2.0}, @@ -208,12 +220,8 @@ def test_two_declared_aliases_share_one_root_and_one_port(self): buffer = obs_value.clone() alias = buffer.detach() - traced_buffer = annotate.input_tensors( - "policy", {"buffer": buffer} - ) - traced_alias = annotate.input_tensors( - "policy", {"alias": alias} - ) + traced_buffer = annotate.input_tensors("policy", {"buffer": buffer}) + traced_alias = annotate.input_tensors("policy", {"alias": alias}) if pass_index == 0: self.assertIs( @@ -240,8 +248,13 @@ def test_two_declared_aliases_share_one_root_and_one_port(self): leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0, - feedback_connections=0) + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=0, + ) self.verify_inference_manager( source_inputs={"policy/buffer": obs_value}, source_outputs={"policy/action": (obs_value + 1.0) * 2.0}, @@ -260,20 +273,27 @@ def test_adopted_declaration_supplies_semantics_the_port_lacks(self): leapp.start(name=self.TEST_GRAPH_NAME) traced_buffer = annotate.input_tensors("policy", {"buffer": buffer}) - annotate.input_tensors("policy", [ - TensorSemantics( - name="alias", ref=alias, kind=InputKindEnum.JOINT_POSITION), - ]) + annotate.input_tensors( + "policy", + [ + TensorSemantics( + name="alias", ref=alias, kind=InputKindEnum.JOINT_POSITION + ), + ], + ) annotate.output_tensors( - "policy", {"action": traced_buffer * 2.0}, export_with="jit") + "policy", {"action": traced_buffer * 2.0}, export_with="jit" + ) node = annotate.get_nodes()["policy"] leapp.stop() ports = [description.name_str for description in node.inputs] self.assertEqual(ports, ["buffer"], "adoption changed the exported port") self.assertEqual( - node.inputs[0].semantics.kind, InputKindEnum.JOINT_POSITION, - "the adopted declaration's semantics were dropped") + node.inputs[0].semantics.kind, + InputKindEnum.JOINT_POSITION, + "the adopted declaration's semantics were dropped", + ) def test_output_alias_with_a_surviving_different_root_faults(self): """A pre-declaration mutation cannot silently escape through an alias. @@ -321,8 +341,13 @@ def test_output_alias_over_a_trimmed_declaration_exports(self): leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0, - feedback_connections=0) + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=0, + ) self.verify_inference_manager( source_inputs={"policy/source": source_value}, source_outputs={"policy/a": source_value}, @@ -352,8 +377,7 @@ def test_from_torch_alias_follows_the_segment(self): leapp.start(name=self.TEST_GRAPH_NAME) for _ in range(2): - obs = annotate.input_tensors( - "node_a", {"obs": obs_value.clone()}) + obs = annotate.input_tensors("node_a", {"obs": obs_value.clone()}) array = wp.from_torch(obs) with annotate.warp_op("node_a", device=self.DEVICE): wp.launch( @@ -366,9 +390,9 @@ def test_from_torch_alias_follows_the_segment(self): self.assertTrue( torch.equal(action.tensor, expected), - f"eager value diverged: got {action.tensor}, expected {expected}") - annotate.output_tensors( - "node_a", {"action": action}, export_with="onnx") + f"eager value diverged: got {action.tensor}, expected {expected}", + ) + annotate.output_tensors("node_a", {"action": action}, export_with="onnx") leapp.stop() leapp.compile_graph(visualize=False) @@ -407,9 +431,9 @@ def test_to_torch_inside_an_open_segment_follows_the_close(self): self.assertTrue( torch.equal(action.tensor, expected), - f"eager value diverged: got {action.tensor}, expected {expected}") - annotate.output_tensors( - "node_a", {"action": action}, export_with="onnx") + f"eager value diverged: got {action.tensor}, expected {expected}", + ) + annotate.output_tensors("node_a", {"action": action}, export_with="onnx") leapp.stop() leapp.compile_graph(visualize=False) @@ -453,9 +477,9 @@ def test_second_warp_handle_on_one_buffer_follows_the_segment(self): self.assertTrue( torch.equal(action.tensor, expected), - f"eager value diverged: got {action.tensor}, expected {expected}") - annotate.output_tensors( - "node_a", {"action": action}, export_with="onnx") + f"eager value diverged: got {action.tensor}, expected {expected}", + ) + annotate.output_tensors("node_a", {"action": action}, export_with="onnx") leapp.stop() leapp.compile_graph(visualize=False) @@ -481,7 +505,8 @@ def test_launch_mutating_two_buffers_keeps_them_separate(self): leapp.start(name=self.TEST_GRAPH_NAME) for _ in range(2): tensor_a, tensor_b = annotate.input_tensors( - "node_a", {"in_a": a_value.clone(), "in_b": b_value.clone()}) + "node_a", {"in_a": a_value.clone(), "in_b": b_value.clone()} + ) array_1 = wp.from_torch(tensor_a) array_2 = wp.from_torch(tensor_b) with annotate.warp_op("node_a", device=self.DEVICE): @@ -493,18 +518,21 @@ def test_launch_mutating_two_buffers_keeps_them_separate(self): ) self.assertIsNot( - tensor_a.proxy.node, tensor_b.proxy.node, - "both buffers were left on one segment output") + tensor_a.proxy.node, + tensor_b.proxy.node, + "both buffers were left on one segment output", + ) out_a = tensor_a + 10.0 out_b = tensor_b * 3.0 self.assertTrue( torch.equal(out_a.tensor, expected_a) and torch.equal(out_b.tensor, expected_b), - f"eager values diverged: got {out_a.tensor} and {out_b.tensor}") + f"eager values diverged: got {out_a.tensor} and {out_b.tensor}", + ) annotate.output_tensors( - "node_a", {"out_a": out_a, "out_b": out_b}, - export_with="onnx") + "node_a", {"out_a": out_a, "out_b": out_b}, export_with="onnx" + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -540,8 +568,7 @@ def test_warp_torch_warp_mutations_chain_on_one_buffer(self): leapp.start(name=self.TEST_GRAPH_NAME) for _ in range(2): - tensor = annotate.input_tensors( - "node_a", {"in_a": source_value.clone()}) + tensor = annotate.input_tensors("node_a", {"in_a": source_value.clone()}) array = wp.from_torch(tensor) with annotate.warp_op("node_a", device=self.DEVICE): @@ -567,10 +594,9 @@ def test_warp_torch_warp_mutations_chain_on_one_buffer(self): action = tensor * 2.0 self.assertTrue( torch.equal(action.tensor, expected), - f"eager value diverged: got {action.tensor}, " - f"expected {expected}") - annotate.output_tensors( - "node_a", {"action": action}, export_with="onnx") + f"eager value diverged: got {action.tensor}, expected {expected}", + ) + annotate.output_tensors("node_a", {"action": action}, export_with="onnx") node = annotate.get_nodes()["node_a"] leapp.stop() @@ -580,7 +606,8 @@ def test_warp_torch_warp_mutations_chain_on_one_buffer(self): # the two would have merged into one segment and lost the middle stage. self.assertEqual( ["warp_segment_0", "warp_segment_1"], - [segment.runner_name for segment in node.warp_segments]) + [segment.runner_name for segment in node.warp_segments], + ) self.verify_inference_manager( source_inputs={"node_a/in_a": source_value}, source_outputs={"node_a/action": expected}, @@ -595,19 +622,22 @@ def test_host_readback_is_a_separate_value(self): leapp.start(name=self.TEST_GRAPH_NAME) array = annotate.input_tensors( "node_a", - {"in_a": wp.array( - [1.0, 2.0, 3.0], dtype=wp.float32, device=self.DEVICE)}, + {"in_a": wp.array([1.0, 2.0, 3.0], dtype=wp.float32, device=self.DEVICE)}, ) host = array.numpy() self.assertIsNot( - array.proxy_view, host.proxy_view, - "a host readback shared a root with its device buffer") + array.proxy_view, + host.proxy_view, + "a host readback shared a root with its device buffer", + ) returned = wp.from_numpy(host, device=self.DEVICE) self.assertIsNot( - host.proxy_view, returned.proxy_view, - "a host-to-device copy shared a root with its host source") + host.proxy_view, + returned.proxy_view, + "a host-to-device copy shared a root with its host source", + ) class TestPersistentAliasAttachment(WarpTestCase, LEAPPFunctionalTestBase): @@ -649,13 +679,18 @@ def test_declared_persistent_alias_receives_the_segment_output(self): # The kernel writes the buffer in place, so each pass starts over. torch_buffer.copy_(source_value) obs, declared_array = annotate.input_tensors( - "node_a", {"obs": torch_buffer, "warp_view": warp_buffer}) + "node_a", {"obs": torch_buffer, "warp_view": warp_buffer} + ) self.assertIs( - declared_array, warp_buffer, - "declaring a raw wp.array did not promote the caller's object") + declared_array, + warp_buffer, + "declaring a raw wp.array did not promote the caller's object", + ) self.assertIs( - obs.proxy_view, warp_buffer.proxy_view, - "the declared aliases did not share one root") + obs.proxy_view, + warp_buffer.proxy_view, + "the declared aliases did not share one root", + ) with annotate.warp_op("node_a", device=self.DEVICE): wp.launch( @@ -668,24 +703,34 @@ def test_declared_persistent_alias_receives_the_segment_output(self): self.assertTrue( torch.equal(action.tensor, expected), - f"eager value diverged: got {action.tensor}, expected {expected}") - annotate.output_tensors( - "node_a", {"action": action}, export_with="onnx") + f"eager value diverged: got {action.tensor}, expected {expected}", + ) + annotate.output_tensors("node_a", {"action": action}, export_with="onnx") node = annotate.get_nodes()["node_a"] leapp.stop() leapp.compile_graph(visualize=False) self.assertTrue( - any("warp_runner" in str(graph_node.target) - for graph_node in node.m.graph.nodes), - "the Warp runner did not survive pruning") + any( + "warp_runner" in str(graph_node.target) + for graph_node in node.m.graph.nodes + ), + "the Warp runner did not survive pruning", + ) self.assertEqual( - node.trimmed_inputs, set(), - "a declared alias produced a port nothing consumed") + node.trimmed_inputs, + set(), + "a declared alias produced a port nothing consumed", + ) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0, - feedback_connections=0) + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=0, + ) self.verify_inference_manager( source_inputs={"node_a/obs": source_value}, source_outputs={"node_a/action": expected}, @@ -711,13 +756,18 @@ def test_declaring_the_warp_side_first_reaches_the_same_root(self): for _ in range(2): torch_buffer.copy_(source_value) declared_array, obs = annotate.input_tensors( - "node_a", {"warp_view": warp_buffer, "obs": torch_buffer}) + "node_a", {"warp_view": warp_buffer, "obs": torch_buffer} + ) self.assertIs( - declared_array, warp_buffer, - "declaring a raw wp.array did not promote the caller's object") + declared_array, + warp_buffer, + "declaring a raw wp.array did not promote the caller's object", + ) self.assertIs( - obs.proxy_view, warp_buffer.proxy_view, - "the declared aliases did not share one root") + obs.proxy_view, + warp_buffer.proxy_view, + "the declared aliases did not share one root", + ) with annotate.warp_op("node_a", device=self.DEVICE): wp.launch( @@ -730,24 +780,34 @@ def test_declaring_the_warp_side_first_reaches_the_same_root(self): self.assertTrue( torch.equal(action.tensor, expected), - f"eager value diverged: got {action.tensor}, expected {expected}") - annotate.output_tensors( - "node_a", {"action": action}, export_with="onnx") + f"eager value diverged: got {action.tensor}, expected {expected}", + ) + annotate.output_tensors("node_a", {"action": action}, export_with="onnx") node = annotate.get_nodes()["node_a"] leapp.stop() leapp.compile_graph(visualize=False) self.assertTrue( - any("warp_runner" in str(graph_node.target) - for graph_node in node.m.graph.nodes), - "the Warp runner did not survive pruning") + any( + "warp_runner" in str(graph_node.target) + for graph_node in node.m.graph.nodes + ), + "the Warp runner did not survive pruning", + ) self.assertEqual( - node.trimmed_inputs, set(), - "a declared alias produced a port nothing consumed") + node.trimmed_inputs, + set(), + "a declared alias produced a port nothing consumed", + ) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0, - feedback_connections=0) + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=0, + ) self.verify_inference_manager( source_inputs={"node_a/warp_view": source_value}, source_outputs={"node_a/action": expected}, @@ -771,10 +831,12 @@ def test_declaring_the_torch_alias_promotes_the_object_the_caller_holds(self): for _ in range(2): torch_buffer.copy_(source_value) annotate.input_tensors( - "node_a", {"obs": torch_buffer, "warp_view": warp_buffer}) + "node_a", {"obs": torch_buffer, "warp_view": warp_buffer} + ) self.assertTrue( torch_buffer.is_tracing, - "declaring a Torch input did not promote the caller's object") + "declaring a Torch input did not promote the caller's object", + ) with annotate.warp_op("node_a", device=self.DEVICE): wp.launch( @@ -788,18 +850,21 @@ def test_declaring_the_torch_alias_promotes_the_object_the_caller_holds(self): self.assertTrue( torch.equal(action.tensor, expected), - f"eager value diverged: got {action.tensor}, expected {expected}") - annotate.output_tensors( - "node_a", {"action": action}, export_with="onnx") + f"eager value diverged: got {action.tensor}, expected {expected}", + ) + annotate.output_tensors("node_a", {"action": action}, export_with="onnx") node = annotate.get_nodes()["node_a"] leapp.stop() leapp.compile_graph(visualize=False) self.assertTrue( - any("warp_runner" in str(graph_node.target) - for graph_node in node.m.graph.nodes), - "the Warp runner did not survive pruning") + any( + "warp_runner" in str(graph_node.target) + for graph_node in node.m.graph.nodes + ), + "the Warp runner did not survive pruning", + ) self.verify_inference_manager( source_inputs={"node_a/obs": source_value}, source_outputs={"node_a/action": expected}, @@ -818,8 +883,7 @@ def test_undeclared_persistent_alias_attaches_to_nothing(self): _, warp_buffer = self._make_persistent_pair(obs_value) leapp.start(name=self.TEST_GRAPH_NAME) - obs = annotate.input_tensors( - "node_a", {"obs": obs_value.clone()}) + obs = annotate.input_tensors("node_a", {"obs": obs_value.clone()}) with annotate.warp_op("node_a", device=self.DEVICE): wp.launch( self.kernels.increment_in_place, @@ -829,11 +893,12 @@ def test_undeclared_persistent_alias_attaches_to_nothing(self): ) self.assertIs( - type(warp_buffer), wp.array, - "an undeclared persistent alias was promoted") + type(warp_buffer), wp.array, "an undeclared persistent alias was promoted" + ) self.assertIsNone( getattr(warp_buffer, "_proxy_view", None), - "an undeclared persistent alias was given provenance") + "an undeclared persistent alias was given provenance", + ) self.assertIsNotNone(obs.proxy, "the declared input lost its provenance") def test_closed_node_output_crosses_the_boundary_untouched(self): @@ -854,8 +919,10 @@ def test_closed_node_output_crosses_the_boundary_untouched(self): consumed = annotate.input_tensors("node_b", {"hidden": hidden}) self.assertIsNot( - consumed.proxy_view, hidden.proxy_view, - "a boundary rewrap shared the producer's root") + consumed.proxy_view, + hidden.proxy_view, + "a boundary rewrap shared the producer's root", + ) class TestNumpySharedMemory(LEAPPFunctionalTestBase): @@ -873,18 +940,23 @@ def test_cpu_numpy_conversion_shares_a_root(self): obs = annotate.input_tensors("policy", {"obs": torch.tensor([1.0, 2.0])}) self.assertIs( - obs.proxy_view, obs.numpy().proxy_view, - "a zero-copy .numpy() did not share the tensor's root") + obs.proxy_view, + obs.numpy().proxy_view, + "a zero-copy .numpy() did not share the tensor's root", + ) def test_identical_layout_numpy_view_shares_a_root(self): """A view cast covers the same bytes at the same layout.""" leapp.start(name=self.TEST_GRAPH_NAME) arr = annotate.input_tensors( - "policy", {"obs": np.array([1.0, 2.0, 3.0], dtype=np.float32)}) + "policy", {"obs": np.array([1.0, 2.0, 3.0], dtype=np.float32)} + ) self.assertIs( - arr.proxy_view, arr.view(type(arr)).proxy_view, - "an identical-layout view did not share its source's root") + arr.proxy_view, + arr.view(type(arr)).proxy_view, + "an identical-layout view did not share its source's root", + ) def test_narrowing_and_allocating_results_are_separate_values(self): """A slice covers fewer bytes and a ufunc result covers other bytes. @@ -894,14 +966,19 @@ def test_narrowing_and_allocating_results_are_separate_values(self): """ leapp.start(name=self.TEST_GRAPH_NAME) arr = annotate.input_tensors( - "policy", {"obs": np.array([1.0, 2.0, 3.0], dtype=np.float32)}) + "policy", {"obs": np.array([1.0, 2.0, 3.0], dtype=np.float32)} + ) self.assertIsNot( - arr.proxy_view, arr[0:2].proxy_view, - "a narrowing slice shared its source's root") + arr.proxy_view, + arr[0:2].proxy_view, + "a narrowing slice shared its source's root", + ) self.assertIsNot( - arr.proxy_view, np.add(arr, 1.0).proxy_view, - "a freshly allocated ufunc result shared its source's root") + arr.proxy_view, + np.add(arr, 1.0).proxy_view, + "a freshly allocated ufunc result shared its source's root", + ) def test_numpy_mutation_reaches_the_torch_carrier(self): """A write through the NumPy carrier moves the Torch carrier with it. @@ -923,7 +1000,8 @@ def test_numpy_mutation_reaches_the_torch_carrier(self): self.assertTrue( torch.equal(action.tensor, expected), - f"eager value diverged: got {action.tensor}, expected {expected}") + f"eager value diverged: got {action.tensor}, expected {expected}", + ) annotate.output_tensors("policy", {"action": action}, export_with="jit") leapp.stop() @@ -942,19 +1020,22 @@ def test_torch_mutation_reaches_the_numpy_carrier(self): property the test above already drives end to end. """ leapp.start(name=self.TEST_GRAPH_NAME) - obs = annotate.input_tensors( - "policy", {"obs": torch.tensor([-1.0, 2.0, -3.0])}) + obs = annotate.input_tensors("policy", {"obs": torch.tensor([-1.0, 2.0, -3.0])}) np_view = obs.numpy() placeholder_proxy = obs.proxy torch.relu_(obs) self.assertIs( - obs.proxy_view, np_view.proxy_view, - "the Torch mutation replaced the root instead of its proxy") + obs.proxy_view, + np_view.proxy_view, + "the Torch mutation replaced the root instead of its proxy", + ) self.assertIsNot( - placeholder_proxy, np_view.proxy, - "the NumPy carrier still reports the pre-mutation proxy") + placeholder_proxy, + np_view.proxy, + "the NumPy carrier still reports the pre-mutation proxy", + ) if __name__ == "__main__": diff --git a/tests/functional_tests/test_state_tracking.py b/tests/functional_tests/test_state_tracking.py index 2145129f..2079bb4a 100644 --- a/tests/functional_tests/test_state_tracking.py +++ b/tests/functional_tests/test_state_tracking.py @@ -26,7 +26,6 @@ from .base import LEAPPFunctionalTestBase - class TestStateTensors(LEAPPFunctionalTestBase): """Tests for the state_tensors API which handles tensors that are both inputs and outputs.""" @@ -34,10 +33,14 @@ def test_state_api_return_shapes_single_multi_and_reentry(self): """state_tensors/update_state should return single tensor for one key, tuple for many keys.""" leapp.start(name=self.TEST_GRAPH_NAME) - obs0 = annotate.input_tensors("policy", {"observation": torch.tensor([1.0, 2.0, 3.0])}) + obs0 = annotate.input_tensors( + "policy", {"observation": torch.tensor([1.0, 2.0, 3.0])} + ) # First trace: state_tensors returns TracedTensor (single) and tuple[TracedTensor] (multi) - single_state = annotate.state_tensors("policy", {"single": torch.tensor([0.1, 0.2, 0.3])}) + single_state = annotate.state_tensors( + "policy", {"single": torch.tensor([0.1, 0.2, 0.3])} + ) self.assertIsInstance(single_state, TracedTensor) multi_state = annotate.state_tensors( @@ -74,7 +77,9 @@ def test_state_api_return_shapes_single_multi_and_reentry(self): leapp.stop() # Passthrough (outside tracing): preserve single-vs-tuple return shape with raw tensors - single_passthrough = annotate.state_tensors("policy", {"single": torch.tensor([0.0, 0.0, 0.0])}) + single_passthrough = annotate.state_tensors( + "policy", {"single": torch.tensor([0.0, 0.0, 0.0])} + ) self.assertIsInstance(single_passthrough, torch.Tensor) multi_passthrough = annotate.state_tensors( @@ -110,15 +115,11 @@ def test_state_tensor_basic(self): obs_value = torch.tensor([1.0, 2.0, 3.0]) # Create regular input - obs = annotate.input_tensors( - "policy", {"observation": obs_value} - ) + obs = annotate.input_tensors("policy", {"observation": obs_value}) # Create state tensor (both input and output) — non-trivial initial value initial_counter = torch.tensor([5.0, -3.0, 0.5]) - counter = annotate.state_tensors( - "policy", {"counter": initial_counter} - ) + counter = annotate.state_tensors("policy", {"counter": initial_counter}) # Use state in computation new_counter = counter + obs @@ -137,8 +138,12 @@ def test_state_tensor_basic(self): # Verify structure: 1 dangling input (observation), 1 dangling output (action), # 1 feedback connection for the recurrent state self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0, - feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) self.verify_all_models_exist("policy") @@ -149,9 +154,11 @@ def test_state_tensor_basic(self): self.assertEqual(len(model_info["outputs"]), 2) # Verify feedback initial values safetensors file - self.verify_feedback_initial_values({ - "policy/counter": initial_counter, - }) + self.verify_feedback_initial_values( + { + "policy/counter": initial_counter, + } + ) self.verify_safetensors_matches_feedback(annotate) self.verify_inference_manager( source_inputs={"policy/observation": obs_value}, @@ -165,12 +172,8 @@ def test_state_tensor_without_update_state_becomes_regular_input(self): obs_value = torch.tensor([1.0, 2.0, 3.0]) initial_counter = torch.tensor([5.0, -3.0, 0.5]) - obs = annotate.input_tensors( - "policy", {"observation": obs_value} - ) - counter = annotate.state_tensors( - "policy", {"counter": initial_counter} - ) + obs = annotate.input_tensors("policy", {"observation": obs_value}) + counter = annotate.state_tensors("policy", {"counter": initial_counter}) action = obs + counter annotate.output_tensors("policy", {"action": action}, export_with="onnx") @@ -179,8 +182,12 @@ def test_state_tensor_without_update_state_becomes_regular_input(self): leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=2, outputs=1, internal_connections=0, - feedback_connections=0 + annotate, + nodes=1, + inputs=2, + outputs=1, + internal_connections=0, + feedback_connections=0, ) self.verify_all_models_exist("policy") self.verify_safetensors_matches_feedback(annotate) @@ -207,7 +214,9 @@ def test_state_tensor_history_buffer(self): ) # History buffer as state [batch, history_len, obs_dim] — non-trivial initial value - initial_history = torch.arange(history_length * obs_dim, dtype=torch.float32).reshape(1, history_length, obs_dim) + initial_history = torch.arange( + history_length * obs_dim, dtype=torch.float32 + ).reshape(1, history_length, obs_dim) history = annotate.state_tensors( "obs_processor", {"observation_history": initial_history} ) @@ -232,15 +241,21 @@ def test_state_tensor_history_buffer(self): # Verify structure: 1 dangling input, 1 dangling output, # 1 feedback connection (observation_history) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0, - feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) self.verify_all_models_exist("obs_processor") # Verify feedback initial values safetensors file - self.verify_feedback_initial_values({ - "obs_processor/observation_history": initial_history, - }) + self.verify_feedback_initial_values( + { + "obs_processor/observation_history": initial_history, + } + ) self.verify_safetensors_matches_feedback(annotate) # Compute expected output for InferenceManager (first iteration uses initial_history) @@ -260,9 +275,7 @@ def test_state_tensor_multiple_states(self): leapp.start(name=self.TEST_GRAPH_NAME) # Input - obs = annotate.input_tensors( - "policy", {"observation": obs_value} - ) + obs = annotate.input_tensors("policy", {"observation": obs_value}) # Multiple state tensors - returns tuple — non-trivial initial values initial_mean = torch.tensor([0.1, -0.2, 0.3]) @@ -302,17 +315,23 @@ def test_state_tensor_multiple_states(self): # Verify structure: 1 dangling input (observation), 1 dangling output (normalized), # 3 feedback connections (running_mean, running_var, step_count) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0, - feedback_connections=3 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=3, ) self.verify_all_models_exist("policy") # Verify feedback initial values safetensors file - self.verify_feedback_initial_values({ - "policy/running_mean": initial_mean, - "policy/running_var": initial_var, - "policy/step_count": initial_count, - }) + self.verify_feedback_initial_values( + { + "policy/running_mean": initial_mean, + "policy/running_var": initial_var, + "policy/step_count": initial_count, + } + ) self.verify_safetensors_matches_feedback(annotate) # Compute expected output for first iteration (uses initial state values) @@ -336,9 +355,7 @@ def test_state_tensor_multi_step_simulation(self): for step in range(5): # Current observation (same value each step for reproducibility) - current_obs = annotate.input_tensors( - "policy", {"current_obs": obs_value} - ) + current_obs = annotate.input_tensors("policy", {"current_obs": obs_value}) # History as state history_state = annotate.state_tensors("policy", {"history": history}) @@ -356,7 +373,9 @@ def test_state_tensor_multi_step_simulation(self): annotate.output_tensors("policy", {"action": action}, export_with="jit") # Update for next step (only affects non-traced execution) - history = new_history.tensor if hasattr(new_history, "tensor") else new_history + history = ( + new_history.tensor if hasattr(new_history, "tensor") else new_history + ) leapp.stop() leapp.compile_graph(visualize=False, validate=True) @@ -365,15 +384,21 @@ def test_state_tensor_multi_step_simulation(self): # 1 dangling input (current_obs), 1 dangling output (action), # 1 feedback connection (history) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0, - feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) self.verify_all_models_exist("policy") # Verify feedback initial values safetensors file - self.verify_feedback_initial_values({ - "policy/history": initial_history, - }) + self.verify_feedback_initial_values( + { + "policy/history": initial_history, + } + ) self.verify_safetensors_matches_feedback(annotate) # Compute expected output (first iteration uses initial_history) @@ -386,7 +411,9 @@ def test_state_tensor_multi_step_simulation(self): source_outputs={"policy/action": expected_action}, ) - def test_state_reentry_reusing_prior_state_output_tensor_preserves_feedback_source(self): + def test_state_reentry_reusing_prior_state_output_tensor_preserves_feedback_source( + self, + ): """Reusing the previous state output tensor should succeed because state feedback shares one source.""" initial_last_action = torch.tensor([0.0, 0.0, 0.0]) first_obs_value = torch.tensor([1.0, 2.0, 3.0]) @@ -415,8 +442,9 @@ def test_state_reentry_reusing_prior_state_output_tensor_preserves_feedback_sour if step_idx == 0: self.assertEqual("last_action", returned_last_action.output_port) - self.assertIs(annotate.nodes["policy"], - returned_last_action.context_obj) + self.assertIs( + annotate.nodes["policy"], returned_last_action.context_obj + ) last_action_value = returned_last_action @@ -425,12 +453,18 @@ def test_state_reentry_reusing_prior_state_output_tensor_preserves_feedback_sour # Should trace one node with one dangling input/output and one feedback edge self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0, - feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, + ) + self.verify_feedback_initial_values( + { + "policy/last_action_in": initial_last_action, + } ) - self.verify_feedback_initial_values({ - "policy/last_action_in": initial_last_action, - }) self.verify_inference_manager( source_inputs={"policy/obs": first_obs_value}, source_outputs={"policy/action": first_obs_value - initial_last_action}, @@ -454,12 +488,12 @@ def test_state_update_also_published_as_regular_output(self): state = annotate.state_tensors("policy", {"acc": state_value}) updated = obs + state state_value = annotate.update_state("policy", {"acc": updated}) - annotate.output_tensors( - "policy", {"action": updated}, export_with="onnx") + annotate.output_tensors("policy", {"action": updated}, export_with="onnx") consumed = annotate.input_tensors("sink", {"action": updated}) annotate.output_tensors( - "sink", {"scaled": consumed * 2.0}, export_with="onnx") + "sink", {"scaled": consumed * 2.0}, export_with="onnx" + ) leapp.stop() leapp.compile_graph(visualize=False, validate=True) @@ -470,15 +504,23 @@ def test_state_update_also_published_as_regular_output(self): # through the regular output, so both ports carry a real edge. The # exported names carry the _in/_out suffixes the backend adds for the # overlapping state name; the ports underneath are unchanged. - data_flow = {source: list(targets) for source, targets - in annotate.detected_pipeline["data_flow"].items()} - feedback_flow = {source: list(targets) for source, targets - in annotate.detected_pipeline["feedback_flow"].items()} + data_flow = { + source: list(targets) + for source, targets in annotate.detected_pipeline["data_flow"].items() + } + feedback_flow = { + source: list(targets) + for source, targets in annotate.detected_pipeline["feedback_flow"].items() + } self.assertEqual({"policy/action": ["sink/action"]}, data_flow) self.assertEqual({"policy/acc_out": ["policy/acc_in"]}, feedback_flow) self.verify_num_connections( - annotate, nodes=2, inputs=1, outputs=1, internal_connections=1, - feedback_connections=1 + annotate, + nodes=2, + inputs=1, + outputs=1, + internal_connections=1, + feedback_connections=1, ) self.verify_feedback_initial_values({"policy/acc_in": initial_state}) self.verify_inference_manager( @@ -506,9 +548,9 @@ def test_nested_state_tensor_raises(self): annotate.input_tensors("policy", {"obs": torch.zeros(3)}) with self.assertRaises(TypeError) as exc: - annotate.state_tensors("policy", { - "state_payload": {"inner": torch.zeros(3)} - }) + annotate.state_tensors( + "policy", {"state_payload": {"inner": torch.zeros(3)}} + ) self.assertIn("does not support nested state structures", str(exc.exception)) leapp.stop() @@ -534,9 +576,7 @@ def test_nested_update_state_raises(self): annotate.state_tensors("policy", {"state": torch.zeros(3)}) with self.assertRaises(TypeError) as exc: - annotate.update_state("policy", { - "state": {"inner": torch.zeros(3)} - }) + annotate.update_state("policy", {"state": {"inner": torch.zeros(3)}}) self.assertIn("does not support nested state structures", str(exc.exception)) leapp.stop() @@ -551,7 +591,9 @@ def test_update_state_non_traced_tensor_raises_clean_error(self): annotate.update_state("policy", {"counter": torch.ones(3)}) msg = str(exc.exception) - self.assertIn("update_state() for node 'policy' received non-traced tensors", msg) + self.assertIn( + "update_state() for node 'policy' received non-traced tensors", msg + ) self.assertNotIn("proxy", msg) annotate.update_state("policy", {"counter": state + obs}) @@ -561,6 +603,7 @@ def test_update_state_non_traced_tensor_raises_clean_error(self): # ── Test models ────────────────────────────────────────────────────────────── + class GRUModel(nn.Module): """GRU with a single hidden state buffer that gets reassigned.""" @@ -650,6 +693,7 @@ def forward(self, obs): # ── Tests ──────────────────────────────────────────────────────────────────── + class TestBufferTracking(LEAPPFunctionalTestBase): """Tests for annotate.module() auto-detection of stateful buffers.""" @@ -665,21 +709,28 @@ def test_single_buffer_gru(self): annotate.module("policy", model) action = model(obs_traced) - annotate.output_tensors("policy", {"action": action}, - export_with="onnx-torchscript") + annotate.output_tensors( + "policy", {"action": action}, export_with="onnx-torchscript" + ) leapp.stop() leapp.compile_graph(visualize=False) # 1 input (obs), 1 output (action), 1 feedback (h_state) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) self.verify_all_models_exist("policy") self.verify_safetensors_matches_feedback(annotate) - self.verify_feedback_initial_values({ - "policy/h_state_in": torch.zeros(1, 1, 8), - }) + self.verify_feedback_initial_values( + { + "policy/h_state_in": torch.zeros(1, 1, 8), + } + ) def test_lstm_two_buffers(self): """LSTM with 2 mutated buffers -> 2 feedback connections.""" @@ -693,22 +744,29 @@ def test_lstm_two_buffers(self): annotate.module("policy", model) action = model(obs_traced) - annotate.output_tensors("policy", {"action": action}, - export_with="onnx-torchscript") + annotate.output_tensors( + "policy", {"action": action}, export_with="onnx-torchscript" + ) leapp.stop() leapp.compile_graph(visualize=False) # 1 input (obs), 1 output (action), 2 feedback (h_state, c_state) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=2 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=2, ) self.verify_all_models_exist("policy") self.verify_safetensors_matches_feedback(annotate) - self.verify_feedback_initial_values({ - "policy/h_state_in": torch.zeros(1, 1, 8), - "policy/c_state_in": torch.zeros(1, 1, 8), - }) + self.verify_feedback_initial_values( + { + "policy/h_state_in": torch.zeros(1, 1, 8), + "policy/c_state_in": torch.zeros(1, 1, 8), + } + ) def test_partial_mutation(self): """3 buffers, 2 mutated -> 2 feedback, 1 regular input (const_mask).""" @@ -730,15 +788,21 @@ def test_partial_mutation(self): # 2 feedback (running_mean, step_count) # const_mask is baked as a constant (not a dynamic input) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=2 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=2, ) self.verify_all_models_exist("policy") self.verify_safetensors_matches_feedback(annotate) - self.verify_feedback_initial_values({ - "policy/running_mean": torch.zeros(1, 4), - "policy/step_count": torch.tensor([0.0]), - }) + self.verify_feedback_initial_values( + { + "policy/running_mean": torch.zeros(1, 4), + "policy/step_count": torch.tensor([0.0]), + } + ) def test_no_buffers_mutated(self): """All buffers read but not reassigned -> 0 feedback, all regular inputs.""" @@ -759,8 +823,12 @@ def test_no_buffers_mutated(self): # 1 dangling input (obs), 1 dangling output, 0 feedback # scale and bias are baked as constants (not dynamic inputs) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=0 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=0, ) self.verify_all_models_exist("policy") self.verify_safetensors_matches_feedback(annotate) @@ -782,8 +850,12 @@ def test_no_buffers(self): leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=0 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=0, ) self.verify_all_models_exist("policy") self.verify_safetensors_matches_feedback(annotate) @@ -801,15 +873,20 @@ def test_buffer_names_filter(self): annotate.module("policy", model, buffer_names=["h_state"]) action = model(obs_traced) - annotate.output_tensors("policy", {"action": action}, - export_with="onnx-torchscript") + annotate.output_tensors( + "policy", {"action": action}, export_with="onnx-torchscript" + ) leapp.stop() leapp.compile_graph(visualize=False) # 1 input (obs), 1 output (action), 1 feedback (h_state only) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) self.verify_all_models_exist("policy") @@ -828,20 +905,27 @@ def test_model_restored_after_tracing(self): annotate.module("policy", model) action = model(obs_traced) - annotate.output_tensors("policy", {"action": action}, - export_with="onnx-torchscript") + annotate.output_tensors( + "policy", {"action": action}, export_with="onnx-torchscript" + ) leapp.stop() leapp.compile_graph(visualize=False) # After tracing, all buffers should be restored restored_buffers = dict(model.named_buffers()) - self.assertEqual(set(original_buffers.keys()), set(restored_buffers.keys()), - "Buffer names should be restored after tracing") + self.assertEqual( + set(original_buffers.keys()), + set(restored_buffers.keys()), + "Buffer names should be restored after tracing", + ) # Restored values should be proper tensors (not TracedTensors) for name, buf in restored_buffers.items(): - self.assertIsInstance(buf, torch.Tensor, - f"Buffer '{name}' should be a torch.Tensor after restore") + self.assertIsInstance( + buf, + torch.Tensor, + f"Buffer '{name}' should be a torch.Tensor after restore", + ) def test_non_mutated_buffers_baked_as_constants(self): """Non-mutated buffers should be baked into the model, not dynamic inputs. @@ -849,6 +933,7 @@ def test_non_mutated_buffers_baked_as_constants(self): Verifies that the exported model produces correct output using the baked buffer values (scale=2, bias=1) without needing them as inputs. """ + class ScaleBiasModel(nn.Module): def __init__(self): super().__init__() @@ -874,8 +959,12 @@ def forward(self, obs): # Only 1 input (obs) — scale and bias baked as constants self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=0 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=0, ) # Load exported model and verify it produces correct values @@ -913,15 +1002,20 @@ def test_nested_submodule_buffers(self): annotate.module("policy", model) action = model(obs_traced) - annotate.output_tensors("policy", {"action": action}, - export_with="onnx-torchscript") + annotate.output_tensors( + "policy", {"action": action}, export_with="onnx-torchscript" + ) leapp.stop() leapp.compile_graph(visualize=False) # GRU's h_state is at rnn.h_state (nested) -> 1 feedback self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) self.verify_all_models_exist("policy") @@ -931,6 +1025,7 @@ class TestBufferTrackingEdgeCases(LEAPPFunctionalTestBase): def test_chained_operations_on_buffer(self): """Buffer used in a chain of operations: reshape, multiply, add.""" + class ChainedOpsModel(nn.Module): def __init__(self): super().__init__() @@ -958,13 +1053,18 @@ def forward(self, obs): leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) self.verify_all_models_exist("policy") def test_buffer_used_in_conditional_arithmetic(self): """Buffer combined with torch.where (element-wise conditional).""" + class ConditionalModel(nn.Module): def __init__(self): super().__init__() @@ -992,12 +1092,17 @@ def forward(self, obs): leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) def test_buffer_with_torch_cat(self): """Buffer concatenated with input using torch.cat.""" + class CatModel(nn.Module): def __init__(self): super().__init__() @@ -1024,12 +1129,17 @@ def forward(self, obs): leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) def test_buffer_with_matmul(self): """Buffer used in matrix multiplication.""" + class MatmulModel(nn.Module): def __init__(self): super().__init__() @@ -1056,12 +1166,17 @@ def forward(self, obs): leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) def test_buffer_with_slicing(self): """Buffer read via slicing (getitem).""" + class SlicingModel(nn.Module): def __init__(self): super().__init__() @@ -1090,12 +1205,17 @@ def forward(self, obs): leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) def test_buffer_with_different_dtypes(self): """Buffers with float32 and int64 dtypes coexist.""" + class MixedDtypeModel(nn.Module): def __init__(self): super().__init__() @@ -1123,12 +1243,17 @@ def forward(self, obs): leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=2 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=2, ) def test_scalar_buffer(self): """Single scalar buffer (0-dimensional after squeeze, 1-element tensor).""" + class ScalarBufferModel(nn.Module): def __init__(self): super().__init__() @@ -1154,12 +1279,17 @@ def forward(self, obs): leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) def test_deeply_nested_submodule_buffer(self): """Buffer three levels deep: model.proj.rnn.h_state.""" + class DeepModel(nn.Module): def __init__(self): super().__init__() @@ -1181,19 +1311,25 @@ def forward(self, obs): annotate.module("policy", model) action = model(obs_traced) - annotate.output_tensors("policy", {"action": action}, - export_with="onnx-torchscript") + annotate.output_tensors( + "policy", {"action": action}, export_with="onnx-torchscript" + ) leapp.stop() leapp.compile_graph(visualize=False) # rnn.rnn.h_state is 3 levels deep -> still detected self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) def test_multiple_outputs_with_buffer(self): """Model returns multiple outputs alongside buffer mutation.""" + class MultiOutputModel(nn.Module): def __init__(self): super().__init__() @@ -1215,18 +1351,24 @@ def forward(self, obs): annotate.module("policy", model) out_a, out_b = model(obs_traced) - annotate.output_tensors("policy", {"action_a": out_a, "action_b": out_b}, - export_with="jit") + annotate.output_tensors( + "policy", {"action_a": out_a, "action_b": out_b}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=2, - internal_connections=0, feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=2, + internal_connections=0, + feedback_connections=1, ) def test_buffer_with_clamp(self): """Buffer updated with torch.clamp — common in RL for bounding state.""" + class ClampModel(nn.Module): def __init__(self): super().__init__() @@ -1253,8 +1395,12 @@ def forward(self, obs): leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, ) @@ -1297,14 +1443,21 @@ def test_stop_restores_buffers_without_output_tensors(self): # Buffers should still be restored by stop() safety net restored_buffer_names = set(dict(model.named_buffers()).keys()) - self.assertEqual(original_buffer_names, restored_buffer_names, - "stop() should restore buffers even without output_tensors()") + self.assertEqual( + original_buffer_names, + restored_buffer_names, + "stop() should restore buffers even without output_tensors()", + ) # Verify they're real tensors, not TracedTensors from leapp.leapp_graph.datatypes import TracedTensor + for name, buf in model.named_buffers(): - self.assertNotIsInstance(buf, TracedTensor, - f"Buffer '{name}' should not be TracedTensor after stop()") + self.assertNotIsInstance( + buf, + TracedTensor, + f"Buffer '{name}' should not be TracedTensor after stop()", + ) def test_model_functional_after_tracing(self): """Model produces correct output after tracing completes (buffers properly restored).""" @@ -1324,8 +1477,9 @@ def test_model_functional_after_tracing(self): obs_traced = annotate.input_tensors("policy", {"obs": obs}) annotate.module("policy", model) action = model(obs_traced) - annotate.output_tensors("policy", {"action": action}, - export_with="onnx-torchscript") + annotate.output_tensors( + "policy", {"action": action}, export_with="onnx-torchscript" + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -1336,7 +1490,8 @@ def test_model_functional_after_tracing(self): self.assertTrue( torch.allclose(ref_output, post_trace_output, atol=1e-6), - "Model should produce identical output before and after tracing") + "Model should produce identical output before and after tracing", + ) def test_buffer_names_filter_nonexistent_name(self): """Filtering by a buffer name that doesn't exist -> no buffers tracked.""" @@ -1357,8 +1512,12 @@ def test_buffer_names_filter_nonexistent_name(self): # No buffers were tracked -> 0 feedback self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, - internal_connections=0, feedback_connections=0 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=0, ) @@ -1375,45 +1534,47 @@ def test_state_tensor_basic(self): leapp.start(name=self.TEST_GRAPH_NAME) obs_value = np.array([1.0, 2.0, 3.0], dtype=np.float32) - obs = annotate.input_tensors( - "policy", {"observation": obs_value} - ) + obs = annotate.input_tensors("policy", {"observation": obs_value}) initial_counter = np.array([5.0, -3.0, 0.5], dtype=np.float32) - counter = annotate.state_tensors( - "policy", {"counter": initial_counter} - ) + counter = annotate.state_tensors("policy", {"counter": initial_counter}) new_counter = counter + obs - returned_state = annotate.update_state( - "policy", {"counter": new_counter}) + returned_state = annotate.update_state("policy", {"counter": new_counter}) self.assertIs(returned_state, new_counter) action = obs * 2.0 - annotate.output_tensors( - "policy", {"action": action}, export_with=None) + annotate.output_tensors("policy", {"action": action}, export_with=None) leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0, - feedback_connections=1 + annotate, + nodes=1, + inputs=1, + outputs=1, + internal_connections=0, + feedback_connections=1, + ) + self.verify_feedback_initial_values( + { + "policy/counter": torch.from_numpy(initial_counter), + } ) - self.verify_feedback_initial_values({ - "policy/counter": torch.from_numpy(initial_counter), - }) self.verify_safetensors_matches_feedback(annotate) def test_register_buffer_with_inplace_assignment(self): """register_buffer with numpy in-place assignment, matching the torch case.""" + class MockModule: def __init__(self): self.values = np.array([1.0, 2.0, 3.0], dtype=np.float32) def run(self, traced_input): self.values = annotate.register_buffer( - "buffer_test", {"values": self.values}) + "buffer_test", {"values": self.values} + ) self.values[:] = traced_input return self.values * 100.0 @@ -1422,21 +1583,20 @@ def run(self, traced_input): leapp.start(name=self.TEST_GRAPH_NAME) input_array = np.array([4.0, 5.0, 6.0], dtype=np.float32) - traced_input = annotate.input_tensors( - "buffer_test", {"input": input_array}) + traced_input = annotate.input_tensors("buffer_test", {"input": input_array}) result = module.run(traced_input) self.assertIsInstance(module.values, TracedNpArray) self.assertIs(module.values.context_obj, annotate.nodes["buffer_test"]) - annotate.output_tensors( - "buffer_test", {"result": result}, export_with=None) + annotate.output_tensors("buffer_test", {"result": result}, export_with=None) leapp.stop() leapp.compile_graph(visualize=False) self.verify_num_connections( - annotate, nodes=1, inputs=1, outputs=1, internal_connections=0) + annotate, nodes=1, inputs=1, outputs=1, internal_connections=0 + ) if __name__ == "__main__": diff --git a/tests/functional_tests/test_unsupported_fail.py b/tests/functional_tests/test_unsupported_fail.py index 59d0e5e4..b4533e7c 100644 --- a/tests/functional_tests/test_unsupported_fail.py +++ b/tests/functional_tests/test_unsupported_fail.py @@ -30,8 +30,7 @@ def _force_boundary_source(source, target): This is what mirror_leapp_tags() does minus its equality check, so a test can build a deliberately incompatible connection. """ - promoted = promote_in_place( - target, source.name, source.context_obj, source.proxy) + promoted = promote_in_place(target, source.name, source.context_obj, source.proxy) promoted.output_port = source.output_port return promoted @@ -40,35 +39,45 @@ def _force_boundary_source(source, target): # Module-level helpers for caller identity tests # --------------------------------------------------------------------------- + def _helper_path_a(name, tensor): - return annotate.input_tensors(name, {'input': tensor}) + return annotate.input_tensors(name, {"input": tensor}) + def _helper_path_b(name, tensor): - return annotate.input_tensors(name, {'input': tensor}) + return annotate.input_tensors(name, {"input": tensor}) + def _shared_helper(name, tensor): - return annotate.input_tensors(name, {'input': tensor}) + return annotate.input_tensors(name, {"input": tensor}) + class _BasePipeline: def annotate_input(self, tensor): - return annotate.input_tensors('func', {'input': tensor}) + return annotate.input_tensors("func", {"input": tensor}) + class _PipelineA(_BasePipeline): def run(self, tensor): return self.annotate_input(tensor) + class _PipelineB(_BasePipeline): def run(self, tensor): return self.annotate_input(tensor) + def _timing_decorator(func): @functools.wraps(func) def wrapper(*args, **kwargs): return func(*args, **kwargs) + return wrapper + def _raw_annotate(tensor): - return annotate.input_tensors('func', {'input': tensor}) + return annotate.input_tensors("func", {"input": tensor}) + _decorated_annotate = _timing_decorator(_raw_annotate) @@ -79,7 +88,7 @@ class TestUnsupportedFail(LEAPPFunctionalTestBase): def test_input_tensors_bare_tensor_fails(self): """Bare tensors must be wrapped in a dict or TensorSemantics.""" with self.assertRaises(TypeError) as exc: - annotate.input_tensors('func', torch.tensor([1.0, 2.0, 3.0])) + annotate.input_tensors("func", torch.tensor([1.0, 2.0, 3.0])) self.assertIn("does not accept a bare tensor", str(exc.exception)) def test_output_tensors_bare_tensor_fails(self): @@ -87,9 +96,11 @@ def test_output_tensors_bare_tensor_fails(self): leapp.start(name=self.TEST_GRAPH_NAME) try: - traced = annotate.input_tensors('func', {'input': torch.tensor([1.0, 2.0, 3.0])}) + traced = annotate.input_tensors( + "func", {"input": torch.tensor([1.0, 2.0, 3.0])} + ) with self.assertRaises(TypeError) as exc: - annotate.output_tensors('func', traced + 1.0, export_with="jit") + annotate.output_tensors("func", traced + 1.0, export_with="jit") self.assertIn("does not accept a bare tensor", str(exc.exception)) finally: leapp.stop() @@ -99,11 +110,13 @@ def test_static_outputs_bare_tensor_fails(self): leapp.start(name=self.TEST_GRAPH_NAME) try: - traced = annotate.input_tensors('func', {'input': torch.tensor([1.0, 2.0, 3.0])}) + traced = annotate.input_tensors( + "func", {"input": torch.tensor([1.0, 2.0, 3.0])} + ) with self.assertRaises(TypeError) as exc: annotate.output_tensors( - 'func', - {'output': traced + 1.0}, + "func", + {"output": traced + 1.0}, static_outputs=torch.tensor([4.0, 5.0, 6.0]), export_with="jit", ) @@ -114,10 +127,13 @@ def test_static_outputs_bare_tensor_fails(self): def test_input_tensors_top_level_raw_list_fails(self): """Top-level raw collections must be named via a dict.""" with self.assertRaises(TypeError) as exc: - annotate.input_tensors('func', [ - torch.tensor([1.0, 2.0, 3.0]), - torch.tensor([4.0, 5.0, 6.0]), - ]) + annotate.input_tensors( + "func", + [ + torch.tensor([1.0, 2.0, 3.0]), + torch.tensor([4.0, 5.0, 6.0]), + ], + ) self.assertIn("expects either a dict of named tensors", str(exc.exception)) def test_output_tensors_top_level_raw_list_fails(self): @@ -125,10 +141,12 @@ def test_output_tensors_top_level_raw_list_fails(self): leapp.start(name=self.TEST_GRAPH_NAME) try: - traced = annotate.input_tensors('func', {'input': torch.tensor([1.0, 2.0, 3.0])}) + traced = annotate.input_tensors( + "func", {"input": torch.tensor([1.0, 2.0, 3.0])} + ) outputs = [traced + 1.0, traced + 2.0] with self.assertRaises(TypeError) as exc: - annotate.output_tensors('func', outputs, export_with="jit") + annotate.output_tensors("func", outputs, export_with="jit") self.assertIn("expects either a dict of named tensors", str(exc.exception)) finally: leapp.stop() @@ -138,15 +156,17 @@ def test_validation_message_uses_sample_index(self): leapp.start(name=self.TEST_GRAPH_NAME, save_path=self.TEST_GRAPH_NAME) for value in (1.0, 2.0): - traced = annotate.input_tensors('func', {'input': torch.tensor([value, value + 1.0])}) - annotate.output_tensors('func', {'output': traced + 1.0}, export_with="jit") + traced = annotate.input_tensors( + "func", {"input": torch.tensor([value, value + 1.0])} + ) + annotate.output_tensors("func", {"output": traced + 1.0}, export_with="jit") leapp.stop() - annotate.nodes['func'].outputs[0].cached_values[0] += 10.0 + annotate.nodes["func"].outputs[0].cached_values[0] += 10.0 results = leapp.compile_graph(visualize=False, validate=True, strict=False) - self.assertFalse(results['func']) + self.assertFalse(results["func"]) log_path = _get_logger().path with open(log_path) as f: @@ -160,12 +180,14 @@ def test_validation_error_hint_when_only_sample_zero_passes(self): leapp.start(name=self.TEST_GRAPH_NAME, save_path=self.TEST_GRAPH_NAME) for value in (1.0, 2.0): - traced = annotate.input_tensors('func', {'input': torch.tensor([value, value + 1.0])}) - annotate.output_tensors('func', {'output': traced + 1.0}, export_with="jit") + traced = annotate.input_tensors( + "func", {"input": torch.tensor([value, value + 1.0])} + ) + annotate.output_tensors("func", {"output": traced + 1.0}, export_with="jit") leapp.stop() - annotate.nodes['func'].outputs[0].cached_values[0] += 10.0 + annotate.nodes["func"].outputs[0].cached_values[0] += 10.0 with self.assertRaises(Exception) as exc: leapp.compile_graph(visualize=False, validate=True, strict=True) @@ -181,8 +203,8 @@ def test_output_tensors_without_prior_input_tensors(self): try: annotate.output_tensors( - 'orphan_node', - {'output': torch.tensor([1.0, 2.0, 3.0])}, + "orphan_node", + {"output": torch.tensor([1.0, 2.0, 3.0])}, export_with="jit", ) self.fail("Expected an exception") @@ -201,11 +223,11 @@ def test_reentry_output_shape_change_fails(self): for idx in range(2): traced = annotate.input_tensors( - 'shape_node', {'input': torch.tensor([1.0, 2.0, 3.0])} + "shape_node", {"input": torch.tensor([1.0, 2.0, 3.0])} ) output = traced + 1.0 if idx == 0 else traced.reshape(1, 3) annotate.output_tensors( - 'shape_node', {'output': output}, export_with="jit" + "shape_node", {"output": output}, export_with="jit" ) leapp.stop() @@ -220,13 +242,14 @@ def test_reentry_output_shape_change_fails(self): def test_same_output_sent_twice_to_one_node_fails(self): """One node output cannot feed two inputs of the same consumer.""" + @annotate.method() def funcA(inputA: torch.Tensor): return inputA @annotate.method() def funcB(inputB: torch.Tensor, inputC: torch.Tensor): - return inputB+inputC + return inputB + inputC leapp.start(name=self.TEST_GRAPH_NAME) retvalA = funcA(torch.tensor([1, 2, 3])) @@ -295,11 +318,14 @@ def funcC(input, detections): try: leapp.compile_graph() except Exception as e: - self.assertEqual(str(e), - "Error requesting input name change for funcC/input: detections is already in use") + self.assertEqual( + str(e), + "Error requesting input name change for funcC/input: detections is already in use", + ) def test_mirror_leapp_tags_data_mismatch(self): """Test mirror_leapp_tags with various data mismatch scenarios""" + @annotate.method(export_with="jit") def funcA(inputA: torch.Tensor): return inputA * 2.0 @@ -309,28 +335,35 @@ def funcB(inputA: torch.Tensor): return inputA, inputA + 1.0 leapp.start(name=self.TEST_GRAPH_NAME) - + # Test 1: Tensors with at least one element that is not equal # This should log an error but not raise an exception out_funcA = funcA(torch.tensor([1.0, 2.0, 3.0])) wrong_buffer = torch.tensor([1.0, 2.0, 4.0]) # Last element is wrong try: - annotate.mirror_leapp_tags(out_funcA, wrong_buffer) # Logs error, doesn't raise + annotate.mirror_leapp_tags( + out_funcA, wrong_buffer + ) # Logs error, doesn't raise self.fail("Expected an exception") except Exception as e: self.assertIn("source and target do not match", str(e)) - + # Test 2: List of tensors with one element wrong # This should log an error but not raise an exception out1, out2 = funcB(torch.tensor([1.0, 2.0, 3.0])) source_list = [out1, out2] - target_list = [out1.clone(), torch.tensor([1.0, 2.0, 99.0])] # Second tensor is wrong - try: - annotate.mirror_leapp_tags(source_list, target_list) # Logs error, doesn't raise + target_list = [ + out1.clone(), + torch.tensor([1.0, 2.0, 99.0]), + ] # Second tensor is wrong + try: + annotate.mirror_leapp_tags( + source_list, target_list + ) # Logs error, doesn't raise self.fail("Expected an exception") except Exception as e: self.assertIn("source and target do not match", str(e)) - + # Test 3: Lists with different number of elements # The exact-match check rejects this before any state is transferred source_list2 = [out1, out2] @@ -340,25 +373,31 @@ def funcB(inputA: torch.Tensor): self.fail("Expected an exception") except Exception as e: self.assertIn("source and target do not match", str(e)) - + # Test 4: Dicts with different keys # This should log an error but not raise an exception - source_dict = {'a': out_funcA, 'b': out_funcA + 1.0} - target_dict = {'a': out_funcA.clone(), 'c': (out_funcA + 1.0).clone()} # 'b' vs 'c' - try: - annotate.mirror_leapp_tags(source_dict, target_dict) # Logs error, doesn't raise + source_dict = {"a": out_funcA, "b": out_funcA + 1.0} + target_dict = { + "a": out_funcA.clone(), + "c": (out_funcA + 1.0).clone(), + } # 'b' vs 'c' + try: + annotate.mirror_leapp_tags( + source_dict, target_dict + ) # Logs error, doesn't raise self.fail("Expected an exception") except Exception as e: self.assertIn("source and target do not match", str(e)) - + leapp.stop() def test_reentrant_tracing_method_inside_method(self): """Test that nested method decorators are properly rejected. - + Calling one @annotate.method from within another should fail because the tracing lock is already acquired. """ + @annotate.method() def inner_func(inputA: torch.Tensor): return inputA * 2.0 @@ -370,7 +409,7 @@ def outer_func(inputA: torch.Tensor): return result leapp.start(name=self.TEST_GRAPH_NAME) - + try: outer_func(torch.tensor([1.0, 2.0, 3.0])) leapp.stop() @@ -379,47 +418,67 @@ def outer_func(inputA: torch.Tensor): leapp.stop() # The error should mention that we're trying to set up a new trace self.assertIn("Mixing active contexts is not allowed", str(e)) - - def test_reentrant_tracing_using_traced_tensors(self): #TODO: these both need to fail + + def test_reentrant_tracing_using_traced_tensors( + self, + ): # TODO: these both need to fail try: leapp.start(name=self.TEST_GRAPH_NAME) - tensors = annotate.input_tensors('func', {'inputA': torch.tensor([1.0, 2.0, 3.0])}) + tensors = annotate.input_tensors( + "func", {"inputA": torch.tensor([1.0, 2.0, 3.0])} + ) tensors += 100 input = torch.tensor([1.0, 2.0, 3.0]) + @annotate.method() def inner_func(inputA: torch.Tensor): return inputA + tensors + output_tensors = inner_func(input) - annotate.output_tensors('func', {'outputA': output_tensors}, export_with="jit") + annotate.output_tensors( + "func", {"outputA": output_tensors}, export_with="jit" + ) leapp.stop() self.fail("Expected an exception") except Exception as e: self.assertIn("detected multiple TracedTensor contexts", str(e)) self.assertIn("output_tensors() to finalize", str(e)) - + def test_passing_traced_tensor_to_method(self): try: + @annotate.method() def inner_func(inputA: torch.Tensor): return inputA + 5 + leapp.start(name=self.TEST_GRAPH_NAME) - tensors = annotate.input_tensors('func', {'inputA': torch.tensor([1.0, 2.0, 3.0])}) + tensors = annotate.input_tensors( + "func", {"inputA": torch.tensor([1.0, 2.0, 3.0])} + ) tensors += 100 output_tensors = inner_func(tensors) - annotate.output_tensors('func', {'outputA': output_tensors}, export_with="jit") + annotate.output_tensors( + "func", {"outputA": output_tensors}, export_with="jit" + ) leapp.stop() self.fail("Expected an exception") except Exception as e: self.assertIn("Mixing active contexts is not allowed", str(e)) - + def test_cross_context_traced_tensor_usage(self): """Basic test: addition of two TracedTensors from different contexts.""" try: leapp.start(name=self.TEST_GRAPH_NAME) - tensor1 = annotate.input_tensors('func1', {'inputA': torch.tensor([1.0, 2.0, 3.0])}) - tensor2 = annotate.input_tensors('func2', {'inputB': torch.tensor([1.0, 2.0, 3.0])}) + tensor1 = annotate.input_tensors( + "func1", {"inputA": torch.tensor([1.0, 2.0, 3.0])} + ) + tensor2 = annotate.input_tensors( + "func2", {"inputB": torch.tensor([1.0, 2.0, 3.0])} + ) output_tensors = tensor1 + tensor2 - annotate.output_tensors('func1', {'outputA': output_tensors}, export_with="jit") + annotate.output_tensors( + "func1", {"outputA": output_tensors}, export_with="jit" + ) leapp.stop() self.fail("Expected an exception") except Exception as e: @@ -430,8 +489,12 @@ def test_cross_context_torch_cat_list(self): """Test torch.cat with a list containing TracedTensors from different contexts.""" try: leapp.start(name=self.TEST_GRAPH_NAME) - tensor1 = annotate.input_tensors('func1', {'inputA': torch.tensor([1.0, 2.0, 3.0])}) - tensor2 = annotate.input_tensors('func2', {'inputB': torch.tensor([4.0, 5.0, 6.0])}) + tensor1 = annotate.input_tensors( + "func1", {"inputA": torch.tensor([1.0, 2.0, 3.0])} + ) + tensor2 = annotate.input_tensors( + "func2", {"inputB": torch.tensor([4.0, 5.0, 6.0])} + ) # torch.cat takes a list of tensors - should detect cross-context in the list output_tensors = torch.cat([tensor1, tensor2], dim=0) leapp.stop() @@ -444,8 +507,12 @@ def test_cross_context_torch_stack(self): """Test torch.stack with tensors from different contexts.""" try: leapp.start(name=self.TEST_GRAPH_NAME) - tensor1 = annotate.input_tensors('func1', {'inputA': torch.tensor([1.0, 2.0, 3.0])}) - tensor2 = annotate.input_tensors('func2', {'inputB': torch.tensor([4.0, 5.0, 6.0])}) + tensor1 = annotate.input_tensors( + "func1", {"inputA": torch.tensor([1.0, 2.0, 3.0])} + ) + tensor2 = annotate.input_tensors( + "func2", {"inputB": torch.tensor([4.0, 5.0, 6.0])} + ) # torch.stack also takes a list output_tensors = torch.stack([tensor1, tensor2], dim=0) leapp.stop() @@ -458,8 +525,12 @@ def test_cross_context_matmul(self): """Test matrix multiplication with TracedTensors from different contexts.""" try: leapp.start(name=self.TEST_GRAPH_NAME) - tensor1 = annotate.input_tensors('func1', {'inputA': torch.tensor([[1.0, 2.0], [3.0, 4.0]])}) - tensor2 = annotate.input_tensors('func2', {'inputB': torch.tensor([[5.0, 6.0], [7.0, 8.0]])}) + tensor1 = annotate.input_tensors( + "func1", {"inputA": torch.tensor([[1.0, 2.0], [3.0, 4.0]])} + ) + tensor2 = annotate.input_tensors( + "func2", {"inputB": torch.tensor([[5.0, 6.0], [7.0, 8.0]])} + ) # Matrix multiplication with two explicit args output_tensors = torch.matmul(tensor1, tensor2) leapp.stop() @@ -472,8 +543,12 @@ def test_cross_context_torch_add_with_alpha(self): """Test torch.add with alpha kwarg - mixing contexts in positional args.""" try: leapp.start(name=self.TEST_GRAPH_NAME) - tensor1 = annotate.input_tensors('func1', {'inputA': torch.tensor([1.0, 2.0, 3.0])}) - tensor2 = annotate.input_tensors('func2', {'inputB': torch.tensor([4.0, 5.0, 6.0])}) + tensor1 = annotate.input_tensors( + "func1", {"inputA": torch.tensor([1.0, 2.0, 3.0])} + ) + tensor2 = annotate.input_tensors( + "func2", {"inputB": torch.tensor([4.0, 5.0, 6.0])} + ) # torch.add(input, other, alpha=1) - tests kwargs handling output_tensors = torch.add(tensor1, tensor2, alpha=2.0) leapp.stop() @@ -486,9 +561,15 @@ def test_cross_context_three_contexts(self): """Test detection with three different contexts mixed together.""" try: leapp.start(name=self.TEST_GRAPH_NAME) - tensor1 = annotate.input_tensors('func1', {'inputA': torch.tensor([1.0, 2.0, 3.0])}) - tensor2 = annotate.input_tensors('func2', {'inputB': torch.tensor([4.0, 5.0, 6.0])}) - tensor3 = annotate.input_tensors('func3', {'inputC': torch.tensor([7.0, 8.0, 9.0])}) + tensor1 = annotate.input_tensors( + "func1", {"inputA": torch.tensor([1.0, 2.0, 3.0])} + ) + tensor2 = annotate.input_tensors( + "func2", {"inputB": torch.tensor([4.0, 5.0, 6.0])} + ) + tensor3 = annotate.input_tensors( + "func3", {"inputC": torch.tensor([7.0, 8.0, 9.0])} + ) # Mix all three contexts output_tensors = torch.cat([tensor1, tensor2, tensor3], dim=0) leapp.stop() @@ -502,9 +583,13 @@ def test_cross_context_torch_where(self): try: leapp.start(name=self.TEST_GRAPH_NAME) # condition from one context - cond_tensor = annotate.input_tensors('cond_ctx', {'cond': torch.tensor([True, False, True])}) + cond_tensor = annotate.input_tensors( + "cond_ctx", {"cond": torch.tensor([True, False, True])} + ) # x from another context - x_tensor = annotate.input_tensors('x_ctx', {'x': torch.tensor([1.0, 2.0, 3.0])}) + x_tensor = annotate.input_tensors( + "x_ctx", {"x": torch.tensor([1.0, 2.0, 3.0])} + ) # y is a regular tensor (should be fine) y_tensor = torch.tensor([10.0, 20.0, 30.0]) # torch.where(condition, x, y) - two TracedTensors from different contexts @@ -514,17 +599,21 @@ def test_cross_context_torch_where(self): except Exception as e: self.assertIn("detected multiple TracedTensor contexts", str(e)) self.assertIn("output_tensors() to finalize", str(e)) - + def test_annotate_multiple_parallel_inputs_with_same_name(self): try: leapp.start(name=self.TEST_GRAPH_NAME) output_tensors = [] for i in range(10): - tensor1 = annotate.input_tensors('func_combined', {'input': torch.tensor([1.0, 2.0, 3.0])}) + tensor1 = annotate.input_tensors( + "func_combined", {"input": torch.tensor([1.0, 2.0, 3.0])} + ) tensor1 += i output_tensors.append(tensor1) - annotate.output_tensors('func_combined', {'outputs': output_tensors}, export_with="jit") + annotate.output_tensors( + "func_combined", {"outputs": output_tensors}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -532,64 +621,72 @@ def test_annotate_multiple_parallel_inputs_with_same_name(self): except Exception as e: self.assertIn("Duplicate i/o name", str(e)) - self.assertIn("Each input/output in the same node must have a unique name", str(e)) + self.assertIn( + "Each input/output in the same node must have a unique name", str(e) + ) def test_output_tensors_with_non_traced_tensor(self): """Test that passing non-TracedTensors to output_tensors raises the correct error. - + This happens when the user doesn't use the returned TracedTensors from input_tensors() in their computations. """ try: leapp.start(name=self.TEST_GRAPH_NAME) - + # Get traced tensors but don't use them - traced_input = annotate.input_tensors('func', {'input': torch.tensor([1.0, 2.0, 3.0])}) - + traced_input = annotate.input_tensors( + "func", {"input": torch.tensor([1.0, 2.0, 3.0])} + ) + # Create a completely new tensor (not derived from traced_input) # This is the user error - they should be using traced_input untraced_output = torch.tensor([4.0, 5.0, 6.0]) - + # This should fail because untraced_output is not a TracedTensor - annotate.output_tensors('func', {'output': untraced_output}, export_with="jit") - + annotate.output_tensors( + "func", {"output": untraced_output}, export_with="jit" + ) + leapp.stop() self.fail("Expected an exception") - + except Exception as e: self.assertIn("non-traced tensors", str(e)) def test_traced_tensor_as_static_output_fails(self): """Test that using a TracedTensor (derived from input) as a static output fails. - + Static outputs should be constant tensors that are NOT derived from inputs. If a user accidentally marks a computed tensor as static, it should error. """ try: leapp.start(name=self.TEST_GRAPH_NAME) - + # Create input and trace it input_tensor = torch.tensor([1.0, 2.0, 3.0]) - traced_input = annotate.input_tensors('func', {'input': input_tensor}) - + traced_input = annotate.input_tensors("func", {"input": input_tensor}) + # Compute a tensor from the traced input computed_tensor = traced_input + 1.0 # This is a TracedTensor - + # Create a proper output proper_output = traced_input * 2.0 - + # User error: trying to use a TracedTensor as a static output # This should fail because static outputs must be raw tensors annotate.output_tensors( - 'func', - {'output': proper_output}, - static_outputs={'bad_static': computed_tensor}, # Error: TracedTensor not allowed - export_with="jit" + "func", + {"output": proper_output}, + static_outputs={ + "bad_static": computed_tensor + }, # Error: TracedTensor not allowed + export_with="jit", ) - + leapp.stop() self.fail("Expected an exception when using TracedTensor as static output") - + except Exception as e: error_msg = str(e) self.assertIn("static output 'bad_static'", error_msg) @@ -599,18 +696,25 @@ def test_input_tensors_after_output_tensors_same_node(self): """Test that calling input_tensors after output_tensors for the same node fails.""" try: leapp.start(name=self.TEST_GRAPH_NAME) - traced = annotate.input_tensors('func', {'input': torch.tensor([1.0, 2.0, 3.0])}) + traced = annotate.input_tensors( + "func", {"input": torch.tensor([1.0, 2.0, 3.0])} + ) result = traced + 1.0 - annotate.output_tensors('func', {'output': result}, export_with="jit") - annotate.input_tensors('func', {'input2': torch.tensor([4.0, 5.0, 6.0])}) + annotate.output_tensors("func", {"output": result}, export_with="jit") + annotate.input_tensors("func", {"input2": torch.tensor([4.0, 5.0, 6.0])}) leapp.stop() - self.fail("Expected an exception for input_tensors after output_tensors on same node") + self.fail( + "Expected an exception for input_tensors after output_tensors on same node" + ) except Exception as e: leapp.stop() - self.assertIn("Cannot reuse a node name from a different call site.", str(e)) + self.assertIn( + "Cannot reuse a node name from a different call site.", str(e) + ) def test_input_tensors_after_method_same_name(self): """Test that calling input_tensors with a name already used by a method fails.""" + @annotate.method(export_with="jit") def funcA(inputA: torch.Tensor): return inputA + 1.0 @@ -618,30 +722,41 @@ def funcA(inputA: torch.Tensor): try: leapp.start(name=self.TEST_GRAPH_NAME) funcA(torch.tensor([1.0, 2.0, 3.0])) - annotate.input_tensors('funcA', {'input2': torch.tensor([4.0, 5.0, 6.0])}) + annotate.input_tensors("funcA", {"input2": torch.tensor([4.0, 5.0, 6.0])}) leapp.stop() - self.fail("Expected an exception for input_tensors reusing a method node name") + self.fail( + "Expected an exception for input_tensors reusing a method node name" + ) except Exception as e: leapp.stop() - self.assertIn("Cannot reuse a node name from a different call site.", str(e)) + self.assertIn( + "Cannot reuse a node name from a different call site.", str(e) + ) def test_expected_fail_connection_shape_mismatch(self): """Shape mismatch should fail at compile_graph edge compatibility validation.""" try: leapp.start(name=self.TEST_GRAPH_NAME) - source_input = annotate.input_tensors('source_node', {'x': torch.tensor([1.0, 2.0, 3.0])}) + source_input = annotate.input_tensors( + "source_node", {"x": torch.tensor([1.0, 2.0, 3.0])} + ) source_output = source_input + 1.0 - annotate.output_tensors('source_node', {'y': source_output}, export_with="jit") + annotate.output_tensors( + "source_node", {"y": source_output}, export_with="jit" + ) # mutate shape after output_tensors and before next input_tensors, # while keeping the source identity so the connection is still established mutated = _force_boundary_source( - source_output, source_output.clone().reshape(1, 3)) + source_output, source_output.clone().reshape(1, 3) + ) - target_input = annotate.input_tensors('target_node', {'y': mutated}) + target_input = annotate.input_tensors("target_node", {"y": mutated}) target_output = target_input + 2.0 - annotate.output_tensors('target_node', {'z': target_output}, export_with="jit") + annotate.output_tensors( + "target_node", {"z": target_output}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -658,18 +773,25 @@ def test_expected_fail_connection_dtype_mismatch(self): try: leapp.start(name=self.TEST_GRAPH_NAME) - source_input = annotate.input_tensors('source_node', {'x': torch.tensor([1.0, 2.0, 3.0])}) + source_input = annotate.input_tensors( + "source_node", {"x": torch.tensor([1.0, 2.0, 3.0])} + ) source_output = source_input + 1.0 - annotate.output_tensors('source_node', {'y': source_output}, export_with="jit") + annotate.output_tensors( + "source_node", {"y": source_output}, export_with="jit" + ) # mutate dtype after output_tensors and before next input_tensors, # while keeping the source identity so the connection is still established mutated = _force_boundary_source( - source_output, source_output.clone().to(torch.int32)) + source_output, source_output.clone().to(torch.int32) + ) - target_input = annotate.input_tensors('target_node', {'y': mutated}) + target_input = annotate.input_tensors("target_node", {"y": mutated}) target_output = target_input + 2 - annotate.output_tensors('target_node', {'z': target_output}, export_with="jit") + annotate.output_tensors( + "target_node", {"z": target_output}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -692,20 +814,23 @@ def test_expected_fail_feedback_connection_shape_mismatch(self): # run twice from identical call sites to establish a feedback edge source_carrier = None for _ in range(2): - obs_traced, fb_traced = annotate.input_tensors('node_a', {'obs': obs, 'fb': feedback_next}) + obs_traced, fb_traced = annotate.input_tensors( + "node_a", {"obs": obs, "fb": feedback_next} + ) a_out = obs_traced + fb_traced - annotate.output_tensors('node_a', {'a_out': a_out}, export_with="jit") + annotate.output_tensors("node_a", {"a_out": a_out}, export_with="jit") - c_in = annotate.input_tensors('node_c', {'a_out': a_out}) + c_in = annotate.input_tensors("node_c", {"a_out": a_out}) c_out = c_in.reshape(1, 3) # source output descriptor shape is [1, 3] - annotate.output_tensors('node_c', {'fb_out': c_out}, export_with="jit") + annotate.output_tensors("node_c", {"fb_out": c_out}, export_with="jit") # mutate value between output_tensors and next input_tensors # keep the source identity to create a feedback connection into node_a/fb if source_carrier is None: source_carrier = c_out feedback_next = _force_boundary_source( - source_carrier, c_out.reshape(3).clone()) + source_carrier, c_out.reshape(3).clone() + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -726,24 +851,26 @@ def test_caller_identity_wrapper_indirection(self): should be detected as different call sites.""" try: leapp.start(name=self.TEST_GRAPH_NAME) - inp = _helper_path_a('func', torch.tensor([1.0, 2.0, 3.0])) + inp = _helper_path_a("func", torch.tensor([1.0, 2.0, 3.0])) result = inp + 1.0 - annotate.output_tensors('func', {'output': result}, export_with="onnx") - _helper_path_b('func', torch.tensor([4.0, 5.0, 6.0])) + annotate.output_tensors("func", {"output": result}, export_with="onnx") + _helper_path_b("func", torch.tensor([4.0, 5.0, 6.0])) leapp.stop() self.fail("Expected an exception for wrapper indirection") except Exception as e: leapp.stop() - self.assertIn("Cannot reuse a node name from a different call site.", str(e)) + self.assertIn( + "Cannot reuse a node name from a different call site.", str(e) + ) def test_caller_identity_same_helper_different_caller(self): """A single helper called from two different lines in user code should now be accepted because the helper itself is the origin.""" leapp.start(name=self.TEST_GRAPH_NAME) - inp = _shared_helper('func', torch.tensor([1.0])) + inp = _shared_helper("func", torch.tensor([1.0])) result = inp + 1.0 - annotate.output_tensors('func', {'output': result}, export_with="onnx") - _shared_helper('func', torch.tensor([2.0])) + annotate.output_tensors("func", {"output": result}, export_with="onnx") + _shared_helper("func", torch.tensor([2.0])) leapp.stop() def test_caller_identity_same_path_in_loop(self): @@ -751,9 +878,9 @@ def test_caller_identity_same_path_in_loop(self): accepted — the stack trace is identical on every iteration.""" leapp.start(name=self.TEST_GRAPH_NAME) for i in range(3): - inp = _shared_helper('func', torch.tensor([float(i)])) + inp = _shared_helper("func", torch.tensor([float(i)])) result = inp + 1.0 - annotate.output_tensors('func', {'output': result}, export_with="onnx") + annotate.output_tensors("func", {"output": result}, export_with="onnx") leapp.stop() def test_caller_identity_conditional_branching(self): @@ -762,36 +889,40 @@ def test_caller_identity_conditional_branching(self): try: leapp.start(name=self.TEST_GRAPH_NAME) if True: - inp = annotate.input_tensors('func', {'input': torch.tensor([1.0])}) + inp = annotate.input_tensors("func", {"input": torch.tensor([1.0])}) else: - inp = annotate.input_tensors('func', {'input': torch.tensor([1.0])}) + inp = annotate.input_tensors("func", {"input": torch.tensor([1.0])}) result = inp + 1.0 - annotate.output_tensors('func', {'output': result}, export_with="onnx") + annotate.output_tensors("func", {"output": result}, export_with="onnx") if False: - annotate.input_tensors('func', {'input': torch.tensor([2.0])}) + annotate.input_tensors("func", {"input": torch.tensor([2.0])}) else: - annotate.input_tensors('func', {'input': torch.tensor([2.0])}) + annotate.input_tensors("func", {"input": torch.tensor([2.0])}) leapp.stop() self.fail("Expected an exception for conditional branching") except Exception as e: leapp.stop() - self.assertIn("Cannot reuse a node name from a different call site.", str(e)) + self.assertIn( + "Cannot reuse a node name from a different call site.", str(e) + ) def test_caller_identity_functools_partial(self): """A functools.partial binding called from two different lines should be detected as different call sites.""" try: leapp.start(name=self.TEST_GRAPH_NAME) - input_func = functools.partial(annotate.input_tensors, 'func') - inp = input_func({'input': torch.tensor([1.0])}) + input_func = functools.partial(annotate.input_tensors, "func") + inp = input_func({"input": torch.tensor([1.0])}) result = inp + 1.0 - annotate.output_tensors('func', {'output': result}, export_with="onnx") - input_func({'input': torch.tensor([2.0])}) + annotate.output_tensors("func", {"output": result}, export_with="onnx") + input_func({"input": torch.tensor([2.0])}) leapp.stop() self.fail("Expected an exception for functools.partial") except Exception as e: leapp.stop() - self.assertIn("Cannot reuse a node name from a different call site.", str(e)) + self.assertIn( + "Cannot reuse a node name from a different call site.", str(e) + ) def test_caller_identity_class_hierarchy(self): """A base-class method called through different subclass methods @@ -801,7 +932,7 @@ def test_caller_identity_class_hierarchy(self): b = _PipelineB() inp = a.run(torch.tensor([1.0])) result = inp + 1.0 - annotate.output_tensors('func', {'output': result}, export_with="onnx") + annotate.output_tensors("func", {"output": result}, export_with="onnx") b.run(torch.tensor([2.0])) leapp.stop() @@ -811,7 +942,7 @@ def test_caller_identity_decorator_adds_frame(self): leapp.start(name=self.TEST_GRAPH_NAME) inp = _raw_annotate(torch.tensor([1.0])) result = inp + 1.0 - annotate.output_tensors('func', {'output': result}, export_with="onnx") + annotate.output_tensors("func", {"output": result}, export_with="onnx") _decorated_annotate(torch.tensor([2.0])) leapp.stop() @@ -820,17 +951,22 @@ def test_caller_identity_lambda_vs_named(self): be detected as different call sites.""" try: leapp.start(name=self.TEST_GRAPH_NAME) - inp = _shared_helper('func', torch.tensor([1.0])) + inp = _shared_helper("func", torch.tensor([1.0])) result = inp + 1.0 - annotate.output_tensors('func', {'output': result}, export_with="onnx") + annotate.output_tensors("func", {"output": result}, export_with="onnx") + def via_lambda(t): - return annotate.input_tensors('func', {'input': t}) + return annotate.input_tensors("func", {"input": t}) + via_lambda(torch.tensor([2.0])) leapp.stop() self.fail("Expected an exception for lambda vs named function") except Exception as e: leapp.stop() - self.assertIn("Cannot reuse a node name from a different call site.", str(e)) + self.assertIn( + "Cannot reuse a node name from a different call site.", str(e) + ) + -if __name__ == '__main__': +if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/functional_tests/test_warp_tracing.py b/tests/functional_tests/test_warp_tracing.py index 118576f5..712c7d5e 100644 --- a/tests/functional_tests/test_warp_tracing.py +++ b/tests/functional_tests/test_warp_tracing.py @@ -27,10 +27,7 @@ from tests.warp_support import WarpTestCase - - class TestWarpOp(WarpTestCase, LEAPPFunctionalTestBase): - def _run_torch_to_warp_to_torch_node(self, device): leapp.start(name=self.TEST_GRAPH_NAME) source = torch.tensor([1.0, 2.0, 3.0], device=device) @@ -118,9 +115,7 @@ def test_warp_write_to_plain_torch_alias_is_traced(self): outputs=[wp.from_torch(out)], device=self.DEVICE, ) - annotate.output_tensors( - "node_a", {"out_a": out * 3.0}, export_with="onnx" - ) + annotate.output_tensors("node_a", {"out_a": out * 3.0}, export_with="onnx") node = annotate.get_nodes()["node_a"] leapp.stop() @@ -129,14 +124,10 @@ def test_warp_write_to_plain_torch_alias_is_traced(self): self.verify_node_io(node, inputs=1, outputs=1) self.verify_inference_manager( source_inputs={ - "node_a/in_a": torch.tensor( - [10.0, 20.0, 30.0], device=self.DEVICE - ) + "node_a/in_a": torch.tensor([10.0, 20.0, 30.0], device=self.DEVICE) }, source_outputs={ - "node_a/out_a": torch.tensor( - [36.0, 66.0, 96.0], device=self.DEVICE - ) + "node_a/out_a": torch.tensor([36.0, 66.0, 96.0], device=self.DEVICE) }, ) @@ -240,7 +231,9 @@ def test_multiple_warp_ops_in_one_node_capture_in_discovery_order(self): [segment.runner_name for segment in node.warp_segments], ["warp_segment_0", "warp_segment_1"], ) - self.assertTrue(all(segment.apic_graph is not None for segment in node.warp_segments)) + self.assertTrue( + all(segment.apic_graph is not None for segment in node.warp_segments) + ) self.verify_all_models_exist("node_a") def test_single_warp_op_can_produce_multiple_outputs(self): @@ -446,7 +439,9 @@ def test_second_pass_with_extra_warp_op_fails(self): out = wp.empty_like(arr) wp.copy(out, arr) if include_extra: - with self.assertRaisesRegex(RuntimeError, "more regions than discovery"): + with self.assertRaisesRegex( + RuntimeError, "more regions than discovery" + ): with annotate.warp_op("node_a"): scaled = wp.empty_like(out) wp.launch( @@ -712,7 +707,6 @@ def test_torch_input_can_be_viewed_as_compound_warp_dtype(self): self.verify_all_models_exist("node_a") - class TestWarpInference(WarpTestCase, LEAPPFunctionalTestBase): NODE_NAME = "node_a" @@ -743,13 +737,12 @@ def test_pt2_runner_requires_custom_op_library(self): lambda values: self._launch_add(values, 1.0), export_with="pt2", ) - with tempfile.TemporaryDirectory() as cache_dir, mock.patch.dict( - os.environ, {"XDG_CACHE_HOME": cache_dir} + with ( + tempfile.TemporaryDirectory() as cache_dir, + mock.patch.dict(os.environ, {"XDG_CACHE_HOME": cache_dir}), ): os.environ.pop("LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY") - with self.assertRaisesRegex( - FileNotFoundError, "leapp-build-warp-runtime" - ): + with self.assertRaisesRegex(FileNotFoundError, "leapp-build-warp-runtime"): leapp.compile_graph(visualize=False) def test_onnx_runner_requires_custom_op_library(self): @@ -759,16 +752,15 @@ def test_onnx_runner_requires_custom_op_library(self): lambda values: self._launch_add(values, 1.0), ) leapp.compile_graph(visualize=False) - with tempfile.TemporaryDirectory() as cache_dir, mock.patch.dict( - os.environ, {"XDG_CACHE_HOME": cache_dir} + with ( + tempfile.TemporaryDirectory() as cache_dir, + mock.patch.dict(os.environ, {"XDG_CACHE_HOME": cache_dir}), ): os.environ.pop("LEAPP_WARP_ONNX_CUSTOM_OP_LIBRARY") manager = InferenceManager( f"{self.TEST_GRAPH_NAME}/{self.TEST_GRAPH_NAME}.yaml" ) - with self.assertRaisesRegex( - FileNotFoundError, "leapp-build-warp-runtime" - ): + with self.assertRaisesRegex(FileNotFoundError, "leapp-build-warp-runtime"): manager.run_policy(manager.get_mock_input()) def test_one_array_read_by_three_launches_is_one_segment_input(self): @@ -779,6 +771,7 @@ def test_one_array_read_by_three_launches_is_one_segment_input(self): post-call walk registers every array it touched, which is what the later launches recognize. """ + def operation(values): first = self._launch_add(values, 1.0) second = self._launch_add(values, 2.0) @@ -798,14 +791,22 @@ def operation(values): segment = node.warp_segments[0] self.assertEqual( - len(segment.input_refs), 1, + len(segment.input_refs), + 1, "one buffer read by three launches produced " - f"{len(segment.input_refs)} segment inputs") + f"{len(segment.input_refs)} segment inputs", + ) self.verify_inference_manager( - source_inputs={f"{self.NODE_NAME}/in_a": torch.tensor( - [1.0, 2.0, 3.0], device=self.DEVICE)}, - source_outputs={f"{self.NODE_NAME}/out_a": torch.tensor( - [2.0, 3.0, 4.0], device=self.DEVICE)}, + source_inputs={ + f"{self.NODE_NAME}/in_a": torch.tensor( + [1.0, 2.0, 3.0], device=self.DEVICE + ) + }, + source_outputs={ + f"{self.NODE_NAME}/out_a": torch.tensor( + [2.0, 3.0, 4.0], device=self.DEVICE + ) + }, ) @@ -973,8 +974,9 @@ def test_torch_view_consumes_pending_warp_output(self): self.verify_inference_manager( source_inputs={f"{self.NODE_NAME}/in_a": values}, source_outputs={ - f"{self.NODE_NAME}/out_a": - (values * 2 + 1).reshape(2, 3)[:, 1:].transpose(0, 1) + f"{self.NODE_NAME}/out_a": (values * 2 + 1) + .reshape(2, 3)[:, 1:] + .transpose(0, 1) }, ) @@ -1077,6 +1079,7 @@ def test_scan_tolerates_a_lazy_import_placeholder(self): which ``getattr``'s default does not absorb because the default only covers ``AttributeError``. """ + class ExceptionWrapper: def __init__(self, exception): object.__setattr__(self, "_exception", exception) @@ -1089,7 +1092,8 @@ def __call__(self, *args, **kwargs): module = types.ModuleType("leapp_test_absent_dependency") module.mesh_loader = ExceptionWrapper( - ModuleNotFoundError("No module named 'networkx'")) + ModuleNotFoundError("No module named 'networkx'") + ) sys.modules[module.__name__] = module self.addCleanup(sys.modules.pop, module.__name__, None) diff --git a/tests/test_examples/base_example_test.py b/tests/test_examples/base_example_test.py index c67f193e..2878e19d 100644 --- a/tests/test_examples/base_example_test.py +++ b/tests/test_examples/base_example_test.py @@ -41,8 +41,7 @@ def setUpClass(cls): # Ensure examples directory exists if not cls.examples_dir.exists(): - raise FileNotFoundError( - f"Examples directory not found: {cls.examples_dir}") + raise FileNotFoundError(f"Examples directory not found: {cls.examples_dir}") # Original working directory cls.original_cwd = os.getcwd() @@ -54,7 +53,7 @@ def setUp(self): def tearDown(self): """Clean up after each test.""" os.chdir(self.original_cwd) - + # Reset ExportManager singleton to avoid stale state between tests # This is necessary because different examples may use the same node names if ExportManager.is_interpret_graph_enabled(): @@ -90,7 +89,7 @@ def _run_example_script(self, script_name, expected_output_dir): cwd=self.repo_root, capture_output=True, text=True, - timeout=300 # 5 minute timeout + timeout=300, # 5 minute timeout ) return result, output_dir_path @@ -106,18 +105,17 @@ def _verify_output_files(self, output_dir_path, expected_files): # Check that output directory was created self.assertTrue( output_dir_path.exists(), - f"Output directory was not created: {output_dir_path}" + f"Output directory was not created: {output_dir_path}", ) # Check that it's actually a directory self.assertTrue( output_dir_path.is_dir(), - f"Output path is not a directory: {output_dir_path}" + f"Output path is not a directory: {output_dir_path}", ) # Get list of actual files - actual_files = [f.name for f in output_dir_path.iterdir() - if f.is_file()] + actual_files = [f.name for f in output_dir_path.iterdir() if f.is_file()] # Check each expected file for expected_file in expected_files: @@ -133,24 +131,25 @@ def _verify_output_files(self, output_dir_path, expected_files): expected_file, actual_files, f"Expected file '{expected_file}' not found in {output_dir_path}. " - f"Actual files: {actual_files}" + f"Actual files: {actual_files}", ) - def _run_and_verify_example(self, script_name, output_dir_name, expected_files, test_description): + def _run_and_verify_example( + self, script_name, output_dir_name, expected_files, test_description + ): """ Complete workflow to run an example and verify its outputs. Args: script_name: Name of the script file - output_dir_name: Expected output directory name + output_dir_name: Expected output directory name expected_files: List of expected output files test_description: Description for logging """ print(f"\n=== {test_description} ===") # Run the example - result, output_dir = self._run_example_script( - script_name, output_dir_name) + result, output_dir = self._run_example_script(script_name, output_dir_name) # Check that script ran successfully if result.returncode != 0: @@ -158,24 +157,24 @@ def _run_and_verify_example(self, script_name, output_dir_name, expected_files, print(f"STDERR:\n{result.stderr}") self.assertEqual( - result.returncode, 0, + result.returncode, + 0, f"{script_name} failed with return code {result.returncode}.\n" f"STDOUT:\n{result.stdout}\n" - f"STDERR:\n{result.stderr}" + f"STDERR:\n{result.stderr}", ) # Verify expected output files exist self._verify_output_files(output_dir, expected_files) print(f"✓ {script_name} completed successfully") - print( - f"✓ Generated {len(expected_files)} expected files in {output_dir}") + print(f"✓ Generated {len(expected_files)} expected files in {output_dir}") self.run_generated_models(output_dir) def run_generated_models(self, output_dir): - '''Run the generated full graph''' - graph_yaml = os.path.split(output_dir)[-1] + '.yaml' + """Run the generated full graph""" + graph_yaml = os.path.split(output_dir)[-1] + ".yaml" graph_yaml_path = os.path.join(output_dir, graph_yaml) # Load the inference manager (validates the graph structure) diff --git a/tests/test_examples/test_function_patch.py b/tests/test_examples/test_function_patch.py index d44dead3..de215c81 100644 --- a/tests/test_examples/test_function_patch.py +++ b/tests/test_examples/test_function_patch.py @@ -23,8 +23,11 @@ def test_function_patch_example(): with tempfile.TemporaryDirectory() as workdir: result = subprocess.run( - [sys.executable, str(EXAMPLE)], cwd=workdir, - capture_output=True, text=True, timeout=120, + [sys.executable, str(EXAMPLE)], + cwd=workdir, + capture_output=True, + text=True, + timeout=120, ) assert result.returncode == 0, f"{result.stdout}\n{result.stderr}" assert "Patch restored: True" in result.stdout @@ -34,7 +37,8 @@ def test_function_patch_example(): manager = InferenceManager(str(graph_dir / "function_patch_example.yaml")) inputs = manager.get_mock_input() assert set(inputs) == { - "camera_preprocess/left_bgr", "camera_preprocess/right_bgr", + "camera_preprocess/left_bgr", + "camera_preprocess/right_bgr", } rng = np.random.default_rng(17) frames = { @@ -43,7 +47,9 @@ def test_function_patch_example(): } # Change each camera independently to catch frozen OpenCV results. for changed_camera in ("left_bgr", "right_bgr"): - frames[changed_camera] = rng.integers(0, 256, size=(8, 12, 3), dtype=np.uint8) + frames[changed_camera] = rng.integers( + 0, 256, size=(8, 12, 3), dtype=np.uint8 + ) expected_views = [] for name, bgr in frames.items(): key = f"camera_preprocess/{name}" diff --git a/tests/test_examples/test_stateful_gru_export.py b/tests/test_examples/test_stateful_gru_export.py index 388bdcc5..321c05fa 100644 --- a/tests/test_examples/test_stateful_gru_export.py +++ b/tests/test_examples/test_stateful_gru_export.py @@ -61,7 +61,9 @@ def _export_gru(self): obs_traced = annotate.input_tensors("policy", {"obs": obs}) annotate.module("policy", model) action = model(obs_traced) - annotate.output_tensors("policy", {"action": action}, export_with="onnx-torchscript") + annotate.output_tensors( + "policy", {"action": action}, export_with="onnx-torchscript" + ) leapp.stop() leapp.compile_graph(visualize=False) @@ -99,21 +101,24 @@ def test_stateful_gru_h_state_updates_between_calls(self): h_initial = manager.value_dict["policy"]["h_state_in"].clone() self.assertTrue( torch.allclose(h_initial, torch.zeros_like(h_initial)), - "Expected h_state to start as zeros") + "Expected h_state to start as zeros", + ) # First call — h_state should change from zeros manager.run_policy(inputs) h_after_first = manager.value_dict["policy"]["h_state_in"].clone() self.assertFalse( torch.allclose(h_after_first, h_initial), - "h_state was not updated after the first call") + "h_state was not updated after the first call", + ) # Second call — h_state should change again (different from after first call) manager.run_policy(inputs) h_after_second = manager.value_dict["policy"]["h_state_in"].clone() self.assertFalse( torch.allclose(h_after_second, h_after_first), - "h_state was not updated between the first and second calls") + "h_state was not updated between the first and second calls", + ) if __name__ == "__main__": diff --git a/tests/test_examples/test_wbc_obj.py b/tests/test_examples/test_wbc_obj.py index 1270e79b..f62e9080 100644 --- a/tests/test_examples/test_wbc_obj.py +++ b/tests/test_examples/test_wbc_obj.py @@ -25,6 +25,7 @@ def run_example(**inputs): from examples.wbc_obj import WBC + wbc = WBC() return wbc.run_model(**inputs) @@ -36,51 +37,48 @@ def test_wbc_obj_execution(self): """Test that wbc_obj.py runs without errors and generates expected files.""" # Expected output files - these should be similar to wbc_plain but for sample_wbc_obj - expected_files = [ - 'wbc_obj.onnx', - 'sample_wbc_obj.png', - 'sample_wbc_obj.yaml' - ] + expected_files = ["wbc_obj.onnx", "sample_wbc_obj.png", "sample_wbc_obj.yaml"] # Run the test self._run_and_verify_example( - script_name='wbc_obj.py', - output_dir_name='sample_wbc_obj', + script_name="wbc_obj.py", + output_dir_name="sample_wbc_obj", expected_files=expected_files, - test_description='Testing wbc_obj.py' + test_description="Testing wbc_obj.py", ) # Test that the exported pipeline produces the same outputs as the original - device = 'cuda' if torch.cuda.is_available() else 'cpu' + device = "cuda" if torch.cuda.is_available() else "cpu" inputs = { - 'joint_pos': torch.randn(19, device=device, dtype=torch.float32), - 'joint_vel': torch.randn(19, device=device, dtype=torch.float32), - 'velocity_commands': torch.randn(3, device=device, dtype=torch.float32), - 'lin_vel_I': torch.randn(3, device=device, dtype=torch.float32), - 'ang_vel_I': torch.randn(3, device=device, dtype=torch.float32), - 'q_IB': torch.randn(4, device=device, dtype=torch.float32), + "joint_pos": torch.randn(19, device=device, dtype=torch.float32), + "joint_vel": torch.randn(19, device=device, dtype=torch.float32), + "velocity_commands": torch.randn(3, device=device, dtype=torch.float32), + "lin_vel_I": torch.randn(3, device=device, dtype=torch.float32), + "ang_vel_I": torch.randn(3, device=device, dtype=torch.float32), + "q_IB": torch.randn(4, device=device, dtype=torch.float32), } exported_pipeline_inputs = { - 'wbc_obj/velocity_commands': inputs['velocity_commands'].clone(), - 'wbc_obj/joint_vel': inputs['joint_vel'].clone(), - 'wbc_obj/joint_pos': inputs['joint_pos'].clone(), - 'wbc_obj/lin_vel_I': inputs['lin_vel_I'].clone(), - 'wbc_obj/ang_vel_I': inputs['ang_vel_I'].clone(), - 'wbc_obj/q_IB': inputs['q_IB'].clone(), + "wbc_obj/velocity_commands": inputs["velocity_commands"].clone(), + "wbc_obj/joint_vel": inputs["joint_vel"].clone(), + "wbc_obj/joint_pos": inputs["joint_pos"].clone(), + "wbc_obj/lin_vel_I": inputs["lin_vel_I"].clone(), + "wbc_obj/ang_vel_I": inputs["ang_vel_I"].clone(), + "wbc_obj/q_IB": inputs["q_IB"].clone(), } - outputs = run_example(**inputs) # Load and run the exported pipeline - exported_example = InferenceManager('sample_wbc_obj/sample_wbc_obj.yaml') + exported_example = InferenceManager("sample_wbc_obj/sample_wbc_obj.yaml") exported_outputs = exported_example.run_policy(exported_pipeline_inputs) - exported_action = exported_outputs['wbc_obj/actions'] + exported_action = exported_outputs["wbc_obj/actions"] try: if torch.cuda.is_available(): torch.cuda.synchronize() # Wait for GPU operations to complete - self.assertTrue(torch.allclose(outputs, exported_action, rtol=1e-6, atol=1e-6)) + self.assertTrue( + torch.allclose(outputs, exported_action, rtol=1e-6, atol=1e-6) + ) except AssertionError as e: print("Outputs do not match:") print(f"Outputs: {outputs}") @@ -88,5 +86,5 @@ def test_wbc_obj_execution(self): raise e -if __name__ == '__main__': +if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/test_examples/test_wbc_plain.py b/tests/test_examples/test_wbc_plain.py index 1f14649e..77413872 100644 --- a/tests/test_examples/test_wbc_plain.py +++ b/tests/test_examples/test_wbc_plain.py @@ -25,9 +25,11 @@ def run_example(**inputs): from examples.wbc_plain import run_model, get_model + model = get_model("models/isaac_velocity_flat_h1_v0.pt") return run_model(model, **inputs) + class TestWBCPlain(BaseExampleTest): """Unit tests for examples/wbc_plain.py""" @@ -36,50 +38,52 @@ def test_wbc_plain_execution(self): # Expected output files based on the sample_wbc_graph directory expected_files = [ - 'wbc_plain.onnx', - 'sample_wbc_graph.png', - 'sample_wbc_graph.yaml' + "wbc_plain.onnx", + "sample_wbc_graph.png", + "sample_wbc_graph.yaml", ] # Run the test self._run_and_verify_example( - script_name='wbc_plain.py', - output_dir_name='sample_wbc_graph', + script_name="wbc_plain.py", + output_dir_name="sample_wbc_graph", expected_files=expected_files, - test_description='Testing wbc_plain.py' + test_description="Testing wbc_plain.py", ) - device = 'cuda' if torch.cuda.is_available() else 'cpu' + device = "cuda" if torch.cuda.is_available() else "cpu" inputs = { - 'joint_pos': torch.randn(19, device=device, dtype=torch.float32), - 'joint_vel': torch.randn(19, device=device, dtype=torch.float32), - 'velocity_commands': torch.randn(3, device=device, dtype=torch.float32), - 'lin_vel_I': torch.randn(3, device=device, dtype=torch.float32), - 'ang_vel_I': torch.randn(3, device=device, dtype=torch.float32), - 'q_IB': torch.randn(4, device=device, dtype=torch.float32), - 'previous_actions': torch.zeros(19, device=device, dtype=torch.float32), + "joint_pos": torch.randn(19, device=device, dtype=torch.float32), + "joint_vel": torch.randn(19, device=device, dtype=torch.float32), + "velocity_commands": torch.randn(3, device=device, dtype=torch.float32), + "lin_vel_I": torch.randn(3, device=device, dtype=torch.float32), + "ang_vel_I": torch.randn(3, device=device, dtype=torch.float32), + "q_IB": torch.randn(4, device=device, dtype=torch.float32), + "previous_actions": torch.zeros(19, device=device, dtype=torch.float32), } exported_pipeline_inputs = { - 'wbc_plain/velocity_commands': inputs['velocity_commands'].clone(), - 'wbc_plain/joint_vel': inputs['joint_vel'].clone(), - 'wbc_plain/joint_pos': inputs['joint_pos'].clone(), - 'wbc_plain/lin_vel_I': inputs['lin_vel_I'].clone(), - 'wbc_plain/ang_vel_I': inputs['ang_vel_I'].clone(), - 'wbc_plain/q_IB': inputs['q_IB'].clone(), + "wbc_plain/velocity_commands": inputs["velocity_commands"].clone(), + "wbc_plain/joint_vel": inputs["joint_vel"].clone(), + "wbc_plain/joint_pos": inputs["joint_pos"].clone(), + "wbc_plain/lin_vel_I": inputs["lin_vel_I"].clone(), + "wbc_plain/ang_vel_I": inputs["ang_vel_I"].clone(), + "wbc_plain/q_IB": inputs["q_IB"].clone(), } run_example(**inputs) outputs, _ = run_example(**inputs) - exported_example = InferenceManager('sample_wbc_graph/sample_wbc_graph.yaml') + exported_example = InferenceManager("sample_wbc_graph/sample_wbc_graph.yaml") exported_outputs = exported_example.run_policy(exported_pipeline_inputs) - exported_action = exported_outputs['wbc_plain/final_actions'] + exported_action = exported_outputs["wbc_plain/final_actions"] try: if torch.cuda.is_available(): torch.cuda.synchronize() # Wait for GPU operations to complete - self.assertTrue(torch.allclose(outputs, exported_action, rtol=1e-3, atol=1e-6)) + self.assertTrue( + torch.allclose(outputs, exported_action, rtol=1e-3, atol=1e-6) + ) except AssertionError as e: print("Outputs do not match:") print(f"Outputs: {outputs}") @@ -87,5 +91,5 @@ def test_wbc_plain_execution(self): raise e -if __name__ == '__main__': +if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/unit_tests/export_format_validation.py b/tests/unit_tests/export_format_validation.py index 4eb5a7b2..de5016a9 100644 --- a/tests/unit_tests/export_format_validation.py +++ b/tests/unit_tests/export_format_validation.py @@ -85,10 +85,7 @@ def verify_onnx( output_names=list(output_names), ) session = ort.InferenceSession(str(onnx_path)) - onnx_inputs = { - name: inp.numpy() - for name, inp in zip(input_names, inputs) - } + onnx_inputs = {name: inp.numpy() for name, inp in zip(input_names, inputs)} output_onnx = session.run(None, onnx_inputs)[0] testcase.assertTrue( np.allclose(output_onnx, expected.detach().numpy(), atol=atol), @@ -168,8 +165,7 @@ def verify_exported_program( f"{test_name}: ExportedProgram output doesn't match expected", ) except Exception as exc: - testcase.fail( - f"{test_name}: ExportedProgram export/execution failed: {exc}") + testcase.fail(f"{test_name}: ExportedProgram export/execution failed: {exc}") def verify_exported_program_on_random_inputs( @@ -198,5 +194,4 @@ def verify_exported_program_on_random_inputs( f"{index + 1}/{num_inputs}", ) except Exception as exc: - testcase.fail( - f"{test_name}: ExportedProgram export/execution failed: {exc}") + testcase.fail(f"{test_name}: ExportedProgram export/execution failed: {exc}") diff --git a/tests/unit_tests/test_conversion.py b/tests/unit_tests/test_conversion.py index f3e376f8..8c4e8f8e 100644 --- a/tests/unit_tests/test_conversion.py +++ b/tests/unit_tests/test_conversion.py @@ -57,9 +57,9 @@ def test_numpy_method_returns_traced_np_array(self): """Test TracedTensor.numpy() returns TracedNpArray when tracing.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - + y = x.numpy() - + self.assertIsInstance(y, TracedNpArray) np.testing.assert_array_almost_equal(y, np.array([1.0, 2.0, 3.0])) @@ -67,10 +67,10 @@ def test_numpy_method_preserves_traced_array_after_compile(self): """Test TracedTensor.numpy() preserves an inactive traced carrier.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - ctx.compile_trace({'output': x * 2}) - + ctx.compile_trace({"output": x * 2}) + y = x.numpy() - + self.assertIsInstance(y, TracedNpArray) self.assertFalse(y.is_tracing) self.assertIs(y.context_obj, ctx) @@ -79,9 +79,9 @@ def test_np_array_returns_traced_np_array(self): """Test np.array(TracedTensor) returns TracedNpArray when patched.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - + y = np.array(x) - + self.assertIsInstance(y, TracedNpArray) np.testing.assert_array_almost_equal(y, np.array([1.0, 2.0, 3.0])) @@ -89,19 +89,21 @@ def test_np_asarray_returns_traced_np_array(self): """Test np.asarray(TracedTensor) returns TracedNpArray when patched.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - + y = np.asarray(x) - + self.assertIsInstance(y, TracedNpArray) np.testing.assert_array_almost_equal(y, np.array([1.0, 2.0, 3.0])) def test_np_array_with_dtype(self): """Test np.array(TracedTensor, dtype=...) preserves tracing when patched.""" ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64), name="x") - + x = ctx.create_input( + torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64), name="x" + ) + y = np.array(x, dtype=np.float32) - + self.assertIsInstance(y, TracedNpArray) self.assertEqual(y.dtype, np.float32) @@ -121,9 +123,9 @@ def test_torch_from_numpy_returns_traced_tensor(self): """Test torch.from_numpy(TracedNpArray) returns TracedTensor when tracing.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(np.array([1.0, 2.0, 3.0], dtype=np.float32), name="x") - + y = torch.from_numpy(x) - + self.assertIsInstance(y, TracedTensor) self.assertTrue(torch.allclose(y.tensor, torch.tensor([1.0, 2.0, 3.0]))) @@ -131,9 +133,9 @@ def test_torch_tensor_returns_traced_tensor(self): """Test torch.tensor(TracedNpArray) returns TracedTensor when tracing.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(np.array([1.0, 2.0, 3.0], dtype=np.float32), name="x") - + y = torch.tensor(x) - + self.assertIsInstance(y, TracedTensor) self.assertTrue(torch.allclose(y.tensor, torch.tensor([1.0, 2.0, 3.0]))) @@ -141,9 +143,9 @@ def test_torch_as_tensor_returns_traced_tensor(self): """Test torch.as_tensor(TracedNpArray) returns TracedTensor when tracing.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(np.array([1.0, 2.0, 3.0], dtype=np.float32), name="x") - + y = torch.as_tensor(x) - + self.assertIsInstance(y, TracedTensor) self.assertTrue(torch.allclose(y.tensor, torch.tensor([1.0, 2.0, 3.0]))) @@ -151,8 +153,8 @@ def test_torch_from_numpy_preserves_traced_tensor_after_compile(self): """Test torch.from_numpy preserves an inactive traced carrier.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(np.array([1.0, 2.0, 3.0], dtype=np.float32), name="x") - ctx.compile_trace({'output': x * 2}) - + ctx.compile_trace({"output": x * 2}) + # After compile, x is inactive but still carries its source context. y = torch.from_numpy(x) @@ -163,7 +165,7 @@ def test_torch_from_numpy_preserves_traced_tensor_after_compile(self): class TestConversionFXGraphCompilation(unittest.TestCase): """Test that conversions work correctly in compiled FX graphs. - + All tests use torch.Tensor as input and output for the compiled graph. """ @@ -177,123 +179,123 @@ def tearDown(self): def test_tensor_to_numpy_and_back(self): """Test: tensor → numpy → operations → tensor → output. - + Verifies that converting to numpy and back preserves tracing. """ ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - + # Convert to numpy, do operation, convert back np_x = x.numpy() np_y = np_x * 2 y = torch.from_numpy(np_y) - + # Compile - ctx.compile_trace({'output': y}) + ctx.compile_trace({"output": y}) graph_module = ctx.m - + # Test with torch tensor input input_tensor = torch.tensor([4.0, 5.0, 6.0]) expected = input_tensor * 2 output = graph_module(input_tensor) - + self.assertIsInstance(output, torch.Tensor) self.assertTrue(torch.allclose(output, expected)) def test_numpy_operations_in_trace(self): """Test: tensor → numpy → numpy ops → tensor → output. - + Verifies numpy operations are recorded correctly. Note: Use keepdims=True for reductions to preserve TracedNpArray (scalar reductions return numpy scalars, losing tracing). """ ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0, 4.0]), name="x") - + # Convert to numpy and do numpy-specific operations np_x = x.numpy() np_y = np.sum(np_x, keepdims=True) # Reduction with keepdims to preserve array y = torch.as_tensor(np_y) - + # Compile - ctx.compile_trace({'output': y}) + ctx.compile_trace({"output": y}) graph_module = ctx.m - + # Test input_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0]) expected = torch.tensor([10.0]) # sum with keepdims output = graph_module(input_tensor) - + self.assertTrue(torch.allclose(output, expected)) def test_numpy_input_to_tensor_output(self): """Test: numpy input → operations → tensor output. - + Graph takes torch tensor, internally uses numpy ops. """ ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - + # Do numpy operation then convert to tensor np_y = x + 10 y = torch.from_numpy(np_y) - + # Compile - ctx.compile_trace({'output': y}) + ctx.compile_trace({"output": y}) graph_module = ctx.m - + # Test with torch tensor input input_tensor = torch.tensor([4.0, 5.0, 6.0]) expected = input_tensor + 10 output = graph_module(input_tensor) - + self.assertIsInstance(output, torch.Tensor) self.assertTrue(torch.allclose(output, expected)) def test_chained_conversions(self): """Test: tensor → numpy → tensor → numpy → tensor. - + Multiple conversions should preserve tracing throughout. """ ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - + # Chain: tensor → numpy → tensor → numpy → tensor np1 = x.numpy() t1 = torch.from_numpy(np1 * 2) np2 = t1.numpy() y = torch.as_tensor(np2 + 1) - + # Compile - ctx.compile_trace({'output': y}) + ctx.compile_trace({"output": y}) graph_module = ctx.m - + # Test input_tensor = torch.tensor([4.0, 5.0, 6.0]) expected = (input_tensor * 2) + 1 output = graph_module(input_tensor) - + self.assertTrue(torch.allclose(output, expected)) def test_np_array_conversion_in_trace(self): """Test np.array() conversion within a trace (with patching).""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - + # Use np.array() for conversion (patched to preserve tracing) np_x = np.array(x) np_y = np_x * 3 y = torch.from_numpy(np_y) - + # Compile - ctx.compile_trace({'output': y}) + ctx.compile_trace({"output": y}) graph_module = ctx.m - + # Test input_tensor = torch.tensor([2.0, 3.0, 4.0]) expected = input_tensor * 3 output = graph_module(input_tensor) - + self.assertTrue(torch.allclose(output, expected)) @@ -311,11 +313,13 @@ def tearDown(self): def test_numpy_with_dtype_change(self): """Test conversion with dtype change preserves tracing.""" ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64), name="x") - + x = ctx.create_input( + torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64), name="x" + ) + # Convert to float32 numpy using np.array with patching np_x = np.array(x, dtype=np.float32) - + self.assertIsInstance(np_x, TracedNpArray) self.assertEqual(np_x.dtype, np.float32) @@ -323,17 +327,17 @@ def test_torch_tensor_with_dtype(self): """Test torch.tensor with dtype preserves tracing.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(np.array([1.0, 2.0, 3.0], dtype=np.float64), name="x") - + # Convert to float32 tensor y = torch.tensor(x, dtype=torch.float32) - + self.assertIsInstance(y, TracedTensor) self.assertEqual(y.dtype, torch.float32) class TestTracedTensorIdentityPreservation(unittest.TestCase): """Test that torch functions preserve TracedTensor identity/type correctly. - + Tests verify that: - torch.as_tensor returns same TracedTensor (no copy) - torch.tensor returns new TracedTensor with same proxy (copy semantics) @@ -351,23 +355,23 @@ def test_torch_as_tensor_preserves_identity(self): """Test torch.as_tensor returns same TracedTensor object.""" ctx = TracedTensorNode(name="test", node_index=0) traced = ctx.create_input(torch.randn(2, 3), name="x") - + result = torch.as_tensor(traced) - + self.assertIsInstance(result, TracedTensor) self.assertIs(result, traced) # Same object (no copy) def test_torch_tensor_creates_new_traced_tensor(self): """Test torch.tensor returns new TracedTensor with same proxy. - + torch.tensor() always creates a copy, so we get a NEW TracedTensor (not the same object), but it preserves the proxy/context for tracing. """ ctx = TracedTensorNode(name="test", node_index=0) traced = ctx.create_input(torch.randn(2, 3), name="x") - + result = torch.tensor(traced) - + self.assertIsInstance(result, TracedTensor) self.assertIsNot(result, traced) # Different object (copy semantics) self.assertIs(result.proxy, traced.proxy) # Same proxy (tracing preserved) @@ -422,7 +426,9 @@ def tearDown(self): def test_wp_to_torch_returns_traced_tensor(self): """Test wp.to_torch(TracedWpArray) returns TracedTensor when tracing.""" ctx = TracedTensorNode(name="test", node_index=0) - x_torch = ctx.create_input(torch.tensor([1.0, 2.0, 3.0], device="cuda"), name="x") + x_torch = ctx.create_input( + torch.tensor([1.0, 2.0, 3.0], device="cuda"), name="x" + ) x = wp.from_torch(x_torch) y = wp.to_torch(x) @@ -470,14 +476,18 @@ def tearDown(self): def test_wp_array_numpy_returns_traced_np_array(self): """Test TracedWpArray.numpy() returns TracedNpArray when tracing.""" ctx = TracedTensorNode(name="test", node_index=0) - x_torch = ctx.create_input(torch.tensor([1.0, 2.0, 3.0], device="cuda"), name="x") + x_torch = ctx.create_input( + torch.tensor([1.0, 2.0, 3.0], device="cuda"), name="x" + ) x = wp.from_torch(x_torch) y = x.numpy() self.assertIsInstance(y, TracedNpArray) self.assertIs(y.proxy, x.proxy) - np.testing.assert_array_almost_equal(y, np.array([1.0, 2.0, 3.0], dtype=np.float32)) + np.testing.assert_array_almost_equal( + y, np.array([1.0, 2.0, 3.0], dtype=np.float32) + ) class TestPatchingBehavior(unittest.TestCase): @@ -504,6 +514,5 @@ def test_patches_can_be_toggled(self): self.assertFalse(patcher.installed) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() - diff --git a/tests/unit_tests/test_export_inference_mode.py b/tests/unit_tests/test_export_inference_mode.py index fbecc883..a38ca467 100644 --- a/tests/unit_tests/test_export_inference_mode.py +++ b/tests/unit_tests/test_export_inference_mode.py @@ -18,16 +18,22 @@ class TestExportInferenceMode(unittest.TestCase): def test_prepare_inputs_escapes_inference_mode(self): for ambient_mode in (False, True): for original_clone in (False, True): - with self.subTest(ambient_mode=ambient_mode, original_clone=original_clone): + with self.subTest( + ambient_mode=ambient_mode, original_clone=original_clone + ): with torch.inference_mode(): - tensor = torch.tensor([1., 2.]) + tensor = torch.tensor([1.0, 2.0]) if original_clone: tensor.original_clone = tensor.clone marker = object() with torch.inference_mode(ambient_mode): - prepared, unchanged = prepare_tensors_for_export((tensor, marker)) + prepared, unchanged = prepare_tensors_for_export( + (tensor, marker) + ) self.assertFalse(prepared.is_inference()) - self.assertEqual(torch.is_inference_mode_enabled(), ambient_mode) + self.assertEqual( + torch.is_inference_mode_enabled(), ambient_mode + ) self.assertTrue(tensor.is_inference()) self.assertNotEqual(prepared.data_ptr(), tensor.data_ptr()) torch.testing.assert_close(prepared, tensor) @@ -41,14 +47,18 @@ def test_compile_reports_inference_tensor_errors(self): ) with tempfile.TemporaryDirectory() as tmpdir: _get_logger().configure(tmpdir, verbose=False) - for backend, target in (("onnx-dynamo", "torch.onnx.export"), - ("pt2", "torch.export.export")): + for backend, target in ( + ("onnx-dynamo", "torch.onnx.export"), + ("pt2", "torch.export.export"), + ): node = LeappNode("inference_test") node.setup_backend(backend, {}) node.m = torch.nn.Identity() for message in messages: for chain in (None, "__cause__", "__context__"): - with self.subTest(backend=backend, message=message, chain=chain): + with self.subTest( + backend=backend, message=message, chain=chain + ): original = RuntimeError(message) error = original if chain: @@ -70,5 +80,10 @@ def test_compile_reports_inference_tensor_errors(self): self.assertIn("FATAL:leapp:", logs.output[0]) self.assertIn(diagnostic, logs.output[0]) else: - self.assertEqual(diagnostic, f"Error compiling model {node.name}: {error}") - self.assertNotIn("torch.inference_mode()", logs.output[0]) + self.assertEqual( + diagnostic, + f"Error compiling model {node.name}: {error}", + ) + self.assertNotIn( + "torch.inference_mode()", logs.output[0] + ) diff --git a/tests/unit_tests/test_module_builder.py b/tests/unit_tests/test_module_builder.py index 3dd6929e..a960be62 100644 --- a/tests/unit_tests/test_module_builder.py +++ b/tests/unit_tests/test_module_builder.py @@ -28,10 +28,10 @@ def test_simple_list_of_tensors(self): t2 = torch.tensor([4.0, 5.0, 6.0]) data = [t1, t2] io_desc, param_fmt = describe_io("inputA", "inputA", data) - + # Verify io_descriptions - self.assertEqual([d.name for d in io_desc], ['inputA_0', 'inputA_1']) - + self.assertEqual([d.name for d in io_desc], ["inputA_0", "inputA_1"]) + # Verify packing expression expected = "inputA = [inputA_0, inputA_1]" self.assertEqual(param_fmt.packed_tensor_expr, expected) @@ -40,67 +40,80 @@ def test_single_tensor_trivial(self): """TEST 2: Single tensor (trivial - no packing needed)""" data = torch.tensor([1.0, 2.0, 3.0]) io_desc, param_fmt = describe_io("single", "single", data) - + # Verify io_descriptions - self.assertEqual([d.name for d in io_desc], ['single']) - + self.assertEqual([d.name for d in io_desc], ["single"]) + # Verify packing expression is empty (trivial case) self.assertEqual(param_fmt.packed_tensor_expr, "") - - - def test_dict_of_tensors(self): """TEST 5: Dict of tensors""" data = {"pose": torch.tensor([1.0, 2.0]), "velocity": torch.tensor([3.0, 4.0])} io_desc, param_fmt = describe_io("state", "state", data) - + # Verify io_descriptions - self.assertEqual([d.name for d in io_desc], ['state_pose', 'state_velocity']) - + self.assertEqual([d.name for d in io_desc], ["state_pose", "state_velocity"]) + # Verify packing expression expected = 'state = {"pose": state_pose, "velocity": state_velocity}' self.assertEqual(param_fmt.packed_tensor_expr, expected) def test_nested_list_of_lists(self): """TEST 6: Nested list of lists""" - data = [[torch.tensor([1.0]), torch.tensor([2.0])], - [torch.tensor([3.0]), torch.tensor([4.0])]] + data = [ + [torch.tensor([1.0]), torch.tensor([2.0])], + [torch.tensor([3.0]), torch.tensor([4.0])], + ] io_desc, param_fmt = describe_io("nested", "nested", data) - + # Verify io_descriptions - self.assertEqual([d.name for d in io_desc], - ['nested_0_0', 'nested_0_1', 'nested_1_0', 'nested_1_1']) - + self.assertEqual( + [d.name for d in io_desc], + ["nested_0_0", "nested_0_1", "nested_1_0", "nested_1_1"], + ) + # Verify packing expression expected = "nested = [[nested_0_0, nested_0_1], [nested_1_0, nested_1_1]]" self.assertEqual(param_fmt.packed_tensor_expr, expected) def test_dict_containing_list(self): """TEST 7: Dict containing list""" - data = {"positions": [torch.tensor([1.0]), torch.tensor([2.0])], - "velocities": [torch.tensor([3.0]), torch.tensor([4.0])]} + data = { + "positions": [torch.tensor([1.0]), torch.tensor([2.0])], + "velocities": [torch.tensor([3.0]), torch.tensor([4.0])], + } io_desc, param_fmt = describe_io("physics", "physics", data) - + # Verify io_descriptions - self.assertEqual([d.name for d in io_desc], - ['physics_positions_0', 'physics_positions_1', - 'physics_velocities_0', 'physics_velocities_1']) - + self.assertEqual( + [d.name for d in io_desc], + [ + "physics_positions_0", + "physics_positions_1", + "physics_velocities_0", + "physics_velocities_1", + ], + ) + # Verify packing expression expected = 'physics = {"positions": [physics_positions_0, physics_positions_1], "velocities": [physics_velocities_0, physics_velocities_1]}' self.assertEqual(param_fmt.packed_tensor_expr, expected) def test_list_of_dicts(self): """TEST 8: List of dicts""" - data = [{"x": torch.tensor([1.0]), "y": torch.tensor([2.0])}, - {"x": torch.tensor([3.0]), "y": torch.tensor([4.0])}] + data = [ + {"x": torch.tensor([1.0]), "y": torch.tensor([2.0])}, + {"x": torch.tensor([3.0]), "y": torch.tensor([4.0])}, + ] io_desc, param_fmt = describe_io("points", "points", data) - + # Verify io_descriptions - self.assertEqual([d.name for d in io_desc], - ['points_0_x', 'points_0_y', 'points_1_x', 'points_1_y']) - + self.assertEqual( + [d.name for d in io_desc], + ["points_0_x", "points_0_y", "points_1_x", "points_1_y"], + ) + # Verify packing expression expected = 'points = [{"x": points_0_x, "y": points_0_y}, {"x": points_1_x, "y": points_1_y}]' self.assertEqual(param_fmt.packed_tensor_expr, expected) @@ -112,11 +125,13 @@ def test_multiple_parameters_with_mixed_transformations(self): io_desc1, param_fmt1 = describe_io("self_data", "self.data", data1) io_desc2, param_fmt2 = describe_io("config", "config", data2) combined_io = io_desc1 + io_desc2 - + # Verify io_descriptions - self.assertEqual([d.name for d in combined_io], - ['self_data_0', 'self_data_1', 'config_a', 'config_b']) - + self.assertEqual( + [d.name for d in combined_io], + ["self_data_0", "self_data_1", "config_a", "config_b"], + ) + # Verify packing expressions for each parameter (uses name_str) expected1 = "self_data = [self_data_0, self_data_1]" expected2 = 'config = {"a": config_a, "b": config_b}' @@ -127,11 +142,12 @@ def test_deeply_nested_structure(self): """TEST 10: Deeply nested structure""" data = {"level1": {"level2": [torch.tensor([1.0]), torch.tensor([2.0])]}} io_desc, param_fmt = describe_io("deep", "deep", data) - + # Verify io_descriptions - self.assertEqual([d.name for d in io_desc], - ['deep_level1_level2_0', 'deep_level1_level2_1']) - + self.assertEqual( + [d.name for d in io_desc], ["deep_level1_level2_0", "deep_level1_level2_1"] + ) + # Verify packing expression expected = 'deep = {"level1": {"level2": [deep_level1_level2_0, deep_level1_level2_1]}}' self.assertEqual(param_fmt.packed_tensor_expr, expected) @@ -146,7 +162,7 @@ def test_simple_list_of_tensors(self): t2 = torch.tensor([4.0, 5.0, 6.0]) data = [t1, t2] io_desc, param_fmt = describe_io("inputA", "inputA", data) - + # Verify unpacking expression expected = "inputA_0 = inputA[0]\ninputA_1 = inputA[1]" self.assertEqual(param_fmt.unpacked_tensor_expr, expected) @@ -155,58 +171,68 @@ def test_single_tensor_trivial(self): """TEST 2: Single tensor (trivial - no unpacking needed)""" data = torch.tensor([1.0, 2.0, 3.0]) io_desc, param_fmt = describe_io("single", "single", data) - + # Verify unpacking expression is empty (trivial case) self.assertEqual(param_fmt.unpacked_tensor_expr, "") - - def test_dict_of_tensors(self): """TEST 5: Dict of tensors""" data = {"pose": torch.tensor([1.0, 2.0]), "velocity": torch.tensor([3.0, 4.0])} io_desc, param_fmt = describe_io("state", "state", data) - + # Verify unpacking expression expected = 'state_pose = state["pose"]\nstate_velocity = state["velocity"]' self.assertEqual(param_fmt.unpacked_tensor_expr, expected) def test_nested_list_of_lists(self): """TEST 6: Nested list of lists""" - data = [[torch.tensor([1.0]), torch.tensor([2.0])], - [torch.tensor([3.0]), torch.tensor([4.0])]] + data = [ + [torch.tensor([1.0]), torch.tensor([2.0])], + [torch.tensor([3.0]), torch.tensor([4.0])], + ] io_desc, param_fmt = describe_io("nested", "nested", data) - + # Verify unpacking expression - expected = ("nested_0_0 = nested[0][0]\n" - "nested_0_1 = nested[0][1]\n" - "nested_1_0 = nested[1][0]\n" - "nested_1_1 = nested[1][1]") + expected = ( + "nested_0_0 = nested[0][0]\n" + "nested_0_1 = nested[0][1]\n" + "nested_1_0 = nested[1][0]\n" + "nested_1_1 = nested[1][1]" + ) self.assertEqual(param_fmt.unpacked_tensor_expr, expected) def test_dict_containing_list(self): """TEST 7: Dict containing list""" - data = {"positions": [torch.tensor([1.0]), torch.tensor([2.0])], - "velocities": [torch.tensor([3.0]), torch.tensor([4.0])]} + data = { + "positions": [torch.tensor([1.0]), torch.tensor([2.0])], + "velocities": [torch.tensor([3.0]), torch.tensor([4.0])], + } io_desc, param_fmt = describe_io("physics", "physics", data) - + # Verify unpacking expression - expected = ('physics_positions_0 = physics["positions"][0]\n' - 'physics_positions_1 = physics["positions"][1]\n' - 'physics_velocities_0 = physics["velocities"][0]\n' - 'physics_velocities_1 = physics["velocities"][1]') + expected = ( + 'physics_positions_0 = physics["positions"][0]\n' + 'physics_positions_1 = physics["positions"][1]\n' + 'physics_velocities_0 = physics["velocities"][0]\n' + 'physics_velocities_1 = physics["velocities"][1]' + ) self.assertEqual(param_fmt.unpacked_tensor_expr, expected) def test_list_of_dicts(self): """TEST 8: List of dicts""" - data = [{"x": torch.tensor([1.0]), "y": torch.tensor([2.0])}, - {"x": torch.tensor([3.0]), "y": torch.tensor([4.0])}] + data = [ + {"x": torch.tensor([1.0]), "y": torch.tensor([2.0])}, + {"x": torch.tensor([3.0]), "y": torch.tensor([4.0])}, + ] io_desc, param_fmt = describe_io("points", "points", data) - + # Verify unpacking expression - expected = ('points_0_x = points[0]["x"]\n' - 'points_0_y = points[0]["y"]\n' - 'points_1_x = points[1]["x"]\n' - 'points_1_y = points[1]["y"]') + expected = ( + 'points_0_x = points[0]["x"]\n' + 'points_0_y = points[0]["y"]\n' + 'points_1_x = points[1]["x"]\n' + 'points_1_y = points[1]["y"]' + ) self.assertEqual(param_fmt.unpacked_tensor_expr, expected) def test_multiple_parameters_with_mixed_transformations(self): @@ -215,7 +241,7 @@ def test_multiple_parameters_with_mixed_transformations(self): data2 = {"a": torch.tensor([3.0]), "b": torch.tensor([4.0])} io_desc1, param_fmt1 = describe_io("self_data", "self.data", data1) io_desc2, param_fmt2 = describe_io("config", "config", data2) - + # Verify unpacking expressions for each parameter (uses name_str) expected1 = "self_data_0 = self_data[0]\nself_data_1 = self_data[1]" expected2 = 'config_a = config["a"]\nconfig_b = config["b"]' @@ -226,13 +252,14 @@ def test_deeply_nested_structure(self): """TEST 10: Deeply nested structure""" data = {"level1": {"level2": [torch.tensor([1.0]), torch.tensor([2.0])]}} io_desc, param_fmt = describe_io("deep", "deep", data) - + # Verify unpacking expression - expected = ('deep_level1_level2_0 = deep["level1"]["level2"][0]\n' - 'deep_level1_level2_1 = deep["level1"]["level2"][1]') + expected = ( + 'deep_level1_level2_0 = deep["level1"]["level2"][0]\n' + 'deep_level1_level2_1 = deep["level1"]["level2"][1]' + ) self.assertEqual(param_fmt.unpacked_tensor_expr, expected) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main(verbosity=2) - diff --git a/tests/unit_tests/test_numpy_compatibility.py b/tests/unit_tests/test_numpy_compatibility.py index ceab16e1..cdd58494 100644 --- a/tests/unit_tests/test_numpy_compatibility.py +++ b/tests/unit_tests/test_numpy_compatibility.py @@ -26,7 +26,7 @@ def test_basic_numpy_function(self): tensor = torch.tensor([1.0, 2.0, 3.0]) leapp.start(name=self.TEST_GRAPH_NAME) - tensor = annotate.input_tensors('basic_numpy_function', {'tensor': tensor}) + tensor = annotate.input_tensors("basic_numpy_function", {"tensor": tensor}) # numpy operations numpy_val = tensor.numpy() @@ -34,16 +34,21 @@ def test_basic_numpy_function(self): numpy_val2 = numpy_val - 1 numpy_val3 = numpy_val1 * numpy_val2 numpy_val4 = numpy_val3 / 2 - torch_tensors = [torch.from_numpy(numpy_val) for numpy_val in [numpy_val1, numpy_val2, numpy_val3, numpy_val4]] - - annotate.output_tensors('basic_numpy_function', {'tensor': torch_tensors}, export_with="jit") + torch_tensors = [ + torch.from_numpy(numpy_val) + for numpy_val in [numpy_val1, numpy_val2, numpy_val3, numpy_val4] + ] + + annotate.output_tensors( + "basic_numpy_function", {"tensor": torch_tensors}, export_with="jit" + ) leapp.stop() leapp.compile_graph(visualize=False) class TestNumpyFunctionsWithPatchesApplied(unittest.TestCase): """Test that regular numpy functions work correctly when patches are applied. - + This verifies that torch.from_numpy, torch.as_tensor, and torch.tensor work correctly with regular numpy arrays between start() and stop(), even though the global patches are applied. @@ -59,92 +64,94 @@ def tearDown(self): def test_numpy_to_torch_between_start_stop(self): """Test torch.from_numpy works with regular numpy arrays during annotation.""" leapp.start(name=self.TEST_GRAPH_NAME) - + # Regular numpy arrays (not TracedTensors) np_array1 = np.array([1.0, 2.0, 3.0], dtype=np.float32) np_array2 = np.array([[1, 2], [3, 4]], dtype=np.float64) - + # These should work normally via the patched functions tensor1 = torch.from_numpy(np_array1) tensor2 = torch.from_numpy(np_array2) - + # Verify the conversions worked correctly self.assertIsInstance(tensor1, torch.Tensor) self.assertIsInstance(tensor2, torch.Tensor) self.assertEqual(tensor1.shape, torch.Size([3])) self.assertEqual(tensor2.shape, torch.Size([2, 2])) self.assertTrue(torch.allclose(tensor1, torch.tensor([1.0, 2.0, 3.0]))) - + leapp.stop() def test_as_tensor_between_start_stop(self): """Test torch.as_tensor works with regular numpy arrays during annotation.""" leapp.start(name=self.TEST_GRAPH_NAME) - + np_array = np.array([1.0, 2.0, 3.0], dtype=np.float32) - + # torch.as_tensor should work normally tensor = torch.as_tensor(np_array) - + self.assertIsInstance(tensor, torch.Tensor) self.assertEqual(tensor.shape, torch.Size([3])) self.assertTrue(torch.allclose(tensor, torch.tensor([1.0, 2.0, 3.0]))) - + # Test with dtype conversion tensor_int = torch.as_tensor(np_array, dtype=torch.int32) self.assertEqual(tensor_int.dtype, torch.int32) - + leapp.stop() def test_torch_tensor_between_start_stop(self): """Test torch.tensor works with regular numpy arrays during annotation.""" leapp.start(name=self.TEST_GRAPH_NAME) - + np_array = np.array([1.0, 2.0, 3.0], dtype=np.float32) - + # torch.tensor should work normally tensor = torch.tensor(np_array) - + self.assertIsInstance(tensor, torch.Tensor) self.assertEqual(tensor.shape, torch.Size([3])) self.assertTrue(torch.allclose(tensor, torch.tensor([1.0, 2.0, 3.0]))) - + # Test with dtype conversion tensor_int = torch.tensor(np_array, dtype=torch.int64) self.assertEqual(tensor_int.dtype, torch.int64) - + leapp.stop() def test_numpy_operations_between_start_stop(self): """Test that standard numpy operations work correctly during annotation.""" leapp.start(name=self.TEST_GRAPH_NAME) - + # Standard numpy operations should work unaffected arr1 = np.array([1.0, 2.0, 3.0]) arr2 = np.array([4.0, 5.0, 6.0]) - + result_add = np.add(arr1, arr2) result_mul = np.multiply(arr1, arr2) result_sin = np.sin(arr1) result_concat = np.concatenate([arr1, arr2]) - + # Verify numpy operations produced correct results np.testing.assert_array_almost_equal(result_add, [5.0, 7.0, 9.0]) np.testing.assert_array_almost_equal(result_mul, [4.0, 10.0, 18.0]) np.testing.assert_array_almost_equal(result_sin, np.sin([1.0, 2.0, 3.0])) - np.testing.assert_array_almost_equal(result_concat, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) - + np.testing.assert_array_almost_equal( + result_concat, [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + ) + # Convert final result to torch tensor = torch.from_numpy(result_concat) self.assertIsInstance(tensor, torch.Tensor) self.assertEqual(tensor.shape, torch.Size([6])) - + leapp.stop() class TestNumpyPatchingDisabled(unittest.TestCase): """Test that global_patching=False correctly disables numpy patching. - + When global_patching=False is passed to start(), TracedTensor.numpy() and np.array(traced_tensor) should return actual numpy arrays, not TracedTensors. """ @@ -159,68 +166,68 @@ def tearDown(self): def test_numpy_method_returns_array_when_patching_disabled(self): """Test that .numpy() returns actual numpy array when global_patching=False.""" from leapp.leapp_graph.datatypes.torch.traced_tensor import TracedTensor - + tensor = torch.tensor([1.0, 2.0, 3.0]) - + # Start with global_patching=False leapp.start(name=self.TEST_GRAPH_NAME, global_patching=False) - traced = annotate.input_tensors('test_node', {'tensor': tensor}) - + traced = annotate.input_tensors("test_node", {"tensor": tensor}) + # Verify we have a TracedTensor self.assertIsInstance(traced, TracedTensor) - + # .numpy() should return actual numpy array, not TracedTensor result = traced.numpy() - + self.assertIsInstance(result, np.ndarray) self.assertNotIsInstance(result, TracedTensor) np.testing.assert_array_almost_equal(result, [1.0, 2.0, 3.0]) - - annotate.output_tensors('test_node', {'tensor': traced}, export_with="jit") + + annotate.output_tensors("test_node", {"tensor": traced}, export_with="jit") leapp.stop() def test_np_array_returns_array_when_patching_disabled(self): """Test that np.array() returns actual numpy array when global_patching=False.""" from leapp.leapp_graph.datatypes.torch.traced_tensor import TracedTensor - + tensor = torch.tensor([4.0, 5.0, 6.0]) - + # Start with global_patching=False leapp.start(name=self.TEST_GRAPH_NAME, global_patching=False) - traced = annotate.input_tensors('test_node', {'tensor': tensor}) - + traced = annotate.input_tensors("test_node", {"tensor": tensor}) + # Verify we have a TracedTensor self.assertIsInstance(traced, TracedTensor) - + # np.array() should return actual numpy array result = np.array(traced) - + self.assertIsInstance(result, np.ndarray) self.assertNotIsInstance(result, TracedTensor) np.testing.assert_array_almost_equal(result, [4.0, 5.0, 6.0]) - - annotate.output_tensors('test_node', {'tensor': traced}, export_with="jit") + + annotate.output_tensors("test_node", {"tensor": traced}, export_with="jit") leapp.stop() def test_np_asarray_returns_array_when_patching_disabled(self): """Test that np.asarray() returns actual numpy array when global_patching=False.""" from leapp.leapp_graph.datatypes.torch.traced_tensor import TracedTensor - + tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) - + # Start with global_patching=False leapp.start(name=self.TEST_GRAPH_NAME, global_patching=False) - traced = annotate.input_tensors('test_node', {'tensor': tensor}) - + traced = annotate.input_tensors("test_node", {"tensor": tensor}) + # Verify we have a TracedTensor self.assertIsInstance(traced, TracedTensor) - + # np.asarray() should return actual numpy array result = np.asarray(traced) - + self.assertIsInstance(result, np.ndarray) self.assertNotIsInstance(result, TracedTensor) np.testing.assert_array_almost_equal(result, [[1.0, 2.0], [3.0, 4.0]]) - - annotate.output_tensors('test_node', {'tensor': traced}, export_with="jit") + + annotate.output_tensors("test_node", {"tensor": traced}, export_with="jit") leapp.stop() diff --git a/tests/unit_tests/test_torch_builtin_compatibility.py b/tests/unit_tests/test_torch_builtin_compatibility.py index a2d1701f..6efe41e4 100644 --- a/tests/unit_tests/test_torch_builtin_compatibility.py +++ b/tests/unit_tests/test_torch_builtin_compatibility.py @@ -33,26 +33,35 @@ from leapp.leapp_graph.traced_node import TracedTensorNode from leapp.leapp_graph.datatypes import TracedTensor -from tests.unit_tests.export_format_validation import verify_exported_program_on_random_inputs - +from tests.unit_tests.export_format_validation import ( + verify_exported_program_on_random_inputs, +) @pytest.mark.filterwarnings("ignore::torch.jit.TracerWarning") class TestFunctional(unittest.TestCase): """Test torch.nn.functional operations trace and compile correctly.""" - + # Number of random inputs to test after compiling NUM_TEST_INPUTS = 5 - def _run_trace_and_compile(self, func, input_shape, test_name, - skip_torchscript=False, skip_onnx=False, - skip_exported_program=False, - onnx_opset=17, atol=1e-5, **func_kwargs): + def _run_trace_and_compile( + self, + func, + input_shape, + test_name, + skip_torchscript=False, + skip_onnx=False, + skip_exported_program=False, + onnx_opset=17, + atol=1e-5, + **func_kwargs, + ): """Helper to trace a function and verify compiled output matches. - + Compiles once, then tests with NUM_TEST_INPUTS random inputs. Tests FX GraphModule, TorchScript, and ONNX exports. - + Args: func: Function to apply (e.g., F.relu, or partial(F.linear, weight=w)) input_shape: Shape of input tensor @@ -63,28 +72,28 @@ def _run_trace_and_compile(self, func, input_shape, test_name, **func_kwargs: Additional kwargs to pass to func """ input_tensor = torch.randn(*input_shape) - + ctx = TracedTensorNode(name=test_name, node_index=0) traced_input = ctx.create_input(input_tensor.clone(), name="x") - + output = func(traced_input, **func_kwargs) self.assertIsInstance(output, TracedTensor) - - ctx.compile_trace({'output': output}) + + ctx.compile_trace({"output": output}) graph_module = ctx.m - + # Test with multiple random inputs across all export formats for i in range(self.NUM_TEST_INPUTS): test_input = torch.randn(*input_shape) expected = func(test_input, **func_kwargs) - + # Test 1: FX GraphModule execution actual_fx = graph_module(test_input) self.assertTrue( torch.allclose(actual_fx, expected, atol=atol), - f"{test_name}: FX output mismatch on input {i+1}/{self.NUM_TEST_INPUTS}" + f"{test_name}: FX output mismatch on input {i + 1}/{self.NUM_TEST_INPUTS}", ) - + # Test 2: TorchScript export and execution (compile once, test with all inputs) if not skip_torchscript: try: @@ -95,11 +104,11 @@ def _run_trace_and_compile(self, func, input_shape, test_name, actual_ts = scripted(test_input) self.assertTrue( torch.allclose(actual_ts, expected, atol=atol), - f"{test_name}: TorchScript output mismatch on input {i+1}" + f"{test_name}: TorchScript output mismatch on input {i + 1}", ) except Exception as e: self.fail(f"{test_name}: TorchScript export/execution failed: {e}") - + # Test 3: ONNX export and execution if not skip_onnx: try: @@ -113,18 +122,20 @@ def _run_trace_and_compile(self, func, input_shape, test_name, dynamo=False, export_params=True, opset_version=onnx_opset, - input_names=['input'], - output_names=['output'], + input_names=["input"], + output_names=["output"], ) session = ort.InferenceSession(str(onnx_path)) - + for i in range(self.NUM_TEST_INPUTS): test_input = torch.randn(*input_shape) expected = func(test_input, **func_kwargs) - output_onnx = session.run(None, {"input": test_input.numpy()})[0] + output_onnx = session.run(None, {"input": test_input.numpy()})[ + 0 + ] self.assertTrue( np.allclose(output_onnx, expected.numpy(), atol=atol), - f"{test_name}: ONNX output mismatch on input {i+1}" + f"{test_name}: ONNX output mismatch on input {i + 1}", ) except Exception as e: self.fail(f"{test_name}: ONNX export/execution failed: {e}") @@ -145,7 +156,7 @@ def _run_trace_and_compile(self, func, input_shape, test_name, # ========================================================================= # Linear # ========================================================================= - + def test_linear(self): """Test F.linear with weight and bias.""" weight = torch.randn(16, 8) @@ -179,8 +190,7 @@ def test_leaky_relu_default(self): def test_leaky_relu_custom_slope(self): """Test F.leaky_relu with custom negative_slope.""" self._run_trace_and_compile( - F.leaky_relu, (4, 8), "test_leaky_relu_custom", - negative_slope=0.2 + F.leaky_relu, (4, 8), "test_leaky_relu_custom", negative_slope=0.2 ) def test_gelu(self): @@ -218,8 +228,7 @@ def test_hardtanh(self): def test_hardtanh_custom_range(self): """Test F.hardtanh with custom min/max values.""" self._run_trace_and_compile( - F.hardtanh, (4, 8), "test_hardtanh_custom", - min_val=-2.0, max_val=2.0 + F.hardtanh, (4, 8), "test_hardtanh_custom", min_val=-2.0, max_val=2.0 ) def test_hardswish(self): @@ -281,15 +290,13 @@ def test_log_softmax(self): def test_normalize_l2_dim1(self): """Test F.normalize with L2 norm along dim=1.""" self._run_trace_and_compile( - F.normalize, (4, 8), "test_normalize_l2", - p=2.0, dim=1 + F.normalize, (4, 8), "test_normalize_l2", p=2.0, dim=1 ) def test_normalize_l1(self): """Test F.normalize with L1 norm.""" self._run_trace_and_compile( - F.normalize, (4, 8), "test_normalize_l1", - p=1.0, dim=1 + F.normalize, (4, 8), "test_normalize_l1", p=1.0, dim=1 ) def test_layer_norm(self): @@ -297,7 +304,9 @@ def test_layer_norm(self): normalized_shape = [8] weight = torch.ones(8) bias = torch.zeros(8) - func = partial(F.layer_norm, normalized_shape=normalized_shape, weight=weight, bias=bias) + func = partial( + F.layer_norm, normalized_shape=normalized_shape, weight=weight, bias=bias + ) self._run_trace_and_compile(func, (4, 8), "test_layer_norm") def test_layer_norm_2d(self): @@ -305,12 +314,14 @@ def test_layer_norm_2d(self): normalized_shape = [8, 8] weight = torch.ones(8, 8) bias = torch.zeros(8, 8) - func = partial(F.layer_norm, normalized_shape=normalized_shape, weight=weight, bias=bias) + func = partial( + F.layer_norm, normalized_shape=normalized_shape, weight=weight, bias=bias + ) self._run_trace_and_compile(func, (4, 8, 8), "test_layer_norm_2d") def test_batch_norm(self): """Test F.batch_norm in eval mode (training=False). - + Uses fixed running_mean/running_var which are required when training=False. Running stats become constants captured in the trace. """ @@ -320,8 +331,12 @@ def test_batch_norm(self): weight = torch.ones(num_features) bias = torch.zeros(num_features) func = partial( - F.batch_norm, running_mean=running_mean, running_var=running_var, - weight=weight, bias=bias, training=False + F.batch_norm, + running_mean=running_mean, + running_var=running_var, + weight=weight, + bias=bias, + training=False, ) self._run_trace_and_compile(func, (2, 3, 8, 8), "test_batch_norm") @@ -333,14 +348,18 @@ def test_batch_norm_1d(self): weight = torch.ones(num_features) bias = torch.zeros(num_features) func = partial( - F.batch_norm, running_mean=running_mean, running_var=running_var, - weight=weight, bias=bias, training=False + F.batch_norm, + running_mean=running_mean, + running_var=running_var, + weight=weight, + bias=bias, + training=False, ) self._run_trace_and_compile(func, (4, 8), "test_batch_norm_1d") def test_instance_norm(self): """Test F.instance_norm. - + Uses default use_input_stats=True (computes stats from input). """ num_features = 3 @@ -357,8 +376,12 @@ def test_instance_norm_with_running_stats(self): weight = torch.ones(num_features) bias = torch.zeros(num_features) func = partial( - F.instance_norm, running_mean=running_mean, running_var=running_var, - weight=weight, bias=bias, use_input_stats=False + F.instance_norm, + running_mean=running_mean, + running_var=running_var, + weight=weight, + bias=bias, + use_input_stats=False, ) self._run_trace_and_compile(func, (2, 3, 8, 8), "test_instance_norm_running") @@ -383,8 +406,7 @@ def test_group_norm_no_affine(self): def test_dropout_eval(self): """Test F.dropout in eval mode (training=False) - should be identity.""" self._run_trace_and_compile( - F.dropout, (4, 8), "test_dropout_eval", - p=0.5, training=False + F.dropout, (4, 8), "test_dropout_eval", p=0.5, training=False ) # ========================================================================= @@ -446,8 +468,12 @@ def test_conv_transpose2d_with_stride(self): """Test F.conv_transpose2d with stride for upsampling.""" weight = torch.randn(3, 16, 4, 4) bias = torch.randn(16) - func = partial(F.conv_transpose2d, weight=weight, bias=bias, stride=2, padding=1) - self._run_trace_and_compile(func, (2, 3, 16, 16), "test_conv_transpose2d_stride") + func = partial( + F.conv_transpose2d, weight=weight, bias=bias, stride=2, padding=1 + ) + self._run_trace_and_compile( + func, (2, 3, 16, 16), "test_conv_transpose2d_stride" + ) # ========================================================================= # Pooling @@ -519,24 +545,30 @@ def test_adaptive_max_pool2d(self): def test_interpolate_nearest(self): """Test F.interpolate with nearest neighbor mode.""" - func = partial(F.interpolate, size=(32, 32), mode='nearest') + func = partial(F.interpolate, size=(32, 32), mode="nearest") self._run_trace_and_compile(func, (2, 3, 16, 16), "test_interpolate_nearest") def test_interpolate_bilinear(self): """Test F.interpolate with bilinear mode.""" - func = partial(F.interpolate, size=(32, 32), mode='bilinear', align_corners=False) + func = partial( + F.interpolate, size=(32, 32), mode="bilinear", align_corners=False + ) self._run_trace_and_compile(func, (2, 3, 16, 16), "test_interpolate_bilinear") def test_interpolate_scale_factor(self): """Test F.interpolate with scale_factor.""" # Note: scale_factor must be float (2.0) not int (2) for TorchScript compatibility - func = partial(F.interpolate, scale_factor=2.0, mode='nearest') + func = partial(F.interpolate, scale_factor=2.0, mode="nearest") self._run_trace_and_compile(func, (2, 3, 16, 16), "test_interpolate_scale") def test_interpolate_trilinear(self): """Test F.interpolate with trilinear mode for 3D.""" - func = partial(F.interpolate, size=(16, 32, 32), mode='trilinear', align_corners=False) - self._run_trace_and_compile(func, (2, 3, 8, 16, 16), "test_interpolate_trilinear") + func = partial( + F.interpolate, size=(16, 32, 32), mode="trilinear", align_corners=False + ) + self._run_trace_and_compile( + func, (2, 3, 8, 16, 16), "test_interpolate_trilinear" + ) # ========================================================================= # Padding @@ -545,23 +577,23 @@ def test_interpolate_trilinear(self): def test_pad_constant(self): """Test F.pad with constant padding.""" # Note: value must be float (0.0) not int (0) for TorchScript compatibility - func = partial(F.pad, pad=(1, 1, 1, 1), mode='constant', value=0.0) + func = partial(F.pad, pad=(1, 1, 1, 1), mode="constant", value=0.0) self._run_trace_and_compile(func, (2, 3, 16, 16), "test_pad_constant") def test_pad_reflect(self): """Test F.pad with reflect padding.""" - func = partial(F.pad, pad=(1, 1, 1, 1), mode='reflect') + func = partial(F.pad, pad=(1, 1, 1, 1), mode="reflect") self._run_trace_and_compile(func, (2, 3, 16, 16), "test_pad_reflect") def test_pad_replicate(self): """Test F.pad with replicate padding.""" - func = partial(F.pad, pad=(1, 1, 1, 1), mode='replicate') + func = partial(F.pad, pad=(1, 1, 1, 1), mode="replicate") self._run_trace_and_compile(func, (2, 3, 16, 16), "test_pad_replicate") def test_pad_1d(self): """Test F.pad on 1D input.""" # Note: value must be float (0.0) not int (0) for TorchScript compatibility - func = partial(F.pad, pad=(2, 2), mode='constant', value=0.0) + func = partial(F.pad, pad=(2, 2), mode="constant", value=0.0) self._run_trace_and_compile(func, (2, 3, 32), "test_pad_1d") # ========================================================================= @@ -581,7 +613,7 @@ def test_unfold_with_params(self): def test_fold(self): """Test F.fold (col2im) - combines sliding local blocks into a tensor. - + Fold is the inverse of unfold. Input shape must match the output of an equivalent unfold operation. Note: aten::col2im requires ONNX opset 18+. @@ -593,18 +625,22 @@ def test_fold(self): def test_unfold_fold_roundtrip(self): """Test unfold → fold roundtrip (with divisor correction). - + Unfold then fold doesn't produce identity without normalizing by the overlap count. This test verifies the ops compose correctly. Note: aten::col2im (fold) requires ONNX opset 18+. """ kernel_size = 3 + def unfold_fold(x): # x: (N, C, H, W) unfolded = F.unfold(x, kernel_size=kernel_size) folded = F.fold(unfolded, output_size=(8, 8), kernel_size=kernel_size) return folded - self._run_trace_and_compile(unfold_fold, (2, 3, 8, 8), "test_unfold_fold", onnx_opset=18) + + self._run_trace_and_compile( + unfold_fold, (2, 3, 8, 8), "test_unfold_fold", onnx_opset=18 + ) # ========================================================================= # Loss Functions @@ -626,9 +662,11 @@ def test_cross_entropy(self): def test_nll_loss(self): """Test F.nll_loss (negative log likelihood).""" target = torch.randint(0, 8, (4,)) + # NLL loss expects log-probabilities, so wrap with log_softmax def nll_with_logsoftmax(x, target): return F.nll_loss(F.log_softmax(x, dim=-1), target) + func = partial(nll_with_logsoftmax, target=target) self._run_trace_and_compile(func, (4, 8), "test_nll_loss") @@ -637,13 +675,14 @@ def test_binary_cross_entropy_with_logits(self): target = torch.rand(4, 8) # Binary targets in [0, 1] func = partial(F.binary_cross_entropy_with_logits, target=target) self._run_trace_and_compile(func, (4, 8), "test_bce_logits") + # ========================================================================= # Embedding # ========================================================================= def test_embedding(self): """Test F.embedding (lookup table). - + Note: Uses a custom test flow because embedding requires integer (LongTensor) inputs, not the float tensors generated by _run_trace_and_compile. """ @@ -652,23 +691,23 @@ def test_embedding(self): weight = torch.randn(vocab_size, embed_dim) input_indices = torch.randint(0, vocab_size, (4, 8)) - + ctx = TracedTensorNode(name="test_embedding", node_index=0) traced_input = ctx.create_input(input_indices.clone(), name="x") - + output = F.embedding(traced_input, weight) self.assertIsInstance(output, TracedTensor) - - ctx.compile_trace({'output': output}) + + ctx.compile_trace({"output": output}) graph_module = ctx.m - + for i in range(self.NUM_TEST_INPUTS): test_input = torch.randint(0, vocab_size, (4, 8)) expected = F.embedding(test_input, weight) actual_fx = graph_module(test_input) self.assertTrue( torch.allclose(actual_fx, expected, atol=1e-5), - f"test_embedding: FX output mismatch on input {i+1}/{self.NUM_TEST_INPUTS}" + f"test_embedding: FX output mismatch on input {i + 1}/{self.NUM_TEST_INPUTS}", ) # ========================================================================= @@ -677,7 +716,7 @@ def test_embedding(self): def test_scaled_dot_product_attention(self): """Test F.scaled_dot_product_attention (PyTorch 2.0+). - + Only query is traced; key and value are fixed via partial. Note: skip_onnx=True because SDPA may not be supported in ONNX opset 17. """ @@ -685,44 +724,46 @@ def test_scaled_dot_product_attention(self): key = torch.randn(2, 4, 8, 16) value = torch.randn(2, 4, 8, 16) func = partial(F.scaled_dot_product_attention, key=key, value=value) - self._run_trace_and_compile( - func, (2, 4, 8, 16), "test_sdpa", - skip_onnx=True - ) + self._run_trace_and_compile(func, (2, 4, 8, 16), "test_sdpa", skip_onnx=True) def test_multi_head_attention(self): """Test F.multi_head_attention_forward (self-attention). - + Wraps the forward function to return only the output tensor (not attention weights). Note: skip_onnx=True because MHA decomposition may not be supported in ONNX opset 17. """ embed_dim = 32 num_heads = 4 - + in_proj_weight = torch.randn(3 * embed_dim, embed_dim) in_proj_bias = torch.randn(3 * embed_dim) out_proj_weight = torch.randn(embed_dim, embed_dim) out_proj_bias = torch.randn(embed_dim) - + def mha_self_attention(query): # Input shape: (seq_len, batch_size, embed_dim) — seq_len first output, _ = F.multi_head_attention_forward( - query, query, query, # self-attention: Q=K=V - embed_dim, num_heads, - in_proj_weight, in_proj_bias, - None, None, # bias_k, bias_v + query, + query, + query, # self-attention: Q=K=V + embed_dim, + num_heads, + in_proj_weight, + in_proj_bias, + None, + None, # bias_k, bias_v False, # add_zero_attn 0.0, # dropout_p - out_proj_weight, out_proj_bias, + out_proj_weight, + out_proj_bias, training=False, need_weights=False, ) return output - + # Shape: (seq_len, batch_size, embed_dim) self._run_trace_and_compile( - mha_self_attention, (8, 2, embed_dim), "test_mha", - skip_onnx=True + mha_self_attention, (8, 2, embed_dim), "test_mha", skip_onnx=True ) # ========================================================================= @@ -731,7 +772,7 @@ def mha_self_attention(query): def test_grid_sample(self): """Test F.grid_sample. - + Grid values are in [-1, 1] range for normalized coordinates. """ # grid shape: (N, H_out, W_out, 2) @@ -742,13 +783,18 @@ def test_grid_sample(self): def test_grid_sample_bilinear(self): """Test F.grid_sample with bilinear interpolation (default).""" grid = torch.rand(2, 8, 8, 2) * 2 - 1 - func = partial(F.grid_sample, grid=grid, mode='bilinear', - padding_mode='zeros', align_corners=False) + func = partial( + F.grid_sample, + grid=grid, + mode="bilinear", + padding_mode="zeros", + align_corners=False, + ) self._run_trace_and_compile(func, (2, 3, 16, 16), "test_grid_sample_bilinear") def test_affine_grid(self): """Test F.affine_grid. - + Generates a sampling grid from an affine transformation matrix (theta). Input theta shape: (N, 2, 3) for 2D spatial transforms. Note: aten::affine_grid_generator requires ONNX opset 20+. @@ -758,16 +804,22 @@ def test_affine_grid(self): def test_affine_grid_and_grid_sample(self): """Test F.affine_grid → F.grid_sample pipeline (common in STN). - + Note: aten::affine_grid_generator requires ONNX opset 20+. """ source = torch.randn(2, 3, 16, 16) - + def stn_transform(theta): grid = F.affine_grid(theta, source.size(), align_corners=True) return F.grid_sample(source, grid, align_corners=True) - - self._run_trace_and_compile(stn_transform, (2, 2, 3), "test_affine_grid_sample", onnx_opset=20, atol=1e-4) + + self._run_trace_and_compile( + stn_transform, + (2, 2, 3), + "test_affine_grid_sample", + onnx_opset=20, + atol=1e-4, + ) # ========================================================================= # Torch Operations (not F.*) @@ -783,24 +835,30 @@ def test_bmm(self): def test_einsum_matmul(self): """Test torch.einsum for matrix multiplication.""" other = torch.randn(8, 16) + def einsum_matmul(x, other): - return torch.einsum('bi,ij->bj', x, other) + return torch.einsum("bi,ij->bj", x, other) + func = partial(einsum_matmul, other=other) self._run_trace_and_compile(func, (4, 8), "test_einsum_matmul") def test_einsum_batched(self): """Test torch.einsum for batched operations.""" other = torch.randn(4, 8, 16) + def einsum_batched(x, other): - return torch.einsum('bik,bkj->bij', x, other) + return torch.einsum("bik,bkj->bij", x, other) + func = partial(einsum_batched, other=other) self._run_trace_and_compile(func, (4, 4, 8), "test_einsum_batched") def test_where(self): """Test torch.where (conditional selection).""" y = torch.randn(4, 8) + def where_positive(x, y): return torch.where(x > 0, x, y) + func = partial(where_positive, y=y) self._run_trace_and_compile(func, (4, 8), "test_where") @@ -815,9 +873,11 @@ def test_abs(self): def test_sqrt(self): """Test torch.sqrt (on positive inputs).""" + # Use abs to ensure positive inputs def sqrt_safe(x): return torch.sqrt(torch.abs(x) + 1e-6) + self._run_trace_and_compile(sqrt_safe, (4, 8), "test_sqrt") def test_exp(self): @@ -826,8 +886,10 @@ def test_exp(self): def test_log(self): """Test torch.log (on positive inputs).""" + def log_safe(x): return torch.log(torch.abs(x) + 1e-6) + self._run_trace_and_compile(log_safe, (4, 8), "test_log") def test_pow(self): @@ -837,38 +899,48 @@ def test_pow(self): def test_sin_cos(self): """Test torch.sin and torch.cos.""" + def sin_cos(x): return torch.sin(x) + torch.cos(x) + self._run_trace_and_compile(sin_cos, (4, 8), "test_sin_cos") def test_stack(self): """Test torch.stack.""" other = torch.randn(4, 8) + def stack_with_other(x, other): return torch.stack([x, other], dim=0) + func = partial(stack_with_other, other=other) self._run_trace_and_compile(func, (4, 8), "test_stack") def test_cat(self): """Test torch.cat.""" other = torch.randn(4, 8) + def cat_with_other(x, other): return torch.cat([x, other], dim=0) + func = partial(cat_with_other, other=other) self._run_trace_and_compile(func, (4, 8), "test_cat") def test_split(self): """Test torch.split.""" + def split_and_sum(x): parts = torch.split(x, 4, dim=1) return parts[0] + parts[1] + self._run_trace_and_compile(split_and_sum, (4, 8), "test_split") def test_chunk(self): """Test torch.chunk.""" + def chunk_and_sum(x): parts = torch.chunk(x, 2, dim=1) return parts[0] + parts[1] + self._run_trace_and_compile(chunk_and_sum, (4, 8), "test_chunk") @@ -904,6 +976,5 @@ def chunk_and_sum(x): # ============================================================================= -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() - diff --git a/tests/unit_tests/test_traced_np_array.py b/tests/unit_tests/test_traced_np_array.py index 44e7591a..fef32093 100644 --- a/tests/unit_tests/test_traced_np_array.py +++ b/tests/unit_tests/test_traced_np_array.py @@ -27,14 +27,10 @@ Array = np.ndarray | TracedNpArray -warnings.filterwarnings( - "ignore", message=".*legacy TorchScript-based ONNX export.*") -warnings.filterwarnings( - "ignore", message=".*on empty non-base types in `__init__`.*") -warnings.filterwarnings( - "ignore", message=".*feature will be removed.*") -warnings.filterwarnings( - "ignore", message=".*Casting complex values to real.*") +warnings.filterwarnings("ignore", message=".*legacy TorchScript-based ONNX export.*") +warnings.filterwarnings("ignore", message=".*on empty non-base types in `__init__`.*") +warnings.filterwarnings("ignore", message=".*feature will be removed.*") +warnings.filterwarnings("ignore", message=".*Casting complex values to real.*") class NpArrayArithmeticFunctions: @@ -47,7 +43,7 @@ def scalar_arithmetic(a: Array) -> Array: c = b - 2 d = c * 3 e = d / 4 - f = e ** 5 + f = e**5 return f @staticmethod @@ -57,7 +53,7 @@ def scalar_arithmetic_reverse(a: Array) -> Array: c = 2 - b d = 3 * c e = 4 / d - f = 5 ** e + f = 5**e return f @staticmethod @@ -77,7 +73,7 @@ def array_arithmetic_reuse(a: Array) -> Array: a = a - 1 a = a * 3 a = a / 4 - a = a ** 5 + a = a**5 return a @staticmethod @@ -362,7 +358,7 @@ def test_pow(self): """Test power.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - y = x ** 2 + y = x**2 self.assertIsInstance(y, TracedNpArray) np.testing.assert_array_almost_equal(y, np.array([1.0, 4.0, 9.0])) @@ -421,7 +417,7 @@ def test_comparison_eq(self): def test_np_sin(self): """Test np.sin ufunc.""" ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([0.0, np.pi/2, np.pi]), name="x") + x = ctx.create_input(np.array([0.0, np.pi / 2, np.pi]), name="x") y = np.sin(x) self.assertIsInstance(y, TracedNpArray) np.testing.assert_array_almost_equal(y, np.array([0.0, 1.0, 0.0]), decimal=5) @@ -429,7 +425,7 @@ def test_np_sin(self): def test_np_cos(self): """Test np.cos ufunc.""" ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([0.0, np.pi/2, np.pi]), name="x") + x = ctx.create_input(np.array([0.0, np.pi / 2, np.pi]), name="x") y = np.cos(x) self.assertIsInstance(y, TracedNpArray) np.testing.assert_array_almost_equal(y, np.array([1.0, 0.0, -1.0]), decimal=5) @@ -743,9 +739,9 @@ def test_method_copy_compiled(self): ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") y = x.copy() * 2 # Use copied array in computation - - ctx.compile_trace({'output': y}) - + + ctx.compile_trace({"output": y}) + # Test compiled graph works input_tensor = torch.tensor([4.0, 5.0, 6.0]) expected = input_tensor.clone() * 2 # clone is torch equivalent of copy @@ -760,9 +756,9 @@ def test_graph_basic_operations(self): x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") y = x + 1 z = y * 2 - + # Compile and verify graph exists - ctx.compile_trace({'output': z}) + ctx.compile_trace({"output": z}) self.assertIsNotNone(ctx.m) self.assertIsNotNone(ctx.graph) @@ -772,9 +768,9 @@ def test_graph_has_nodes(self): x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") y = x + 1 z = y * 2 - - ctx.compile_trace({'output': z}) - + + ctx.compile_trace({"output": z}) + # Check graph has nodes nodes = list(ctx.graph.nodes) self.assertGreater(len(nodes), 0) @@ -785,8 +781,8 @@ def test_graph_multiple_inputs(self): x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") y = ctx.create_input(np.array([4.0, 5.0, 6.0]), name="y") z = x + y - - ctx.compile_trace({'output': z}) + + ctx.compile_trace({"output": z}) self.assertIsNotNone(ctx.m) # ==================== Operations After Compile Tests ==================== @@ -797,7 +793,7 @@ def test_operations_after_compile_return_native_arrays(self): x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") y = x * 2 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) # After compilation, is_tracing should be False self.assertFalse(ctx.is_tracing) @@ -857,7 +853,7 @@ def test_conversion_copy_property(self): ) self.assertTrue(np.array_equal(result.view(np.ndarray), [1.0, 2.0, 3.0])) - ctx.compile_trace({'aliased': aliased, 'copied': copied}) + ctx.compile_trace({"aliased": aliased, "copied": copied}) # Each copy is recorded, while the alias reuses its source proxy targets = [node.target for node in ctx.graph.nodes] @@ -881,18 +877,19 @@ def test_conversion_dtype_swap_property(self): finally: patcher.uninstall() - for result, expected in ((positional, np.float32), - (keyword, np.float32), - (widened, np.int64)): + for result, expected in ( + (positional, np.float32), + (keyword, np.float32), + (widened, np.int64), + ): self.assertIsInstance(result, TracedNpArray) self.assertEqual(result.dtype, expected) - ctx.compile_trace({'positional': positional, 'widened': widened}) + ctx.compile_trace({"positional": positional, "widened": widened}) # Each output cast is a real node rather than a reuse of the source proxy casts = [node for node in ctx.graph.nodes if node.target == "to"] - self.assertEqual([node.args[1] for node in casts], - [torch.float32, torch.int64]) + self.assertEqual([node.args[1] for node in casts], [torch.float32, torch.int64]) # A float64 input proves the recorded cast changes dtype rather than # coincidentally matching torch's float32 default. @@ -906,9 +903,9 @@ def test_getitem_after_compile_returns_native_array(self): x = ctx.create_input(np.array([1.0, 2.0, 3.0, 4.0, 5.0]), name="x") # Use slice to get array output (scalar indexing returns np.float64, not array) y = x[0:2] - - ctx.compile_trace({'y': y}) - + + ctx.compile_trace({"y": y}) + z = x[2:4] self.assertNotIsInstance(z, TracedNpArray) self.assertIsInstance(z, np.ndarray) @@ -920,10 +917,10 @@ def test_multiple_contexts_raises_error(self): """Test that mixing contexts raises an error.""" ctx1 = TracedTensorNode(name="test1", node_index=0) ctx2 = TracedTensorNode(name="test2", node_index=1) - + x = ctx1.create_input(np.array([1.0, 2.0, 3.0]), name="x") y = ctx2.create_input(np.array([4.0, 5.0, 6.0]), name="y") - + # Mixing arrays from different contexts should raise with self.assertRaises(Exception): _ = x + y @@ -934,9 +931,9 @@ def test_chained_operations(self): """Test chained operations.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - + result = ((x + 1) * 2 - 3) / 4 - + self.assertIsInstance(result, TracedNpArray) expected = ((np.array([1.0, 2.0, 3.0]) + 1) * 2 - 3) / 4 np.testing.assert_array_almost_equal(result, expected) @@ -945,12 +942,12 @@ def test_complex_computation(self): """Test complex computation chain.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - + # Multiple operations y = np.sin(x) z = np.cos(x) - w = y ** 2 + z ** 2 # Should be ~1.0 - + w = y**2 + z**2 # Should be ~1.0 + self.assertIsInstance(w, TracedNpArray) np.testing.assert_array_almost_equal(w, np.ones(3), decimal=5) @@ -967,18 +964,18 @@ def _test_function(self, func, name): """Helper to test a function produces correct results.""" # Get expected result from regular numpy expected = func(self.test_data.copy()) - + # Create traced version ctx = TracedTensorNode(name=f"test_{name}", node_index=0) traced_input = ctx.create_input(self.test_data.copy(), name="input") traced_result = func(traced_input) - + # Compare results if isinstance(traced_result, TracedNpArray): actual = traced_result else: actual = np.array(traced_result) - + if isinstance(expected, np.ndarray): np.testing.assert_array_almost_equal(actual, expected, decimal=5) else: @@ -986,15 +983,22 @@ def _test_function(self, func, name): def test_scalar_arithmetic(self): """Test scalar arithmetic function.""" - self._test_function(NpArrayArithmeticFunctions.scalar_arithmetic, "scalar_arithmetic") + self._test_function( + NpArrayArithmeticFunctions.scalar_arithmetic, "scalar_arithmetic" + ) def test_scalar_arithmetic_reverse(self): """Test reverse scalar arithmetic function.""" - self._test_function(NpArrayArithmeticFunctions.scalar_arithmetic_reverse, "scalar_arithmetic_reverse") + self._test_function( + NpArrayArithmeticFunctions.scalar_arithmetic_reverse, + "scalar_arithmetic_reverse", + ) def test_array_arithmetic(self): """Test array arithmetic function.""" - self._test_function(NpArrayArithmeticFunctions.array_arithmetic, "array_arithmetic") + self._test_function( + NpArrayArithmeticFunctions.array_arithmetic, "array_arithmetic" + ) def test_sum_method(self): """Test sum method.""" @@ -1018,7 +1022,9 @@ def test_reshape_method(self): def test_reshape_function(self): """Test reshape function.""" - self._test_function(NpArrayArithmeticFunctions.reshape_function, "reshape_function") + self._test_function( + NpArrayArithmeticFunctions.reshape_function, "reshape_function" + ) # Suppress TracerWarning from setitem tests where torch.tensor/torch.as_tensor constants @@ -1027,7 +1033,7 @@ def test_reshape_function(self): @pytest.mark.filterwarnings("ignore::torch.jit.TracerWarning") class TestCompiledGraphExecution(unittest.TestCase): """Test that traced numpy operations compile and execute correctly. - + These tests verify that: 1. The FX graph module produces correct output 2. TorchScript export works and produces correct output @@ -1035,10 +1041,16 @@ class TestCompiledGraphExecution(unittest.TestCase): 4. ExportedProgram export works and produces correct output """ - def _test_operation(self, op_name: str, np_func, input_data: np.ndarray, - expected_output: np.ndarray, atol: float = 1e-5): + def _test_operation( + self, + op_name: str, + np_func, + input_data: np.ndarray, + expected_output: np.ndarray, + atol: float = 1e-5, + ): """Helper to test a single operation through all export formats. - + Args: op_name: Name of the operation (for error messages) np_func: Function that takes TracedNpArray and returns result @@ -1050,44 +1062,47 @@ def _test_operation(self, op_name: str, np_func, input_data: np.ndarray, ctx = TracedTensorNode(name=f"test_{op_name}", node_index=0) traced_input = ctx.create_input(input_data.copy(), name="input") traced_result = np_func(traced_input) - - self.assertIsInstance(traced_result, TracedNpArray, - f"{op_name}: result should be TracedNpArray, got {type(traced_result)}") - + + self.assertIsInstance( + traced_result, + TracedNpArray, + f"{op_name}: result should be TracedNpArray, got {type(traced_result)}", + ) + # Compile the trace - ctx.compile_trace({'output': traced_result}) + ctx.compile_trace({"output": traced_result}) graph_module = ctx.m - + # Convert input to torch tensor for execution input_tensor = torch.from_numpy(input_data.copy()).float() - + # Test 1: FX GraphModule execution try: graph_output = graph_module(input_tensor) graph_output_np = graph_output.detach().numpy() np.testing.assert_allclose( - graph_output_np.flatten(), - expected_output.flatten(), + graph_output_np.flatten(), + expected_output.flatten(), atol=atol, - err_msg=f"{op_name}: FX GraphModule output doesn't match expected" + err_msg=f"{op_name}: FX GraphModule output doesn't match expected", ) except Exception as e: self.fail(f"{op_name}: FX GraphModule execution failed: {e}") - + # Test 2: TorchScript export and execution try: scripted = torch.jit.script(graph_module) script_output = scripted(input_tensor) script_output_np = script_output.detach().numpy() np.testing.assert_allclose( - script_output_np.flatten(), - expected_output.flatten(), + script_output_np.flatten(), + expected_output.flatten(), atol=atol, - err_msg=f"{op_name}: TorchScript output doesn't match expected" + err_msg=f"{op_name}: TorchScript output doesn't match expected", ) except Exception as e: self.fail(f"{op_name}: TorchScript export/execution failed: {e}") - + # Test 3: ONNX export and execution try: with tempfile.TemporaryDirectory() as tmpdir: @@ -1100,18 +1115,20 @@ def _test_operation(self, op_name: str, np_func, input_data: np.ndarray, export_params=True, opset_version=17, do_constant_folding=True, - input_names=['input'], - output_names=['output'], + input_names=["input"], + output_names=["output"], ) - + session = ort.InferenceSession(str(onnx_path)) - onnx_output = session.run(None, {"input": input_data.astype(np.float32)})[0] - + onnx_output = session.run( + None, {"input": input_data.astype(np.float32)} + )[0] + np.testing.assert_allclose( - onnx_output.flatten(), - expected_output.flatten(), + onnx_output.flatten(), + expected_output.flatten(), atol=atol, - err_msg=f"{op_name}: ONNX output doesn't match expected" + err_msg=f"{op_name}: ONNX output doesn't match expected", ) except Exception as e: self.fail(f"{op_name}: ONNX export/execution failed: {e}") @@ -1156,21 +1173,24 @@ def test_mul_np_float32(self): input_data = np.array([1.0, 2.0, 3.0], dtype=np.float32) expected = input_data * np.float32(2.0) self._test_operation( - "mul_np_float32", lambda x: x * np.float32(2.0), input_data, expected) + "mul_np_float32", lambda x: x * np.float32(2.0), input_data, expected + ) def test_mul_np_float64(self): """Test: x * np.float64(2.0) — numpy scalars must not leak into the FX graph.""" input_data = np.array([1.0, 2.0, 3.0]) expected = input_data * np.float64(2.0) self._test_operation( - "mul_np_float64", lambda x: x * np.float64(2.0), input_data, expected) + "mul_np_float64", lambda x: x * np.float64(2.0), input_data, expected + ) def test_div_norm(self): """Test: x / np.linalg.norm(x) — unmapped norm returns a numpy scalar.""" input_data = np.array([3.0, 4.0], dtype=np.float32) expected = input_data / np.linalg.norm(input_data) self._test_operation( - "div_norm", lambda x: x / np.linalg.norm(x), input_data, expected) + "div_norm", lambda x: x / np.linalg.norm(x), input_data, expected + ) def test_div_scalar(self): """Test: x / 2""" @@ -1181,8 +1201,8 @@ def test_div_scalar(self): def test_pow_scalar(self): """Test: x ** 2""" input_data = np.array([1.0, 2.0, 3.0]) - expected = input_data ** 2 - self._test_operation("pow_scalar", lambda x: x ** 2, input_data, expected) + expected = input_data**2 + self._test_operation("pow_scalar", lambda x: x**2, input_data, expected) def test_neg(self): """Test: -x""" @@ -1272,7 +1292,9 @@ def test_np_mean_scalar(self): """Test: np.mean(x) returning a scalar — must stay traced.""" input_data = np.array([1.0, 2.0, 3.0]) expected = np.asarray(np.mean(input_data)) - self._test_operation("np_mean_scalar", lambda x: np.mean(x), input_data, expected) + self._test_operation( + "np_mean_scalar", lambda x: np.mean(x), input_data, expected + ) def test_np_var_scalar(self): """Test: np.var(x) returning a scalar — must stay traced.""" @@ -1290,13 +1312,17 @@ def test_np_prod_scalar(self): """Test: np.prod(x) returning a scalar — must stay traced.""" input_data = np.array([1.0, 2.0, 3.0]) expected = np.asarray(np.prod(input_data)) - self._test_operation("np_prod_scalar", lambda x: np.prod(x), input_data, expected) + self._test_operation( + "np_prod_scalar", lambda x: np.prod(x), input_data, expected + ) def test_np_prod_axis(self): """Test: np.prod(x, axis=1) — axis must be renamed to dim for aten.""" input_data = np.array([[1.0, 2.0], [3.0, 4.0]]) expected = np.prod(input_data, axis=1) - self._test_operation("np_prod_axis", lambda x: np.prod(x, axis=1), input_data, expected) + self._test_operation( + "np_prod_axis", lambda x: np.prod(x, axis=1), input_data, expected + ) def test_np_max_scalar(self): """Test: np.max(x) returning a scalar — must stay traced.""" @@ -1314,72 +1340,95 @@ def test_np_sum_keepdims(self): """Test: np.sum(x, keepdims=True) - tests keepdims→keepdim conversion""" input_data = np.array([1.0, 2.0, 3.0]) expected = np.sum(input_data, keepdims=True) - self._test_operation("np_sum_keepdims", lambda x: np.sum(x, keepdims=True), input_data, expected) + self._test_operation( + "np_sum_keepdims", lambda x: np.sum(x, keepdims=True), input_data, expected + ) def test_np_mean_keepdims(self): """Test: np.mean(x, keepdims=True) - tests keepdims→keepdim conversion""" input_data = np.array([1.0, 2.0, 3.0]) expected = np.mean(input_data, keepdims=True) - self._test_operation("np_mean_keepdims", lambda x: np.mean(x, keepdims=True), input_data, expected) + self._test_operation( + "np_mean_keepdims", + lambda x: np.mean(x, keepdims=True), + input_data, + expected, + ) def test_np_max_axis(self): """Test: np.max(x, axis=0)""" input_data = np.array([[1.0, 3.0, 2.0], [4.0, 2.0, 5.0]]) expected = np.max(input_data, axis=0) - self._test_operation("np_max_axis", lambda x: np.max(x, axis=0), input_data, expected) + self._test_operation( + "np_max_axis", lambda x: np.max(x, axis=0), input_data, expected + ) def test_np_min_axis(self): """Test: np.min(x, axis=0)""" input_data = np.array([[1.0, 3.0, 2.0], [4.0, 2.0, 5.0]]) expected = np.min(input_data, axis=0) - self._test_operation("np_min_axis", lambda x: np.min(x, axis=0), input_data, expected) + self._test_operation( + "np_min_axis", lambda x: np.min(x, axis=0), input_data, expected + ) def test_np_stack_axis(self): """Test: np.stack([x, x], axis=0) — axis must be renamed to dim for aten.""" input_data = np.array([1.0, 2.0, 3.0]) expected = np.stack([input_data, input_data], axis=0) self._test_operation( - "np_stack_axis", lambda x: np.stack([x, x], axis=0), input_data, expected) + "np_stack_axis", lambda x: np.stack([x, x], axis=0), input_data, expected + ) def test_np_split_axis(self): """Test: np.split(x, 2, axis=0) — axis must be renamed to dim for aten.""" input_data = np.array([[1.0, 2.0], [3.0, 4.0]]) expected = np.split(input_data, 2, axis=0)[0] self._test_operation( - "np_split_axis", lambda x: np.split(x, 2, axis=0)[0], input_data, expected) + "np_split_axis", lambda x: np.split(x, 2, axis=0)[0], input_data, expected + ) def test_np_array_split_axis(self): """Test: np.array_split(x, 2, axis=0) — axis must be renamed to dim for aten.""" input_data = np.array([[1.0, 2.0], [3.0, 4.0]]) expected = np.array_split(input_data, 2, axis=0)[0] self._test_operation( - "np_array_split_axis", lambda x: np.array_split(x, 2, axis=0)[0], - input_data, expected) + "np_array_split_axis", + lambda x: np.array_split(x, 2, axis=0)[0], + input_data, + expected, + ) def test_np_flip_axis(self): """Test: np.flip(x, axis=1) — axis becomes dims as a tuple.""" input_data = np.array([[1.0, 2.0], [3.0, 4.0]]) expected = np.ascontiguousarray(np.flip(input_data, axis=1)) - self._test_operation("np_flip_axis", lambda x: np.flip(x, axis=1), input_data, expected) + self._test_operation( + "np_flip_axis", lambda x: np.flip(x, axis=1), input_data, expected + ) def test_np_flip_positional(self): """Test: np.flip(x, 1) — integer dims must be wrapped for torch.flip.""" input_data = np.array([[1.0, 2.0], [3.0, 4.0]]) expected = np.ascontiguousarray(np.flip(input_data, 1)) - self._test_operation("np_flip_positional", lambda x: np.flip(x, 1), input_data, expected) + self._test_operation( + "np_flip_positional", lambda x: np.flip(x, 1), input_data, expected + ) def test_np_flip_all_axes(self): """Test: np.flip(x) — omitted axis reverses every dimension.""" input_data = np.array([1.0, 2.0, 3.0]) expected = np.ascontiguousarray(np.flip(input_data)) - self._test_operation("np_flip_all_axes", lambda x: np.flip(x), input_data, expected) + self._test_operation( + "np_flip_all_axes", lambda x: np.flip(x), input_data, expected + ) def test_np_roll_axis(self): """Test: np.roll(x, 1, axis=1) — axis becomes dims, not dim.""" input_data = np.array([[1.0, 2.0], [3.0, 4.0]]) expected = np.roll(input_data, 1, axis=1) self._test_operation( - "np_roll_axis", lambda x: np.roll(x, 1, axis=1), input_data, expected) + "np_roll_axis", lambda x: np.roll(x, 1, axis=1), input_data, expected + ) # ==================== Shape Operations ==================== @@ -1387,20 +1436,29 @@ def test_reshape(self): """Test: np.reshape(x, (2, 3))""" input_data = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) expected = np.reshape(input_data, (2, 3)) - self._test_operation("reshape", lambda x: np.reshape(x, (2, 3)), input_data, expected) + self._test_operation( + "reshape", lambda x: np.reshape(x, (2, 3)), input_data, expected + ) def test_transpose_no_axes(self): """Test: np.transpose(x) - tests automatic axes reversal for torch.permute""" input_data = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) # No axes specified - numpy reverses all dims, framework should handle this expected = np.transpose(input_data) - self._test_operation("transpose_no_axes", lambda x: np.transpose(x), input_data, expected) + self._test_operation( + "transpose_no_axes", lambda x: np.transpose(x), input_data, expected + ) def test_transpose_with_axes(self): """Test: np.transpose(x, axes) - explicit axes""" input_data = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) expected = np.transpose(input_data, (1, 0)) - self._test_operation("transpose_with_axes", lambda x: np.transpose(x, (1, 0)), input_data, expected) + self._test_operation( + "transpose_with_axes", + lambda x: np.transpose(x, (1, 0)), + input_data, + expected, + ) def test_squeeze(self): """Test: np.squeeze(x)""" @@ -1414,19 +1472,25 @@ def test_chained_arithmetic(self): """Test: (x + 1) * 2 - 3""" input_data = np.array([1.0, 2.0, 3.0]) expected = (input_data + 1) * 2 - 3 - self._test_operation("chained_arithmetic", lambda x: (x + 1) * 2 - 3, input_data, expected) + self._test_operation( + "chained_arithmetic", lambda x: (x + 1) * 2 - 3, input_data, expected + ) def test_chained_ufuncs(self): """Test: np.exp(np.sin(x))""" input_data = np.array([0.0, 0.5, 1.0]) expected = np.exp(np.sin(input_data)) - self._test_operation("chained_ufuncs", lambda x: np.exp(np.sin(x)), input_data, expected) + self._test_operation( + "chained_ufuncs", lambda x: np.exp(np.sin(x)), input_data, expected + ) def test_complex_computation(self): """Test: np.sqrt(x**2 + 1)""" input_data = np.array([1.0, 2.0, 3.0]) expected = np.sqrt(input_data**2 + 1) - self._test_operation("complex_computation", lambda x: np.sqrt(x**2 + 1), input_data, expected) + self._test_operation( + "complex_computation", lambda x: np.sqrt(x**2 + 1), input_data, expected + ) # ==================== Indexing Operations ==================== @@ -1442,8 +1506,9 @@ def test_greater_than(self): """Test: x > 2""" input_data = np.array([1.0, 2.0, 3.0]) expected = (input_data > 2).astype(np.float32) - self._test_operation("greater_than", lambda x: (x > 2).astype(np.float32), - input_data, expected) + self._test_operation( + "greater_than", lambda x: (x > 2).astype(np.float32), input_data, expected + ) # ==================== Clip Operation ==================== @@ -1451,21 +1516,25 @@ def test_np_clip(self): """Test: np.clip(x, 1.5, 2.5)""" input_data = np.array([1.0, 2.0, 3.0]) expected = np.clip(input_data, 1.5, 2.5) - self._test_operation("np_clip", lambda x: np.clip(x, 1.5, 2.5), input_data, expected) + self._test_operation( + "np_clip", lambda x: np.clip(x, 1.5, 2.5), input_data, expected + ) def test_np_maximum_scalar(self): """Test: np.maximum(x, 3.0) — torch.maximum requires tensor operands.""" input_data = np.array([1.0, 4.0, 2.0]) expected = np.maximum(input_data, 3.0) self._test_operation( - "np_maximum_scalar", lambda x: np.maximum(x, 3.0), input_data, expected) + "np_maximum_scalar", lambda x: np.maximum(x, 3.0), input_data, expected + ) def test_np_minimum_scalar(self): """Test: np.minimum(x, 3.0) — torch.minimum requires tensor operands.""" input_data = np.array([1.0, 4.0, 2.0]) expected = np.minimum(input_data, 3.0) self._test_operation( - "np_minimum_scalar", lambda x: np.minimum(x, 3.0), input_data, expected) + "np_minimum_scalar", lambda x: np.minimum(x, 3.0), input_data, expected + ) def test_np_round(self): """Test: np.round(x) — dispatches as a function, not a ufunc.""" @@ -1480,14 +1549,14 @@ def test_np_matmul(self): input_data = np.array([[1.0, 2.0], [3.0, 4.0]]) expected = np.matmul(input_data, input_data.T) self._test_operation( - "np_matmul", lambda x: np.matmul(x, x.T), input_data, expected) + "np_matmul", lambda x: np.matmul(x, x.T), input_data, expected + ) def test_matmul_operator(self): """Test: x @ x.T — the operator form routes through the same ufunc.""" input_data = np.array([[1.0, 2.0], [3.0, 4.0]]) expected = input_data @ input_data.T - self._test_operation( - "matmul_operator", lambda x: x @ x.T, input_data, expected) + self._test_operation("matmul_operator", lambda x: x @ x.T, input_data, expected) def test_np_tensordot_axes(self): """Test: np.tensordot(x, x, axes=...) — torch.tensordot takes dims.""" @@ -1495,7 +1564,10 @@ def test_np_tensordot_axes(self): expected = np.tensordot(input_data, input_data, axes=([1], [1])) self._test_operation( "np_tensordot_axes", - lambda x: np.tensordot(x, x, axes=([1], [1])), input_data, expected) + lambda x: np.tensordot(x, x, axes=([1], [1])), + input_data, + expected, + ) # ==================== Sorting and Searching Operations ==================== @@ -1509,44 +1581,66 @@ def test_np_sort_axis(self): """Test: np.sort(x, axis=1)""" input_data = np.array([[3.0, 1.0, 2.0], [6.0, 4.0, 5.0]]) expected = np.sort(input_data, axis=1) - self._test_operation("np_sort_axis", lambda x: np.sort(x, axis=1), input_data, expected) + self._test_operation( + "np_sort_axis", lambda x: np.sort(x, axis=1), input_data, expected + ) def test_np_argsort(self): """Test: np.argsort(x)""" input_data = np.array([3.0, 1.0, 4.0, 1.0, 5.0]) - expected = np.argsort(input_data).astype(np.float32) # Convert to float for comparison + expected = np.argsort(input_data).astype( + np.float32 + ) # Convert to float for comparison # Note: argsort returns indices as int64, but we compare as float for consistency - self._test_operation("np_argsort", lambda x: np.argsort(x).astype(np.float32), - input_data, expected) + self._test_operation( + "np_argsort", + lambda x: np.argsort(x).astype(np.float32), + input_data, + expected, + ) def test_np_where_compiled(self): """Test: np.where(x > 2, x, 0)""" input_data = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) expected = np.where(input_data > 2, input_data, 0.0) - self._test_operation("np_where", lambda x: np.where(x > 2, x, 0.0), input_data, expected) + self._test_operation( + "np_where", lambda x: np.where(x > 2, x, 0.0), input_data, expected + ) def test_np_nonzero_rows(self): """Test: np.nonzero(x)[0] — numpy returns one index array per dimension.""" input_data = np.array([[0.0, 0.0, 1.0], [1.0, 0.0, 1.0]]) expected = np.nonzero(input_data)[0].astype(np.float32) self._test_operation( - "np_nonzero_rows", lambda x: np.nonzero(x)[0].astype(np.float32), - input_data, expected) + "np_nonzero_rows", + lambda x: np.nonzero(x)[0].astype(np.float32), + input_data, + expected, + ) def test_np_nonzero_cols(self): """Test: np.nonzero(x)[1] — the second element indexes the last dimension.""" input_data = np.array([[0.0, 0.0, 1.0], [1.0, 0.0, 1.0]]) expected = np.nonzero(input_data)[1].astype(np.float32) self._test_operation( - "np_nonzero_cols", lambda x: np.nonzero(x)[1].astype(np.float32), - input_data, expected) + "np_nonzero_cols", + lambda x: np.nonzero(x)[1].astype(np.float32), + input_data, + expected, + ) # ==================== Setitem Operations ==================== - def _test_setitem(self, test_name: str, input_data: np.ndarray, - setitem_func, expected_output: np.ndarray, atol: float = 1e-5): + def _test_setitem( + self, + test_name: str, + input_data: np.ndarray, + setitem_func, + expected_output: np.ndarray, + atol: float = 1e-5, + ): """Helper to test setitem operations through all export formats. - + Args: test_name: Name of the test (for error messages and ONNX filename) input_data: Input numpy array @@ -1558,36 +1652,36 @@ def _test_setitem(self, test_name: str, input_data: np.ndarray, ctx = TracedTensorNode(name=f"test_{test_name}", node_index=0) traced_input = ctx.create_input(input_data.copy(), name="input") traced_result = setitem_func(traced_input) - + # Compile the trace - ctx.compile_trace({'output': traced_result}) + ctx.compile_trace({"output": traced_result}) graph_module = ctx.m - + # Convert input to torch tensor for execution torch_input = torch.from_numpy(input_data.copy()).float() torch_expected = torch.from_numpy(expected_output).float() - + # Test 1: FX GraphModule execution try: output = graph_module(torch_input) self.assertTrue( torch.allclose(output, torch_expected, atol=atol), - f"{test_name}: FX GraphModule execution failed - output doesn't match" + f"{test_name}: FX GraphModule execution failed - output doesn't match", ) except Exception as e: self.fail(f"{test_name}: FX GraphModule execution failed: {e}") - + # Test 2: TorchScript export try: scripted = torch.jit.script(graph_module) output_ts = scripted(torch_input) self.assertTrue( torch.allclose(output_ts, torch_expected, atol=atol), - f"{test_name}: TorchScript output doesn't match" + f"{test_name}: TorchScript output doesn't match", ) except Exception as e: self.fail(f"{test_name}: TorchScript export/execution failed: {e}") - + # Test 3: ONNX export try: with tempfile.TemporaryDirectory() as tmpdir: @@ -1599,14 +1693,14 @@ def _test_setitem(self, test_name: str, input_data: np.ndarray, dynamo=False, export_params=True, opset_version=17, - input_names=['input'], - output_names=['output'], + input_names=["input"], + output_names=["output"], ) session = ort.InferenceSession(str(onnx_path)) output_onnx = session.run(None, {"input": torch_input.numpy()})[0] self.assertTrue( np.allclose(output_onnx, expected_output, atol=atol), - f"{test_name}: ONNX output doesn't match" + f"{test_name}: ONNX output doesn't match", ) except Exception as e: self.fail(f"{test_name}: ONNX export/execution failed: {e}") @@ -1624,48 +1718,47 @@ def _test_setitem(self, test_name: str, input_data: np.ndarray, def test_setitem_single_index(self): """Test: x[0] = 10""" input_data = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) - + def setitem_func(x): x[0] = 10.0 return x * 2 - + expected = np.array([20.0, 4.0, 6.0, 8.0, 10.0]) self._test_setitem("setitem_single", input_data, setitem_func, expected) def test_setitem_slice(self): """Test: x[1:3] = [10, 20]""" input_data = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) - + def setitem_func(x): x[1:3] = np.array([10.0, 20.0]) return x * 2 - + expected = np.array([2.0, 20.0, 40.0, 8.0, 10.0]) self._test_setitem("setitem_slice", input_data, setitem_func, expected) def test_setitem_full_slice(self): """Test: x[:] = constant""" input_data = np.array([1.0, 2.0, 3.0]) - + def setitem_func(x): x[:] = np.array([10.0, 20.0, 30.0]) return x * 2 - + expected = np.array([20.0, 40.0, 60.0]) self._test_setitem("setitem_full", input_data, setitem_func, expected) def test_setitem_step_slice(self): """Test: x[::2] = [10, 30, 50] with step""" input_data = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) - + def setitem_func(x): x[::2] = np.array([10.0, 30.0, 50.0]) # Assigns to indices 0, 2, 4 return x * 2 - + expected = np.array([20.0, 4.0, 60.0, 8.0, 100.0]) self._test_setitem("setitem_step", input_data, setitem_func, expected) - def _validate_dynamic_setitem(self, graph_module, inputs, expected, test_name): """Validate export-first backends for traced NumPy index inputs.""" actual = graph_module(*inputs) @@ -1675,9 +1768,7 @@ def _validate_dynamic_setitem(self, graph_module, inputs, expected, test_name): scripted_actual = scripted(*inputs) self.assertTrue(torch.allclose(scripted_actual, expected, atol=1e-5)) - verify_onnx_dynamo( - self, graph_module, inputs, expected, test_name=test_name - ) + verify_onnx_dynamo(self, graph_module, inputs, expected, test_name=test_name) verify_exported_program( self, graph_module, inputs, expected, test_name=test_name ) @@ -1685,12 +1776,8 @@ def _validate_dynamic_setitem(self, graph_module, inputs, expected, test_name): def test_setitem_traced_integer_index_augmented_assignment(self): """NumPy x[:, indices] += source uses shared getitem/setitem lowering.""" ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input( - np.arange(8, dtype=np.float32).reshape(2, 4), name="x" - ) - indices = ctx.create_input( - np.array([0, 2], dtype=np.int64), name="indices" - ) + x = ctx.create_input(np.arange(8, dtype=np.float32).reshape(2, 4), name="x") + indices = ctx.create_input(np.array([0, 2], dtype=np.int64), name="indices") source = ctx.create_input( np.array([[10.0, 20.0], [30.0, 40.0]], dtype=np.float32), name="source", @@ -1717,12 +1804,8 @@ def test_setitem_traced_integer_index_augmented_assignment(self): def test_setitem_traced_boolean_mask_augmented_assignment(self): """NumPy x[:, mask] *= scalar uses shared getitem/setitem lowering.""" ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input( - np.arange(8, dtype=np.float32).reshape(2, 4), name="x" - ) - mask = ctx.create_input( - np.array([True, False, True, False]), name="mask" - ) + x = ctx.create_input(np.arange(8, dtype=np.float32).reshape(2, 4), name="x") + mask = ctx.create_input(np.array([True, False, True, False]), name="mask") x[:, mask] *= 3.0 result = x + 1.0 @@ -1744,12 +1827,8 @@ def test_setitem_traced_boolean_mask_augmented_assignment(self): def test_setitem_traced_index_plain_numpy_source(self): """A traced NumPy destination accepts a constant ndarray source.""" ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input( - np.arange(8, dtype=np.float32).reshape(2, 4), name="x" - ) - indices = ctx.create_input( - np.array([0, 2], dtype=np.int64), name="indices" - ) + x = ctx.create_input(np.arange(8, dtype=np.float32).reshape(2, 4), name="x") + indices = ctx.create_input(np.array([0, 2], dtype=np.int64), name="indices") source = np.array([[11.0, 12.0], [13.0, 14.0]], dtype=np.float32) x[:, indices] = source @@ -1772,4 +1851,3 @@ def test_setitem_traced_index_plain_numpy_source(self): if __name__ == "__main__": unittest.main() - diff --git a/tests/unit_tests/test_traced_tensor.py b/tests/unit_tests/test_traced_tensor.py index 40544067..f28090b5 100644 --- a/tests/unit_tests/test_traced_tensor.py +++ b/tests/unit_tests/test_traced_tensor.py @@ -26,12 +26,9 @@ Tensor = torch.Tensor | TracedTensor -warnings.filterwarnings( - "ignore", message=".*legacy TorchScript-based ONNX export.*") -warnings.filterwarnings( - "ignore", message=".*on empty non-base types in `__init__`.*") -warnings.filterwarnings( - "ignore", message=".*feature will be removed.*") +warnings.filterwarnings("ignore", message=".*legacy TorchScript-based ONNX export.*") +warnings.filterwarnings("ignore", message=".*on empty non-base types in `__init__`.*") +warnings.filterwarnings("ignore", message=".*feature will be removed.*") class TensorArithmeticFunctions: @@ -527,7 +524,7 @@ def validate_export( verify_dynamo_onnx=False, ): """Validate graph_module execution across FX, TorchScript, and ONNX. - + Args: graph_module: The compiled FX GraphModule to test inputs: Tuple of input tensors @@ -537,25 +534,25 @@ def validate_export( # Ensure inputs is a tuple if not isinstance(inputs, tuple): inputs = (inputs,) - + # Generate input names based on count input_names = [f"input_{i}" for i in range(len(inputs))] - + # Test 1: FX GraphModule execution output = graph_module(*inputs) self.assertTrue( torch.allclose(output, expected, atol=1e-5), - f"{test_name}: FX GraphModule output doesn't match expected" + f"{test_name}: FX GraphModule output doesn't match expected", ) - + # Test 2: TorchScript export scripted = torch.jit.script(graph_module) output_ts = scripted(*inputs) self.assertTrue( torch.allclose(output_ts, expected, atol=1e-5), - f"{test_name}: TorchScript output doesn't match expected" + f"{test_name}: TorchScript output doesn't match expected", ) - + # Test 3: legacy ONNX export if verify_legacy_onnx: with tempfile.TemporaryDirectory() as tmpdir: @@ -594,7 +591,8 @@ def validate_export( # Test 5: ExportedProgram export verify_exported_program( - self, graph_module, inputs, expected, test_name=test_name) + self, graph_module, inputs, expected, test_name=test_name + ) def test_tensor_metadata_properties(self): """Tensor metadata access matches native Tensor behavior.""" @@ -656,10 +654,11 @@ def test_regular_tensor_div_traced_tensor(self): self.assertTrue(torch.allclose(result.tensor, expected)) # Verify it compiles to TorchScript - ctx.compile_trace({'result': result}) + ctx.compile_trace({"result": result}) scripted = torch.jit.script(ctx.m) - self.assertTrue(torch.allclose( - scripted(torch.tensor([2.0, 4.0, 6.0])), expected)) + self.assertTrue( + torch.allclose(scripted(torch.tensor([2.0, 4.0, 6.0])), expected) + ) def test_tensor_norm_method_rewrites_torch_tensor_function_target(self): """Test tensor.norm() becomes a call_method node before export.""" @@ -672,20 +671,24 @@ def test_tensor_norm_method_rewrites_torch_tensor_function_target(self): expected = torch.tensor([5.0, 13.0]) self.assertTrue(torch.allclose(result.tensor, expected)) - ctx.compile_trace({'result': result}) + ctx.compile_trace({"result": result}) norm_nodes = [ - node for node in ctx.m.graph.nodes + node + for node in ctx.m.graph.nodes if node.op == "call_method" and node.target == "norm" ] - self.assertTrue(norm_nodes, "Expected tensor.norm() to be rewritten to call_method") + self.assertTrue( + norm_nodes, "Expected tensor.norm() to be rewritten to call_method" + ) traced = torch.jit.trace( ctx.m, (torch.tensor([[3.0, 4.0], [5.0, 12.0]]),), ) - self.assertTrue(torch.allclose( - traced(torch.tensor([[3.0, 4.0], [5.0, 12.0]])), expected)) + self.assertTrue( + torch.allclose(traced(torch.tensor([[3.0, 4.0], [5.0, 12.0]])), expected) + ) # ==================== is_tracing=False Tests ==================== @@ -695,7 +698,7 @@ def test_operations_after_compile_return_native_tensors(self): x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") y = x * 2 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) # After compilation, is_tracing should be False self.assertFalse(ctx.is_tracing) @@ -715,7 +718,7 @@ def test_getitem_after_compile_returns_native_tensor(self): x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") y = x[0] - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) z = x[1] self.assertNotIsInstance(z, TracedTensor) @@ -729,7 +732,7 @@ def test_to_after_compile_returns_native_tensor(self): x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") y = x.to(torch.float64) - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) z = x.to(torch.float32) self.assertNotIsInstance(z, TracedTensor) @@ -749,11 +752,9 @@ def test_clone_preserves_type_and_compiles(self): self.assertTrue(torch.allclose(cloned.tensor, x.tensor)) output = cloned * 2 - ctx.compile_trace({'output': output}) + ctx.compile_trace({"output": output}) runtime_input = torch.tensor([4.0, 5.0, 6.0]) - self.assertTrue(torch.allclose( - ctx.m(runtime_input), runtime_input.clone() * 2 - )) + self.assertTrue(torch.allclose(ctx.m(runtime_input), runtime_input.clone() * 2)) def test_contiguous(self): """Test contiguous operation.""" @@ -777,11 +778,11 @@ def test_detach_preserves_type_and_compiles(self): self.assertFalse(detached.requires_grad) output = detached * 2 - ctx.compile_trace({'output': output}) + ctx.compile_trace({"output": output}) runtime_input = torch.tensor([4.0, 5.0, 6.0]) - self.assertTrue(torch.allclose( - ctx.m(runtime_input), runtime_input.detach() * 2 - )) + self.assertTrue( + torch.allclose(ctx.m(runtime_input), runtime_input.detach() * 2) + ) def test_detach_in_place_returns_the_receiver(self): """``detach_`` keeps its receiver so chained in-place idioms survive. @@ -807,7 +808,7 @@ def test_cpu(self): self.assertIsInstance(y, TracedTensor) self.assertTrue(torch.allclose(y.tensor, x.tensor)) - self.assertEqual(y.device.type, 'cpu') + self.assertEqual(y.device.type, "cpu") # ==================== Additional Math Operations ==================== @@ -844,7 +845,7 @@ def test_exp_log(self): def test_sin_cos(self): """Test sin and cos operations.""" ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor([0.0, 3.14159/2, 3.14159]), name="x") + x = ctx.create_input(torch.tensor([0.0, 3.14159 / 2, 3.14159]), name="x") sin_x = torch.sin(x) cos_x = torch.cos(x) @@ -1043,11 +1044,11 @@ def test_comparison_operators(self): def test_inplace_operators_preserve_receiver_and_compile(self): """In-place operators mutate one receiver and replay functionally.""" cases = ( - ("add", operator.iadd, torch.tensor([1., 2., 3.]), 2.0), - ("sub", operator.isub, torch.tensor([5., 6., 7.]), 2.0), - ("mul", operator.imul, torch.tensor([1., 2., 3.]), 3.0), - ("div", operator.itruediv, torch.tensor([4., 8., 12.]), 2.0), - ("pow", operator.ipow, torch.tensor([2., 3., 4.]), 2.0), + ("add", operator.iadd, torch.tensor([1.0, 2.0, 3.0]), 2.0), + ("sub", operator.isub, torch.tensor([5.0, 6.0, 7.0]), 2.0), + ("mul", operator.imul, torch.tensor([1.0, 2.0, 3.0]), 3.0), + ("div", operator.itruediv, torch.tensor([4.0, 8.0, 12.0]), 2.0), + ("pow", operator.ipow, torch.tensor([2.0, 3.0, 4.0]), 2.0), ) for name, operation, input_tensor, operand in cases: @@ -1062,10 +1063,8 @@ def test_inplace_operators_preserve_receiver_and_compile(self): self.assertIsInstance(receiver, TracedTensor) self.assertTrue(torch.allclose(receiver.tensor, expected)) - ctx.compile_trace({'x': receiver}) - self.assertTrue(torch.allclose( - ctx.m(input_tensor.clone()), expected - )) + ctx.compile_trace({"x": receiver}) + self.assertTrue(torch.allclose(ctx.m(input_tensor.clone()), expected)) def test_inplace_method_add_(self): """Test in-place method add_().""" @@ -1082,7 +1081,7 @@ def test_inplace_method_add_(self): expected = torch.tensor([3.0, 4.0, 5.0]) self.assertTrue(torch.allclose(x.tensor, expected)) - ctx.compile_trace({'x': x}) + ctx.compile_trace({"x": x}) graph_result = ctx.m(torch.tensor([1.0, 2.0, 3.0])) self.assertTrue(torch.allclose(graph_result, expected)) @@ -1098,7 +1097,7 @@ def test_inplace_method_mul_(self): expected = torch.tensor([3.0, 6.0, 9.0]) self.assertTrue(torch.allclose(x.tensor, expected)) - ctx.compile_trace({'x': x}) + ctx.compile_trace({"x": x}) graph_result = ctx.m(torch.tensor([1.0, 2.0, 3.0])) self.assertTrue(torch.allclose(graph_result, expected)) @@ -1111,7 +1110,7 @@ def test_inplace_relu_on_inference_tensor(self): y = x * 2 torch.nn.functional.relu(y, inplace=True) - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) targets = [node.target for node in ctx.m.graph.nodes] self.assertIn(torch.nn.functional.relu, targets) @@ -1133,7 +1132,7 @@ def test_equivalent_copies_preserve_port_before_inference_mutation(self): x = ctx.create_input(torch.tensor([-1.0, 2.0, -3.0]), name="x") with mode(): y = x * 2 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) ctx.publish_output_port(y, "y") self.assertFalse(y.is_tracing) @@ -1156,19 +1155,19 @@ def test_inplace_index_add_preserves_receiver_and_exports(self): with self.subTest(traced_index=traced_index): ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.zeros(3), name="x") - source = ctx.create_input(torch.tensor([10., 20.]), name="source") + source = ctx.create_input(torch.tensor([10.0, 20.0]), name="source") index = torch.tensor([0, 2]) if traced_index: index = ctx.create_input(index, name="index") self.assertIs(x.index_add_(0, index, source), x) - self.assertTrue(torch.equal(x.tensor, torch.tensor([10., 0., 20.]))) + self.assertTrue(torch.equal(x.tensor, torch.tensor([10.0, 0.0, 20.0]))) x.index_add_(0, index, source, alpha=2) - self.assertTrue(torch.equal(x.tensor, torch.tensor([30., 0., 60.]))) - ctx.compile_trace({'out': x}) + self.assertTrue(torch.equal(x.tensor, torch.tensor([30.0, 0.0, 60.0]))) + ctx.compile_trace({"out": x}) - runtime_x = torch.tensor([1., 2., 3.]) - runtime_source = torch.tensor([4., 5.]) + runtime_x = torch.tensor([1.0, 2.0, 3.0]) + runtime_source = torch.tensor([4.0, 5.0]) runtime_index = torch.tensor([1, 1]) if traced_index else index expected = runtime_x.clone() expected.index_add_(0, runtime_index, runtime_source) @@ -1177,7 +1176,9 @@ def test_inplace_index_add_preserves_receiver_and_exports(self): if traced_index: inputs += (runtime_index,) self.validate_export( - ctx.m, inputs, expected, + ctx.m, + inputs, + expected, test_name="inplace_index_add", verify_legacy_onnx=False, verify_dynamo_onnx=True, @@ -1188,8 +1189,8 @@ def test_inplace_index_add_transit_invalidates_port(self): ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.zeros(2, 3), name="x") index = ctx.create_input(torch.tensor([0, 2]), name="index") - source = ctx.create_input(torch.tensor([[1., 2.], [3., 4.]]), name="source") - ctx.compile_trace({'x': x, 'index': index, 'source': source}) + source = ctx.create_input(torch.tensor([[1.0, 2.0], [3.0, 4.0]]), name="source") + ctx.compile_trace({"x": x, "index": index, "source": source}) ctx.publish_output_port(x, "x") self.assertFalse(x.is_tracing) self.assertEqual(x.output_port, "x") @@ -1211,7 +1212,7 @@ def test_inplace_method_fill_discarded_result(self): expected = torch.tensor([3.0, 3.0, 3.0]) self.assertTrue(torch.allclose(x.tensor, expected)) - ctx.compile_trace({'x': x}) + ctx.compile_trace({"x": x}) graph_result = ctx.m(torch.tensor([1.0, 2.0, 3.0])) self.assertTrue(torch.allclose(graph_result, expected)) @@ -1227,7 +1228,7 @@ def test_inplace_method_copy_discarded_result(self): self.assertEqual(id(x), original_id) self.assertTrue(torch.allclose(x.tensor, source)) - ctx.compile_trace({'x': x}) + ctx.compile_trace({"x": x}) graph_result = ctx.m(torch.tensor([1.0, 2.0, 3.0])) self.assertTrue(torch.allclose(graph_result, source)) @@ -1247,7 +1248,7 @@ def test_inplace_chained_operations(self): expected = torch.tensor([1.0, 3.0, 5.0]) self.assertTrue(torch.allclose(x.tensor, expected)) - ctx.compile_trace({'x': x}) + ctx.compile_trace({"x": x}) result = ctx.m(torch.tensor([1.0, 2.0, 3.0])) self.assertTrue(torch.allclose(result, expected)) @@ -1262,7 +1263,7 @@ def test_inplace_with_tensor_operand(self): expected = torch.tensor([11.0, 22.0, 33.0]) self.assertTrue(torch.allclose(x.tensor, expected)) - ctx.compile_trace({'x': x}) + ctx.compile_trace({"x": x}) result = ctx.m(torch.tensor([1.0, 2.0, 3.0])) self.assertTrue(torch.allclose(result, expected)) @@ -1276,7 +1277,7 @@ def test_inplace_after_compile_preserves_traced_tensor(self): x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") y = x * 2 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) result = x result += 1 @@ -1295,7 +1296,7 @@ def test_inplace_after_compile_all_operators(self): # Test += ctx = TracedTensorNode(name="test_add", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - ctx.compile_trace({'y': x * 2}) + ctx.compile_trace({"y": x * 2}) x += 10 self.assertIsInstance(x, TracedTensor) self.assertTrue(torch.allclose(x.tensor, torch.tensor([11.0, 12.0, 13.0]))) @@ -1303,7 +1304,7 @@ def test_inplace_after_compile_all_operators(self): # Test -= ctx = TracedTensorNode(name="test_sub", node_index=0) x = ctx.create_input(torch.tensor([10.0, 20.0, 30.0]), name="x") - ctx.compile_trace({'y': x * 2}) + ctx.compile_trace({"y": x * 2}) x -= 5 self.assertIsInstance(x, TracedTensor) self.assertTrue(torch.allclose(x.tensor, torch.tensor([5.0, 15.0, 25.0]))) @@ -1311,7 +1312,7 @@ def test_inplace_after_compile_all_operators(self): # Test *= ctx = TracedTensorNode(name="test_mul", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - ctx.compile_trace({'y': x + 1}) + ctx.compile_trace({"y": x + 1}) x *= 3 self.assertIsInstance(x, TracedTensor) self.assertTrue(torch.allclose(x.tensor, torch.tensor([3.0, 6.0, 9.0]))) @@ -1319,7 +1320,7 @@ def test_inplace_after_compile_all_operators(self): # Test /= ctx = TracedTensorNode(name="test_div", node_index=0) x = ctx.create_input(torch.tensor([10.0, 20.0, 30.0]), name="x") - ctx.compile_trace({'y': x + 1}) + ctx.compile_trace({"y": x + 1}) x /= 2 self.assertIsInstance(x, TracedTensor) self.assertTrue(torch.allclose(x.tensor, torch.tensor([5.0, 10.0, 15.0]))) @@ -1327,7 +1328,7 @@ def test_inplace_after_compile_all_operators(self): # Test **= ctx = TracedTensorNode(name="test_pow", node_index=0) x = ctx.create_input(torch.tensor([2.0, 3.0, 4.0]), name="x") - ctx.compile_trace({'y': x + 1}) + ctx.compile_trace({"y": x + 1}) x **= 2 self.assertIsInstance(x, TracedTensor) self.assertTrue(torch.allclose(x.tensor, torch.tensor([4.0, 9.0, 16.0]))) @@ -1336,13 +1337,13 @@ def test_inplace_after_compile_chained(self): """Test chained in-place operations after compile.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - ctx.compile_trace({'y': x * 2}) + ctx.compile_trace({"y": x * 2}) # Chain multiple in-place operations - x += 1 # [2, 3, 4] - x *= 2 # [4, 6, 8] - x -= 3 # [1, 3, 5] - x /= 2 # [0.5, 1.5, 2.5] + x += 1 # [2, 3, 4] + x *= 2 # [4, 6, 8] + x -= 3 # [1, 3, 5] + x /= 2 # [0.5, 1.5, 2.5] self.assertIsInstance(x, TracedTensor) self.assertFalse(x.is_tracing) @@ -1353,7 +1354,7 @@ def test_inplace_methods_after_compile(self): # Test add_() ctx = TracedTensorNode(name="test_add", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - ctx.compile_trace({'y': x * 2}) + ctx.compile_trace({"y": x * 2}) result = x.add_(10) self.assertIs(result, x) self.assertTrue(torch.allclose(result.tensor, torch.tensor([11.0, 12.0, 13.0]))) @@ -1361,7 +1362,7 @@ def test_inplace_methods_after_compile(self): # Test sub_() ctx = TracedTensorNode(name="test_sub", node_index=0) x = ctx.create_input(torch.tensor([10.0, 20.0, 30.0]), name="x") - ctx.compile_trace({'y': x * 2}) + ctx.compile_trace({"y": x * 2}) result = x.sub_(5) self.assertIs(result, x) self.assertTrue(torch.allclose(result.tensor, torch.tensor([5.0, 15.0, 25.0]))) @@ -1369,7 +1370,7 @@ def test_inplace_methods_after_compile(self): # Test mul_() ctx = TracedTensorNode(name="test_mul", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - ctx.compile_trace({'y': x + 1}) + ctx.compile_trace({"y": x + 1}) result = x.mul_(5) self.assertIs(result, x) self.assertTrue(torch.allclose(result.tensor, torch.tensor([5.0, 10.0, 15.0]))) @@ -1377,7 +1378,7 @@ def test_inplace_methods_after_compile(self): # Test div_() ctx = TracedTensorNode(name="test_div", node_index=0) x = ctx.create_input(torch.tensor([10.0, 20.0, 30.0]), name="x") - ctx.compile_trace({'y': x + 1}) + ctx.compile_trace({"y": x + 1}) result = x.div_(2) self.assertIs(result, x) self.assertTrue(torch.allclose(result.tensor, torch.tensor([5.0, 10.0, 15.0]))) @@ -1385,7 +1386,7 @@ def test_inplace_methods_after_compile(self): # Test pow_() ctx = TracedTensorNode(name="test_pow", node_index=0) x = ctx.create_input(torch.tensor([2.0, 3.0, 4.0]), name="x") - ctx.compile_trace({'y': x + 1}) + ctx.compile_trace({"y": x + 1}) result = x.pow_(2) self.assertIs(result, x) self.assertTrue(torch.allclose(result.tensor, torch.tensor([4.0, 9.0, 16.0]))) @@ -1393,7 +1394,7 @@ def test_inplace_methods_after_compile(self): # Test copy_() ctx = TracedTensorNode(name="test_copy", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - ctx.compile_trace({'y': x * 2}) + ctx.compile_trace({"y": x * 2}) x.copy_(torch.tensor([100.0, 200.0, 300.0])) self.assertIsInstance(x, TracedTensor) self.assertTrue(torch.allclose(x.tensor, torch.tensor([100.0, 200.0, 300.0]))) @@ -1402,18 +1403,22 @@ def test_inplace_with_broadcasting(self): """Test in-place operations with broadcasting (2D tensor + 1D tensor).""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), name="x") - + # Add a 1D tensor that broadcasts across rows x += torch.tensor([10.0, 20.0, 30.0]) - + self.assertIsInstance(x, TracedTensor) expected = torch.tensor([[11.0, 22.0, 33.0], [14.0, 25.0, 36.0]]) self.assertTrue(torch.allclose(x.tensor, expected)) - + # Verify it works after compile too - ctx.compile_trace({'x': x}) + ctx.compile_trace({"x": x}) result = ctx.m(torch.tensor([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]])) - self.assertTrue(torch.allclose(result, torch.tensor([[11.0, 21.0, 31.0], [12.0, 22.0, 32.0]]))) + self.assertTrue( + torch.allclose( + result, torch.tensor([[11.0, 21.0, 31.0], [12.0, 22.0, 32.0]]) + ) + ) # ==================== Advanced Indexing Tests ==================== @@ -1431,7 +1436,7 @@ def test_indexing_with_none_newaxis(self): self.assertTrue(torch.allclose(y.tensor, expected)) # Test compilation - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) result = ctx.m(torch.tensor([1.0, 2.0, 3.0])) self.assertTrue(torch.allclose(result, expected)) @@ -1463,7 +1468,7 @@ def test_indexing_with_ellipsis(self): self.assertEqual(y.shape, torch.Size([2, 3])) # Test compilation - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) input_tensor = torch.randn(2, 3, 4) result = ctx.m(input_tensor) expected = input_tensor[..., 0] @@ -1491,15 +1496,14 @@ def test_indexing_empty_slice(self): self.assertEqual(y.shape, torch.Size([0])) # Test compilation - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) result = ctx.m(torch.tensor([1.0, 2.0, 3.0])) self.assertEqual(result.shape, torch.Size([0])) def test_indexing_with_list_advanced(self): """Test advanced integer list indexing.""" ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(torch.tensor( - [10.0, 20.0, 30.0, 40.0, 50.0]), name="x") + x = ctx.create_input(torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0]), name="x") # Index with a list y = x[[0, 2, 4]] @@ -1510,9 +1514,8 @@ def test_indexing_with_list_advanced(self): self.assertTrue(torch.allclose(y.tensor, expected)) # Test compilation - ctx.compile_trace({'y': y}) - result = ctx.m( - torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0])) + ctx.compile_trace({"y": y}) + result = ctx.m(torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0])) self.assertTrue(torch.allclose(result, expected)) def test_indexing_with_list_2d(self): @@ -1540,7 +1543,7 @@ def test_torch_index_select_with_traced_indices(self): expected = torch.tensor([10.0, 30.0, 50.0]) self.assertTrue(torch.allclose(y.tensor, expected)) - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) self.validate_export( ctx.m, (input_tensor, input_indices), @@ -1560,8 +1563,7 @@ def test_indexing_negative_indices_advanced(self): # Negative slice ctx2 = TracedTensorNode(name="test2", node_index=1) - x2 = ctx2.create_input(torch.tensor( - [1.0, 2.0, 3.0, 4.0, 5.0]), name="x") + x2 = ctx2.create_input(torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]), name="x") z = x2[-3:-1] self.assertEqual(z.shape, torch.Size([2])) expected = torch.tensor([3.0, 4.0]) @@ -1579,7 +1581,7 @@ def test_indexing_combination_advanced(self): self.assertEqual(y.shape, torch.Size([3, 2])) # Test compilation - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) input_tensor = torch.randn(2, 3, 4, 5) result = ctx.m(input_tensor) expected = input_tensor[0, :, 1:3, -2] @@ -1607,7 +1609,7 @@ def test_variable_reassignment(self): self.assertTrue(torch.allclose(z.tensor, expected)) # Test compilation - ctx.compile_trace({'z': z}) + ctx.compile_trace({"z": z}) result = ctx.m(torch.tensor([1.0, 2.0, 3.0])) self.assertTrue(torch.allclose(result, expected)) @@ -1630,7 +1632,7 @@ def test_multiple_outputs_with_shared_intermediate(self): out3 = y * 3 # [6, 12, 18] # Compile with multiple outputs - ctx.compile_trace({'out1': out1, 'out2': out2, 'out3': out3}) + ctx.compile_trace({"out1": out1, "out2": out2, "out3": out3}) # Test execution result = ctx.m(torch.tensor([1.0, 2.0, 3.0])) @@ -1638,12 +1640,9 @@ def test_multiple_outputs_with_shared_intermediate(self): # When multiple outputs, result should be a tuple if isinstance(result, tuple): self.assertEqual(len(result), 3) - self.assertTrue(torch.allclose( - result[0], torch.tensor([3.0, 5.0, 7.0]))) - self.assertTrue(torch.allclose( - result[1], torch.tensor([1.0, 3.0, 5.0]))) - self.assertTrue(torch.allclose( - result[2], torch.tensor([6.0, 12.0, 18.0]))) + self.assertTrue(torch.allclose(result[0], torch.tensor([3.0, 5.0, 7.0]))) + self.assertTrue(torch.allclose(result[1], torch.tensor([1.0, 3.0, 5.0]))) + self.assertTrue(torch.allclose(result[2], torch.tensor([6.0, 12.0, 18.0]))) def test_diamond_dependency_graph(self): """Test diamond-shaped dependency graph (one input, multiple paths, one output).""" @@ -1665,7 +1664,7 @@ def test_diamond_dependency_graph(self): self.assertTrue(torch.allclose(c.tensor, expected)) # Test compilation - ctx.compile_trace({'c': c}) + ctx.compile_trace({"c": c}) result = ctx.m(torch.tensor([1.0, 2.0, 3.0])) self.assertTrue(torch.allclose(result, expected)) @@ -1684,7 +1683,7 @@ def test_deep_computation_chain(self): self.assertTrue(torch.allclose(result.tensor, expected, atol=1e-5)) # Test compilation - ctx.compile_trace({'result': result}) + ctx.compile_trace({"result": result}) output = ctx.m(torch.tensor([1.0])) self.assertTrue(torch.allclose(output, expected, atol=1e-5)) @@ -1700,7 +1699,7 @@ def test_unused_intermediate_values(self): # Only use x directly result = x + 1 - ctx.compile_trace({'result': result}) + ctx.compile_trace({"result": result}) # Check that the graph doesn't contain the unused operations graph_str = str(ctx.m.graph) @@ -1708,7 +1707,7 @@ def test_unused_intermediate_values(self): # The graph should have: placeholder (x), add (x + 1), output # Count the number of operations (excluding placeholder and output) nodes = list(ctx.m.graph.nodes) - call_function_nodes = [n for n in nodes if n.op == 'call_function'] + call_function_nodes = [n for n in nodes if n.op == "call_function"] # Should only have 1 call_function (the add for result) self.assertEqual(len(call_function_nodes), 1) @@ -1731,7 +1730,7 @@ def test_complex_branching_graph(self): self.assertIsInstance(result, TracedTensor) # Test compilation - ctx.compile_trace({'result': result}) + ctx.compile_trace({"result": result}) output = ctx.m(torch.tensor([1.0, 2.0, 3.0])) # Verify result matches @@ -1797,8 +1796,8 @@ def test_iter_tracedtensor_1d(self): for elem in elements: # Could be either TracedTensor or torch.Tensor self.assertTrue( - isinstance(elem, (TracedTensor, torch.Tensor)) or - isinstance(elem, (int, float)) + isinstance(elem, (TracedTensor, torch.Tensor)) + or isinstance(elem, (int, float)) ) except TypeError as e: # If iteration is not supported, that's okay @@ -1857,19 +1856,13 @@ def test_reversed_tracedtensor(self): # Check if order is reversed # Note: elements might be TracedTensors or tensors if isinstance(reversed_x[0], TracedTensor): - self.assertTrue(torch.allclose( - reversed_x[0].tensor, torch.tensor(3.0) - )) - self.assertTrue(torch.allclose( - reversed_x[-1].tensor, torch.tensor(1.0) - )) + self.assertTrue(torch.allclose(reversed_x[0].tensor, torch.tensor(3.0))) + self.assertTrue( + torch.allclose(reversed_x[-1].tensor, torch.tensor(1.0)) + ) elif isinstance(reversed_x[0], torch.Tensor): - self.assertTrue(torch.allclose( - reversed_x[0], torch.tensor(3.0) - )) - self.assertTrue(torch.allclose( - reversed_x[-1], torch.tensor(1.0) - )) + self.assertTrue(torch.allclose(reversed_x[0], torch.tensor(3.0))) + self.assertTrue(torch.allclose(reversed_x[-1], torch.tensor(1.0))) except TypeError as e: # If reversed is not supported, check the error self.fail(f"reversed() failed: {e}") @@ -1931,7 +1924,7 @@ def test_boolean_indexing_with_traced_mask_lowers_to_masked_select(self): expected = torch.tensor([2.0, 3.0]) self.assertTrue(torch.allclose(y.tensor, expected)) - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) self.validate_export( ctx.m, (input_x,), @@ -1953,7 +1946,7 @@ def test_advanced_indexing_with_traced_indices_lowers_to_index_select(self): expected = torch.tensor([10.0, 30.0, 50.0]) self.assertTrue(torch.allclose(y.tensor, expected)) - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) self.validate_export( ctx.m, (input_x, input_indices), @@ -1989,7 +1982,7 @@ def test_nd_advanced_indexing_with_traced_indices_lowers_to_index_select(self): self.assertEqual(y.shape, expected.shape) self.assertTrue(torch.allclose(y.tensor, expected)) - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) self.validate_export( ctx.m, (input_x, input_indices), @@ -2014,10 +2007,9 @@ def test_all_arithmetic_functions(self): expected = func(input_tensor) ctx = TracedTensorNode(name="test", node_index=0) - input_tracked = ctx.create_input( - input_tensor, name="my_test_input") + input_tracked = ctx.create_input(input_tensor, name="my_test_input") result_tracked = func(input_tracked) - ctx.compile_trace({'result_tracked': result_tracked}) + ctx.compile_trace({"result_tracked": result_tracked}) graph_module = ctx.m self.assertIsInstance( @@ -2065,8 +2057,7 @@ def test_all_arithmetic_functions(self): ) # Test 4: ExportedProgram export - exported = torch.export.export( - graph_module, (input_tensor.clone(),)) + exported = torch.export.export(graph_module, (input_tensor.clone(),)) result_exported = exported.module()(input_tensor) self.assertTrue( torch.allclose(result_exported, expected), @@ -2091,13 +2082,14 @@ def test_all_arithmetic_functions(self): export_params=True, opset_version=17, do_constant_folding=True, - input_names=['my_test_input'], - output_names=['result_tracked'], + input_names=["my_test_input"], + output_names=["result_tracked"], ) session = ort.InferenceSession(str(onnx_path)) result_onnx = session.run( - None, {"my_test_input": input_tensor.numpy()})[0] + None, {"my_test_input": input_tensor.numpy()} + )[0] self.assertTrue( torch.allclose( @@ -2120,7 +2112,7 @@ def test_all_arithmetic_functions(self): def test_setitem_full_slice_with_traced_tensor(self): """Test full slice assignment where value is another TracedTensor. - + buffer[:] = y uses index_put(buffer, indices, y), so the graph references both buffer and x (through y = x * 2). """ @@ -2141,8 +2133,8 @@ def test_setitem_full_slice_with_traced_tensor(self): # Test compilation - graph uses index_put(buffer, indices, y) # which references both buffer and x (via y) - ctx.compile_trace({'buffer': buffer}) - + ctx.compile_trace({"buffer": buffer}) + # Run GraphModule with both inputs (x and buffer) input_x = torch.tensor([1.0, 2.0, 3.0]) input_buffer = torch.zeros(3) @@ -2153,7 +2145,7 @@ def test_setitem_full_slice_with_traced_tensor(self): scripted = torch.jit.script(ctx.m) result_scripted = scripted(input_x, input_buffer) self.assertTrue(torch.allclose(result_scripted, expected)) - + # Test with different x input to ensure it's actually computing from x result_different = ctx.m(torch.tensor([10.0, 20.0, 30.0]), input_buffer) expected_different = torch.tensor([20.0, 40.0, 60.0]) @@ -2185,7 +2177,7 @@ def test_full_assignment_into_raw_buffer_keeps_buffer_memory(self): self.assertEqual(buffer.data_ptr(), buffer_data_ptr) self.assertTrue(torch.allclose(buffer.tensor, expected)) - ctx.compile_trace({'result': buffer}) + ctx.compile_trace({"result": buffer}) output = ctx.m(torch.tensor([1.0, 2.0, 3.0])) self.assertTrue(torch.allclose(output, expected)) @@ -2195,7 +2187,7 @@ def test_setitem_after_compile_no_trace(self): x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") y = x * 2 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) # After compile, setitem should still work but not trace x[:] = torch.tensor([10.0, 20.0, 30.0]) @@ -2219,12 +2211,10 @@ def test_setitem_single_index_fx_torchscript_onnx(self): x[0] = 10.0 self.assertIs(x, receiver) self.assertIsInstance(x, TracedTensor) - self.assertTrue(torch.equal( - x.tensor, torch.tensor([10.0, 2.0, 3.0, 4.0, 5.0]) - )) + self.assertTrue(torch.equal(x.tensor, torch.tensor([10.0, 2.0, 3.0, 4.0, 5.0]))) y = x * 2 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) input_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]) expected = torch.tensor([20.0, 4.0, 6.0, 8.0, 10.0]) @@ -2239,12 +2229,12 @@ def test_setitem_slice_fx_torchscript_onnx(self): x[1:3] = torch.tensor([10.0, 20.0]) self.assertIs(x, receiver) self.assertIsInstance(x, TracedTensor) - self.assertTrue(torch.equal( - x.tensor, torch.tensor([1.0, 10.0, 20.0, 4.0, 5.0]) - )) + self.assertTrue( + torch.equal(x.tensor, torch.tensor([1.0, 10.0, 20.0, 4.0, 5.0])) + ) y = x * 2 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) input_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]) expected = torch.tensor([2.0, 20.0, 40.0, 8.0, 10.0]) @@ -2259,12 +2249,10 @@ def test_setitem_full_slice_constant_fx_torchscript_onnx(self): x[:] = torch.tensor([10.0, 20.0, 30.0]) self.assertIs(x, receiver) self.assertIsInstance(x, TracedTensor) - self.assertTrue(torch.equal( - x.tensor, torch.tensor([10.0, 20.0, 30.0]) - )) + self.assertTrue(torch.equal(x.tensor, torch.tensor([10.0, 20.0, 30.0]))) y = x * 2 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) input_tensor = torch.tensor([1.0, 2.0, 3.0]) expected = torch.tensor([20.0, 40.0, 60.0]) @@ -2274,12 +2262,12 @@ def test_setitem_step_slice_fx_torchscript_onnx(self): """Test: x[::2] = [10, 20, 30] with step exports correctly.""" ctx = TracedTensorNode(name="test_setitem_step_export", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]), name="x") - + x[::2] = torch.tensor([10.0, 30.0, 50.0]) # Assigns to indices 0, 2, 4 y = x * 2 - - ctx.compile_trace({'y': y}) - + + ctx.compile_trace({"y": y}) + input_tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]) expected = torch.tensor([20.0, 4.0, 60.0, 8.0, 100.0]) self.validate_export(ctx.m, (input_tensor,), expected, "setitem_step") @@ -2288,30 +2276,38 @@ def test_setitem_with_traced_tensor_slice_fx_torchscript_onnx(self): """Test traced integer and slice sources in one assignment chain exports correctly.""" ctx = TracedTensorNode(name="test_setitem_cross_tensor", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]), name="x") - other = ctx.create_input(torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0]), name="other") + other = ctx.create_input( + torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0]), name="other" + ) x[0] = other[4] x[2:4] = other[1:3] y = x * 2 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) input_x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]) input_other = torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0]) # x after mutations: [50.0, 2.0, 20.0, 30.0, 5.0] expected = torch.tensor([100.0, 4.0, 40.0, 60.0, 10.0]) - self.validate_export(ctx.m, (input_x, input_other), expected, "setitem_cross_tensor") + self.validate_export( + ctx.m, (input_x, input_other), expected, "setitem_cross_tensor" + ) def test_setitem_multidim_partial_slice_fx_torchscript_dynamo_onnx(self): """Test tuple slice setitem for export-first backends.""" - ctx = TracedTensorNode(name="test_setitem_multidim_partial_export", node_index=0) + ctx = TracedTensorNode( + name="test_setitem_multidim_partial_export", node_index=0 + ) x = ctx.create_input(torch.arange(36.0).reshape(1, 6, 6), name="x") - other = ctx.create_input(torch.arange(36.0, 72.0).reshape(1, 6, 6), name="other") + other = ctx.create_input( + torch.arange(36.0, 72.0).reshape(1, 6, 6), name="other" + ) x[:, 0:3, 0:3] = other[:, 0:3, 0:3] y = x[:, 0:3, 0:3] * 2.0 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) self.assertNotIn("__setitem__", str(ctx.m.graph)) @@ -2340,7 +2336,7 @@ def test_setitem_multidim_partial_imul_fx_torchscript_dynamo_onnx(self): x[:, 7:13] *= 100.0 y = x + 1.0 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) self.assertNotIn("__setitem__", str(ctx.m.graph)) @@ -2364,7 +2360,7 @@ def test_setitem_integer_step_slice_scalar_fx_torchscript_dynamo_onnx(self): x[1, 1:5:2] = 9.0 y = x * 2.0 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) self.assertNotIn("__setitem__", str(ctx.m.graph)) @@ -2486,7 +2482,9 @@ def test_setitem_with_traced_integer_index(self): ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.arange(8.0).reshape(2, 4), name="x") index = ctx.create_input(torch.tensor([1, 3]), name="index") - source = ctx.create_input(torch.tensor([[20.0, 30.0], [40.0, 50.0]]), name="source") + source = ctx.create_input( + torch.tensor([[20.0, 30.0], [40.0, 50.0]]), name="source" + ) x[:, index] = source y = x * 2.0 @@ -2534,7 +2532,7 @@ def test_setitem_with_traced_boolean_mask(self): verify_legacy_onnx=False, verify_dynamo_onnx=True, ) - + def test_plain_destination_with_traced_index(self): """A traced key alone is enough to promote a plain destination.""" ctx = TracedTensorNode(name="test", node_index=0) @@ -2576,7 +2574,7 @@ def test_class_swap_full_slice(self): y = buf * 2 self.assertIsInstance(y, TracedTensor) - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) input_tensor = torch.tensor([4.0, 5.0, 6.0]) expected = input_tensor * 2 @@ -2593,7 +2591,7 @@ def test_class_swap_ellipsis(self): self.assertIsInstance(buf, TracedTensor) y = buf + 1.0 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) input_tensor = torch.tensor([4.0, 5.0, 6.0]) expected = input_tensor + 1.0 @@ -2610,7 +2608,7 @@ def test_class_swap_copy_(self): self.assertIsInstance(buf, TracedTensor) y = buf - 1.0 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) input_tensor = torch.tensor([1.0, 2.0, 3.0]) expected = input_tensor - 1.0 @@ -2674,24 +2672,18 @@ def _validate_plain_inplace_method( ) destination = initial.clone() - result = getattr(destination, method)( - *traced_operands, *extra_args, **kwargs - ) + result = getattr(destination, method)(*traced_operands, *extra_args, **kwargs) self.assertIs(result, destination) self.assertIsInstance(destination, TracedTensor) eager_expected = initial.clone() - getattr(eager_expected, method)( - *trace_operands, *extra_args, **kwargs - ) + getattr(eager_expected, method)(*trace_operands, *extra_args, **kwargs) self.assertEqual(destination.dtype, eager_expected.dtype) self.assertTrue(torch.allclose(destination.tensor, eager_expected)) ctx.compile_trace({"out": destination}) runtime_expected = initial.clone() - getattr(runtime_expected, method)( - *runtime_operands, *extra_args, **kwargs - ) + getattr(runtime_expected, method)(*runtime_operands, *extra_args, **kwargs) self.validate_export( ctx.m, runtime_operands, @@ -2881,9 +2873,7 @@ def test_registered_inplace_method_mutates_traced_receiver(self): result = receiver.remainder_(other) self.assertIs(result, receiver) - self.assertTrue( - torch.allclose(receiver.tensor, torch.tensor([1.0, 3.0, 4.0])) - ) + self.assertTrue(torch.allclose(receiver.tensor, torch.tensor([1.0, 3.0, 4.0]))) ctx.compile_trace({"out": receiver}) runtime_receiver = torch.tensor([8.0, 11.0, 14.0]) runtime_other = torch.tensor([3.0, 5.0, 6.0]) @@ -2957,7 +2947,7 @@ def test_inplace_masked_fill_exports_functionally(self): self.assertIs(result, receiver) self.assertTrue( - torch.equal(receiver.tensor, torch.tensor([-1., 2., -1., 4.])) + torch.equal(receiver.tensor, torch.tensor([-1.0, 2.0, -1.0, 4.0])) ) ctx.compile_trace({"out": receiver}) runtime_input = torch.tensor([10.0, 20.0, 30.0, 40.0]) @@ -2989,9 +2979,7 @@ def test_inplace_index_copy_exports_functionally(self): runtime_input = torch.full((4, 2), 5.0) runtime_index = torch.tensor([1, 3]) runtime_source = torch.tensor([[10.0, 20.0], [30.0, 40.0]]) - runtime_expected = runtime_input.index_copy( - 0, runtime_index, runtime_source - ) + runtime_expected = runtime_input.index_copy(0, runtime_index, runtime_source) self.validate_export( ctx.m, (runtime_input, runtime_index, runtime_source), @@ -3030,18 +3018,14 @@ def test_inplace_schema_discovery_is_conservative(self): } for method, expected_spec in expected.items(): with self.subTest(method=method): - functional, mutates_metadata = ( - TracedTensor._functional_inplace_spec(method) - ) - self.assertEqual( - (functional.__name__, mutates_metadata), expected_spec + functional, mutates_metadata = TracedTensor._functional_inplace_spec( + method ) + self.assertEqual((functional.__name__, mutates_metadata), expected_spec) for method in ("normal_", "polygamma_", "resize_"): with self.subTest(method=method): - self.assertIsNone( - TracedTensor._functional_inplace_spec(method) - ) + self.assertIsNone(TracedTensor._functional_inplace_spec(method)) def test_schema_discovered_comparison_preserves_receiver_dtype(self): """Functional comparison output is cast to the receiver's dtype.""" @@ -3058,9 +3042,7 @@ def test_schema_discovered_comparison_preserves_receiver_dtype(self): ctx.compile_trace({"out": receiver}) runtime_input = torch.tensor([4.0, 16.0, 25.0]) runtime_other = torch.tensor([4.0, 5.0, 25.0]) - runtime_expected = runtime_input.eq(runtime_other).to( - runtime_input.dtype - ) + runtime_expected = runtime_input.eq(runtime_other).to(runtime_input.dtype) self.validate_export( ctx.m, (runtime_input, runtime_other), @@ -3193,7 +3175,7 @@ def test_class_swap_multidim(self): self.assertIsInstance(buf, TracedTensor) y = torch.relu(buf) - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) input_tensor = torch.randn(3, 4) expected = torch.relu(input_tensor) @@ -3211,7 +3193,7 @@ def test_class_swap_chained(self): buf2[:] = intermediate y = buf2 + 10.0 - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) input_tensor = torch.tensor([4.0, 5.0, 6.0]) expected = input_tensor * 3.0 + 10.0 @@ -3253,7 +3235,7 @@ def test_class_swap_does_not_fire_for_traced_target(self): self.assertIsInstance(x, TracedTensor) z = x * 2 - ctx.compile_trace({'z': z}) + ctx.compile_trace({"z": z}) input_x = torch.tensor([0.0, 0.0, 0.0]) input_y = torch.tensor([7.0, 8.0, 9.0]) @@ -3266,9 +3248,7 @@ def test_tuple_traced_boolean_mask_augmented_assignment(self): """x[:, mask] *= scalar lowers as getitem, multiply, then setitem.""" ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.arange(8.0).reshape(2, 4), name="x") - mask = ctx.create_input( - torch.tensor([True, False, True, False]), name="mask" - ) + mask = ctx.create_input(torch.tensor([True, False, True, False]), name="mask") x[:, mask] *= 3.0 y = x + 1.0 @@ -3332,12 +3312,12 @@ def test_tuple_unsupported_traced_index_raises_not_implemented(self): # ==================== TorchScript Module Interaction Tests ==================== def test_tracing_through_torchscript_module(self): """Test that TorchScript module with parameters raises error with TracedTensor.""" + # Create a TorchScript module with parameters class LinearModule(torch.nn.Module): def __init__(self): super().__init__() - self.weight = torch.nn.Parameter( - torch.tensor([[2.0, 3.0, 4.0]])) + self.weight = torch.nn.Parameter(torch.tensor([[2.0, 3.0, 4.0]])) def forward(self, x): return torch.matmul(self.weight, x) @@ -3349,7 +3329,7 @@ def forward(self, x): ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") y = scripted_module(x) - ctx.compile_trace({'y': y}) + ctx.compile_trace({"y": y}) torch.manual_seed(42) input_tensor = torch.randn((3,)) @@ -3359,6 +3339,7 @@ def forward(self, x): def test_torchscript_with_normal_normal_torch_operations(self): """Test that TorchScript module nested with traced ops raises error.""" + class ReLUModule(torch.nn.Module): def forward(self, x): return torch.relu(x) @@ -3370,22 +3351,24 @@ def forward(self, x): def operation_chain(x): y = x * 2 - z = torch.abs(y)-1 + z = torch.abs(y) - 1 z = scripted_relu(z) return z + z = operation_chain(x) - ctx.compile_trace({'z': z}) + ctx.compile_trace({"z": z}) output = ctx.m(torch.tensor([-1.0, 0.0, 1.0, 2.0])) torch.manual_seed(42) - input_tensor = torch.randn((4,))*2 + input_tensor = torch.randn((4,)) * 2 expected = operation_chain(input_tensor) output = ctx.m(input_tensor) self.assertTrue(torch.allclose(output, expected)) def test_torchscript_with_multiple_inputs(self): """Test that TorchScript module with multiple TracedTensor inputs raises error.""" + class AddModule(torch.nn.Module): def forward(self, x, y): return x + y * 2 @@ -3396,10 +3379,9 @@ def forward(self, x, y): x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") y = ctx.create_input(torch.tensor([4.0, 5.0, 6.0]), name="y") - z = scripted_add(x, y) - ctx.compile_trace({'z': z}) + ctx.compile_trace({"z": z}) input_x = torch.randn((3,)) input_y = torch.randn((3,)) @@ -3409,13 +3391,15 @@ def forward(self, x, y): def test_torchscript_module_with_state_and_untraced_tensor_input(self): """Test that TorchScript module with state raises error with TracedTensor.""" + class StatefulModule(torch.nn.Module): def __init__(self): super().__init__() self.linear = torch.nn.Linear(3, 2) # Set fixed weights for reproducibility self.linear.weight.data = torch.tensor( - [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] + ) self.linear.bias.data = torch.tensor([0.0, 0.0]) def forward(self, x, y): @@ -3426,10 +3410,10 @@ def forward(self, x, y): ctx = TracedTensorNode(name="test", node_index=0) x = ctx.create_input(torch.tensor([1.0, 2.0, 3.0]), name="x") - y = torch.tensor([4.0, 5.0]) #<-- this should stay constant for the tests too + y = torch.tensor([4.0, 5.0]) # <-- this should stay constant for the tests too z = scripted_module(x, y) - ctx.compile_trace({'z': z}) + ctx.compile_trace({"z": z}) input_x = torch.randn((3,)) # we expect the other input is actually inlined into the graph diff --git a/tests/unit_tests/test_two_pass_planning.py b/tests/unit_tests/test_two_pass_planning.py index b095e8d1..1dc13e18 100644 --- a/tests/unit_tests/test_two_pass_planning.py +++ b/tests/unit_tests/test_two_pass_planning.py @@ -81,7 +81,9 @@ def test_warp_op_exit_without_begin_is_fatal(self): capture=False, ) - with self.assertRaisesRegex(RuntimeError, "before the Warp operation was initialized"): + with self.assertRaisesRegex( + RuntimeError, "before the Warp operation was initialized" + ): warp_op.__exit__(None, None, None) def test_backend_call_stack_match_uses_open_call_stack(self): @@ -119,9 +121,7 @@ def wrapper(): def install_patch(): backend._patches.install(module, "target", original, wrapper) - oracle_path = ( - "leapp.leapp_graph.datatypes.warp.patching.WarpCudaOracle" - ) + oracle_path = "leapp.leapp_graph.datatypes.warp.patching.WarpCudaOracle" with ( patch("torch.cuda.is_available", return_value=False), patch(oracle_path) as oracle_class, diff --git a/tests/unit_tests/test_user_patching.py b/tests/unit_tests/test_user_patching.py index 25ca3ef4..5d8e4136 100644 --- a/tests/unit_tests/test_user_patching.py +++ b/tests/unit_tests/test_user_patching.py @@ -213,10 +213,13 @@ def test_missing_runtime_library_leaves_warp_functions_untouched(self): self.skipTest("warp-lang is not installed") original_rand_init = wp.rand_init - with tempfile.TemporaryDirectory() as temp_dir, mock.patch.dict( - os.environ, - {"XDG_CACHE_HOME": temp_dir}, - clear=False, + with ( + tempfile.TemporaryDirectory() as temp_dir, + mock.patch.dict( + os.environ, + {"XDG_CACHE_HOME": temp_dir}, + clear=False, + ), ): for env_name in patching_module.WARP_RUNTIME_ENVIRONMENTS.values(): os.environ.pop(env_name, None) @@ -240,10 +243,13 @@ def test_missing_cupti_leaves_warp_functions_untouched(self): from leapp.leapp_graph.datatypes.warp import cupti_oracle original_rand_init = wp.rand_init - with tempfile.TemporaryDirectory() as temp_dir, mock.patch.dict( - os.environ, - {"XDG_CACHE_HOME": temp_dir}, - clear=False, + with ( + tempfile.TemporaryDirectory() as temp_dir, + mock.patch.dict( + os.environ, + {"XDG_CACHE_HOME": temp_dir}, + clear=False, + ), ): for env_name in patching_module.WARP_RUNTIME_ENVIRONMENTS.values(): os.environ.pop(env_name, None) @@ -251,11 +257,14 @@ def test_missing_cupti_leaves_warp_functions_untouched(self): runtime_library.parent.mkdir(parents=True) runtime_library.touch() - with mock.patch.object( - warp_context, - "runtime", - object(), - ), mock.patch.object(cupti_oracle, "CUPTI_AVAILABLE", False): + with ( + mock.patch.object( + warp_context, + "runtime", + object(), + ), + mock.patch.object(cupti_oracle, "CUPTI_AVAILABLE", False), + ): patcher = TracingPatcher() patcher.install() try: diff --git a/tests/unit_tests/test_warp_runtime.py b/tests/unit_tests/test_warp_runtime.py index 2f93841e..9add6b53 100644 --- a/tests/unit_tests/test_warp_runtime.py +++ b/tests/unit_tests/test_warp_runtime.py @@ -13,16 +13,12 @@ from leapp import warp_runtime _BUILD_SCRIPT_PATH = ( - Path(__file__).resolve().parents[2] - / "scripts" - / "leapp-build-warp-runtime" + Path(__file__).resolve().parents[2] / "scripts" / "leapp-build-warp-runtime" ) _BUILD_SCRIPT_LOADER = SourceFileLoader( "leapp_build_warp_runtime", str(_BUILD_SCRIPT_PATH) ) -_BUILD_SCRIPT_SPEC = spec_from_loader( - _BUILD_SCRIPT_LOADER.name, _BUILD_SCRIPT_LOADER -) +_BUILD_SCRIPT_SPEC = spec_from_loader(_BUILD_SCRIPT_LOADER.name, _BUILD_SCRIPT_LOADER) assert _BUILD_SCRIPT_SPEC is not None _BUILD_SCRIPT = module_from_spec(_BUILD_SCRIPT_SPEC) _BUILD_SCRIPT_LOADER.exec_module(_BUILD_SCRIPT) @@ -37,9 +33,7 @@ def _create_build_resources(tmp_path, windows=False): ort_dir = tmp_path / "onnxruntime" (ort_dir / "capi").mkdir(parents=True) - ort_library = ( - "onnxruntime.dll" if windows else "libonnxruntime.so.1.2.3" - ) + ort_library = "onnxruntime.dll" if windows else "libonnxruntime.so.1.2.3" (ort_dir / "capi" / ort_library).touch() torch_dir = tmp_path / "torch" @@ -53,9 +47,7 @@ def _create_build_resources(tmp_path, windows=False): __version__="1.2.3", ), "torch": SimpleNamespace( - utils=SimpleNamespace( - cmake_prefix_path=str(torch_dir / "share" / "cmake") - ) + utils=SimpleNamespace(cmake_prefix_path=str(torch_dir / "share" / "cmake")) ), } return modules @@ -88,8 +80,7 @@ def test_windows_cache_and_artifact_names(tmp_path, monkeypatch): def test_environment_override_wins_over_cached_artifact(tmp_path, monkeypatch): monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) cache_library = ( - warp_runtime.warp_runtime_build_dir() - / "libleapp_wrp_onnx_custom_op.so" + warp_runtime.warp_runtime_build_dir() / "libleapp_wrp_onnx_custom_op.so" ) cache_library.parent.mkdir(parents=True) cache_library.touch() @@ -100,24 +91,17 @@ def test_environment_override_wins_over_cached_artifact(tmp_path, monkeypatch): assert warp_runtime.resolve_warp_runtime_library("onnx") == override -def test_missing_environment_override_does_not_fall_back( - tmp_path, monkeypatch -): +def test_missing_environment_override_does_not_fall_back(tmp_path, monkeypatch): monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache")) cached_library = ( - warp_runtime.warp_runtime_build_dir() - / "libleapp_wrp_torch_custom_op.so" + warp_runtime.warp_runtime_build_dir() / "libleapp_wrp_torch_custom_op.so" ) cached_library.parent.mkdir(parents=True) cached_library.touch() missing_override = tmp_path / "missing.so" - monkeypatch.setenv( - "LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY", str(missing_override) - ) + monkeypatch.setenv("LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY", str(missing_override)) - with pytest.raises( - FileNotFoundError, match="LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY" - ): + with pytest.raises(FileNotFoundError, match="LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY"): warp_runtime.resolve_warp_runtime_library("pt2") @@ -164,18 +148,14 @@ def test_missing_warp_dependency_recommends_cuda_extra(monkeypatch): def missing_module(name): raise ImportError(name) - monkeypatch.setattr( - _BUILD_SCRIPT.importlib, "import_module", missing_module - ) + monkeypatch.setattr(_BUILD_SCRIPT.importlib, "import_module", missing_module) monkeypatch.setattr( _BUILD_SCRIPT.shutil, "which", lambda name: "/usr/bin/cmake" if name == "cmake" else None, ) - with pytest.raises( - RuntimeError, match="never installs dependencies" - ) as exc_info: + with pytest.raises(RuntimeError, match="never installs dependencies") as exc_info: _BUILD_SCRIPT._validate_build_resources() message = str(exc_info.value) @@ -183,9 +163,7 @@ def missing_module(name): assert "leapp[warp-cu13]" in message -def test_windows_preflight_reports_missing_warp_import_library( - tmp_path, monkeypatch -): +def test_windows_preflight_reports_missing_warp_import_library(tmp_path, monkeypatch): modules = _create_build_resources(tmp_path, windows=True) monkeypatch.setattr(_BUILD_SCRIPT.sys, "platform", "win32") monkeypatch.delenv("LEAPP_WARP_IMPORT_LIBRARY", raising=False) @@ -217,14 +195,10 @@ def test_build_configures_both_adapters(tmp_path, monkeypatch): def run(command, check): commands.append(command) if "--build" in command: - for path in warp_runtime.warp_runtime_artifact_paths( - build_dir - ).values(): + for path in warp_runtime.warp_runtime_artifact_paths(build_dir).values(): path.touch() - monkeypatch.setattr( - _BUILD_SCRIPT, "_runtime_source_dir", lambda: source_dir - ) + monkeypatch.setattr(_BUILD_SCRIPT, "_runtime_source_dir", lambda: source_dir) monkeypatch.setattr( _BUILD_SCRIPT, "_validate_build_resources", @@ -241,15 +215,12 @@ def run(command, check): assert not stale_file.exists() -def test_discovery_reports_cache_and_missing_artifacts( - tmp_path, monkeypatch, capsys -): +def test_discovery_reports_cache_and_missing_artifacts(tmp_path, monkeypatch, capsys): monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) monkeypatch.delenv("LEAPP_WARP_ONNX_CUSTOM_OP_LIBRARY", raising=False) monkeypatch.delenv("LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY", raising=False) onnx_library = ( - warp_runtime.warp_runtime_build_dir() - / "libleapp_wrp_onnx_custom_op.so" + warp_runtime.warp_runtime_build_dir() / "libleapp_wrp_onnx_custom_op.so" ) onnx_library.parent.mkdir(parents=True) onnx_library.touch() diff --git a/tests/unit_tests/test_warp_torch_boundaries.py b/tests/unit_tests/test_warp_torch_boundaries.py index b539ef4a..6a82f1cd 100644 --- a/tests/unit_tests/test_warp_torch_boundaries.py +++ b/tests/unit_tests/test_warp_torch_boundaries.py @@ -51,18 +51,31 @@ def pending_output(): session.close_warp_segment(requester=session.owner_token) -@pytest.mark.parametrize("consume", [ - lambda x: x.reshape(3, 2), - lambda x: x.view(3, 2), - lambda x: x.transpose(0, 1), - lambda x: x[:, 1:], - lambda x: x + 2, - lambda x: torch.cat([x, x], dim=0), - lambda x: x.detach(), - lambda x: x[:], - lambda x: x.contiguous(), -], ids=["reshape", "view", "transpose", "slice", "add", "nested_args", - "detach", "full_slice", "contiguous"]) +@pytest.mark.parametrize( + "consume", + [ + lambda x: x.reshape(3, 2), + lambda x: x.view(3, 2), + lambda x: x.transpose(0, 1), + lambda x: x[:, 1:], + lambda x: x + 2, + lambda x: torch.cat([x, x], dim=0), + lambda x: x.detach(), + lambda x: x[:], + lambda x: x.contiguous(), + ], + ids=[ + "reshape", + "view", + "transpose", + "slice", + "add", + "nested_args", + "detach", + "full_slice", + "contiguous", + ], +) def test_torch_consumer_uses_finalized_warp_output(pending_output, consume): node, tensor, segment, session = pending_output result = consume(tensor) @@ -76,22 +89,26 @@ def test_torch_consumer_uses_finalized_warp_output(pending_output, consume): node.graph.lint() -@pytest.mark.parametrize("mutate", [ - lambda x: x.copy_(torch.zeros(2, 3)), - lambda x: x.__setitem__((slice(None), 0), 0.0), -], ids=["copy", "setitem"]) +@pytest.mark.parametrize( + "mutate", + [ + lambda x: x.copy_(torch.zeros(2, 3)), + lambda x: x.__setitem__((slice(None), 0), 0.0), + ], + ids=["copy", "setitem"], +) def test_torch_write_closes_segment_before_mutating_storage(pending_output, mutate): _, tensor, _, session = pending_output operation = session.active_warp_op terminate = operation.terminate def check_before_close(*args, **kwargs): - torch.testing.assert_close( - tensor.tensor, torch.arange(6.0).reshape(2, 3) + 7 - ) + torch.testing.assert_close(tensor.tensor, torch.arange(6.0).reshape(2, 3) + 7) return terminate(*args, **kwargs) - with mock.patch.object(operation, "terminate", side_effect=check_before_close) as close: + with mock.patch.object( + operation, "terminate", side_effect=check_before_close + ) as close: mutate(tensor) close.assert_called_once() diff --git a/tests/warp_support.py b/tests/warp_support.py index 45f9e17f..48f70794 100644 --- a/tests/warp_support.py +++ b/tests/warp_support.py @@ -212,12 +212,12 @@ def scale_and_split( @classmethod def setUpClass(cls): - os.environ[ - "LEAPP_WARP_ONNX_CUSTOM_OP_LIBRARY" - ] = ensure_warp_onnx_custom_op_library() - os.environ[ - "LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY" - ] = ensure_warp_pt2_custom_op_library() + os.environ["LEAPP_WARP_ONNX_CUSTOM_OP_LIBRARY"] = ( + ensure_warp_onnx_custom_op_library() + ) + os.environ["LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY"] = ( + ensure_warp_pt2_custom_op_library() + ) # Most Warp entry points initialize the runtime on first use, but the # Torch conversions read it directly, so a test whose first Warp call is # ``wp.from_torch`` would fail on an uninitialized runtime. diff --git a/uv.lock b/uv.lock index 8572b86a..94ed234a 100644 --- a/uv.lock +++ b/uv.lock @@ -78,59 +78,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0d/45/c7cd8d36d3b506bbd02db5066fae3340284781168f0d08dac25deef5f69d/ast_serialize-0.9.0-cp39-abi3-win_arm64.whl", hash = "sha256:74473258a5c55855d5306c864a5c799fbff03a0f0ea1197346b2b5cc5b4ea48a", size = 1128237, upload-time = "2026-09-02T15:50:43.496Z" }, ] -[[package]] -name = "black" -version = "26.5.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "click" }, - { name = "mypy-extensions" }, - { name = "packaging" }, - { name = "pathspec" }, - { name = "platformdirs" }, - { name = "pytokens" }, - { name = "tomli", marker = "python_full_version < '3.11' or (extra == 'extra-5-leapp-warp-cu12' and extra == 'extra-5-leapp-warp-cu13')" }, - { name = "typing-extensions", marker = "python_full_version < '3.11' or (extra == 'extra-5-leapp-warp-cu12' and extra == 'extra-5-leapp-warp-cu13')" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c0/37/5628dd55bf2b34257fc7603f0fe97c40e3aaf24265f416a9c85c95ca1436/black-26.5.1.tar.gz", hash = "sha256:dd321f668053961824bcc1be1cc1df748b2d7e4fa28086b08331e577b0100a73", size = 679439, upload-time = "2026-05-18T16:53:36.107Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/be/84/b3f55026206a9e8820a91503308075ca48eadc515e436731ca01dbe043b3/black-26.5.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:9942db8888e06943c5dde66ca0037dcff82a2a4ec1ad0ada9e0d2ee9d9823893", size = 1987719, upload-time = "2026-05-18T17:05:02.757Z" }, - { url = "https://files.pythonhosted.org/packages/c6/34/7db312c5e5783d6e76cffd9d5ac8972a32badae4c6e3288dac0eed8d3bed/black-26.5.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:89c93167a74d3a75dfaa38a5c7cca015537d5820dd7f17d63267d674a61cae90", size = 1810083, upload-time = "2026-05-18T17:05:04.302Z" }, - { url = "https://files.pythonhosted.org/packages/33/e2/e0101e73c2c8727634e2efcb35e2b34bd23ad70dfa673789f5773a591b21/black-26.5.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:22f2cd76d069cc54c71f10360744ba8983fbb616903b4304a85b734915c8e1b4", size = 1860633, upload-time = "2026-05-18T17:05:06.391Z" }, - { url = "https://files.pythonhosted.org/packages/b0/4c/e15c0c5b23cf3651035fe5addcce90e283af3548a3f91bb03d81b83106ab/black-26.5.1-cp310-cp310-win_amd64.whl", hash = "sha256:87ed5c6f450580a2f6790bc7cbfb016dfc73bc750249762268a3695361315eef", size = 1477886, upload-time = "2026-05-18T17:05:07.96Z" }, - { url = "https://files.pythonhosted.org/packages/9f/3f/59d43ade98d2ce5c8dc34a4e46cbecd177e6d55d7d4092969c6003ccc655/black-26.5.1-cp310-cp310-win_arm64.whl", hash = "sha256:58b4bd92cf88aacf83d88479c8f9caee044b1ec55f2451a337354a7ea2590a22", size = 1277111, upload-time = "2026-05-18T17:05:09.473Z" }, - { url = "https://files.pythonhosted.org/packages/4b/96/3c3e09f09f44a37aac36b178a279cd19aa7001bd796187a7b162a294c81f/black-26.5.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:96ae2c733b2aabdd9986e2c5df628ff3473676cd1c5faded1ff496cf6d74083c", size = 1970639, upload-time = "2026-05-18T17:05:11.461Z" }, - { url = "https://files.pythonhosted.org/packages/83/ea/5ad117b9ee3ecd933c712bcbae610006e5b7cc9f41c526cd7ed3b6c4124c/black-26.5.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:0e48b87e03bf109288e55cfceadcfa15ff5470aca2851a851950ed2926f450d7", size = 1792130, upload-time = "2026-05-18T17:05:12.983Z" }, - { url = "https://files.pythonhosted.org/packages/06/3a/7c448bc623fcdfa96672531beb5a616ea5e64f6975955254d7731ffb0ad9/black-26.5.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5119fa92ae61f786e8c3662fd60aece1d0a2dd5cca5d0c79417a95e7a4272a59", size = 1846134, upload-time = "2026-05-18T17:05:14.506Z" }, - { url = "https://files.pythonhosted.org/packages/a1/5b/0b39b3a5917f0657ac014ad2edb58c139553a478adfe7f817abf1622ff6e/black-26.5.1-cp311-cp311-win_amd64.whl", hash = "sha256:30d3c14661f2792e9142cce3eeeb1cbc175b3eb5f733be0c8eeb99651e52b0c3", size = 1478883, upload-time = "2026-05-18T17:05:16.542Z" }, - { url = "https://files.pythonhosted.org/packages/4c/48/dc222692e0f95030db1bbfb6c857e76858bad09058221ea7aae815255327/black-26.5.1-cp311-cp311-win_arm64.whl", hash = "sha256:1ef92b76f7733f282fd096ea406200b5a286c42947412b0eaff3a74e3616cefe", size = 1277776, upload-time = "2026-05-18T17:05:18.029Z" }, - { url = "https://files.pythonhosted.org/packages/24/99/7744b906703228264ef73bdd534df88ec1ef3de45c4e78f6d31b9e32d0c9/black-26.5.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:4ad6fa01f941920f54f2bbb35f3df7673428a0ef98a0b0840c2eaef3b110efa8", size = 2012518, upload-time = "2026-05-18T17:05:20.108Z" }, - { url = "https://files.pythonhosted.org/packages/b7/c0/c5a3b1636dfd09c42534f2b3cf33506814f6d3e066fb0879ffa16c1ae860/black-26.5.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:3915f256e75a2d7cf88d8953d37f780455dc586cc72dee059c528fe77f581217", size = 1816016, upload-time = "2026-05-18T17:05:21.84Z" }, - { url = "https://files.pythonhosted.org/packages/1f/0e/36044316b65ca471d3bb6d3703fd06fb50c6b727c3562f6a5a3153634f88/black-26.5.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9d98d4137277c75dfb898ec8d846c4fd68ba1e9cf77f95e2865c203dc18f4c3d", size = 1884150, upload-time = "2026-05-18T17:05:23.546Z" }, - { url = "https://files.pythonhosted.org/packages/b3/33/dafc5808c2af43672912111d7c3354af1615f7e2be3bed7a878461abbe4d/black-26.5.1-cp312-cp312-win_amd64.whl", hash = "sha256:a1dca32d9f1784af512a13410ec204c6f7f0aa9797a111c42e1c03449821c264", size = 1486825, upload-time = "2026-05-18T17:05:25.004Z" }, - { url = "https://files.pythonhosted.org/packages/82/14/b965ee6ad2a311f28bdbf692def3ee9848d2ae289dab28b27657fcee3e78/black-26.5.1-cp312-cp312-win_arm64.whl", hash = "sha256:1037d5ac7b7b310b2632ad867ec8d0e4c4819dcdb0b820f63135da746a24e418", size = 1288646, upload-time = "2026-05-18T17:05:26.477Z" }, - { url = "https://files.pythonhosted.org/packages/3f/5c/c384363980e11e25ca6b93205949bb331fbf35f4e0dbec376dfa6326cec8/black-26.5.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:2b36cf2ddf5566e205f6535f782a62194a184d33e175b64ae8c40b1737522be3", size = 2009020, upload-time = "2026-05-18T17:05:28.132Z" }, - { url = "https://files.pythonhosted.org/packages/0b/df/9f31c5e0babbfed77d505fc5d120beb98b21b33feaeded3924ea941fe360/black-26.5.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1f7ea64ebfa01b50f693508fc39f875e264446d3b097088f84f203b9d09618a0", size = 1813335, upload-time = "2026-05-18T17:05:31.266Z" }, - { url = "https://files.pythonhosted.org/packages/fb/24/8e7b9a2fa61b0afd82209efe937557d180a1fa055bd7f6161eb9defc3719/black-26.5.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ecb3e624844c798144e9bd986954e0adc81d8911a1f30f375e1252fe26e8c294", size = 1881614, upload-time = "2026-05-18T17:05:32.718Z" }, - { url = "https://files.pythonhosted.org/packages/49/ad/b4e0d9365ba8ac34f6bbab62a4b1b2dd5d618fac3fa1b8db968c844201b5/black-26.5.1-cp313-cp313-win_amd64.whl", hash = "sha256:e1a26503279b6b310669fb0b219c39e4820b77e8189fe80f522bb511f247db0a", size = 1488925, upload-time = "2026-05-18T17:05:34.259Z" }, - { url = "https://files.pythonhosted.org/packages/a1/4b/652b859bf5df88a751c30451b09338f7fd26a77d1271c666992f836b7711/black-26.5.1-cp313-cp313-win_arm64.whl", hash = "sha256:5c34b25da232ead53a6f335b76dbea124f4d152ad568b9080d6f944bc2b34b52", size = 1289883, upload-time = "2026-05-18T17:05:36.019Z" }, - { url = "https://files.pythonhosted.org/packages/a6/16/a8da8eb208c51c7f4ce74609a45d0dcc6d8a2141e45e81ee5289d1bb0d59/black-26.5.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:e88976690a64b0af98312ca958415849cb42423423c5f2ee74af4b49a97a2168", size = 2004800, upload-time = "2026-05-18T17:05:38.182Z" }, - { url = "https://files.pythonhosted.org/packages/11/8a/a479296a19e383b70a725882a6cf3d786540601ff03cabbaaf1cce864c5a/black-26.5.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:32d5ea7f6c8bdfa6e648326ebca1f02b0764e2a029edc6f8dce2627e19d468c3", size = 1815576, upload-time = "2026-05-18T17:05:40.309Z" }, - { url = "https://files.pythonhosted.org/packages/81/6b/cfaf3d39f25132c156a068f6b805576c9103a84086019507c70e1911ee7d/black-26.5.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ea8d16dc41655aa113cd64665e7219446cd7e4ff2248d7178eaa905190c86b18", size = 1877927, upload-time = "2026-05-18T17:05:42.463Z" }, - { url = "https://files.pythonhosted.org/packages/66/76/302e313964bcff7e28df329d39f84f5270095730d85ff0acc260610a0d82/black-26.5.1-cp314-cp314-win_amd64.whl", hash = "sha256:577f21094ea469ef92ec1adaf2c9441a226d2144d01a5be2fa823cecf6543e50", size = 1511860, upload-time = "2026-05-18T17:05:43.943Z" }, - { url = "https://files.pythonhosted.org/packages/27/4e/a3827e35e0e567f9f9ee59e2a0ab979267dca98718f25547ca8c6733afd4/black-26.5.1-cp314-cp314-win_arm64.whl", hash = "sha256:ed1a20af114c301a0269bf01163d51dbef72737fd65f850001e7cbe7f3c7abae", size = 1316632, upload-time = "2026-05-18T17:05:45.521Z" }, - { url = "https://files.pythonhosted.org/packages/94/51/f975cae76d44274cc2868dc9040ac5d58d464784610234455b4e7b19c6ef/black-26.5.1-py3-none-any.whl", hash = "sha256:4ed7f7da04046d2e488437170797d3b4a4ad83906683bcb7dfc68b673bbce5e2", size = 213693, upload-time = "2026-05-18T16:53:33.964Z" }, -] - -[[package]] -name = "click" -version = "8.5.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/c7/0e/7fa0ef50764b67090eca4114772a2abf8b6148198475e54c660b97caeee6/click-8.5.0.tar.gz", hash = "sha256:ba0d2089de75ea0310e2dde03160e6ca10009947fb95a182f9b54021bb272e34", size = 382235, upload-time = "2026-08-26T13:33:14.56Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/58/50/6c0d534c5f134586a8e1ba4e330569e32f057e33372ae556463212fb4cd3/click-8.5.0-py3-none-any.whl", hash = "sha256:255bc9599cf7748b4b1a446ccc735421bd08a2ae529a8b88597d3de5664ee360", size = 125251, upload-time = "2026-08-26T13:33:12.928Z" }, -] - [[package]] name = "colorama" version = "0.4.6" @@ -449,10 +396,10 @@ dependencies = [ [package.optional-dependencies] dev = [ - { name = "black" }, { name = "flake8" }, { name = "mypy" }, { name = "pytest" }, + { name = "ruff" }, ] test = [ { name = "opencv-python-headless" }, @@ -472,7 +419,6 @@ warp-cu13 = [ [package.metadata] requires-dist = [ - { name = "black", marker = "extra == 'dev'" }, { name = "cupti-python", marker = "sys_platform == 'linux' and extra == 'warp-cu12'", specifier = ">=12.8,<13" }, { name = "cupti-python", marker = "sys_platform == 'linux' and extra == 'warp-cu13'", specifier = ">=13,<14" }, { name = "fast-sugiyama", marker = "python_full_version >= '3.11'", specifier = ">=0.5.3" }, @@ -489,6 +435,7 @@ requires-dist = [ { name = "pillow", marker = "python_full_version >= '3.11'", specifier = ">=10.0.0" }, { name = "pytest", marker = "extra == 'dev'" }, { name = "pyyaml", specifier = ">=6.0" }, + { name = "ruff", marker = "extra == 'dev'", specifier = "==0.17.1" }, { name = "safetensors", specifier = ">=0.5.0" }, { name = "torch", specifier = ">=2.6.0" }, { name = "warp-lang", marker = "extra == 'test'", specifier = ">=1.16" }, @@ -2033,15 +1980,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/36/54/0169bc772ec491108b62f644f8ecf1fe5d8ae5ebafde2ee2142210166903/pillow-12.3.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:04f01d28a6aaff387bf842a13be313df23ba0597a44f1a976c9feb3c6ff4711a", size = 7231786, upload-time = "2026-07-01T11:56:35.046Z" }, ] -[[package]] -name = "platformdirs" -version = "4.11.7" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/69/b7/802a56eca9f2fac455b8bab5375a2647b0f0e14a2cd63ef077de3c4a7658/platformdirs-4.11.7.tar.gz", hash = "sha256:4f41487eeeeeb07f3a6625e61d9bc0ae6809f92d3386dbd74392fbb76108104d", size = 35127, upload-time = "2026-09-01T13:35:10.502Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/27/6e/80993e10a0482f630cef528635789233224f36b1ffd11592aa15d13ff9ce/platformdirs-4.11.7-py3-none-any.whl", hash = "sha256:8a02cb259042c79d1cd0450facc2fe6dc9d303ae7901afbe33bf8ea0b188cef6", size = 23938, upload-time = "2026-09-01T13:35:09.02Z" }, -] - [[package]] name = "pluggy" version = "1.6.0" @@ -2120,45 +2058,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, ] -[[package]] -name = "pytokens" -version = "0.4.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/b6/34/b4e015b99031667a7b960f888889c5bd34ef585c85e1cb56a594b92836ac/pytokens-0.4.1.tar.gz", hash = "sha256:292052fe80923aae2260c073f822ceba21f3872ced9a68bb7953b348e561179a", size = 23015, upload-time = "2026-01-30T01:03:45.924Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/42/24/f206113e05cb8ef51b3850e7ef88f20da6f4bf932190ceb48bd3da103e10/pytokens-0.4.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2a44ed93ea23415c54f3face3b65ef2b844d96aeb3455b8a69b3df6beab6acc5", size = 161522, upload-time = "2026-01-30T01:02:50.393Z" }, - { url = "https://files.pythonhosted.org/packages/d4/e9/06a6bf1b90c2ed81a9c7d2544232fe5d2891d1cd480e8a1809ca354a8eb2/pytokens-0.4.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:add8bf86b71a5d9fb5b89f023a80b791e04fba57960aa790cc6125f7f1d39dfe", size = 246945, upload-time = "2026-01-30T01:02:52.399Z" }, - { url = "https://files.pythonhosted.org/packages/69/66/f6fb1007a4c3d8b682d5d65b7c1fb33257587a5f782647091e3408abe0b8/pytokens-0.4.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:670d286910b531c7b7e3c0b453fd8156f250adb140146d234a82219459b9640c", size = 259525, upload-time = "2026-01-30T01:02:53.737Z" }, - { url = "https://files.pythonhosted.org/packages/04/92/086f89b4d622a18418bac74ab5db7f68cf0c21cf7cc92de6c7b919d76c88/pytokens-0.4.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:4e691d7f5186bd2842c14813f79f8884bb03f5995f0575272009982c5ac6c0f7", size = 262693, upload-time = "2026-01-30T01:02:54.871Z" }, - { url = "https://files.pythonhosted.org/packages/b4/7b/8b31c347cf94a3f900bdde750b2e9131575a61fdb620d3d3c75832262137/pytokens-0.4.1-cp310-cp310-win_amd64.whl", hash = "sha256:27b83ad28825978742beef057bfe406ad6ed524b2d28c252c5de7b4a6dd48fa2", size = 103567, upload-time = "2026-01-30T01:02:56.414Z" }, - { url = "https://files.pythonhosted.org/packages/3d/92/790ebe03f07b57e53b10884c329b9a1a308648fc083a6d4a39a10a28c8fc/pytokens-0.4.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d70e77c55ae8380c91c0c18dea05951482e263982911fc7410b1ffd1dadd3440", size = 160864, upload-time = "2026-01-30T01:02:57.882Z" }, - { url = "https://files.pythonhosted.org/packages/13/25/a4f555281d975bfdd1eba731450e2fe3a95870274da73fb12c40aeae7625/pytokens-0.4.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a58d057208cb9075c144950d789511220b07636dd2e4708d5645d24de666bdc", size = 248565, upload-time = "2026-01-30T01:02:59.912Z" }, - { url = "https://files.pythonhosted.org/packages/17/50/bc0394b4ad5b1601be22fa43652173d47e4c9efbf0044c62e9a59b747c56/pytokens-0.4.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b49750419d300e2b5a3813cf229d4e5a4c728dae470bcc89867a9ad6f25a722d", size = 260824, upload-time = "2026-01-30T01:03:01.471Z" }, - { url = "https://files.pythonhosted.org/packages/4e/54/3e04f9d92a4be4fc6c80016bc396b923d2a6933ae94b5f557c939c460ee0/pytokens-0.4.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:d9907d61f15bf7261d7e775bd5d7ee4d2930e04424bab1972591918497623a16", size = 264075, upload-time = "2026-01-30T01:03:04.143Z" }, - { url = "https://files.pythonhosted.org/packages/d1/1b/44b0326cb5470a4375f37988aea5d61b5cc52407143303015ebee94abfd6/pytokens-0.4.1-cp311-cp311-win_amd64.whl", hash = "sha256:ee44d0f85b803321710f9239f335aafe16553b39106384cef8e6de40cb4ef2f6", size = 103323, upload-time = "2026-01-30T01:03:05.412Z" }, - { url = "https://files.pythonhosted.org/packages/41/5d/e44573011401fb82e9d51e97f1290ceb377800fb4eed650b96f4753b499c/pytokens-0.4.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:140709331e846b728475786df8aeb27d24f48cbcf7bcd449f8de75cae7a45083", size = 160663, upload-time = "2026-01-30T01:03:06.473Z" }, - { url = "https://files.pythonhosted.org/packages/f0/e6/5bbc3019f8e6f21d09c41f8b8654536117e5e211a85d89212d59cbdab381/pytokens-0.4.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6d6c4268598f762bc8e91f5dbf2ab2f61f7b95bdc07953b602db879b3c8c18e1", size = 255626, upload-time = "2026-01-30T01:03:08.177Z" }, - { url = "https://files.pythonhosted.org/packages/bf/3c/2d5297d82286f6f3d92770289fd439956b201c0a4fc7e72efb9b2293758e/pytokens-0.4.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:24afde1f53d95348b5a0eb19488661147285ca4dd7ed752bbc3e1c6242a304d1", size = 269779, upload-time = "2026-01-30T01:03:09.756Z" }, - { url = "https://files.pythonhosted.org/packages/20/01/7436e9ad693cebda0551203e0bf28f7669976c60ad07d6402098208476de/pytokens-0.4.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:5ad948d085ed6c16413eb5fec6b3e02fa00dc29a2534f088d3302c47eb59adf9", size = 268076, upload-time = "2026-01-30T01:03:10.957Z" }, - { url = "https://files.pythonhosted.org/packages/2e/df/533c82a3c752ba13ae7ef238b7f8cdd272cf1475f03c63ac6cf3fcfb00b6/pytokens-0.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:3f901fe783e06e48e8cbdc82d631fca8f118333798193e026a50ce1b3757ea68", size = 103552, upload-time = "2026-01-30T01:03:12.066Z" }, - { url = "https://files.pythonhosted.org/packages/cb/dc/08b1a080372afda3cceb4f3c0a7ba2bde9d6a5241f1edb02a22a019ee147/pytokens-0.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8bdb9d0ce90cbf99c525e75a2fa415144fd570a1ba987380190e8b786bc6ef9b", size = 160720, upload-time = "2026-01-30T01:03:13.843Z" }, - { url = "https://files.pythonhosted.org/packages/64/0c/41ea22205da480837a700e395507e6a24425151dfb7ead73343d6e2d7ffe/pytokens-0.4.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5502408cab1cb18e128570f8d598981c68a50d0cbd7c61312a90507cd3a1276f", size = 254204, upload-time = "2026-01-30T01:03:14.886Z" }, - { url = "https://files.pythonhosted.org/packages/e0/d2/afe5c7f8607018beb99971489dbb846508f1b8f351fcefc225fcf4b2adc0/pytokens-0.4.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:29d1d8fb1030af4d231789959f21821ab6325e463f0503a61d204343c9b355d1", size = 268423, upload-time = "2026-01-30T01:03:15.936Z" }, - { url = "https://files.pythonhosted.org/packages/68/d4/00ffdbd370410c04e9591da9220a68dc1693ef7499173eb3e30d06e05ed1/pytokens-0.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:970b08dd6b86058b6dc07efe9e98414f5102974716232d10f32ff39701e841c4", size = 266859, upload-time = "2026-01-30T01:03:17.458Z" }, - { url = "https://files.pythonhosted.org/packages/a7/c9/c3161313b4ca0c601eeefabd3d3b576edaa9afdefd32da97210700e47652/pytokens-0.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:9bd7d7f544d362576be74f9d5901a22f317efc20046efe2034dced238cbbfe78", size = 103520, upload-time = "2026-01-30T01:03:18.652Z" }, - { url = "https://files.pythonhosted.org/packages/8f/a7/b470f672e6fc5fee0a01d9e75005a0e617e162381974213a945fcd274843/pytokens-0.4.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4a14d5f5fc78ce85e426aa159489e2d5961acf0e47575e08f35584009178e321", size = 160821, upload-time = "2026-01-30T01:03:19.684Z" }, - { url = "https://files.pythonhosted.org/packages/80/98/e83a36fe8d170c911f864bfded690d2542bfcfacb9c649d11a9e6eb9dc41/pytokens-0.4.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:97f50fd18543be72da51dd505e2ed20d2228c74e0464e4262e4899797803d7fa", size = 254263, upload-time = "2026-01-30T01:03:20.834Z" }, - { url = "https://files.pythonhosted.org/packages/0f/95/70d7041273890f9f97a24234c00b746e8da86df462620194cef1d411ddeb/pytokens-0.4.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dc74c035f9bfca0255c1af77ddd2d6ae8419012805453e4b0e7513e17904545d", size = 268071, upload-time = "2026-01-30T01:03:21.888Z" }, - { url = "https://files.pythonhosted.org/packages/da/79/76e6d09ae19c99404656d7db9c35dfd20f2086f3eb6ecb496b5b31163bad/pytokens-0.4.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f66a6bbe741bd431f6d741e617e0f39ec7257ca1f89089593479347cc4d13324", size = 271716, upload-time = "2026-01-30T01:03:23.633Z" }, - { url = "https://files.pythonhosted.org/packages/79/37/482e55fa1602e0a7ff012661d8c946bafdc05e480ea5a32f4f7e336d4aa9/pytokens-0.4.1-cp314-cp314-win_amd64.whl", hash = "sha256:b35d7e5ad269804f6697727702da3c517bb8a5228afa450ab0fa787732055fc9", size = 104539, upload-time = "2026-01-30T01:03:24.788Z" }, - { url = "https://files.pythonhosted.org/packages/30/e8/20e7db907c23f3d63b0be3b8a4fd1927f6da2395f5bcc7f72242bb963dfe/pytokens-0.4.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:8fcb9ba3709ff77e77f1c7022ff11d13553f3c30299a9fe246a166903e9091eb", size = 168474, upload-time = "2026-01-30T01:03:26.428Z" }, - { url = "https://files.pythonhosted.org/packages/d6/81/88a95ee9fafdd8f5f3452107748fd04c24930d500b9aba9738f3ade642cc/pytokens-0.4.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:79fc6b8699564e1f9b521582c35435f1bd32dd06822322ec44afdeba666d8cb3", size = 290473, upload-time = "2026-01-30T01:03:27.415Z" }, - { url = "https://files.pythonhosted.org/packages/cf/35/3aa899645e29b6375b4aed9f8d21df219e7c958c4c186b465e42ee0a06bf/pytokens-0.4.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d31b97b3de0f61571a124a00ffe9a81fb9939146c122c11060725bd5aea79975", size = 303485, upload-time = "2026-01-30T01:03:28.558Z" }, - { url = "https://files.pythonhosted.org/packages/52/a0/07907b6ff512674d9b201859f7d212298c44933633c946703a20c25e9d81/pytokens-0.4.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:967cf6e3fd4adf7de8fc73cd3043754ae79c36475c1c11d514fc72cf5490094a", size = 306698, upload-time = "2026-01-30T01:03:29.653Z" }, - { url = "https://files.pythonhosted.org/packages/39/2a/cbbf9250020a4a8dd53ba83a46c097b69e5eb49dd14e708f496f548c6612/pytokens-0.4.1-cp314-cp314t-win_amd64.whl", hash = "sha256:584c80c24b078eec1e227079d56dc22ff755e0ba8654d8383b2c549107528918", size = 116287, upload-time = "2026-01-30T01:03:30.912Z" }, - { url = "https://files.pythonhosted.org/packages/c6/78/397db326746f0a342855b81216ae1f0a32965deccfd7c830a2dbc66d2483/pytokens-0.4.1-py3-none-any.whl", hash = "sha256:26cef14744a8385f35d0e095dc8b3a7583f6c953c2e3d269c7f82484bf5ad2de", size = 13729, upload-time = "2026-01-30T01:03:45.029Z" }, -] - [[package]] name = "pywin32" version = "312" @@ -2248,6 +2147,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" }, ] +[[package]] +name = "ruff" +version = "0.17.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e7/e6d5749983b020e6cb1055abec0a6f361a8e94f0184cdcf56d4760995f02/ruff-0.17.1.tar.gz", hash = "sha256:5bb796c5112e9fb9527f2ef2faa0a130d6312d19de92bcfaddd70dbb2d41020f", size = 5323349, upload-time = "2026-10-15T15:23:18.64Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/78/90c46ca6250709f7c37e397b60d06061adee8d0f7d13959ed6a8e49b14fd/ruff-0.17.1-py3-none-linux_armv6l.whl", hash = "sha256:1796c41eafec66f4635125dd1e074a7aa27b8d20c187d484360acf11d247fc4a", size = 10147528, upload-time = "2026-10-15T15:22:33.667Z" }, + { url = "https://files.pythonhosted.org/packages/aa/d2/a2e28bc35b083ae51fa07dfba5280a667c8b01ecaee7802fc8849dcd5c46/ruff-0.17.1-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:8eadb09a788010be29c895418bdfa513e4e1c65e64b94e5b19e4cd381f6cc7fb", size = 10225974, upload-time = "2026-10-15T15:22:36.328Z" }, + { url = "https://files.pythonhosted.org/packages/8e/71/7a9029a81c8206ee01a98f6cfeeb56e6c095af86857581a37045c46d7e94/ruff-0.17.1-py3-none-macosx_11_0_arm64.whl", hash = "sha256:29b454f18ca3eeaf37fbf140841dcf27d36109e29dee9aae8acecb68c26c6f74", size = 9270361, upload-time = "2026-10-15T15:22:38.651Z" }, + { url = "https://files.pythonhosted.org/packages/c0/8c/e075e9c917cb0575297880c96768304e881240171c870cbd8d9b7b54bec2/ruff-0.17.1-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:872e03dd231744d6c5ec5c8af6be52b71323fc6ee096e4c8b34d9f844ee3b978", size = 9938300, upload-time = "2026-10-15T15:22:41.284Z" }, + { url = "https://files.pythonhosted.org/packages/87/27/0b4bb926a49de2383f161272b4d863571a4c7ecdbf595a6bc58813a13367/ruff-0.17.1-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4c0d6cb5f8fbac77eec9af49f4fc0eaf856b05545e372b0694d24b4e76b1c640", size = 9948464, upload-time = "2026-10-15T15:22:43.681Z" }, + { url = "https://files.pythonhosted.org/packages/14/27/518d57e0e02aecb8ea4acac6b57b51a53173d25fa734952bfdd7d531e88e/ruff-0.17.1-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:78beced4a3c94dfedbee4b999b95b6a27b35820fdc0afbacc6d37a8da26c1cd1", size = 10638639, upload-time = "2026-10-15T15:22:46.081Z" }, + { url = "https://files.pythonhosted.org/packages/01/82/d89a4e498ade59e6cee52d67f8da569019372212b87468c1fdeae884601c/ruff-0.17.1-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7feefc63921202894f0e158b0e32f45a5950d70a7fb693ccb9264816c25b1c05", size = 11337238, upload-time = "2026-10-15T15:22:48.894Z" }, + { url = "https://files.pythonhosted.org/packages/a2/7c/3b0fdef9352656317c977086fe1d802902406e4f2a0d8578da90c0e6d1c5/ruff-0.17.1-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:684154861a819bf1f433669d2a0ce9c3e54677c65b8ebb866cc2a7ec86345785", size = 10896327, upload-time = "2026-10-15T15:22:51.45Z" }, + { url = "https://files.pythonhosted.org/packages/62/86/5566290219708ce2928f3a29c4f843a1185b341ec6111c4631a33c5ceced/ruff-0.17.1-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a6ef98e51de8ef38082f90653ceeb84944bd7673783eea13e1647c5b9b064189", size = 10350811, upload-time = "2026-10-15T15:22:54.401Z" }, + { url = "https://files.pythonhosted.org/packages/ad/d2/ca8462dabb3c413990d49703ba306c30c862195a75448d080fa8fcce9775/ruff-0.17.1-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:8f9aa5b5bd5fa1937803ea609c8f153b8de39bf3e0e149897ae34b0c5b52f467", size = 10569994, upload-time = "2026-10-15T15:22:57.142Z" }, + { url = "https://files.pythonhosted.org/packages/e3/ad/dd1864c1f1db17c8cfb297ce70b352388e6977620f66f7182c52855407df/ruff-0.17.1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:80e53b0e5faa3c5667f035daaa032893ae42b7cfe708c886d39ed131e54688ab", size = 10059672, upload-time = "2026-10-15T15:23:00.084Z" }, + { url = "https://files.pythonhosted.org/packages/65/45/3f56cf1d2168c2165cdd5ac6f9668aee66a4cfdd4caf2b2cc5ccb2b1d568/ruff-0.17.1-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:056c5d5fa0ac790051ae165d7035c7bf956d02de2db557f2e9b31d877a505fca", size = 9947803, upload-time = "2026-10-15T15:23:03.132Z" }, + { url = "https://files.pythonhosted.org/packages/aa/4b/5457cb2be4c1d1ea5e919558131a0f197a4a454a0530d9e30820e8206e4c/ruff-0.17.1-py3-none-musllinux_1_2_i686.whl", hash = "sha256:bd8f5a55b10f371642eacd0e8665a42a8b97da15c9a7187c79e697b77b5e0ce4", size = 10351807, upload-time = "2026-10-15T15:23:05.628Z" }, + { url = "https://files.pythonhosted.org/packages/49/fb/db15b4f8fa24a1e36b99e1401f42ef6df16ee299a48670d5d262b771e7d7/ruff-0.17.1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:ce05e6f3973611d9550a600e04bb531061e368baba31610eb57c1509a8f21a0a", size = 10765716, upload-time = "2026-10-15T15:23:08.092Z" }, + { url = "https://files.pythonhosted.org/packages/0a/e8/1e186578076aff7a80735f625005fde678412ab609cd323ed182eb7d6b87/ruff-0.17.1-py3-none-win32.whl", hash = "sha256:1de0030a8aa78c65c5fecb84b2e5c390ac684c0452875a02afc7c083d49a496a", size = 9895093, upload-time = "2026-10-15T15:23:10.582Z" }, + { url = "https://files.pythonhosted.org/packages/6a/2e/da3d9589bae8e8b58f3a48bd5bc0b980232fb14a83bf02c3be7660f87c1b/ruff-0.17.1-py3-none-win_amd64.whl", hash = "sha256:5df76f16580fab58d0f6c98a3ac254a4abceb7b2489b8b678335869cff77b3d7", size = 10097955, upload-time = "2026-10-15T15:23:13.309Z" }, + { url = "https://files.pythonhosted.org/packages/df/c9/0f2a4dea06a5e12c1191ee3d96b1e1233797d0d5e18f97a2b909a95a9836/ruff-0.17.1-py3-none-win_arm64.whl", hash = "sha256:42400b9a6ff8515ff27796f07f92e0195be011c86e81e381606f58ea5502fef4", size = 9853542, upload-time = "2026-10-15T15:23:16.05Z" }, +] + [[package]] name = "safetensors" version = "0.8.0" From da03834f6d263817937ae6fe64713a6df0ac34c1 Mon Sep 17 00:00:00 2001 From: Frank Lai NV Date: Sat, 17 Oct 2026 12:51:56 -0700 Subject: [PATCH 09/21] docs: clarify whole-pipeline export and optional stage boundaries (#29) * docs: clarify whole-pipeline export and optional stage boundaries Signed-off-by: Frank Lai * docs: streamline pipeline pitch and single-model quick start Signed-off-by: Frank Lai * docs: correct capture scope and getting started references Signed-off-by: Frank Lai * docs: focus README on general pipeline export workflow Signed-off-by: Frank Lai * docs: clarify stateful cached validation and sequential replay Signed-off-by: Frank Lai * docs: remove em dashes from README Signed-off-by: Frank Lai * docs: describe capture flexibility as a capability Signed-off-by: Frank Lai * docs: remove redundant supported computation overview Signed-off-by: Frank Lai * docs: restore concise YAML artifact summary Signed-off-by: Frank Lai --------- Signed-off-by: Frank Lai --- README.md | 57 +++++++----- SKILL.md | 2 + .../_static/images/getting_started_graph.png | 4 +- docs/source/api/index.rst | 12 ++- docs/source/getting_started.rst | 90 +++++++------------ docs/source/guides/debugging.rst | 46 +++++++++- docs/source/guides/graph.rst | 6 ++ docs/source/guides/nodes.rst | 15 ++++ docs/source/index.rst | 40 +++++---- docs/source/tensor_libraries/limitations.rst | 15 +++- examples/getting_started.py | 32 +++---- tests/test_examples/test_getting_started.py | 3 +- 12 files changed, 189 insertions(+), 133 deletions(-) diff --git a/README.md b/README.md index 18f29ad7..bb531d73 100644 --- a/README.md +++ b/README.md @@ -3,50 +3,67 @@ SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All SPDX-License-Identifier: Apache-2.0 --> -# LEAPP — Lightweight Export Annotations for Policy Pipelines +# LEAPP: Lightweight Export Annotations for Policy Pipelines -LEAPP is a Python package for tracing and exporting multi-step PyTorch computational graphs. Annotate your existing code with lightweight markers, and LEAPP captures the graph structure, exports each stage as an individual model, and generates a deployment-ready YAML specification. +[Full LEAPP docs](https://nvidia-isaac.github.io/leapp/) -LEAPP is designed for pipelines that chain multiple PyTorch models or processing stages together — where you need to export the whole pipeline, not just a single model. It traces real execution of your code, records which tensors flow between stages, exports each stage independently, and writes a YAML that describes how to wire everything back together at inference time. +**Export the policy code you already run: preprocessing, model inference, and action processing together.** -## Installation +LEAPP traces supported PyTorch and NumPy computation across existing functions and objects. Export the complete path as **one ONNX, TorchScript, or PT2 model**, or choose separate stages when useful. + +Input and output annotations capture computation as your application runs. A single model can capture computation across multiple input call sites, helper functions, and class methods. The annotations work directly in application code without a prescribed RL environment interface. [See annotation patterns](https://nvidia-isaac.github.io/leapp/guides/nodes.html). + +- **Keep processing with the policy.** Capture supported NumPy preprocessing, conversions into PyTorch, and action transforms together. [Warp kernel regions](https://nvidia-isaac.github.io/leapp/tensor_libraries/warp.html) can also be captured into ONNX or PT2 with an additional runtime adapter. +- **Adapt unsupported calls with [FunctionPatch](https://nvidia-isaac.github.io/leapp/guides/patching.html).** Register a faithful replacement built from supported operations; existing module-level callers can stay in place. +- **Describe the deployment interface.** Attach tensor roles, joint names, and timing information with [semantic metadata](https://nvidia-isaac.github.io/leapp/semantics/usage.html), so downstream applications can interpret model inputs and outputs. +- **Inspect and validate the export.** Compare exported outputs, including state updates, against cached samples from consecutive Python iterations. Configure sample coverage, inspect the graph, and use [sequential replay](https://nvidia-isaac.github.io/leapp/guides/debugging.html#stateful-transitions-and-sequential-replay) to check accumulated drift and runtime feedback. + +## Quick start ```bash pip install leapp ``` -LEAPP requires Python 3.10+ and PyTorch 2.6.0+. +Requires Python 3.10+ and PyTorch 2.6.0+. -## Quick example +This example exports NumPy preprocessing, a small PyTorch policy, and action scaling and clipping together: ```python +import numpy as np import torch import leapp from leapp import annotate -leapp.start(name="my_pipeline") +def make_observation(position, velocity): + return np.concatenate([np.clip(position, -1.0, 1.0), velocity * 0.1]) + +policy = torch.nn.Sequential( + torch.nn.Linear(12, 32), torch.nn.Tanh(), torch.nn.Linear(32, 6), +).eval() + +leapp.start(name="robot_policy") +position, velocity = annotate.input_tensors("control", { + "position": np.zeros(6, dtype=np.float32), + "velocity": np.ones(6, dtype=np.float32), +}) -x = annotate.input_tensors("preprocess", {"obs": torch.randn(8)}) -features = (x - x.mean()) / (x.std() + 1e-6) -annotate.output_tensors("preprocess", {"features": features}, export_with="jit") +with torch.no_grad(): + observation = torch.from_numpy(make_observation(position, velocity)) + targets = torch.clamp(policy(observation) * 0.25, -1.0, 1.0) -features = annotate.input_tensors("policy", {"features": features}) -action = torch.tanh(features @ torch.randn(8, 4)) -annotate.output_tensors("policy", {"action": action}, export_with="onnx") +annotate.output_tensors("control", {"targets": targets}, export_with="onnx") leapp.stop() -leapp.compile_graph() +leapp.compile_graph(validate=True) ``` -Running this produces a `my_pipeline/` directory containing the exported models, a YAML pipeline spec, and, on Python 3.11 or later, a PNG graph visualization. On Python 3.10, LEAPP warns and skips visualization while completing the rest of the export. +Produces **`robot_policy/control.onnx`**, a YAML interface specification, and a graph PNG on Python 3.11+. This stateless model runs directly in ONNX Runtime; [InferenceManager](https://nvidia-isaac.github.io/leapp/leapp_runtime.html) loads the bundle for Python replay. Other applications can consume the exported artifacts through their chosen backend runtime. -## Documentation +## Capture scope -Full guides and API reference live at **[https://nvidia-isaac.github.io/leapp/](https://nvidia-isaac.github.io/leapp/)**: +LEAPP captures the executed tensor path. Data-dependent Python branches are not preserved dynamically, and unsupported operations or scalar extraction can freeze values. Check tracing warnings and validate on fresh inputs. See [PyTorch](https://nvidia-isaac.github.io/leapp/tensor_libraries/torch_limitations.html) and [NumPy](https://nvidia-isaac.github.io/leapp/tensor_libraries/limitations.html) limitations. -- [Getting started](https://nvidia-isaac.github.io/leapp/getting_started.html) — install, first pipeline, generated output -- [Guides](https://nvidia-isaac.github.io/leapp/guides/index.html) — node patterns, export options, feedback graphs, semantic data, runtime validation -- [API reference](https://nvidia-isaac.github.io/leapp/api/index.html) +[Getting started](https://nvidia-isaac.github.io/leapp/getting_started.html) · [Export options](https://nvidia-isaac.github.io/leapp/guides/export.html) · [API reference](https://nvidia-isaac.github.io/leapp/api/index.html) ## License diff --git a/SKILL.md b/SKILL.md index 698681e6..5526ebbe 100644 --- a/SKILL.md +++ b/SKILL.md @@ -180,6 +180,8 @@ Warp is part of this skill. **Before integrating Warp, read [the complete Warp g ## Validate the whole exported pipeline +`compile_graph(validate=True)` compares all cached node outputs, including exported state updates. Consecutive samples validate the original pipeline's stateful transitions, not just its actions. Each sample supplies recorded reference inputs/state; sequential runtime replay instead feeds exported state forward and checks accumulated numerical drift and runtime feedback/reset handling. Cache size controls transition coverage, not feedback replay. See [stateful validation](https://nvidia-isaac.github.io/leapp/guides/debugging.html#stateful-transitions-and-sequential-replay). + 1. Compare the original application and annotated execution using the same inputs, weights, state, execution mode, and random seeds where relevant. Check function adapters against original functions separately. 2. Run `compile_graph(validate=True, strict=True)` with justified tolerances. Inspect warnings and generated YAML: expected external inputs, final outputs, forward/feedback edges, initial values, backend types, and existing artifact paths. Explain any unexpected dangling input. 3. Load the bundle with `InferenceManager` and compare **fresh representative inputs not used during tracing** against the original pipeline. A mock-input smoke test alone is insufficient. Include realistic value ranges and relevant branches/shapes; for stateful pipelines compare a sequence from matching initial/reset state. diff --git a/docs/source/_static/images/getting_started_graph.png b/docs/source/_static/images/getting_started_graph.png index 6b194875..0e14c707 100644 --- a/docs/source/_static/images/getting_started_graph.png +++ b/docs/source/_static/images/getting_started_graph.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:e01b3c6be57ebb3458236a70cf68ae8c4eb7656405164e5f9a3a4c86b185e6e8 -size 59696 +oid sha256:53a3e55d9e207dd969b1f9ea897c6253c7d8b47b81a050d21fb6c1d08403eee3 +size 43904 diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst index e8756f0b..58efe11d 100644 --- a/docs/source/api/index.rst +++ b/docs/source/api/index.rst @@ -58,8 +58,10 @@ Parameters * ``non_traced`` (list[str], optional): Node names to exclude from export. They are still traced, still capture I/O, contribute to graph connectivity, and appear in the YAML. Defaults to ``None``. -* ``max_cached_io`` (int, optional): How many re-entry I/O examples - LEAPP caches per node for multi-example validation. Defaults to ``5``. +* ``max_cached_io`` (int, optional): Maximum number of captured input/output + samples per node, including the initial trace. Repeated calls can record + consecutive stateful transitions, including state inputs and updates. + Defaults to ``5``. * ``global_patching`` (bool, optional): Enable all process-global patches, including user-defined replacements. Defaults to ``True``. * ``patching`` (sequence[FunctionPatch] | None, optional): User-defined @@ -197,8 +199,10 @@ Parameters visualization errors stop compilation. * ``verbose`` (bool | None, optional): Override verbose logging for the compile step. ``None`` leaves the current setting unchanged. -* ``validate`` (bool, optional): Validate exported models against - captured traced outputs. Defaults to ``True``. +* ``validate`` (bool, optional): Validate exported models against all cached + input/output samples, including exported state updates. Each sample uses + its recorded inputs and state; see :ref:`stateful-validation` for the + distinction from sequential runtime replay. Defaults to ``True``. * ``rtol`` (float, optional): Relative tolerance for ``torch.allclose()``. Defaults to ``1e-3``. * ``atol`` (float, optional): Absolute tolerance for diff --git a/docs/source/getting_started.rst b/docs/source/getting_started.rst index 0878a144..e4341220 100644 --- a/docs/source/getting_started.rst +++ b/docs/source/getting_started.rst @@ -11,8 +11,8 @@ Installation LEAPP requires Python 3.10+ and PyTorch 2.6.0+. Graph visualization requires Python 3.11+; on Python 3.10, LEAPP warns and skips PNG generation while -completing the rest of the export. The full dependency list is -in the :doc:`/index`. +completing the rest of the export. The full dependency list is in +`pyproject.toml `_. .. note:: @@ -20,13 +20,13 @@ in the :doc:`/index`. is required. See :doc:`tensor_libraries/installation`. Welcome to LEAPP! This guide walks you through the basics of using LEAPP -to trace and export computational graphs from PyTorch code. +to export observation processing, policy inference, and action processing +together as one ONNX model. You'll learn how to: -* Use traced tensors to annotate node inputs and outputs -* Build multi-node graphs by connecting one node's outputs to a later node's - inputs +* Mark the inputs and final outputs of an existing pipeline +* Capture computation across helpers and multiple input call sites in one model * Understand where LEAPP writes models, YAML, and graph visualizations .. note:: @@ -39,8 +39,9 @@ You'll learn how to: Example: a robot sensor pipeline ================================ -This pipeline preprocesses robot observations into a feature vector, then -runs a small policy stage on those features. +This pipeline normalizes joint state, projects gravity, runs a small policy, +and scales and clips its actions. All of that computation is captured in one +``control.onnx`` model. The helper functions keep their existing organization. .. code-block:: python @@ -63,7 +64,7 @@ runs a small policy stage on those features. # Simulates reading live joint sensors each time this function is called. JOINT_POS = torch.randn(6) JOINT_VEL = torch.randn(6) - return annotate.input_tensors("obs_processor", { + return annotate.input_tensors("control", { "joint_pos": JOINT_POS, "joint_vel": JOINT_VEL, }) @@ -72,7 +73,7 @@ runs a small policy stage on those features. # Simulates reading live orientation and command context each call. ORIENTATION = torch.tensor([1.0, 0.0, 0.0, 0.0]) CMD_VEL = torch.tensor([0.5, 0.0, 0.1]) - return annotate.input_tensors("obs_processor", { + return annotate.input_tensors("control", { "orientation": ORIENTATION, "cmd_vel": CMD_VEL, }) @@ -96,7 +97,7 @@ runs a small policy stage on those features. def main(): leapp.start(name="sample_pipeline") - # NODE 1: observation preprocessing. + # One node captures the complete computation. # Multiple input_tensors() calls can contribute to the same node. pos, vel = capture_joint_observations() quat, cmd = capture_context() @@ -105,26 +106,17 @@ runs a small policy stage on those features. gravity_vec = project_gravity(quat) obs_features = torch.cat([pos_norm, vel_norm, gravity_vec, cmd]) - annotate.output_tensors( - "obs_processor", - {"obs_features": obs_features}, - export_with="jit", - ) - - # NODE 2: policy. - feat = annotate.input_tensors("policy", - {"obs_features": obs_features}) - raw_action = feat @ _W + _b + raw_action = obs_features @ _W + _b joint_targets = scale_and_clip(raw_action) annotate.output_tensors( - "policy", + "control", {"joint_targets": joint_targets}, - export_with="jit", + export_with="onnx", ) leapp.stop() - leapp.compile_graph() + leapp.compile_graph(validate=True) if __name__ == "__main__": main() @@ -150,7 +142,7 @@ Traced tensors annotate.output_tensors('feature_extractor', { 'result': result, - }, export_with="jit") + }, export_with="onnx") This pair gives you the most flexible capture surface: @@ -168,11 +160,8 @@ Graph flow joint_pos, joint_vel, orientation, cmd_vel | - [obs_processor] - | - obs_features - | - [policy] + [control] + normalize -> policy -> scale/clip | joint_targets @@ -197,11 +186,12 @@ artifacts and the metadata needed to wire them together. .. grid:: 1 1 2 2 :gutter: 3 - .. grid-item-card:: Exported node models + .. grid-item-card:: One exported model :class-card: sd-rounded-3 - Per-node TorchScript, ExportedProgram (``.pt2``), or ONNX artifacts - containing traced compute, constants, and model weights. + ``control.onnx`` contains preprocessing, policy weights and inference, + and action processing. Other export formats include TorchScript and + ExportedProgram (``.pt2``). .. grid-item-card:: Pipeline specification :class-card: sd-rounded-3 @@ -224,7 +214,8 @@ artifacts and the metadata needed to wire them together. Try it yourself =============== -1. Run the maintained example: ``python examples/getting_started.py`` +1. In a repository checkout, run ``python examples/getting_started.py``; + with a pip-only installation, save and run the code above instead. 2. Open the generated ``sample_pipeline/`` directory. 3. Inspect the files to see your exported pipeline. @@ -234,7 +225,7 @@ Graph visualization On Python 3.11+, LEAPP writes a PNG graph visualization showing exported nodes, graph inputs and outputs, and data-flow connections between nodes. Python 3.10 warns and skips this artifact. Use it to verify that LEAPP -detected the node boundaries and cross-node connections you intended. +captured the external inputs and final outputs you intended. .. image:: _static/images/getting_started_graph.png :alt: Getting started graph @@ -248,37 +239,20 @@ LEAPP also generates a complete specification of the pipeline. More details in .. code-block:: yaml - models: - obs_processor: - inputs: - - dtype: float32 - name: joint_pos - shape: [6] - type: tensor - outputs: - - dtype: float32 - name: obs_features - shape: [18] - type: tensor - parameters: - backend: jit - model_path: obs_processor.pt - pipeline: inputs: - obs_processor: [joint_pos, joint_vel, orientation, cmd_vel] + control: [joint_pos, joint_vel, orientation, cmd_vel] outputs: - policy: [joint_targets] - data_flow: - obs_processor/obs_features: [policy/obs_features] + control: [joint_targets] + data_flow: {} feedback_flow: {} -This YAML contains: +The complete YAML contains: * Complete node specifications with input/output tensor descriptions. * Shape and dtype information for all tensors. * Graph structure metadata ready for deployment systems. -* Export configuration showing how each node should be compiled. +* Model paths and export backend settings. Next steps ========== @@ -286,5 +260,5 @@ Next steps * Take a look at the :doc:`semantics/usage` guide to learn how to add semantic metadata to your pipeline. * Learn graph annotation patterns in :doc:`guides/nodes` and exported-model runtime details in :doc:`generated_configs`. -* Learn how to interpret and run the generated bundle in :doc:`generated_configs`. +* Add recurrent state and feedback with :doc:`guides/graph`. * Browse the full :doc:`api/index`. diff --git a/docs/source/guides/debugging.rst b/docs/source/guides/debugging.rst index afb2fdd6..284be016 100644 --- a/docs/source/guides/debugging.rst +++ b/docs/source/guides/debugging.rst @@ -22,9 +22,11 @@ correct. Confirm the export ================== -``compile_graph(validate=True)`` checks that each exported node matches -the traced Python on the inputs captured during tracing. LEAPP runs the -exported model, compares outputs with +``compile_graph(validate=True)`` checks each exported node against all its +cached Python input/output samples, including exported state inputs and +updates. Repeated execution can capture consecutive stateful transitions; +``max_cached_io`` controls how many samples are retained per node. LEAPP runs +the exported model on each sample, compares outputs with ``torch.allclose(..., rtol=..., atol=...)``, logs any deviation, and returns a dict mapping node names to results. @@ -80,6 +82,44 @@ Re-entry also checks names, shapes, dtypes, and connectivity against the first step. Those failures show up during tracing, not at ``compile_graph()``. +.. _stateful-validation: + +Stateful transitions and sequential replay +------------------------------------------ + +When consecutive calls carry state forward, cached samples preserve the +original trajectory, including each state's next value: + +.. code-block:: text + + (observation_0, state_0) -> (action_0, state_1) + (observation_1, state_1) -> (action_1, state_2) + (observation_2, state_2) -> (action_2, state_3) + +LEAPP supplies each recorded observation and state to the exported node and +compares both its action and its next state with the recorded outputs. This +validates the state-transition computation across the captured trajectory. +State must be included in the exported interface as described in +:doc:`graph`; unannotated Python state is outside this check. + +For deterministic computation with all relevant state represented, exact +agreement on every transition also implies agreement when those transitions +are chained with the same observations, initial state, and reset boundaries: +each exported next state equals the next recorded input state. + +Validation accepts numerical differences within ``rtol`` and ``atol``. +Cached validation supplies the recorded reference state again at each step, +so small accepted errors do not accumulate between samples. **Sequential +runtime replay feeds the exported state outputs into subsequent calls** and +can reveal whether those errors grow, as well as problems in feedback wiring, +initialization, or reset handling. Increasing the cache size expands recorded +transition coverage; it does not change how validation supplies state. + +Use both checks when assessing stateful deployment: capture representative +consecutive transitions, then compare a sequence through +:doc:`../leapp_runtime` against the original pipeline from matching initial +state. Include episode resets and fresh observations beyond the capture set. + Inspect what was captured ========================= diff --git a/docs/source/guides/graph.rst b/docs/source/guides/graph.rst index f03c99c5..f5c156a8 100644 --- a/docs/source/guides/graph.rst +++ b/docs/source/guides/graph.rst @@ -19,6 +19,12 @@ State capture has two complementary paths: LEAPP can also infer **feedback connections** when outputs from later nodes are fed back into earlier nodes across repeated graph execution. +``compile_graph(validate=True)`` checks exported state updates alongside other +outputs on every cached sample. Capture consecutive calls to validate stateful +transitions across the original trajectory. Sequential runtime replay +additionally checks accumulated numerical drift and feedback handling; see +:ref:`stateful-validation`. + Explicit state capture ====================== diff --git a/docs/source/guides/nodes.rst b/docs/source/guides/nodes.rst index d2f92002..1066af56 100644 --- a/docs/source/guides/nodes.rst +++ b/docs/source/guides/nodes.rst @@ -7,6 +7,21 @@ This guide covers advanced node patterns. The primary interface is :func:`~leapp.annotate.output_tensors`; use :func:`~leapp.annotate.method` when a node fits cleanly inside one function and you want a shorthand. +One model for the whole pipeline +================================ + +A node is an export boundary, not a restriction on how you organize Python +code. Use the same node name from the first input annotation through the final +output annotation to export preprocessing, model inference, and action +processing together as one model. The captured computation can cross helper +functions and class methods, including supported NumPy-to-PyTorch conversions. +There is no requirement to wrap the pipeline in a single ``nn.Module`` or an +RL environment adapter. + +Introduce additional node boundaries when you want separate model artifacts, +different export backends, or stage interfaces you can inspect independently. +LEAPP exports one model per node; splitting the pipeline is optional. + Distributed inputs: one node, many call sites ============================================= diff --git a/docs/source/index.rst b/docs/source/index.rst index 5dd6099a..957591d7 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -25,11 +25,12 @@ LEAPP Export the whole policy pipeline, not just the model weights. -LEAPP helps teams move learned policies and multi-stage PyTorch pipelines -from research code into deployable artifacts. It traces real execution across -preprocessing, model inference, postprocessing, constants, and tensor library -glue, then exports per-node models and a YAML pipeline specification that -downstream runtimes can consume. +LEAPP exports the policy code you already run into ONNX, TorchScript, or +PyTorch ExportedProgram models. Capture preprocessing, model inference, and +action processing together in **one model**, even when the computation spans +helper functions, class methods, and supported PyTorch and NumPy operations. +Choose separate models only when your deployment benefits from independent +stages. LEAPP is designed for the full policy path, including the logic that surrounds the learned model: normalization, feature transforms, recurrent @@ -40,7 +41,10 @@ into deployment. With LEAPP, you mark the input and output boundaries of the compute you already run. Execute the annotated pipeline with sample inputs, and LEAPP compiles the -traced graph into portable artifacts. +traced graph and validates exported outputs against captured Python outputs. +Your code does not need to implement an RL environment interface or fit into a +single ``forward()`` method. Isaac Lab also provides a native LEAPP export +integration; see :doc:`ecosystem`. .. button-ref:: getting_started :ref-type: doc @@ -83,11 +87,10 @@ checkpoint. It is a strong fit when you need to: Keep preprocessing, postprocessing, constants, and tensor operations around the learned model. - .. grid-item-card:: Avoid specialized export-only implementations + .. grid-item-card:: Reuse existing computation :class-card: leapp-fit-card sd-rounded-3 - Eliminate the need for a separate export-only implementation of the same - computation. + Reuse supported computation directly from your application. .. grid-item-card:: Move between runtimes :class-card: leapp-fit-card sd-rounded-3 @@ -95,11 +98,11 @@ checkpoint. It is a strong fit when you need to: Carry policy logic into another runtime, simulator, application, or deployment stack. - .. grid-item-card:: Export modular policy graphs + .. grid-item-card:: Choose one model or several :class-card: leapp-fit-card sd-rounded-3 - Represent a policy as connected nodes, so each stage can be inspected, - exported, and wired independently. + Export the full computation as one model, or choose separate stages + with independent export backends when needed. .. grid-item-card:: Describe tensor meaning :class-card: leapp-fit-card sd-rounded-3 @@ -114,10 +117,11 @@ How LEAPP Fits In .. grid-item:: - LEAPP traces the code path you actually execute. It records tensor flow - between annotated stages, exports each stage independently, and writes a - YAML specification describing how the exported pieces connect at - inference time. + LEAPP traces the code path you actually execute. One annotated input/output + boundary can capture the entire policy pipeline into one model file. + Multiple boundaries let you export separate models instead. A YAML + specification describes the model interfaces, semantic metadata, and + any state feedback or connections between models. The result is an export bundle that can be inspected, validated, and handed to downstream systems without requiring those systems to @@ -135,8 +139,7 @@ How LEAPP Fits In
export bundle
- node1.pt - node2.onnx + control.onnx pipeline.yaml visualization.png
@@ -144,7 +147,6 @@ How LEAPP Fits In
- Documentation Map ================= diff --git a/docs/source/tensor_libraries/limitations.rst b/docs/source/tensor_libraries/limitations.rst index c8c2f928..bf01433b 100644 --- a/docs/source/tensor_libraries/limitations.rst +++ b/docs/source/tensor_libraries/limitations.rst @@ -71,7 +71,14 @@ extensions such as **SciPy**, **OpenCV**, and **numba**, NumPy's own compiled ro including ``np.fft``, ``np.linalg.inv``, ``np.interp``, and ``np.pad``, and plain Python iteration over array elements. -Where a rewrite in traced operations is not practical, split the node: +When a faithful replacement can be expressed in supported operations, use +:doc:`../guides/patching` to register it for the unsupported module-level +function. Existing helpers can keep their calls to that function; the adapter's +operations become part of the exported model. Replacements must be pure and +match the behavior used by the application. Validate them against the original +function on fresh inputs. + +If a faithful function patch is not practical, split the node: call :func:`~leapp.annotate.output_tensors` before the untraceable step and :func:`~leapp.annotate.input_tensors` after it, so the gap becomes an explicit graph boundary that the deploying runtime has to fill. @@ -133,9 +140,9 @@ Extracting Python values captured at trace time and drop out of the graph. Each logs a warning. Using an array in a boolean context, such as ``if array.any():``, records only -the branch taken during tracing and logs an error. An exported graph is a -static structure and cannot hold a data-dependent branch; use ``np.where`` for -element-wise choices, or split the node around the decision. +the branch taken during tracing and logs an error. LEAPP's capture does not +preserve data-dependent Python branches; use ``np.where`` for element-wise +choices, or split the node around the decision. Casts are limited to the shared dtypes ====================================== diff --git a/examples/getting_started.py b/examples/getting_started.py index 06c4ad6e..5272f515 100644 --- a/examples/getting_started.py +++ b/examples/getting_started.py @@ -10,7 +10,7 @@ 2. multiple input_tensors() calls contributing inputs to the same node Pipeline: - [obs_processor] -> [policy] + joint state + orientation + command -> [control] -> joint targets Run: python examples/getting_started.py @@ -35,22 +35,22 @@ def normalize_joints(pos: torch.Tensor, vel: torch.Tensor): def capture_joint_observations(): - """First input_tensors() call for the obs_processor node.""" + """First input_tensors() call for the control node.""" # Simulates reading live joint sensors each time this function is called. JOINT_POS = torch.randn(6) JOINT_VEL = torch.randn(6) - return annotate.input_tensors("obs_processor", { + return annotate.input_tensors("control", { "joint_pos": JOINT_POS, "joint_vel": JOINT_VEL, }) def capture_context(): - """Second input_tensors() call for the same obs_processor node.""" + """Second input_tensors() call for the same control node.""" # Simulates reading live orientation and command context each call. ORIENTATION = torch.tensor([1.0, 0.0, 0.0, 0.0]) CMD_VEL = torch.tensor([0.5, 0.0, 0.1]) - return annotate.input_tensors("obs_processor", { + return annotate.input_tensors("control", { "orientation": ORIENTATION, "cmd_vel": CMD_VEL, }) @@ -82,35 +82,25 @@ def scale_and_clip(raw: torch.Tensor) -> torch.Tensor: def main(): leapp.start(name="sample_pipeline") - # Node 1: observation preprocessing. - # These two calls contribute inputs to the same traced node. + # One node captures preprocessing, policy, and action processing. + # These two calls contribute inputs at different call sites. pos, vel = capture_joint_observations() quat, cmd = capture_context() pos_norm, vel_norm = normalize_joints(pos, vel) gravity_vec = project_gravity(quat) # helper function calls are traced too obs_features = torch.cat([pos_norm, vel_norm, gravity_vec, cmd]) - annotate.output_tensors( - "obs_processor", - {"obs_features": obs_features}, - export_with="jit", - ) - - # Node 2: tiny policy + post-processing. - feat = annotate.input_tensors("policy", { - "obs_features": obs_features, - }) - raw_action = feat @ _W + _b + raw_action = obs_features @ _W + _b joint_targets = scale_and_clip(raw_action) annotate.output_tensors( - "policy", + "control", {"joint_targets": joint_targets}, - export_with="jit", + export_with="onnx", ) leapp.stop() - leapp.compile_graph(visualize=True) + leapp.compile_graph(visualize=True, validate=True) out_dir = "sample_pipeline" print(f"\nExported to {out_dir}/") diff --git a/tests/test_examples/test_getting_started.py b/tests/test_examples/test_getting_started.py index 6ec2808a..4a907073 100644 --- a/tests/test_examples/test_getting_started.py +++ b/tests/test_examples/test_getting_started.py @@ -27,8 +27,7 @@ class TestGettingStarted(BaseExampleTest): def test_getting_started_execution(self): """Test that getting_started.py runs and generates expected artifacts.""" expected_files = [ - "obs_processor.pt", - "policy.pt", + "control.onnx", "sample_pipeline.png", "sample_pipeline.yaml", ] From 698ea6c8c1c9f38b082ed0400efe00a6fa312515 Mon Sep 17 00:00:00 2001 From: Frank Lai NV Date: Sat, 17 Oct 2026 14:12:26 -0700 Subject: [PATCH 10/21] docs: identify Isaac LEAPP and publish agent documentation index (#31) Signed-off-by: Frank Lai --- README.md | 4 +-- docs/source/_templates/layout.html | 32 ++++++++++++++++++++++++ docs/source/conf.py | 7 ++++-- docs/source/index.rst | 15 +++++++---- docs/source/llms.txt | 40 ++++++++++++++++++++++++++++++ 5 files changed, 89 insertions(+), 9 deletions(-) create mode 100644 docs/source/_templates/layout.html create mode 100644 docs/source/llms.txt diff --git a/README.md b/README.md index bb531d73..f1603a29 100644 --- a/README.md +++ b/README.md @@ -5,11 +5,11 @@ SPDX-License-Identifier: Apache-2.0 # LEAPP: Lightweight Export Annotations for Policy Pipelines -[Full LEAPP docs](https://nvidia-isaac.github.io/leapp/) +[Full Isaac LEAPP docs](https://nvidia-isaac.github.io/leapp/) **Export the policy code you already run: preprocessing, model inference, and action processing together.** -LEAPP traces supported PyTorch and NumPy computation across existing functions and objects. Export the complete path as **one ONNX, TorchScript, or PT2 model**, or choose separate stages when useful. +Isaac LEAPP is NVIDIA's library for machine learning pipeline export. It traces supported PyTorch and NumPy computation across existing functions and objects. Export the complete path as **one ONNX, TorchScript, or PT2 model**, or choose separate stages when useful. Input and output annotations capture computation as your application runs. A single model can capture computation across multiple input call sites, helper functions, and class methods. The annotations work directly in application code without a prescribed RL environment interface. [See annotation patterns](https://nvidia-isaac.github.io/leapp/guides/nodes.html). diff --git a/docs/source/_templates/layout.html b/docs/source/_templates/layout.html new file mode 100644 index 00000000..c2c71524 --- /dev/null +++ b/docs/source/_templates/layout.html @@ -0,0 +1,32 @@ +{# +SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 +#} +{% extends "!layout.html" %} + +{% block extrahead %} +{{ super() }} + +{% if pagename == 'index' %} + +{% endif %} +{% endblock %} diff --git a/docs/source/conf.py b/docs/source/conf.py index abee0867..84c8e886 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -11,7 +11,7 @@ # -- Project information ----------------------------------------------------- -project = 'LEAPP' +project = 'Isaac LEAPP' author = 'NVIDIA Corporation' copyright = f'{datetime.now().year}, NVIDIA Corporation' @@ -81,6 +81,9 @@ # -- HTML output ------------------------------------------------------------- html_theme = 'nvidia_sphinx_theme' +html_title = 'Isaac LEAPP: machine learning pipeline export' +html_short_title = 'Isaac LEAPP' +html_baseurl = 'https://nvidia-isaac.github.io/leapp/' html_theme_options = { 'copyright_override': { @@ -91,7 +94,7 @@ } html_static_path = ['_static'] -html_extra_path = ['_redirects'] +html_extra_path = ['_redirects', 'llms.txt'] html_css_files = ['css/leapp.css'] html_last_updated_fmt = '%b %d, %Y' diff --git a/docs/source/index.rst b/docs/source/index.rst index 957591d7..29a1c702 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -1,12 +1,15 @@ -===== -LEAPP -===== +=========== +Isaac LEAPP +=========== + +.. meta:: + :description: Isaac LEAPP is NVIDIA's library for machine learning pipeline export. Capture supported preprocessing, inference, and postprocessing in ONNX, TorchScript, or PT2 models. .. raw:: html
-

LEAPP

+

Isaac LEAPP

Lightweight Export Annotations
for Policy Pipelines

@@ -25,7 +28,9 @@ LEAPP Export the whole policy pipeline, not just the model weights. -LEAPP exports the policy code you already run into ONNX, TorchScript, or +Isaac LEAPP is NVIDIA's library for machine learning pipeline export. +LEAPP stands for Lightweight Export Annotations for Policy Pipelines. +It exports the computation you already run into ONNX, TorchScript, or PyTorch ExportedProgram models. Capture preprocessing, model inference, and action processing together in **one model**, even when the computation spans helper functions, class methods, and supported PyTorch and NumPy operations. diff --git a/docs/source/llms.txt b/docs/source/llms.txt new file mode 100644 index 00000000..51509328 --- /dev/null +++ b/docs/source/llms.txt @@ -0,0 +1,40 @@ +# Isaac LEAPP + +> NVIDIA's library for machine learning pipeline export. LEAPP stands for Lightweight Export Annotations for Policy Pipelines. Capture supported preprocessing, model inference, and postprocessing from existing Python code in deployable models. + +The Python package and import name are `leapp`: install with `pip install leapp` and use `import leapp`. + +One exported ONNX, TorchScript, or PT2 model can span existing helper functions, class methods, and multiple input call sites. Separate models are optional. Capture supports PyTorch and supported NumPy operations, including conversions between them. Warp kernel regions require an additional runtime adapter. + +LEAPP is a standalone library with ecosystem integrations, including Isaac Lab and Isaac ROS Deploy. Its Python `InferenceManager` provides bundle replay; exported artifacts can also run in compatible backend runtimes, including C++ applications. Runtime operator support and any custom adapters must match the export. + +Capture follows the executed tensor path. Data-dependent Python branches are not preserved dynamically, and unsupported operations or scalar extraction can freeze values. Cached validation includes exported state updates; sequential replay checks accumulated drift and runtime feedback. Check the installed version against the documentation before using an API. + +## Agent integration instructions + +- [Integration skill](https://raw.githubusercontent.com/nvidia-isaac/leapp/main/SKILL.md): Markdown workflow for annotating existing code, exporting, and validating a pipeline. +- [Skill bundle](https://nvidia-isaac.github.io/leapp/_downloads/leapp-integrate.zip): Download the skill with its supporting Markdown references. +- [Project overview](https://github.com/nvidia-isaac/leapp): Source repository, installation, and a complete single-model example. + +## Capture and export + +- [Getting started](https://nvidia-isaac.github.io/leapp/getting_started.html): Export preprocessing, inference, and postprocessing as one ONNX model. +- [Annotation patterns](https://nvidia-isaac.github.io/leapp/guides/nodes.html): Define boundaries across existing application code. +- [Export configuration](https://nvidia-isaac.github.io/leapp/guides/export.html): Select model formats and backend settings. +- [Function patches](https://nvidia-isaac.github.io/leapp/guides/patching.html): Supply faithful replacements for unsupported module-level functions. +- [State and feedback](https://nvidia-isaac.github.io/leapp/guides/graph.html): Expose state inputs, updates, and feedback connections. +- [API reference](https://nvidia-isaac.github.io/leapp/api/index.html): Public signatures and parameters. + +## Validation and deployment + +- [Validation guide](https://nvidia-isaac.github.io/leapp/guides/debugging.html): Cached samples, fresh inputs, sequential replay, and tracing diagnostics. +- [Generated configuration](https://nvidia-isaac.github.io/leapp/generated_configs.html): Model paths, tensor interfaces, and graph connections in the YAML bundle. +- [Python bundle runtime](https://nvidia-isaac.github.io/leapp/leapp_runtime.html): Load and replay bundles with InferenceManager. +- [Deployment semantics](https://nvidia-isaac.github.io/leapp/semantics/usage.html): Describe tensor roles, element names, and timing for downstream integrations. +- [Ecosystem integrations](https://nvidia-isaac.github.io/leapp/ecosystem.html): Isaac Lab export and Isaac ROS Deploy runtime documentation. + +## Capture limits + +- [PyTorch limitations](https://nvidia-isaac.github.io/leapp/tensor_libraries/torch_limitations.html): Control flow, scalar extraction, and other tracing restrictions. +- [NumPy limitations](https://nvidia-isaac.github.io/leapp/tensor_libraries/limitations.html): Supported computation boundaries and trace-loss risks. +- [Warp capture](https://nvidia-isaac.github.io/leapp/tensor_libraries/warp.html): Kernel capture and runtime adapter requirements. From 3dbe6fc75aac94d6cb0d224e19d30b1419ee8780 Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Mon, 19 Oct 2026 12:21:10 -0700 Subject: [PATCH 11/21] docs: describe current bundle and lifecycle behavior Signed-off-by: Marc Alwan --- packages/leapp-runtime-rust/README.md | 56 +++++++------------ .../bindings/common/README.md | 21 +++---- .../bindings/common/tests/lifecycle.rs | 2 +- packages/leapp-runtime-rust/tests/bundle.rs | 4 +- 4 files changed, 32 insertions(+), 51 deletions(-) diff --git a/packages/leapp-runtime-rust/README.md b/packages/leapp-runtime-rust/README.md index b7653bd9..e128c130 100644 --- a/packages/leapp-runtime-rust/README.md +++ b/packages/leapp-runtime-rust/README.md @@ -5,16 +5,12 @@ SPDX-License-Identifier: Apache-2.0 # LEAPP Rust bundle API and shared lifecycle sources -This is the standalone common foundation extracted from the Triton-based -and native runtime implementations. It loads and validates -exported LEAPP bundles and exposes their public inspection API in Rust, Python, -and C++. The original common runtime/session lifecycle is retained unchanged -for integration with either engine. Bundle loading does not initialize an -inference backend. - -The package path and library names match the runtime branches to make this a -base for either implementation. This worktree does not refactor those branches -to depend on a new crate. The two Python distributions share the `leapp_runtime` +This package loads and validates exported LEAPP bundles and exposes their public +inspection API in Rust, Python, and C++. It also provides shared runtime/session +lifecycle sources for ownership, error handling, and shutdown. Bundle loading +does not initialize an inference backend. + +The inspection and full-runtime Python distributions share the `leapp_runtime` import name and should be installed in separate environments. ## Included code @@ -23,18 +19,18 @@ import name and should be installed in separate environments. | --- | --- | | Rust core | `Bundle`, YAML schema parsing, graph and artifact validation, SHA-256 checks, safetensors initial-state validation | | Public inspection types | `PortRef`, `PortSpec`, `DataType`, `Metadata`, `MetadataValue`, `BundleError`, `ErrorCode` | -| Common lifecycle sources | Original `RuntimeHandle`, `SessionHandle`, and lifecycle errors; three focused regression tests | +| Common lifecycle sources | `RuntimeHandle`, `SessionHandle`, and lifecycle errors; three focused regression tests | | Python | PyO3 bundle bindings, detached metadata and port descriptors, typed stubs and structured exceptions | | C++ | C++20 bundle facade, typed metadata, status and data-type headers, private cxx transport, installable CMake target | -| Tests | Existing Rust bundle regression suite and independent Python/C++ inspection and lifetime tests | +| Tests | Rust bundle regression suite and independent Python/C++ inspection and lifetime tests | Concrete engine runtimes, sessions, results, tensor/DLPack code, graph planners, -the Triton converter, and backend adapters remain in the runtime branches. +and backend adapters are supplied by the full runtime packages. Python and C++ share the core bundle implementation. The `bindings/common` lifecycle source depends on concrete execution types, so it is not a separate -compiled crate in this bundle-only checkout. See the +compiled crate in this inspection-only package. See the [common lifecycle sources](bindings/common/README.md) for integration and testing. -The extraction has no Triton, CUDA, ONNX Runtime, PyTorch, or NumPy dependency. +This package has no Triton, CUDA, ONNX Runtime, PyTorch, or NumPy dependency. The C++ binding currently supports Unix platforms (Linux and macOS), using native path bytes. @@ -61,13 +57,15 @@ execution. Bundle loading validates graph ports and connections, artifacts, and optional initial-value files. It does not check whether an inference engine can execute -the graph. Declared external inputs can remain unresolved, matching the source -implementations' inspection contract. Python and C++ represent missing input -descriptors as `None` and `std::nullopt`, respectively. +the graph. Declared external inputs can remain unresolved. Python and C++ +represent missing input descriptors as `None` and `std::nullopt`, respectively. The loader normalizes known backend aliases, accepts custom nonempty backend names, and leaves adapter availability to execution. +The loader preserves YAML metadata, including custom tags, duplicate-key +handling, and YAML merges. + ## Build and test Run from this directory with a Rust toolchain: @@ -87,22 +85,8 @@ also requires a C++20 compiler and Python 3.10 or newer. See [C++ bindings](bindings/cxx/README.md) for installation and binding tests. All inspection tests create opaque artifact bytes; no model generation is needed. -## Extraction provenance - -The shared sources, loader, and tests were extracted from the existing runtime -implementations for backend-independent inspection. - -Execution feature gates and dependencies were removed. The C++ bundle facade -was split out of `runtime.cc`, and `DataType` moved from `tensor.h` into -`data_type.h`. Binding manifests, Python exports and stubs, and CMake packaging -now expose only implemented inspection operations. - -The shared `bindings/common/src/lib.rs` and `runtime.rs` from the same source -commit are retained byte-for-byte. The lifecycle regression target compiles the -original implementation against a small test-only runtime API. No generic engine -interface or lifecycle cleanup changes are introduced by this extraction. +## Shared lifecycle -Binding regression tests also exposed a pre-existing tagged-metadata parsing -failure. The loader now reads YAML tags directly and separates unknown metadata -from serde's schema-field buffer before restoring it to the snapshot. Tests cover -tags at every inspection level, nested values, duplicate keys, and YAML merges. +`bindings/common/src/runtime.rs` manages session ownership, serializes calls, and +coordinates explicit shutdown. The lifecycle regression target compiles these +sources against a small test-only runtime API, without inference SDKs. diff --git a/packages/leapp-runtime-rust/bindings/common/README.md b/packages/leapp-runtime-rust/bindings/common/README.md index e4f0eb50..d25b5e9c 100644 --- a/packages/leapp-runtime-rust/bindings/common/README.md +++ b/packages/leapp-runtime-rust/bindings/common/README.md @@ -3,24 +3,21 @@ SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All SPDX-License-Identifier: Apache-2.0 --> -# Original shared lifecycle sources +# Shared lifecycle sources -`src/lib.rs` and `src/runtime.rs` are unchanged copies of the common lifecycle -used by both original runtime implementations. They retain session -ownership, locking, explicit closure, and the original error handling. +`src/lib.rs` and `src/runtime.rs` provide session ownership, locking, and explicit +closure. This source depends on the concrete `Runtime`, `Session`, tensor, options, and error types in `leapp_runtime`. Those execution types are not part of this -bundle-only checkout, so this directory is not a separate Cargo crate here. +inspection-only package, so this directory is not a separate Cargo crate here. The Python and C++ inspection bindings use the core bundle crate directly. -When integrating with either runtime branch, use these common sources with that -branch's existing common-crate manifest and execution feature. +Runtime packages compile these sources with their execution types and features. -The root Cargo test target includes the original `runtime.rs` against a small -runtime API defined only in the test binary. Three tests verify that sessions -retain their runtime, explicit close invalidates all handles while preserving -bundle info, and close waits for an active run. These test lifecycle behavior, -not model execution. No adapter traits or cleanup behavior changes are included. +The root Cargo test target includes `runtime.rs` against a small runtime API +defined only in the test binary. Three tests cover session ownership, close +invalidation and retained bundle info, and waiting for an active run. Execution +behavior is tested in the full runtime integration. ```sh cargo test --locked --test binding_lifecycle diff --git a/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs b/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs index dc6b9637..eebd3bf9 100644 --- a/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs +++ b/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -// Compile the original lifecycle source against a test-only runtime API. +// Compile the shared lifecycle source against a test-only runtime API. // The real dependency still supplies Bundle and ErrorCode; no production shim. extern crate leapp_runtime as bundle_api; extern crate self as leapp_runtime; diff --git a/packages/leapp-runtime-rust/tests/bundle.rs b/packages/leapp-runtime-rust/tests/bundle.rs index c23705be..67d850fc 100644 --- a/packages/leapp-runtime-rust/tests/bundle.rs +++ b/packages/leapp-runtime-rust/tests/bundle.rs @@ -482,7 +482,7 @@ fn yaml_and_artifacts_use_the_same_normalized_path() { } #[test] -fn cpp_bundle_fixtures_load_through_the_new_api() { +fn cpp_bundle_fixtures_load() { let f = Fixture::new(); fs::create_dir(f.directory.path().join("artifacts")).unwrap(); fs::write( @@ -507,7 +507,7 @@ fn cpp_bundle_fixtures_load_through_the_new_api() { } #[test] -fn retains_yaml12_scalar_resolution_from_the_prototype() { +fn uses_yaml12_scalar_resolution() { let f = Fixture::new(); fs::write( &f.path, From 732d9a7f05dd76bfdccc198baa29557a539829db Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Mon, 19 Oct 2026 11:54:45 -0700 Subject: [PATCH 12/21] refactor: organize bundle inspection modules Signed-off-by: Marc Alwan --- packages/leapp-runtime-rust/src/bundle.rs | 130 ++++++++++++++++++ .../src/{ => bundle}/load.rs | 8 +- .../src/{ => bundle}/schema.rs | 0 .../src/{ => bundle}/validate.rs | 2 +- packages/leapp-runtime-rust/src/lib.rs | 125 +---------------- 5 files changed, 137 insertions(+), 128 deletions(-) create mode 100644 packages/leapp-runtime-rust/src/bundle.rs rename packages/leapp-runtime-rust/src/{ => bundle}/load.rs (98%) rename packages/leapp-runtime-rust/src/{ => bundle}/schema.rs (100%) rename packages/leapp-runtime-rust/src/{ => bundle}/validate.rs (98%) diff --git a/packages/leapp-runtime-rust/src/bundle.rs b/packages/leapp-runtime-rust/src/bundle.rs new file mode 100644 index 00000000..74ebe263 --- /dev/null +++ b/packages/leapp-runtime-rust/src/bundle.rs @@ -0,0 +1,130 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Immutable bundle inspection, with parsing and validation in child modules. + +mod load; +mod schema; +mod validate; + +use crate::BundleError; +use indexmap::IndexMap; +use schema::{BundleDescription, InitialValueSpec}; +pub use schema::{DataType, Metadata, MetadataValue, PortRef, PortSpec}; +use std::{ + path::{Path, PathBuf}, + sync::Arc, +}; + +#[derive(Debug)] +struct BundleData { + yaml_path: PathBuf, + description: BundleDescription, + inputs: Vec, + outputs: Vec, + // Validated initial descriptors are retained for later bundle conversion. + #[allow(dead_code)] + initial_values: IndexMap, +} + +/// Shared ownership of an immutable, validated bundle snapshot. +/// +/// Cloning shares the snapshot. Inspection borrows are tied to the bundle; +/// concurrent reads are safe. Artifact files must remain available and unchanged +/// for any subsequent execution that uses the resolved paths. +#[derive(Debug, Clone)] +pub struct Bundle(Arc); + +impl Bundle { + /// Parse and validate a YAML bundle, resolving artifacts beside its path. + /// + /// The path is made absolute and lexically normalized before reading YAML + /// and resolving relative artifact paths. Symlinks are not canonicalized. + /// Artifacts are sought beside the YAML first, then at the supplied path. + /// + /// Returns a new bundle on success and an error with a phase and diagnostic + /// on failure. Existing bundles are never modified by an open attempt. + pub fn open(path: impl AsRef) -> Result { + load::read_bundle(path.as_ref()).map(|data| Self(Arc::new(data))) + } + + /// Absolute YAML path, without resolving symlinks. + pub fn yaml_path(&self) -> &Path { + &self.0.yaml_path + } + + /// Model names in declaration order, borrowed from this snapshot. + pub fn model_names(&self) -> Vec<&str> { + self.0 + .description + .models + .keys() + .map(String::as_str) + .collect() + } + + /// Declared input references in order; these need not identify existing ports. + pub fn pipeline_inputs(&self) -> &[PortRef] { + &self.0.inputs + } + + /// Public output references in declaration order. + pub fn pipeline_outputs(&self) -> &[PortRef] { + &self.0.outputs + } + + /// Pipeline configs; an empty mapping when absent. + pub fn pipeline_configs(&self) -> &Metadata { + &self.0.description.pipeline.configs + } + + /// Additional descriptive fields at pipeline level. + pub fn pipeline_metadata(&self) -> &Metadata { + &self.0.description.pipeline.metadata + } + + /// Export provenance from the `system information` mapping. + pub fn system_information(&self) -> &Metadata { + &self.0.description.system_information + } + + /// Additional descriptive fields at the bundle root. + pub fn metadata(&self) -> &Metadata { + &self.0.description.metadata + } + + /// Original exporter parameters, including unknown fields and original paths. + /// Returns `None` when the model is absent. + pub fn model_parameters(&self, node: &str) -> Option<&Metadata> { + Some(&self.0.description.models.get(node)?.original_parameters) + } + + /// Additional model metadata, or `None` when the model is absent. + pub fn model_metadata(&self, node: &str) -> Option<&Metadata> { + Some(&self.0.description.models.get(node)?.metadata) + } + + /// The first input descriptor with this name, including internal ports. + /// Returns `None` if absent; duplicate declarations are permitted. + pub fn input_info(&self, node: &str, port: &str) -> Option<&PortSpec> { + self.0 + .description + .models + .get(node)? + .inputs + .iter() + .find(|p| p.name == port) + } + + /// The first output descriptor with this name, including internal ports. + /// Returns `None` if absent; duplicate declarations are permitted. + pub fn output_info(&self, node: &str, port: &str) -> Option<&PortSpec> { + self.0 + .description + .models + .get(node)? + .outputs + .iter() + .find(|p| p.name == port) + } +} diff --git a/packages/leapp-runtime-rust/src/load.rs b/packages/leapp-runtime-rust/src/bundle/load.rs similarity index 98% rename from packages/leapp-runtime-rust/src/load.rs rename to packages/leapp-runtime-rust/src/bundle/load.rs index 2ad31691..401bda86 100644 --- a/packages/leapp-runtime-rust/src/load.rs +++ b/packages/leapp-runtime-rust/src/bundle/load.rs @@ -1,11 +1,11 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use crate::{ - error::in_phase, +use super::{ schema::{BundleDescription, NodeParameters}, - validate, BundleData, BundleError, Metadata, MetadataValue, PortRef, + validate, BundleData, }; +use crate::{error::in_phase, BundleError, Metadata, MetadataValue, PortRef}; use anyhow::{ensure, Context, Result}; use indexmap::IndexMap; use serde::{ @@ -113,7 +113,7 @@ pub(crate) fn read_bundle(path: &Path) -> Result { Ok(Default::default()) } })?; - let boundary = |map: &crate::schema::BoundaryPorts| { + let boundary = |map: &super::schema::BoundaryPorts| { map.iter() .flat_map(|(node, ports)| { ports.iter().map(|port| PortRef { diff --git a/packages/leapp-runtime-rust/src/schema.rs b/packages/leapp-runtime-rust/src/bundle/schema.rs similarity index 100% rename from packages/leapp-runtime-rust/src/schema.rs rename to packages/leapp-runtime-rust/src/bundle/schema.rs diff --git a/packages/leapp-runtime-rust/src/validate.rs b/packages/leapp-runtime-rust/src/bundle/validate.rs similarity index 98% rename from packages/leapp-runtime-rust/src/validate.rs rename to packages/leapp-runtime-rust/src/bundle/validate.rs index 59246c61..59840e14 100644 --- a/packages/leapp-runtime-rust/src/validate.rs +++ b/packages/leapp-runtime-rust/src/bundle/validate.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use crate::schema::{BundleDescription, InitialValueSpec, NodeSpec, PortRef}; +use super::schema::{BundleDescription, InitialValueSpec, NodeSpec, PortRef}; use anyhow::{ensure, Context, Result}; use indexmap::IndexMap; use safetensors::SafeTensors; diff --git a/packages/leapp-runtime-rust/src/lib.rs b/packages/leapp-runtime-rust/src/lib.rs index b12214f5..0abd9730 100644 --- a/packages/leapp-runtime-rust/src/lib.rs +++ b/packages/leapp-runtime-rust/src/lib.rs @@ -20,129 +20,8 @@ //! ``` #![forbid(unsafe_code)] +mod bundle; mod error; -mod load; -mod schema; -mod validate; +pub use bundle::{Bundle, DataType, Metadata, MetadataValue, PortRef, PortSpec}; pub use error::{BundleError, ErrorCode}; -use indexmap::IndexMap; -use schema::{BundleDescription, InitialValueSpec}; -pub use schema::{DataType, Metadata, MetadataValue, PortRef, PortSpec}; -use std::{ - path::{Path, PathBuf}, - sync::Arc, -}; - -#[derive(Debug)] -struct BundleData { - yaml_path: PathBuf, - description: BundleDescription, - inputs: Vec, - outputs: Vec, - // Validated initial descriptors are retained for later bundle conversion. - #[allow(dead_code)] - initial_values: IndexMap, -} - -/// Shared ownership of an immutable, validated bundle snapshot. -/// -/// Cloning shares the snapshot. Inspection borrows are tied to the bundle; -/// concurrent reads are safe. Artifact files must remain available and unchanged -/// for any subsequent execution that uses the resolved paths. -#[derive(Debug, Clone)] -pub struct Bundle(Arc); - -impl Bundle { - /// Parse and validate a YAML bundle, resolving artifacts beside its path. - /// - /// The path is made absolute and lexically normalized before reading YAML - /// and resolving relative artifact paths. Symlinks are not canonicalized. - /// Artifacts are sought beside the YAML first, then at the supplied path. - /// - /// Returns a new bundle on success and an error with a phase and diagnostic - /// on failure. Existing bundles are never modified by an open attempt. - pub fn open(path: impl AsRef) -> Result { - load::read_bundle(path.as_ref()).map(|data| Self(Arc::new(data))) - } - - /// Absolute YAML path, without resolving symlinks. - pub fn yaml_path(&self) -> &Path { - &self.0.yaml_path - } - - /// Model names in declaration order, borrowed from this snapshot. - pub fn model_names(&self) -> Vec<&str> { - self.0 - .description - .models - .keys() - .map(String::as_str) - .collect() - } - - /// Declared input references in order; these need not identify existing ports. - pub fn pipeline_inputs(&self) -> &[PortRef] { - &self.0.inputs - } - - /// Public output references in declaration order. - pub fn pipeline_outputs(&self) -> &[PortRef] { - &self.0.outputs - } - - /// Pipeline configs; an empty mapping when absent. - pub fn pipeline_configs(&self) -> &Metadata { - &self.0.description.pipeline.configs - } - - /// Additional descriptive fields at pipeline level. - pub fn pipeline_metadata(&self) -> &Metadata { - &self.0.description.pipeline.metadata - } - - /// Export provenance from the `system information` mapping. - pub fn system_information(&self) -> &Metadata { - &self.0.description.system_information - } - - /// Additional descriptive fields at the bundle root. - pub fn metadata(&self) -> &Metadata { - &self.0.description.metadata - } - - /// Original exporter parameters, including unknown fields and original paths. - /// Returns `None` when the model is absent. - pub fn model_parameters(&self, node: &str) -> Option<&Metadata> { - Some(&self.0.description.models.get(node)?.original_parameters) - } - - /// Additional model metadata, or `None` when the model is absent. - pub fn model_metadata(&self, node: &str) -> Option<&Metadata> { - Some(&self.0.description.models.get(node)?.metadata) - } - - /// The first input descriptor with this name, including internal ports. - /// Returns `None` if absent; duplicate declarations are permitted. - pub fn input_info(&self, node: &str, port: &str) -> Option<&PortSpec> { - self.0 - .description - .models - .get(node)? - .inputs - .iter() - .find(|p| p.name == port) - } - - /// The first output descriptor with this name, including internal ports. - /// Returns `None` if absent; duplicate declarations are permitted. - pub fn output_info(&self, node: &str, port: &str) -> Option<&PortSpec> { - self.0 - .description - .models - .get(node)? - .outputs - .iter() - .find(|p| p.name == port) - } -} From 7b88a73afca6db23143fe9c6bbc2f934306231c5 Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Mon, 19 Oct 2026 12:06:41 -0700 Subject: [PATCH 13/21] refactor: remove custom YAML tag support Signed-off-by: Marc Alwan --- packages/leapp-runtime-rust/README.md | 5 +- .../leapp-runtime-rust/bindings/cxx/README.md | 3 +- .../bindings/cxx/include/leapp/metadata.h | 9 +- .../bindings/cxx/src/bridge.rs | 1 - .../bindings/cxx/src/bundle.cc | 3 - .../bindings/cxx/src/bundle.rs | 7 +- .../bindings/cxx/tests/bundle_test.cc | 10 +- .../bindings/python/README.md | 6 +- .../python/python/leapp_runtime/__init__.py | 2 - .../python/python/leapp_runtime/_native.pyi | 7 - .../bindings/python/src/bundle.rs | 15 +- .../bindings/python/src/lib.rs | 1 - .../bindings/python/tests/test_bundle.py | 6 +- .../leapp-runtime-rust/src/bundle/load.rs | 144 +----------------- packages/leapp-runtime-rust/tests/bundle.rs | 43 ------ 15 files changed, 23 insertions(+), 239 deletions(-) diff --git a/packages/leapp-runtime-rust/README.md b/packages/leapp-runtime-rust/README.md index e128c130..28100ff8 100644 --- a/packages/leapp-runtime-rust/README.md +++ b/packages/leapp-runtime-rust/README.md @@ -63,8 +63,9 @@ represent missing input descriptors as `None` and `std::nullopt`, respectively. The loader normalizes known backend aliases, accepts custom nonempty backend names, and leaves adapter availability to execution. -The loader preserves YAML metadata, including custom tags, duplicate-key -handling, and YAML merges. +The loader preserves ordinary YAML metadata, duplicate-key handling, and YAML +merges. Custom YAML tags are unsupported and fail during parsing; LEAPP exports +do not generate them. ## Build and test diff --git a/packages/leapp-runtime-rust/bindings/cxx/README.md b/packages/leapp-runtime-rust/bindings/cxx/README.md index 0d0439c4..d1233fdb 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/README.md +++ b/packages/leapp-runtime-rust/bindings/cxx/README.md @@ -19,7 +19,8 @@ cmake --build target/consumer --parallel Consumers use `find_package(leapp_runtime CONFIG REQUIRED)` and link `leapp::runtime`. Include `leapp/bundle.h` for inspection or `leapp/data_type.h` -for scalar type declarations. +for scalar type declarations. Custom YAML tags are unsupported; `Bundle::Open` +returns an invalid-argument status in the `parse` phase. ```cpp leapp::Bundle bundle; diff --git a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/metadata.h b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/metadata.h index 342262e9..2b45b4c0 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/metadata.h +++ b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/metadata.h @@ -5,7 +5,6 @@ #define LEAPP_METADATA_H_ #include -#include #include #include #include @@ -17,13 +16,12 @@ namespace leapp { /// keys. struct Metadata { struct Entry; - struct Tagged; using Sequence = std::vector; using Mapping = std::vector; using Value = std::variant>; + Sequence, Mapping>; Value value; ///< Defaults to null. @@ -36,11 +34,6 @@ struct Metadata::Entry { Metadata value; }; -struct Metadata::Tagged { - std::string tag; - Metadata value; -}; - } // namespace leapp #endif // LEAPP_METADATA_H_ diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs b/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs index a4e9a7a6..a5852262 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs +++ b/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs @@ -36,7 +36,6 @@ pub mod ffi { String, Sequence, Mapping, - Tagged, } // Preorder nodes carry immediate child counts; mappings alternate key/value nodes. diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc index bd2262f1..b9ea0bc7 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc +++ b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc @@ -107,9 +107,6 @@ Metadata ReadMetadata(const rust::Vec& nodes, } return {std::move(items)}; } - case Kind::Tagged: - return {std::make_shared( - Metadata::Tagged{String(node.text), ReadMetadata(nodes, index)})}; } // An unknown kind violates the private generated bridge contract. std::abort(); diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs index 7e2b2280..5cb77e97 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs +++ b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs @@ -103,11 +103,7 @@ fn append(value: &MetadataValue, nodes: &mut Vec) { n.kind = K::Mapping; n.children = v.len(); } - MetadataValue::Tagged(v) => { - n.kind = K::Tagged; - n.text = v.tag.to_string(); - n.children = 1; - } + MetadataValue::Tagged(_) => unreachable!("bundle loading rejects YAML tags"), } nodes.push(n); match value { @@ -122,7 +118,6 @@ fn append(value: &MetadataValue, nodes: &mut Vec) { append(v, nodes); } } - MetadataValue::Tagged(v) => append(&v.value, nodes), _ => {} } } diff --git a/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc b/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc index 288b11b7..de7ee8f1 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc +++ b/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc @@ -34,7 +34,7 @@ int main(int argc, char** argv) { std::filesystem::copy_file(argv[1], path); std::ofstream(directory / "model.bin", std::ios::binary) << "abc"; std::ofstream(path, std::ios::app) - << "tagged: !sensor {axes: [x, y]}\n" + << "sensor: {axes: [x, y]}\n" "sequence_keys: {[x, y]: coordinates}\n" "unsigned: 18446744073709551615\n"; @@ -66,12 +66,8 @@ int main(int argc, char** argv) { Check(bundle.model_parameters("absent") == nullptr, "absent parameters"); const auto* frequency = bundle.pipeline_configs().Find("frequency"); Check(frequency && std::get(frequency->value) == 50, "typed config"); - const auto* tag = bundle.metadata().Find("tagged"); - Check(tag != nullptr, "tag metadata"); - const auto& tagged = - std::get>(tag->value); - Check(tagged->tag == "!sensor" && tagged->value.Find("axes"), - "tag preserved"); + const auto* sensor = bundle.metadata().Find("sensor"); + Check(sensor && sensor->Find("axes"), "nested mapping metadata"); const auto* mapping = bundle.metadata().Find("sequence_keys"); Check(mapping != nullptr, "mapping metadata"); const auto& entries = std::get(mapping->value); diff --git a/packages/leapp-runtime-rust/bindings/python/README.md b/packages/leapp-runtime-rust/bindings/python/README.md index d8e3c9f0..e805da9a 100644 --- a/packages/leapp-runtime-rust/bindings/python/README.md +++ b/packages/leapp-runtime-rust/bindings/python/README.md @@ -20,9 +20,9 @@ for name, spec in bundle.pipeline_outputs().items(): ``` Metadata and port descriptors are detached from the immutable Rust snapshot. -Editing returned metadata does not change later inspection results. Unknown -YAML tags become `TaggedValue`; mappings with unhashable or colliding Python -keys become `MetadataMapping` so no entries are lost. +Editing returned metadata does not change later inspection results. Mappings +with unhashable or colliding Python keys become `MetadataMapping` so no entries +are lost. Custom YAML tags are unsupported and raise `BundleError` during parsing. `BundleError` exposes `code` (an `ErrorCode`), `phase`, `path`, and `message`. Only bundle inspection is exported. This distribution uses the same diff --git a/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/__init__.py b/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/__init__.py index 18f73185..fbaf68d9 100644 --- a/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/__init__.py +++ b/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/__init__.py @@ -11,7 +11,6 @@ MetadataMapping, PortRef, PortSpec, - TaggedValue, ) @@ -31,5 +30,4 @@ class ErrorCode(IntEnum): "MetadataMapping", "PortRef", "PortSpec", - "TaggedValue", ] diff --git a/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi b/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi index 567f4b28..19c20e78 100644 --- a/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi +++ b/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi @@ -35,13 +35,6 @@ class PortSpec: @property def metadata(self) -> dict[str, Any]: ... -class TaggedValue: - @property - def tag(self) -> str: ... - - @property - def value(self) -> Any: ... - class MetadataMapping: @property def entries(self) -> list[tuple[Any, Any]]: ... diff --git a/packages/leapp-runtime-rust/bindings/python/src/bundle.rs b/packages/leapp-runtime-rust/bindings/python/src/bundle.rs index 4e388718..cad524c7 100644 --- a/packages/leapp-runtime-rust/bindings/python/src/bundle.rs +++ b/packages/leapp-runtime-rust/bindings/python/src/bundle.rs @@ -34,12 +34,6 @@ pub struct PortSpec { metadata: Py, } -#[pyclass(frozen, get_all, module = "leapp_runtime")] -pub struct TaggedValue { - tag: String, - value: Py, -} - /// Lossless fallback for YAML mappings containing keys Python cannot hash. #[pyclass(frozen, get_all, module = "leapp_runtime")] pub struct MetadataMapping { @@ -95,14 +89,7 @@ pub fn value(py: Python<'_>, v: &MetadataValue) -> PyResult> { } } } - MetadataValue::Tagged(v) => Py::new( - py, - TaggedValue { - tag: v.tag.to_string(), - value: value(py, &v.value)?, - }, - )? - .into_any(), + MetadataValue::Tagged(_) => unreachable!("bundle loading rejects YAML tags"), }) } diff --git a/packages/leapp-runtime-rust/bindings/python/src/lib.rs b/packages/leapp-runtime-rust/bindings/python/src/lib.rs index 7f03cf8b..6c609e24 100644 --- a/packages/leapp-runtime-rust/bindings/python/src/lib.rs +++ b/packages/leapp-runtime-rust/bindings/python/src/lib.rs @@ -31,7 +31,6 @@ fn _native(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; - m.add_class::()?; m.add_class::()?; Ok(()) } diff --git a/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py b/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py index 2ebb9de3..1a32f006 100644 --- a/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py +++ b/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py @@ -70,16 +70,14 @@ def test_lossless_metadata(bundle_path): bundle_path.write_text( YAML + """ -tagged: !sensor {axes: [x, y], enabled: true} +sensor: {axes: [x, y], enabled: true} sequence_keys: {[x, y]: coordinates} colliding_keys: {true: boolean, 1: integer} unsigned: 18446744073709551615 """ ) metadata = lr.Bundle.open(bundle_path).metadata() - assert isinstance(metadata["tagged"], lr.TaggedValue) - assert metadata["tagged"].tag == "!sensor" - assert metadata["tagged"].value == {"axes": ["x", "y"], "enabled": True} + assert metadata["sensor"] == {"axes": ["x", "y"], "enabled": True} assert isinstance(metadata["sequence_keys"], lr.MetadataMapping) assert metadata["sequence_keys"].entries == [(["x", "y"], "coordinates")] assert isinstance(metadata["colliding_keys"], lr.MetadataMapping) diff --git a/packages/leapp-runtime-rust/src/bundle/load.rs b/packages/leapp-runtime-rust/src/bundle/load.rs index 401bda86..e376dcee 100644 --- a/packages/leapp-runtime-rust/src/bundle/load.rs +++ b/packages/leapp-runtime-rust/src/bundle/load.rs @@ -5,17 +5,13 @@ use super::{ schema::{BundleDescription, NodeParameters}, validate, BundleData, }; -use crate::{error::in_phase, BundleError, Metadata, MetadataValue, PortRef}; +use crate::{error::in_phase, BundleError, MetadataValue, PortRef}; use anyhow::{ensure, Context, Result}; use indexmap::IndexMap; -use serde::{ - de::{EnumAccess, MapAccess, SeqAccess, VariantAccess, Visitor}, - Deserialize, Deserializer, -}; -use serde_yaml_ng::value::{Tag, TaggedValue}; +use serde::Deserialize; use sha2::{Digest, Sha256}; use std::{ - fmt, fs, io, + fs, io, path::{Component, Path, PathBuf}, }; @@ -30,39 +26,8 @@ pub(crate) fn read_bundle(path: &Path) -> Result { let mut description = in_phase(path, "parse", || { // PyYAML keeps the last duplicate mapping entry and expands YAML merges. let value = serde_yaml_ng::from_str::(&yaml_text)?.into_value()?; - // Serde's flatten buffer cannot carry YAML tags. Deserialize known - // schema fields separately, then attach the original metadata values. - let mut schema = value.clone(); - let metadata = take_metadata(Some(&mut schema), ROOT_FIELDS); - let pipeline_metadata = take_metadata(schema.get_mut("pipeline"), PIPELINE_FIELDS); - if let Some(models) = schema - .get_mut("models") - .and_then(MetadataValue::as_mapping_mut) - { - for node in models.values_mut() { - take_metadata(Some(node), MODEL_FIELDS); - for field in ["inputs", "outputs"] { - if let Some(ports) = - node.get_mut(field).and_then(MetadataValue::as_sequence_mut) - { - for port in ports { - take_metadata(Some(port), PORT_FIELDS); - } - } - } - } - } - let mut description: BundleDescription = serde_path_to_error::deserialize(schema)?; - description.metadata = metadata; - description.pipeline.metadata = pipeline_metadata; + let mut description: BundleDescription = serde_path_to_error::deserialize(value.clone())?; for (name, node) in &mut description.models { - let mut original = value["models"][name].clone(); - node.metadata = take_metadata(Some(&mut original), MODEL_FIELDS); - for (field, ports) in [("inputs", &mut node.inputs), ("outputs", &mut node.outputs)] { - for (index, port) in ports.iter_mut().enumerate() { - port.metadata = take_metadata(original[field].get_mut(index), PORT_FIELDS); - } - } node.original_parameters = serde_yaml_ng::from_value(value["models"][name]["parameters"].clone())?; } @@ -132,36 +97,6 @@ pub(crate) fn read_bundle(path: &Path) -> Result { }) } -const ROOT_FIELDS: &[&str] = &["models", "pipeline", "system information"]; -const PIPELINE_FIELDS: &[&str] = &[ - "data_flow", - "feedback_flow", - "inputs", - "outputs", - "initial_values", - "configs", -]; -const MODEL_FIELDS: &[&str] = &["inputs", "outputs", "parameters"]; -const PORT_FIELDS: &[&str] = &["name", "dtype", "shape", "type"]; - -fn take_metadata(value: Option<&mut MetadataValue>, fields: &[&str]) -> Metadata { - let Some(mapping) = value.and_then(MetadataValue::as_mapping_mut) else { - // Leave malformed structure to the typed deserializer's path diagnostics. - return Metadata::new(); - }; - let mut metadata = Metadata::new(); - mapping.retain(|key, value| { - if let Some(key) = key.as_str() { - if !fields.contains(&key) { - metadata.insert(key.into(), value.clone()); - return false; - } - } - true - }); - metadata -} - // Inspection normalizes known exporter aliases without requiring an adapter. pub(crate) fn backend_name(name: &str) -> &str { match name { @@ -206,74 +141,13 @@ fn normalize(path: &Path) -> PathBuf { } // Decode mappings before Value so duplicate keys follow Python's last-value rule. +// Serde's untagged deserializer rejects custom YAML tags, including nested keys. +#[derive(Deserialize)] +#[serde(untagged)] enum YamlValue { Mapping(IndexMap), Sequence(Vec), Scalar(MetadataValue), - Tagged(String, Box), -} - -// An untagged serde enum cannot accept YAML tags. Read each value directly so -// metadata retains tags while mappings still use Python's last-value rule. -impl<'de> Deserialize<'de> for YamlValue { - fn deserialize>(deserializer: D) -> Result { - struct YamlVisitor; - - impl<'de> Visitor<'de> for YamlVisitor { - type Value = YamlValue; - - fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str("a YAML value") - } - - fn visit_bool(self, value: bool) -> Result { - Ok(YamlValue::Scalar(MetadataValue::Bool(value))) - } - - fn visit_i64(self, value: i64) -> Result { - Ok(YamlValue::Scalar(MetadataValue::Number(value.into()))) - } - - fn visit_u64(self, value: u64) -> Result { - Ok(YamlValue::Scalar(MetadataValue::Number(value.into()))) - } - - fn visit_f64(self, value: f64) -> Result { - Ok(YamlValue::Scalar(MetadataValue::Number(value.into()))) - } - - fn visit_str(self, value: &str) -> Result { - Ok(YamlValue::Scalar(MetadataValue::String(value.into()))) - } - - fn visit_unit(self) -> Result { - Ok(YamlValue::Scalar(MetadataValue::Null)) - } - - fn visit_seq>(self, mut sequence: A) -> Result { - let mut values = Vec::new(); - while let Some(value) = sequence.next_element()? { - values.push(value); - } - Ok(YamlValue::Sequence(values)) - } - - fn visit_map>(self, mut mapping: A) -> Result { - let mut values = IndexMap::new(); - while let Some((key, value)) = mapping.next_entry()? { - values.insert(key, value); - } - Ok(YamlValue::Mapping(values)) - } - - fn visit_enum>(self, tagged: A) -> Result { - let (tag, value) = tagged.variant::()?; - Ok(YamlValue::Tagged(tag, Box::new(value.newtype_variant()?))) - } - } - - deserializer.deserialize_any(YamlVisitor) - } } impl YamlValue { @@ -296,10 +170,6 @@ impl YamlValue { .collect::>() .map(MetadataValue::Sequence), Self::Scalar(value) => Ok(value), - Self::Tagged(tag, value) => Ok(MetadataValue::Tagged(Box::new(TaggedValue { - tag: Tag::new(tag), - value: value.into_value()?, - }))), } } } diff --git a/packages/leapp-runtime-rust/tests/bundle.rs b/packages/leapp-runtime-rust/tests/bundle.rs index 67d850fc..71260793 100644 --- a/packages/leapp-runtime-rust/tests/bundle.rs +++ b/packages/leapp-runtime-rust/tests/bundle.rs @@ -637,46 +637,3 @@ fn nested_yaml_merges_expand_before_overrides() { assert_eq!(b.metadata()["combined"]["frequency"], value("100")); assert_eq!(b.metadata()["combined"]["label"], "explicit"); } - -#[test] -fn tagged_metadata_preserves_nested_values_duplicates_and_merges() { - let f = Fixture::new(); - fs::write(&f.path, format!("{YAML}\nbase: &base {{label: inherited}}\ntagged: !sensor {{<<: *base, label: old, label: current, samples: [!units radians, null, 1.25]}}\n")).unwrap(); - let b = Bundle::open(&f.path).unwrap(); - let Value::Tagged(tagged) = &b.metadata()["tagged"] else { - panic!("metadata tag was lost"); - }; - assert_eq!(tagged.tag, "sensor"); - assert_eq!(tagged.value["label"], "current"); - assert_eq!( - tagged.value["samples"], - value("[!units radians, null, 1.25]") - ); -} - -#[test] -fn tagged_metadata_is_preserved_at_every_inspection_level() { - let mut f = Fixture::new(); - let tag = value("!description {units: !unit radians}"); - f.value["description"] = tag.clone(); - f.value["pipeline"]["description"] = tag.clone(); - f.value["pipeline"]["configs"]["description"] = tag.clone(); - f.value["system information"]["description"] = tag.clone(); - f.value["models"]["a_policy"]["description"] = tag.clone(); - f.value["models"]["a_policy"]["parameters"]["description"] = tag.clone(); - f.value["models"]["a_policy"]["inputs"][0]["description"] = tag.clone(); - f.value["models"]["a_policy"]["outputs"][0]["description"] = tag.clone(); - let b = f.open(); - for metadata in [ - b.metadata(), - b.pipeline_metadata(), - b.pipeline_configs(), - b.system_information(), - b.model_metadata("a_policy").unwrap(), - b.model_parameters("a_policy").unwrap(), - &b.input_info("a_policy", "features").unwrap().metadata, - &b.output_info("a_policy", "actions").unwrap().metadata, - ] { - assert_eq!(metadata["description"], tag); - } -} From 8ced418a6db9daea4e3231e29c4e289cde32f4c8 Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Mon, 19 Oct 2026 12:11:32 -0700 Subject: [PATCH 14/21] feat: expose model input and output names Signed-off-by: Marc Alwan --- packages/leapp-runtime-rust/README.md | 6 +++++ .../leapp-runtime-rust/bindings/cxx/README.md | 12 +++++++++ .../bindings/cxx/include/leapp/bundle.h | 12 +++++++++ .../bindings/cxx/src/bridge.rs | 7 +++++ .../bindings/cxx/src/bundle.cc | 20 ++++++++++++++ .../bindings/cxx/src/bundle.rs | 14 ++++++++++ .../bindings/cxx/tests/bundle_test.cc | 8 ++++++ .../bindings/cxx/tests/consumer/main.cc | 1 + .../bindings/python/README.md | 9 +++++++ .../python/python/leapp_runtime/_native.pyi | 2 ++ .../bindings/python/src/bundle.rs | 11 ++++++++ .../bindings/python/tests/test_bundle.py | 11 ++++++++ packages/leapp-runtime-rust/src/bundle.rs | 17 ++++++++++++ packages/leapp-runtime-rust/tests/bundle.rs | 27 +++++++++++++++++++ 14 files changed, 157 insertions(+) diff --git a/packages/leapp-runtime-rust/README.md b/packages/leapp-runtime-rust/README.md index 28100ff8..9a572faf 100644 --- a/packages/leapp-runtime-rust/README.md +++ b/packages/leapp-runtime-rust/README.md @@ -47,6 +47,12 @@ for output in bundle.pipeline_outputs() { # Ok::<(), leapp_runtime::BundleError>(()) ``` +`model_io(node)` returns `Option<(Vec<&str>, Vec<&str>)>`: input names followed +by output names, including internal ports. Names are unqualified and borrowed +from the bundle, with declaration order and duplicates preserved. An absent +model returns `None`; a model with no ports returns two empty lists. Use these +names with `input_info(node, port)` and `output_info(node, port)` for descriptors. + `Bundle` is an immutable, reference-counted snapshot. Clones share the snapshot; inspection is safe across threads. Getters preserve declaration order and unknown metadata. Exporter parameters retain their original paths and backend diff --git a/packages/leapp-runtime-rust/bindings/cxx/README.md b/packages/leapp-runtime-rust/bindings/cxx/README.md index d1233fdb..2234466f 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/README.md +++ b/packages/leapp-runtime-rust/bindings/cxx/README.md @@ -32,8 +32,20 @@ if (!status.ok()) { for (const auto& [name, spec] : bundle.pipeline_outputs()) { // Inspect the optional PortSpec and its typed Metadata. } +for (const auto& model : bundle.model_names()) { + const auto* io = bundle.model_io(model); + for (const auto& name : io->inputs) { + const auto* spec = bundle.input_info(model, name); + // Inspect the internal input descriptor. + } +} ``` +`model_io(node)` returns a retained `const ModelIO*` with `inputs` and `outputs` +name lists, including internal ports. Names are unqualified; declaration order +and duplicates are preserved. An absent model returns `nullptr`, while a model +with no ports has empty lists. + `Bundle` copies retain a shared immutable snapshot. Getter references remain valid while any copy owns the snapshot. Concurrent reads are safe; assignment or destruction of the same handle requires synchronization. `Open` requires an diff --git a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/bundle.h b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/bundle.h index 74e0a9c6..19d8c4fd 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/bundle.h +++ b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/bundle.h @@ -37,6 +37,12 @@ struct PortSpec { /// declaration. using PortMap = std::vector>>; +/// @brief Unqualified model port names in declaration order, including duplicates. +struct ModelIO { + std::vector inputs; + std::vector outputs; +}; + /// @brief Shared immutable bundle snapshot and retained C++ inspection values. /// Getter references survive while any copy retains the snapshot. Concurrent /// reads are safe; assignment or destruction of the same handle needs @@ -52,6 +58,12 @@ class Bundle { const std::filesystem::path& yaml_path() const; const std::vector& model_names() const; + + /// @brief Borrow input/output names, including internal ports. An absent model + /// returns nullptr; a model with no ports has empty lists. The pointer survives + /// while any Bundle copy retains the snapshot. + const ModelIO* model_io(std::string_view node) const; + const PortMap& pipeline_inputs() const; const PortMap& pipeline_outputs() const; const Metadata& pipeline_configs() const; diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs b/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs index a5852262..307de97f 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs +++ b/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs @@ -63,6 +63,12 @@ pub mod ffi { port: String, } + struct ModelIoData { + present: bool, + inputs: Vec, + outputs: Vec, + } + enum MetadataField { Root, System, @@ -83,6 +89,7 @@ pub mod ffi { fn bundle_open(path: &[u8], output: &mut BundleHandle) -> StatusData; fn bundle_path(bundle: &BundleHandle) -> Vec; fn bundle_models(bundle: &BundleHandle) -> Vec; + fn bundle_model_io(bundle: &BundleHandle, node: &[u8]) -> ModelIoData; fn bundle_ports(bundle: &BundleHandle, inputs: bool) -> Vec; fn bundle_port(bundle: &BundleHandle, node: &[u8], port: &[u8], input: bool) -> PortData; fn bundle_metadata( diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc index b9ea0bc7..8dd650ca 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc +++ b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc @@ -154,6 +154,17 @@ struct Bundle::Data { parameters.emplace(name, Field(*handle, wire::MetadataField::Parameters, name)); models.emplace(name, Field(*handle, wire::MetadataField::Model, name)); + const auto io = wire::bundle_model_io(*handle, Bytes(name)); + if (io.present) { + ModelIO ports; + for (const auto& input : io.inputs) { + ports.inputs.push_back(String(input)); + } + for (const auto& output : io.outputs) { + ports.outputs.push_back(String(output)); + } + model_ports.emplace(name, std::move(ports)); + } } inputs = Ports(true); outputs = Ports(false); @@ -209,6 +220,7 @@ struct Bundle::Data { Metadata pipeline; std::map> parameters; std::map> models; + std::map> model_ports; // Map nodes never move; returned descriptor addresses survive later lookups. mutable std::mutex mutex; mutable std::map, @@ -239,6 +251,14 @@ const std::vector& Bundle::model_names() const { return data_ ? data_->names : empty; } +const ModelIO* Bundle::model_io(std::string_view node) const { + if (!data_) { + return nullptr; + } + const auto it = data_->model_ports.find(node); + return it == data_->model_ports.end() ? nullptr : &it->second; +} + const PortMap& Bundle::pipeline_inputs() const { static const PortMap empty; return data_ ? data_->inputs : empty; diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs index 5cb77e97..5bd9685b 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs +++ b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs @@ -40,6 +40,20 @@ pub fn bundle_models(bundle: &BundleHandle) -> Vec { .unwrap_or_default() } +pub fn bundle_model_io(bundle: &BundleHandle, node: &[u8]) -> ffi::ModelIoData { + let io = bundle + .0 + .as_ref() + .and_then(|bundle| bundle.model_io(std::str::from_utf8(node).ok()?)); + let present = io.is_some(); + let (inputs, outputs) = io.unwrap_or_default(); + ffi::ModelIoData { + present, + inputs: inputs.into_iter().map(String::from).collect(), + outputs: outputs.into_iter().map(String::from).collect(), + } +} + pub fn bundle_ports(bundle: &BundleHandle, inputs: bool) -> Vec { let Some(b) = &bundle.0 else { return Vec::new(); diff --git a/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc b/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc index de7ee8f1..b4e7e7b7 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc +++ b/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc @@ -41,6 +41,7 @@ int main(int argc, char** argv) { leapp::Bundle empty; Check(!empty.is_valid() && empty.model_names().empty(), "empty handle"); Check(empty.input_info("missing", "port") == nullptr, "empty descriptor"); + Check(empty.model_io("missing") == nullptr, "empty model I/O"); Check(!leapp::Bundle::Open(path, nullptr).ok(), "null destination"); auto status = leapp::Bundle::Open(directory / "absent.yaml", &empty); Check(status.code() == leapp::ErrorCode::kNotFound, "missing YAML category"); @@ -54,6 +55,11 @@ int main(int argc, char** argv) { Check(bundle.model_names() == std::vector({"z_observation", "a_policy"}), "model declaration order"); + const auto* io = bundle.model_io("a_policy"); + Check(io && io->inputs == std::vector({"features", "state"}) && + io->outputs == std::vector({"actions", "next_state"}), + "model I/O includes internal ports in declaration order"); + Check(bundle.model_io("absent") == nullptr, "absent model I/O"); Check(bundle.pipeline_inputs().size() == 1, "pipeline inputs"); Check(bundle.pipeline_outputs().size() == 3, "pipeline outputs"); Check(bundle.pipeline_outputs()[0].first == "a_policy/next_state", @@ -85,6 +91,8 @@ int main(int argc, char** argv) { auto retained = bundle; bundle = {}; std::filesystem::remove(path); + Check(retained.model_io("a_policy") == io && io->inputs.front() == "features", + "model I/O lifetime follows the retained snapshot"); std::vector readers; for (int i = 0; i < 4; ++i) { readers.emplace_back([retained, port] { diff --git a/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/main.cc b/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/main.cc index 2c2f2e68..ab42c01f 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/main.cc +++ b/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/main.cc @@ -7,6 +7,7 @@ int main() { leapp::Bundle bundle; const auto status = leapp::Bundle::Open("", &bundle); return !status.ok() && !bundle.is_valid() && + bundle.model_io("absent") == nullptr && status.code() == leapp::ErrorCode::kInvalidArgument ? 0 : 1; diff --git a/packages/leapp-runtime-rust/bindings/python/README.md b/packages/leapp-runtime-rust/bindings/python/README.md index e805da9a..fd08da6f 100644 --- a/packages/leapp-runtime-rust/bindings/python/README.md +++ b/packages/leapp-runtime-rust/bindings/python/README.md @@ -17,8 +17,17 @@ from leapp_runtime import Bundle, BundleError bundle = Bundle.open("export/policy.yaml") for name, spec in bundle.pipeline_outputs().items(): print(name, spec.dtype, spec.shape) +for model in bundle.model_names(): + inputs, outputs = bundle.model_io(model) + for name in inputs: + print(model, name, bundle.input_info(model, name).shape) ``` +`model_io(model_name)` returns a tuple of detached input and output name lists, +including internal ports. Names are unqualified; declaration order and duplicates +are preserved. An absent model returns `None`, while a model with no ports returns +`([], [])`. + Metadata and port descriptors are detached from the immutable Rust snapshot. Editing returned metadata does not change later inspection results. Mappings with unhashable or colliding Python keys become `MetadataMapping` so no entries diff --git a/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi b/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi index 19c20e78..9f91a0d9 100644 --- a/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi +++ b/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi @@ -47,6 +47,8 @@ class Bundle: def model_names(self) -> list[str]: ... + def model_io(self, model_name: str) -> tuple[list[str], list[str]] | None: ... + def pipeline_inputs(self) -> dict[str, PortSpec | None]: ... def pipeline_outputs(self) -> dict[str, PortSpec | None]: ... diff --git a/packages/leapp-runtime-rust/bindings/python/src/bundle.rs b/packages/leapp-runtime-rust/bindings/python/src/bundle.rs index cad524c7..e21391b2 100644 --- a/packages/leapp-runtime-rust/bindings/python/src/bundle.rs +++ b/packages/leapp-runtime-rust/bindings/python/src/bundle.rs @@ -151,6 +151,17 @@ impl PyBundle { .collect() } + /// Detached input/output name lists in declaration order, or None if absent. + /// Includes internal ports and preserves duplicate, unqualified names. + fn model_io(&self, model_name: &str) -> Option<(Vec, Vec)> { + self.inner.model_io(model_name).map(|(inputs, outputs)| { + ( + inputs.into_iter().map(String::from).collect(), + outputs.into_iter().map(String::from).collect(), + ) + }) + } + /// Ordered mapping of qualified input names to detached port descriptors. fn pipeline_inputs<'py>(&self, py: Python<'py>) -> PyResult> { let result = PyDict::new(py); diff --git a/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py b/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py index 1a32f006..8e31c240 100644 --- a/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py +++ b/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py @@ -55,6 +55,17 @@ def test_ordered_inspection_and_detached_values(bundle_path): assert spec.shape == [3] +def test_model_io_returns_detached_internal_port_names(bundle_path): + bundle = lr.Bundle.open(bundle_path) + expected = (["features", "state"], ["actions", "next_state"]) + inputs, outputs = bundle.model_io(model_name="a_policy") + assert (inputs, outputs) == expected + assert bundle.model_io("absent") is None + inputs.clear() + outputs.append("extra") + assert bundle.model_io("a_policy") == expected + + def test_custom_backend_and_unresolved_input(bundle_path): bundle_path.write_text( YAML.replace("backend: pt2", "backend: custom").replace( diff --git a/packages/leapp-runtime-rust/src/bundle.rs b/packages/leapp-runtime-rust/src/bundle.rs index 74ebe263..d6961d94 100644 --- a/packages/leapp-runtime-rust/src/bundle.rs +++ b/packages/leapp-runtime-rust/src/bundle.rs @@ -63,6 +63,23 @@ impl Bundle { .collect() } + /// Input and output names for a model, including internal ports. + /// + /// Names are unqualified and borrowed from this snapshot. Both lists preserve + /// declaration order and duplicates. An absent model returns `None`; a model + /// with no ports returns two empty lists. + pub fn model_io(&self, node: &str) -> Option<(Vec<&str>, Vec<&str>)> { + let model = self.0.description.models.get(node)?; + Some(( + model.inputs.iter().map(|port| port.name.as_str()).collect(), + model + .outputs + .iter() + .map(|port| port.name.as_str()) + .collect(), + )) + } + /// Declared input references in order; these need not identify existing ports. pub fn pipeline_inputs(&self) -> &[PortRef] { &self.0.inputs diff --git a/packages/leapp-runtime-rust/tests/bundle.rs b/packages/leapp-runtime-rust/tests/bundle.rs index 71260793..9e14dbdc 100644 --- a/packages/leapp-runtime-rust/tests/bundle.rs +++ b/packages/leapp-runtime-rust/tests/bundle.rs @@ -99,6 +99,33 @@ fn inspection_preserves_order_and_metadata() { ); } +#[test] +fn model_io_preserves_internal_port_order_and_distinguishes_absent_models() { + let mut f = Fixture::new(); + let model = &mut f.value["models"]["a_policy"]; + for direction in ["inputs", "outputs"] { + let ports = model[direction].as_sequence_mut().unwrap(); + ports.push(ports[0].clone()); + } + let mut empty = model.clone(); + empty["inputs"] = value("[]"); + empty["outputs"] = value("[]"); + f.value["models"]["empty"] = empty; + + let b = f.open(); + let (inputs, outputs) = b.model_io("a_policy").unwrap(); + assert_eq!(inputs, ["features", "state", "features"]); + assert_eq!(outputs, ["actions", "next_state", "actions"]); + for name in inputs { + assert!(b.input_info("a_policy", name).is_some()); + } + for name in outputs { + assert!(b.output_info("a_policy", name).is_some()); + } + assert_eq!(b.model_io("empty"), Some((vec![], vec![]))); + assert!(b.model_io("absent").is_none()); +} + #[test] fn getters_preserve_unknown_fields_at_each_level() { let mut f = Fixture::new(); From 5f509b2e918a9652bf0f5e1c24ce4c5db45f2909 Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Mon, 19 Oct 2026 12:23:29 -0700 Subject: [PATCH 15/21] fix: preserve error chains and complete shutdown Signed-off-by: Marc Alwan --- packages/leapp-runtime-rust/README.md | 12 +-- .../bindings/common/README.md | 30 +++++++- .../bindings/common/src/runtime.rs | 47 +++++++++-- .../bindings/common/tests/lifecycle.rs | 77 ++++++++++++++++++- 4 files changed, 146 insertions(+), 20 deletions(-) diff --git a/packages/leapp-runtime-rust/README.md b/packages/leapp-runtime-rust/README.md index 9a572faf..2e4941e0 100644 --- a/packages/leapp-runtime-rust/README.md +++ b/packages/leapp-runtime-rust/README.md @@ -19,7 +19,7 @@ import name and should be installed in separate environments. | --- | --- | | Rust core | `Bundle`, YAML schema parsing, graph and artifact validation, SHA-256 checks, safetensors initial-state validation | | Public inspection types | `PortRef`, `PortSpec`, `DataType`, `Metadata`, `MetadataValue`, `BundleError`, `ErrorCode` | -| Common lifecycle sources | `RuntimeHandle`, `SessionHandle`, and lifecycle errors; three focused regression tests | +| Common lifecycle sources | `RuntimeHandle`, `SessionHandle`, error-source preservation, and cleanup after poisoned locks; five focused regression tests | | Python | PyO3 bundle bindings, detached metadata and port descriptors, typed stubs and structured exceptions | | C++ | C++20 bundle facade, typed metadata, status and data-type headers, private cxx transport, installable CMake target | | Tests | Rust bundle regression suite and independent Python/C++ inspection and lifetime tests | @@ -84,7 +84,7 @@ cargo clippy --locked --workspace --all-targets -- -D warnings cargo fmt --all --check ``` -The default workspace member is the core crate, including the three common +The default workspace member is the core crate, including the five common lifecycle tests. These tests require no engine or language-binding SDK. Building the complete workspace also requires a C++20 compiler and Python 3.10 or newer. See @@ -94,6 +94,8 @@ All inspection tests create opaque artifact bytes; no model generation is needed ## Shared lifecycle -`bindings/common/src/runtime.rs` manages session ownership, serializes calls, and -coordinates explicit shutdown. The lifecycle regression target compiles these -sources against a small test-only runtime API, without inference SDKs. +`bindings/common/src/runtime.rs` preserves runtime error sources and completes +shutdown cleanup even when locks are poisoned. The lifecycle regression target +compiles this implementation against a small test-only runtime API, without +inference SDKs. See the [session reset contract](bindings/common/README.md#session-reset) +for state ownership and reset behavior. diff --git a/packages/leapp-runtime-rust/bindings/common/README.md b/packages/leapp-runtime-rust/bindings/common/README.md index d25b5e9c..07e6466c 100644 --- a/packages/leapp-runtime-rust/bindings/common/README.md +++ b/packages/leapp-runtime-rust/bindings/common/README.md @@ -6,7 +6,10 @@ SPDX-License-Identifier: Apache-2.0 # Shared lifecycle sources `src/lib.rs` and `src/runtime.rs` provide session ownership, locking, and explicit -closure. +closure. Wrapped runtime errors expose their source chain, and close finishes +releasing every session and the runtime even when locks are poisoned. +Close reports the lock error after cleanup; subsequent +close calls succeed. Normal run/reset operations still fail on poisoned locks. This source depends on the concrete `Runtime`, `Session`, tensor, options, and error types in `leapp_runtime`. Those execution types are not part of this @@ -15,10 +18,29 @@ The Python and C++ inspection bindings use the core bundle crate directly. Runtime packages compile these sources with their execution types and features. The root Cargo test target includes `runtime.rs` against a small runtime API -defined only in the test binary. Three tests cover session ownership, close -invalidation and retained bundle info, and waiting for an active run. Execution -behavior is tested in the full runtime integration. +defined only in the test binary. Five tests cover session ownership, close +invalidation and retained bundle info, waiting for an active run, error-source +preservation, and cleanup after poisoned locks. Execution behavior is tested in +the full runtime integration. ```sh cargo test --locked --test binding_lifecycle ``` + +## Session reset + +`SessionHandle::reset()` delegates to the execution session. State belongs to +each session, so reset restores its feedback and default inputs to their captured +initial values. Those values are loaded when the runtime opens: exported +initial tensors are used when present, otherwise zeros are used for statically +shaped state. Dynamic state requires an exported initial tensor with a concrete +shape. + +Reset serializes with runs on the same session. It preserves other sessions and +retained results, and it does not reread files, reload models, reset +backend-internal mutable state, or flush shared allocation pools. The shared +wrapper rejects reset after explicit runtime closure. + +This inspection-only package provides the reset wrapper as source; it does not +export an executable runtime/session API. Tensor-state restoration requires the +execution engine and is outside the lifecycle tests here. diff --git a/packages/leapp-runtime-rust/bindings/common/src/runtime.rs b/packages/leapp-runtime-rust/bindings/common/src/runtime.rs index 865821e9..3a28982d 100644 --- a/packages/leapp-runtime-rust/bindings/common/src/runtime.rs +++ b/packages/leapp-runtime-rust/bindings/common/src/runtime.rs @@ -11,6 +11,7 @@ use std::{ }; /// Binding lifecycle failures preserve native execution diagnostics. +/// Wrapped runtime failures remain accessible through [`std::error::Error::source`]. #[derive(Debug)] pub enum Error { Closed, @@ -47,7 +48,14 @@ impl fmt::Display for Error { } } -impl std::error::Error for Error {} +impl std::error::Error for Error { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Runtime(error) => Some(error), + Self::Closed | Self::Lock => None, + } + } +} type SessionSlot = Mutex>; @@ -104,17 +112,40 @@ impl RuntimeHandle { }) } - /// Wait for active native calls and close all sessions. Repeated calls succeed. + /// Wait for active native calls and close all sessions. + /// Poisoned locks are recovered only for cleanup, then reported as `Error::Lock`. + /// Once cleanup completes, repeated calls succeed. pub fn close(&self) -> Result<(), Error> { - let mut runtime = self.0.runtime.write().map_err(|_| Error::Lock)?; - let mut sessions = self.0.sessions.lock().map_err(|_| Error::Lock)?; + let mut poisoned = false; + let mut runtime = self.0.runtime.write().unwrap_or_else(|error| { + poisoned = true; + error.into_inner() + }); + let mut sessions = self.0.sessions.lock().unwrap_or_else(|error| { + poisoned = true; + error.into_inner() + }); + // The runtime write guard excludes active calls. Recover poisoned state + // only to drop it, and visit every session before reporting the failure. for weak in sessions.drain(..) { if let Some(session) = weak.upgrade() { - session.lock().map_err(|_| Error::Lock)?.take(); + let mut slot = session.lock().unwrap_or_else(|error| { + poisoned = true; + error.into_inner() + }); + slot.take(); + session.clear_poison(); } } runtime.take(); - Ok(()) + // Empty, closed state is valid again. Clear poison while holding the guards. + self.0.sessions.clear_poison(); + self.0.runtime.clear_poison(); + if poisoned { + Err(Error::Lock) + } else { + Ok(()) + } } } @@ -138,7 +169,9 @@ impl SessionHandle { .map_err(Error::Runtime) } - /// Restore the session's explicit feedback state. + /// Restore captured initial feedback and default inputs for this session. + /// Retained results are unaffected; files and models are not reloaded. + /// Backend-internal mutable state is not reset. pub fn reset(&self) -> Result<(), Error> { let guard = self.state.runtime.read().map_err(|_| Error::Lock)?; if guard.is_none() { diff --git a/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs b/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs index eebd3bf9..28cffcbb 100644 --- a/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs +++ b/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs @@ -14,8 +14,51 @@ use std::{ time::Duration, }; -#[path = "../src/runtime.rs"] -mod lifecycle; +mod lifecycle { + include!("../src/runtime.rs"); + + #[test] + fn close_recovers_poisoned_locks_and_releases_every_session() { + use std::panic::{catch_unwind, AssertUnwindSafe}; + + for poisoned in ["session", "registry", "runtime"] { + let (runtime, owner) = crate::runtime(); + let sessions: Vec<_> = (0..3).map(|_| runtime.create_session().unwrap()).collect(); + assert!(catch_unwind(AssertUnwindSafe(|| match poisoned { + "session" => { + // Cleanup must visit sessions on both sides of the poisoned one. + let _guard = sessions[1].slot.lock().unwrap(); + panic!("simulate session panic"); + } + "registry" => { + let _guard = runtime.0.sessions.lock().unwrap(); + panic!("simulate registry panic"); + } + "runtime" => { + let _guard = runtime.0.runtime.write().unwrap(); + panic!("simulate runtime panic"); + } + _ => unreachable!(), + })) + .is_err()); + + assert!(matches!(runtime.close(), Err(Error::Lock)), "{poisoned}"); + assert!(!runtime.is_open().unwrap()); + assert!(owner.upgrade().is_none()); + assert!(runtime.0.sessions.lock().unwrap().is_empty()); + for session in &sessions { + assert!(session.slot.lock().unwrap().is_none()); + assert!(matches!(session.reset(), Err(Error::Closed))); + assert!(matches!( + session.run(&[], &RunOptions::default()), + Err(Error::Closed) + )); + } + assert!(matches!(runtime.create_session(), Err(Error::Closed))); + runtime.close().unwrap(); + } + } +} use lifecycle::{Error, RuntimeHandle}; pub type RuntimeOptions = Arc<()>; @@ -30,7 +73,7 @@ pub struct Runtime(Arc<()>); pub struct Session(Weak<()>); #[derive(Debug)] -pub struct RuntimeError; +pub struct RuntimeError(std::io::Error); impl RuntimeError { pub fn code(&self) -> ErrorCode { @@ -44,7 +87,13 @@ impl RuntimeError { impl fmt::Display for RuntimeError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.write_str("test runtime error") + write!(f, "test runtime error: {}", self.0) + } +} + +impl std::error::Error for RuntimeError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(&self.0) } } @@ -100,6 +149,26 @@ fn runtime() -> (RuntimeHandle, Weak<()>) { ) } +#[test] +fn runtime_errors_preserve_the_source_chain_and_diagnostics() { + use std::error::Error as _; + + let runtime = RuntimeError(std::io::Error::other("test backend error")); + let message = runtime.to_string(); + let error = Error::Runtime(runtime); + let source = error + .source() + .unwrap() + .downcast_ref::() + .unwrap(); + assert_eq!(error.code(), source.code()); + assert_eq!(error.phase(), source.phase()); + assert_eq!(error.to_string(), message); + assert_eq!(source.source().unwrap().to_string(), "test backend error"); + assert!(Error::Closed.source().is_none()); + assert!(Error::Lock.source().is_none()); +} + #[test] fn sessions_keep_runtime_alive_after_the_last_runtime_handle_is_dropped() { let (runtime, owner) = runtime(); From c694c2754e1a08c4cae62af5cfda976098c3d4f1 Mon Sep 17 00:00:00 2001 From: Frank Lai Date: Mon, 19 Oct 2026 15:56:58 -0700 Subject: [PATCH 16/21] docs: use repository skill references in agent instructions Signed-off-by: Frank Lai --- AGENTS.md | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 561368df..628001b7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -16,17 +16,22 @@ Python 3.10+ and PyTorch 2.6.0+. Visualization extras install only on ## Load this first +All instruction and skill references below are relative to the repository +root and tracked in Git. Do not require machine-local or editor-specific +skill installations. + | If you are changing… | Read this before editing | |---|---| -| Tracing, `TracedData` / `TracedTensor` / `TracedNpArray`, indexing, mutation, node-boundary identity, FX proxies, NumPy dispatch | `.cursor/skills/leap_agent/SKILL.md` | -| Warp tracing (`TracedWpArray`, APIC, `warp_op`) | `warp_support.md` | +| Tracing, `TracedData` / `TracedTensor` / `TracedNpArray`, indexing, mutation, node-boundary identity, FX proxies, NumPy dispatch | [SKILL.md](SKILL.md), especially "Preserve graph connections and computation" and "Diagnose trace loss before changing exporters"; implementation in `leapp/leapp_graph/` | +| Warp tracing (`TracedWpArray`, APIC, `warp_op`) | [Warp guide](references/warp.md) | | Public annotate/lifecycle API or docs examples | `docs/source/api/index.rst` plus the matching guide under `docs/source/` | | Graph PNG layout | `packages/leapp-visualization/` | | Export backends | `leapp/backends/` | | YAML / `InferenceManager` | `leapp/leapp.py` (`compile_graph`), `leapp/inference_manager.py`, `docs/source/generated_configs.rst` | -Do not mix the skill's current `(context_obj, output_port)` identity with -legacy `leapp_tag` helpers. The skill states which checkout you have. +Do not mix the current `(context_obj, output_port)` identity with +legacy `leapp_tag` helpers. Check the implementation in this repository +when changing tracing internals. ## Tree (start here) @@ -72,7 +77,7 @@ Narrower (use these while iterating): | Annotate / I/O / state | `pytest tests/functional_tests/ -v` | | Examples | `pytest tests/test_examples/ -v` | | Visualization | `pytest packages/leapp-visualization/tests/ -v` (Python 3.11+) | -| Warp tracing | see `warp_support.md` | +| Warp tracing | see [Warp guide](references/warp.md) | ### Python formatting @@ -187,7 +192,7 @@ User-facing integration recipes belong here, not in this file. lifecycle stays in `warp/session.py`, `warp_op.py`, and `warp/warp_segment.py`. - After shared indexing changes: both torch and NumPy unit suites, plus - FX / jit / exported-program coverage as the skill describes. + FX / jit / exported-program coverage in `tests/`. `onnx-dynamo` for modern index cases; do not require legacy ONNX if that exporter cannot lower the op. - Do not invent graph edges for data a node did not publish. Clearing From ba891b394b06d6d779c2563e1e22cd33215865e8 Mon Sep 17 00:00:00 2001 From: Frank Lai NV Date: Tue, 20 Oct 2026 16:03:53 -0700 Subject: [PATCH 17/21] docs: update Isaac Lab policy export URL (#38) Signed-off-by: Frank Lai NV --- docs/source/ecosystem.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/ecosystem.rst b/docs/source/ecosystem.rst index 2884b322..9eb04f9b 100644 --- a/docs/source/ecosystem.rst +++ b/docs/source/ecosystem.rst @@ -11,7 +11,7 @@ handling by hand. :gutter: 3 .. grid-item-card:: Isaac Lab - :link: https://isaac-sim.github.io/IsaacLab/release/3.0.0/source/policy_deployment/05_leapp/exporting_policies_with_leapp.html + :link: https://isaac-sim.github.io/IsaacLab/main/source/policy_deployment/05_leapp/exporting_policies_with_leapp.html :link-type: url :img-bottom: _static/images/ecosystem_isaac_lab.gif :img-alt: Isaac Lab From 1a749d30409d3196f95558063f200f373fa9aab6 Mon Sep 17 00:00:00 2001 From: Frank Lai NV Date: Tue, 20 Oct 2026 23:30:29 -0700 Subject: [PATCH 18/21] Capture NumPy functions as ATen and support reverse slices (#32) * feat: lower NumPy calls through ATen and support reverse slices Signed-off-by: Frank Lai * docs: remove NumPy ATen experiment report Signed-off-by: Frank Lai * test: consolidate NumPy tracing coverage Signed-off-by: Frank Lai * docs: explain NumPy ATen capture and supported limits Signed-off-by: Frank Lai * docs: describe NumPy capture without historical comparisons Signed-off-by: Frank Lai * fix: defer recoverable NumPy capture errors until output validation Signed-off-by: Frank Lai * perf: validate NumPy numerics at exported node outputs Signed-off-by: Frank Lai * test: fold NumPy comparison probe into replay coverage Signed-off-by: Frank Lai * docs: simplify NumPy capture explanation Signed-off-by: Frank Lai * test: use portable cosine dtypes and remove duplicate NumPy checks Signed-off-by: Frank Lai --------- Signed-off-by: Frank Lai --- AGENTS.md | 13 + SKILL.md | 7 +- docs/source/spelling_wordlist.txt | 5 + docs/source/tensor_libraries/how_it_works.rst | 149 ++- docs/source/tensor_libraries/limitations.rst | 185 +-- docs/source/tensor_libraries/numpy.rst | 8 +- docs/source/tensor_libraries/torch.rst | 6 +- .../tensor_libraries/warp_how_it_works.rst | 6 +- .../datatypes/numpy/aten_adapter.py | 248 ++++ .../leapp_graph/datatypes/numpy/overrides.py | 76 ++ .../datatypes/numpy/traced_np_array.py | 479 +------ leapp/leapp_graph/datatypes/traced_data.py | 86 ++ leapp/leapp_graph/traced_node.py | 10 + tests/unit_tests/test_numpy_compatibility.py | 76 ++ tests/unit_tests/test_numpy_slicing.py | 134 ++ tests/unit_tests/test_traced_np_array.py | 1182 +++++++---------- 16 files changed, 1306 insertions(+), 1364 deletions(-) create mode 100644 leapp/leapp_graph/datatypes/numpy/aten_adapter.py create mode 100644 leapp/leapp_graph/datatypes/numpy/overrides.py create mode 100644 tests/unit_tests/test_numpy_slicing.py diff --git a/AGENTS.md b/AGENTS.md index 628001b7..6978293c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -184,6 +184,19 @@ User-facing integration recipes belong here, not in this file. ## Guardrails when changing behavior +- Tracing can enter branches that are not needed to produce the declared + outputs. Report recoverable capture or translation failures with + `_get_logger().error(...)`, which logs without raising; do not abort the + entire trace merely because such a branch cannot be captured. Keep this + distinction from `_get_logger().fatal(...)`, which raises. Test that an + unsupported operation on an unused branch does not prevent exporting the + supported output graph. Preserve failed-operation dependencies until output + pruning, and reject export if a declared output depends on a failed capture; + test direct and downstream dependencies separately. +- NumPy capture checks graph structure and mutation safety. Leave numerical + comparisons to final exported-node validation; do not add per-call eager + parity checks to the adapter. Retain shape, dtype, value, and changed-input + parity coverage in tests, including final validation rejecting wrong outputs. - Prefer extending `TracedData` once (`_map_structure`, index lowering, port preserve/clear) over copying logic into tensor, NumPy, and Warp classes. Warp value semantics live on `TracedWpArray`; `WarpPatchBackend` diff --git a/SKILL.md b/SKILL.md index 5526ebbe..2d18bce0 100644 --- a/SKILL.md +++ b/SKILL.md @@ -155,10 +155,11 @@ Calls in NumPy, SciPy, scikit-learn, scikit-image, OpenCV, or other libraries ** | Hazard | What the agent should do | |--------|--------------------------| -| Unmapped NumPy/native operation | Treat the trace warning as a correctness failure. A plain result mixed back into traced data can become a constant. Patch the operation and compare new inputs against the original. | +| Unsupported NumPy/native operation | NumPy capture checks graph structure; numerical comparison happens during final exported-node validation. Recoverable capture failures log errors and continue eager execution. Unused failures are pruned; a failed dependency in the declared output graph blocks export. Unsupported mutations remain immediate errors. Calls that bypass the hooks can still return plain values that become constants. Patch unsupported computation and compare new inputs against the original. | | `.item()`, `.tolist()`, `float()`, `int()` | Audit explicitly: Torch extraction can be silent. Keep dynamic data in traced operations; logs alone cannot prove correctness. | -| NumPy indexing down to a scalar or NumPy scalar operands | Check for loss of array tracing or unsupported operand types. Use a faithful adapter preserving the needed array computation. | -| NumPy negative-stride slicing, unsupported `astype` dtype | Eager and exported behavior can disagree. Check the online limitations and implement supported equivalent operations within the adapter. | +| NumPy indexing down to a scalar | Keep an array dimension to preserve runtime dependence. Constant NumPy scalar operands are accepted, but cannot restore a dependency lost by indexing. | +| NumPy negative-stride slicing | Static basic reads use `flip` or normalized indices with `index_select`. Runtime bounds, mixed advanced indices, reverse assignment, and mutation through reversed views need separate support. | +| Unsupported `astype` dtype | Eager and exported dtypes can disagree. Use a supported dtype and check replay against the original. | | Tensor-dependent Python `if`/`while` or iteration counts | Only the executed path is captured. Use a faithful traceable expression in an adapter where possible; otherwise document the supported domain or external control-flow boundary. Another exporter cannot recover a branch already lost. | Supported NumPy/Torch conversions such as `torch.from_numpy()` and `.numpy()` can preserve tracing with global patching enabled. Do not rewrite all NumPy code or treat every conversion as a trace break. diff --git a/docs/source/spelling_wordlist.txt b/docs/source/spelling_wordlist.txt index e064964a..e6c61747 100644 --- a/docs/source/spelling_wordlist.txt +++ b/docs/source/spelling_wordlist.txt @@ -267,3 +267,8 @@ isaac canonicalize canonicalized canonicalizing +ATen +ufunc +namespaces +functionalizes +TopK diff --git a/docs/source/tensor_libraries/how_it_works.rst b/docs/source/tensor_libraries/how_it_works.rst index 6cd8dc00..0b6ffc74 100644 --- a/docs/source/tensor_libraries/how_it_works.rst +++ b/docs/source/tensor_libraries/how_it_works.rst @@ -1,21 +1,46 @@ -============= +============ How it works -============= - -``TracedNpArray`` subclasses ``np.ndarray``, so it holds real array data and -behaves like an array everywhere LEAPP is not involved. It implements NumPy's -two dispatch hooks, ``__array_ufunc__`` and ``__array_function__``, which lets -it see every NumPy call made on it before NumPy executes it. - -Each intercepted call does two things: - -#. Runs the original NumPy operation eagerly on the underlying buffer, so the - value you get back is exactly what NumPy would have produced. -#. Looks the NumPy function up in a translation table and records the - **equivalent torch operation** into the FX graph. - -So ``np.clip`` on a traced array returns a clipped NumPy array and records -``torch.clamp``. Given this node: +============ + +``TracedNpArray`` subclasses ``np.ndarray``, so it holds real array data. +Its ``__array_ufunc__`` and ``__array_function__`` hooks intercept NumPy calls +that participate in those protocols. During active tracing, LEAPP captures +those calls through PyTorch's tensor-backed NumPy compatibility layer, +``torch._numpy``. + +From a NumPy call to ATen +========================= + +For each intercepted function or direct ufunc call, LEAPP: + +#. Uses PyTorch's NumPy compatibility layer to express the call as tensor + operations. +#. Captures those operations using tensor versions of the inputs and removes + temporary in-place operations from the graph. +#. Checks that the graph contains supported functional operations and tensor + outputs that can be connected to the eager result. It does not execute the + captured graph again to compare intermediate shapes, dtypes, or values. +#. Inserts the captured ATen operations into the current LEAPP FX graph and + wraps the eager NumPy result with the corresponding graph references. + +ATen is PyTorch's tensor operator layer. One NumPy call can contribute several +ATen operations. The exported graph contains tensor operations, with no +``torch._numpy`` arrays or calls needed during replay. LEAPP still executes +real NumPy for the eager result, preserving its values and view/storage +behavior during tracing. + +Numerical validation happens at the exported node outputs. +``compile_graph()`` defaults to ``validate=True, strict=True``: it compares +exported outputs with cached eager results and raises on validation failure. +Collect representative, varied input samples; see :doc:`/guides/debugging`. + +A recoverable capture or graph-structure failure logs an error and continues eager +execution. LEAPP retains the failed operation's dependencies until it knows +which computations produce the declared outputs. Unused failures are pruned; +a failure needed by an output prevents export. After the owning node finishes +tracing, its arrays use native NumPy behavior again. + +For example, this preprocessing node stays in NumPy at the call site: .. code-block:: python @@ -27,73 +52,55 @@ So ``np.clip`` on a traced array returns a clipped NumPy array and records state_norm = np.clip((state - STATE_MEAN) / STATE_STD, -5.0, 5.0) obs = np.concatenate([state_norm, velocity]) -LEAPP records a graph that mentions no NumPy at all: +Its recorded computation contains ATen subtraction, division, clipping, and +concatenation. Schematically, omitting constant-copy nodes: .. code-block:: text - %state = placeholder[target=state] - %velocity = placeholder[target=velocity] - %_tensor_constant0 = get_attr[target=_tensor_constant0] - %sub = call_function[target=torch.sub](args = (%state, %_tensor_constant0)) - %_tensor_constant1 = get_attr[target=_tensor_constant1] - %div = call_function[target=torch.div](args = (%sub, %_tensor_constant1)) - %clamp = call_function[target=torch.clamp](args = (%div, -5.0, 5.0)) - %cat = call_function[target=torch.cat](args = ([%clamp, %velocity],)) + state, velocity = placeholders + mean, std, lower, upper = graph constants + sub = torch.ops.aten.sub(state, mean) + div = torch.ops.aten.div(sub, std) + clamp = torch.ops.aten.clamp(div, lower, upper) + cat = torch.ops.aten.cat([clamp, velocity]) -Two details in that graph matter for :doc:`limitations`. Plain NumPy arrays -used as operands, here ``STATE_MEAN`` and ``STATE_STD``, become frozen -``get_attr`` constants. And each recorded node comes from a lookup, so a -NumPy call with no entry in the table records nothing. +Plain arrays such as ``STATE_MEAN`` and ``STATE_STD`` become graph constants. +If they must vary at runtime, declare them as inputs. Exact ATen sequences can +vary with the installed PyTorch version. What gets traced ================ -Tracing is driven by an explicit NumPy-to-torch table, so support is a fixed -list rather than a general rule. +Coverage depends on the callable, arguments, dtype, installed PyTorch version, +and export backend. Finding a compatibility implementation is only the first +step; that particular call must also pass capture and structural checks. +Examples exercised by the tests include: .. list-table:: :header-rows: 1 :widths: 26 74 * - Category - - Covered - * - Arithmetic and math - - ``+ - * / // % **``, ``np.add`` and the other arithmetic ufuncs, - ``np.abs``, ``np.sqrt``, ``np.square``, ``np.exp``, ``np.log``, - trigonometric and hyperbolic functions, ``np.floor``, ``np.ceil``, - ``np.round`` - * - Comparison and logic - - ``== != < <= > >=``, ``np.maximum``, ``np.minimum``, the logical and - bitwise ufuncs, ``np.isnan``, ``np.isinf``, ``np.isfinite`` + - Examples + * - Arithmetic and comparison + - Arithmetic ufuncs, ``np.sin``, ``np.sqrt``, ``np.cbrt``, and + ``np.isclose`` feeding ``np.where`` * - Reductions - - ``np.sum``, ``np.mean``, ``np.std``, ``np.var``, ``np.min``, ``np.max``, - ``np.argmin``, ``np.argmax``, ``np.cumsum``, ``np.all``, ``np.any``, - with ``axis`` and ``keepdims`` - * - Shape and layout - - ``np.reshape``, ``np.transpose``, ``.T``, ``np.squeeze``, - ``np.expand_dims``, ``np.concatenate``, ``np.stack``, ``np.split``, - ``np.swapaxes``, ``np.moveaxis``, ``np.roll`` - * - Selection - - ``np.where``, ``np.clip``, ``np.sort``, ``np.argsort``, ``np.nonzero``, - slicing, boolean masks, index arrays, indexed assignment + - ``np.sum``, ``np.mean``, ``np.nanmean``, ``np.var``, and + ``np.quantile`` with supported axis and dtype arguments + * - Shape and selection + - ``np.reshape``, ``np.transpose``, ``np.concatenate``, ``np.diff``, + ``np.broadcast_to``, constant padding, and sorting * - Linear algebra - - ``np.matmul``, ``@``, ``np.dot``, ``np.einsum``, ``np.tensordot``, - ``np.trace``, ``np.tril``, ``np.triu`` - * - Array creation - - ``np.zeros_like``, ``np.ones_like``, ``np.full_like``, ``np.zeros``, - ``np.ones``, ``np.arange``, ``np.linspace``, ``np.eye`` - * - Conversion - - ``astype``, ``.copy()``, ``np.asarray``, ``torch.from_numpy``, - ``torch.as_tensor``, ``TracedTensor.numpy()`` - -Anything outside the table runs normally and is skipped by the tracer, with a -warning naming the call: - -.. code-block:: text - - No torch equivalent for numpy function . Operation will not be traced. - -Treat that warning as an error. See :doc:`limitations`. After tracing, -run ``compile_graph(validate=True)`` so a later sample can catch an -untraced result that was frozen as a constant. The workflow is on -:doc:`/guides/debugging`. + - ``np.matmul``, ``np.linalg.norm``, ``np.linalg.inv``, and ``np.cross`` + +Array methods, indexing, and NumPy/Torch conversions retain their own tracing +paths. Static basic reverse slices use ``torch.flip`` for full reversal and +normalized indices with ``torch.index_select`` for other negative strides. +Calls that never reach a traced operand, such as standalone array creation, +are not made traceable by compatibility-layer discovery. + +These examples are not a promise of every NumPy signature or backend. The +adapter uses private PyTorch APIs, and agreement on one input cannot prove +correctness for other values. Validate exported results on changed inputs as +well as the capture input; see :doc:`/guides/debugging` and :doc:`limitations`. diff --git a/docs/source/tensor_libraries/limitations.rst b/docs/source/tensor_libraries/limitations.rst index bf01433b..8f0c6430 100644 --- a/docs/source/tensor_libraries/limitations.rst +++ b/docs/source/tensor_libraries/limitations.rst @@ -2,61 +2,82 @@ Limitations =========== -NumPy tracing records a torch graph from a lookup of equivalent operations. -Calls that are not in that lookup, and values that leave the NumPy array type, -are not represented in the exported model. This page covers the cases that -follow from that. +NumPy function capture reuses ``torch._numpy`` implementations and checks their +captured ATen graph structure. Numerical comparison with real NumPy happens at +the exported node outputs. This expands function coverage, but +does not intercept every way that Python or a native extension can consume an +array. Calls that fail capture and values that escape the tracing hooks need +different handling. -Untraced results become baked-in constants -========================================== - -This is the failure to watch for, because nothing stops it. - -When a NumPy call is not traced, its result is an ordinary ``np.ndarray`` with -no link to the graph. If that value is later combined with traced data, LEAPP -records it the way it records any other plain array: as a **constant**. The -exported model then returns the value captured during tracing, for every input. - -Unmapped calls log a warning and then continue: - -.. code-block:: text - - No torch equivalent for numpy function . Operation will not be traced. +Capture failures are checked against outputs +============================================ -Treat that warning as an error. +A valid NumPy call can fail ATen capture because its implementation is missing, +its signature is unsupported, or its graph structure is unsupported. LEAPP logs a +recoverable error and continues eager execution. Tracing can therefore pass +through work that is not needed by the declared outputs: .. code-block:: python state = annotate.input_tensors("node", {"state": raw_state}) - - # np.interp has no torch equivalent, so this line is not recorded. - calibrated = np.interp(state, CALIB_IN, CALIB_OUT) - - result = state * 0.1 + calibrated - + spacing = np.spacing(state) # logs a capture error; NumPy still runs + result = np.sin(state) annotate.output_tensors("node", {"result": result}, export_with="jit") -That node exports without an error. Replay validation passes as well, because -validation replays the tracing inputs, which are the inputs the constant was -captured from. The graph is wrong for every other input. - -LEAPP only raises when the untraced value is a node output itself: - -.. code-block:: text - - output_tensors() for node 'node' received non-traced tensors: {'ndarray'} +The unused ``spacing`` computation is pruned. If ``spacing`` itself, or a +traced computation that uses it, is declared as an output, graph finalization +raises ``NotImplementedError`` with the original capture diagnostic. Retaining +that dependency prevents a failed operation from silently becoming a constant +in an otherwise supported output graph. + +The diagnostic identifies the NumPy operation and node. Ordinary arrays and +inactive traced arrays continue to use native NumPy behavior. Invalid NumPy +calls still raise their native exceptions. Unsupported mutations still stop +capture: a write can affect a required input even if its return value is unused. + +Support is specific to the arguments and dtypes being captured: + +* ``out=`` mutation and ``overwrite_input=True`` are rejected before modifying + caller storage. General in-place NumPy functions are not supported by this + adapter. Indexed assignment has a separate tracing path. +* Padding accepts static widths, constant mode, and a scalar fill value. + Stable or structured sorting cannot be captured. +* Direct ufunc calls are supported through the adapter; methods such as + ``np.add.reduce`` are not. Use the supported reduction function where + appropriate, such as ``np.sum``. +* Intermediate shapes, dtypes, and values are not compared against NumPy during + capture. Compatibility-layer differences can propagate into later operations; + a matching function signature alone does not establish correct output behavior. + Final validation checks the exported node results, so an intermediate + difference that disappears before the output need not prevent export. + +A successful capture does not establish support for every runtime value, +shape, device, or export backend. ``compile_graph()`` enables numerical +validation and strict failure handling by default. Keep those checks enabled +and collect varied input samples; validating only the capture sample can miss +frozen values. Exported output interfaces must still meet the application's +shape and dtype requirements; removing intermediate comparisons does not +relax that contract. See :doc:`/guides/debugging`. + +Untraced values can still become constants +========================================== -.. warning:: +The error policy applies to calls that reach the dispatch hooks. A function +that accesses the raw buffer or otherwise bypasses those hooks can still +return an ordinary array with no graph reference. If that result is mixed +back into traced data, it can be recorded as a constant. Standalone array +constructors and writes into plain ``np.ndarray`` destinations are not made +traceable by the ATen adapter. - Read the tracing log. An untraced NumPy call in the middle of a node is - reported once, as a warning, and never again. +Declare runtime-dependent values as inputs, keep their computation in traced +operations, and check replay on different inputs. Intentional plain-array +constants, such as normalization statistics, remain valid graph operands. Third-party functions ===================== -The rule is not "NumPy yes, everything else no". It is whether the function -decomposes into operations LEAPP can see. A helper written in Python on top of -traced NumPy calls traces fine, even from another package: +A helper written in Python on top of supported traced operations can be +captured even when it comes from another package: .. code-block:: python @@ -65,11 +86,10 @@ traced NumPy calls traces fine, even from another package: scale = np.sqrt(np.mean(centered * centered, axis=-1, keepdims=True)) return centered / (scale + 1e-6) -Anything that reaches the raw buffer instead is invisible to the tracer, and -its result feeds the constant-folding case above. That covers compiled -extensions such as **SciPy**, **OpenCV**, and **numba**, NumPy's own compiled routines -including ``np.fft``, ``np.linalg.inv``, ``np.interp``, and ``np.pad``, and -plain Python iteration over array elements. +Code that accesses the raw buffer, including some SciPy, OpenCV, or numba +routines, may bypass tracing. Being implemented in native code is not itself +a reason a NumPy call fails: functions such as ``np.linalg.inv`` and supported +``np.pad`` calls are intercepted and captured through the compatibility layer. When a faithful replacement can be expressed in supported operations, use :doc:`../guides/patching` to register it for the unsupported module-level @@ -86,58 +106,48 @@ graph boundary that the deploying runtime has to fill. Indexing down to a scalar breaks the chain ========================================== -NumPy returns a scalar, not an array, when indexing removes every axis. That -scalar is a plain ``np.float32``, so tracing stops there. Torch does not have -this problem, because ``tensor[0]`` is still a tensor. +NumPy returns a scalar when indexing removes every axis. LEAPP still returns +that plain scalar, losing its runtime dependency. This remains a separate +limitation from function capture: .. code-block:: python - row = x[0] # 2-D input: TracedNpArray, one axis left - corner = x[0, 0] # 2-D input: np.float32, no longer traced - - corner = x[0:1, 0:1] # TracedNpArray, shape (1, 1) - -Only the keys that reduce to a scalar are affected, so this bites hardest on -1-D arrays, where ``x[0]`` is already a scalar. Keep at least one axis in the -key. Slicing that leaves an array behind, including steps, boolean masks, -index arrays, ``None``, and ``Ellipsis``, is traced normally. - -Reversed slicing disagrees with the graph -========================================= + row = x[0] # 2-D input: TracedNpArray, one axis left + corner = x[0, 0] # plain NumPy scalar, no longer traced + result = x + corner # corner can become a capture-time constant -``x[::-1]`` is ordinary NumPy but has no torch equivalent, since torch does not -support negative strides. The eager value is reversed, the recorded graph is -not, and nothing reports it at trace time. Use an index array instead: + corner = x[0:1, 0:1] # retains an array and its graph reference -.. code-block:: python - - flipped = x[::-1] # graph is not reversed - flipped = x[np.arange(x.shape[0] - 1, -1, -1)] # traced correctly - -``leapp.compile_graph(validate=True)``, the default, does catch this one: -the exported model disagrees with the traced Python output on the same input. +Keep at least one axis in the index when its value must remain dynamic. +In contrast, supported reductions such as ``np.mean(x)`` or +``np.linalg.norm(x)`` can return traced zero-dimensional arrays. Constant +NumPy scalar operands such as ``np.float32(2.0)`` are also accepted; accepting +a scalar operand does not restore a dependency already lost by indexing. -NumPy scalar types are not valid operands -========================================= +Reverse slices require static basic indices +=========================================== -Python scalars and array constants are recorded as graph values; NumPy's own -scalar types are not. +Negative-stride reads are supported for fixed input shapes and static basic +indices: .. code-block:: python - y = x * 2.0 # ok - y = x * np.array(2.0, np.float32) # ok - y = x * np.float32(2.0) # NotImplementedError during tracing + flipped = x[::-1] # full reversal lowers to torch.flip + selected = x[5:0:-2] # normalized indices lower to torch.index_select -This also arrives indirectly, since many NumPy calls return scalar types. For -example ``np.linalg.norm(x)`` without an ``axis`` returns a ``np.float32``, so -``x / np.linalg.norm(x)`` fails on the division rather than on the norm. +Basic tuples can combine these slices with integers, ``None``, and +``Ellipsis``. Runtime slice bounds and advanced indices combined with negative +strides are rejected during capture. Reverse assignment and replay of +mutations through reversed views are not supported by this read lowering. +The eager NumPy result still has NumPy view/storage behavior; that alone does +not establish alias-mutation behavior in the exported graph. Extracting Python values ======================== ``.item()``, ``.tolist()``, ``float()``, and ``int()`` return concrete values -captured at trace time and drop out of the graph. Each logs a warning. +captured at trace time and drop out of the graph. ``.item()`` and ``.tolist()`` +log warnings; do not rely on warnings to detect every scalar conversion. Using an array in a boolean context, such as ``if array.any():``, records only the branch taken during tracing and logs an error. LEAPP's capture does not @@ -150,12 +160,5 @@ Casts are limited to the shared dtypes ``astype`` is recorded only for dtypes that exist in both libraries: ``float16``, ``float32``, ``float64``, ``int8``, ``int16``, ``int32``, ``int64``, ``uint8``, and ``bool``. Any other target dtype casts eagerly, warns, -and leaves the graph at the original dtype. - -Guidance -======== - -Keep the NumPy portion of a pipeline to the element-wise, reduction, and shape -work close to the data source, and cross into torch as early as the pipeline -allows. That is where NumPy support pays for itself, and it keeps the surface -that has to be checked small. +and leaves the graph at the original dtype. This method-level behavior is +separate from function capture and final numerical validation. diff --git a/docs/source/tensor_libraries/numpy.rst b/docs/source/tensor_libraries/numpy.rst index 4d766d6f..8798a792 100644 --- a/docs/source/tensor_libraries/numpy.rst +++ b/docs/source/tensor_libraries/numpy.rst @@ -8,9 +8,11 @@ recorded into the same graph as the surrounding PyTorch code. A pipeline can keep the NumPy preprocessing it already has instead of rewriting it in torch before it can be exported. -The exported artifact is always a torch graph. NumPy is a frontend to that -graph, not a second export backend. See :doc:`how_it_works` for the -interception model and :doc:`limitations` for the cases that follow from that. +Supported NumPy functions are captured through PyTorch's tensor-backed NumPy +compatibility layer. The exported artifact contains torch operations; NumPy +is a frontend to that graph, not a second export backend. See +:doc:`how_it_works` for the interception model and :doc:`limitations` for its +supported scope. Example: a dataset-driven policy pipeline ========================================= diff --git a/docs/source/tensor_libraries/torch.rst b/docs/source/tensor_libraries/torch.rst index 608fd5e0..10848d63 100644 --- a/docs/source/tensor_libraries/torch.rst +++ b/docs/source/tensor_libraries/torch.rst @@ -8,9 +8,9 @@ PyTorch is the path LEAPP is built around. Tensors handed to and every torch operation applied to one is recorded into the graph that gets exported. -Nothing is translated along the way. :doc:`NumPy ` tracing looks each -call up in a table of torch equivalents and :doc:`Warp ` hands whole -regions to APIC capture, but a torch call is recorded as itself. See +A torch call is recorded as itself. :doc:`NumPy ` function tracing +captures tensor-backed compatibility implementations as ATen operations, while +:doc:`Warp ` hands whole regions to APIC capture. See :doc:`torch_how_it_works` for the interception model and :doc:`torch_limitations` for what it cannot represent. diff --git a/docs/source/tensor_libraries/warp_how_it_works.rst b/docs/source/tensor_libraries/warp_how_it_works.rst index 115e52cb..6f2590b8 100644 --- a/docs/source/tensor_libraries/warp_how_it_works.rst +++ b/docs/source/tensor_libraries/warp_how_it_works.rst @@ -5,9 +5,9 @@ How it works LEAPP does not rewrite Warp kernels as torch operations. It hands the work to Warp's own APIC capture and stores the result, so LEAPP exports a subset of what APIC already supports: if APIC cannot capture and replay a piece of Warp -code, LEAPP cannot export it either. This is in contrast to -:doc:`NumPy tracing `, where each call is looked up and recorded -as an equivalent torch operation. +code, LEAPP cannot export it either. :doc:`NumPy tracing ` +captures each supported function through a tensor-backed compatibility +implementation and inserts its ATen operations into the graph. The unit LEAPP captures is a **segment**: a consecutive run of Warp calls on values belonging to one LEAPP node. A segment is neither a single launch nor diff --git a/leapp/leapp_graph/datatypes/numpy/aten_adapter.py b/leapp/leapp_graph/datatypes/numpy/aten_adapter.py new file mode 100644 index 00000000..d319b87f --- /dev/null +++ b/leapp/leapp_graph/datatypes/numpy/aten_adapter.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Capture NumPy operations through PyTorch's private compatibility layer. + +A small override registry handles compatibility gaps; other calls use PyTorch's +NumPy implementation. Each call is captured independently and structurally checked, +then inlined as ATen into the owning node. No torch._numpy objects survive capture. +Numerical validation happens on the exported node outputs. +""" + +import inspect +import operator +from functools import lru_cache + +import numpy as np +import torch +import torch._numpy as tnp +from torch.fx.experimental.proxy_tensor import make_fx +from torch.fx.node import map_arg + +from leapp.utils.logging import _get_logger + +from ..traced_data import TracedData +from .overrides import NUMPY_OVERRIDES + + +@lru_cache(maxsize=1) +def _implementations(): + mappings = {} + # Match actual function identity within known namespaces, not just __name__. + # Random/stateful functions are intentionally outside this adapter. + for numpy_module, torch_module in ( + (np, tnp), + (np.linalg, tnp.linalg), + (np.fft, tnp.fft), + ): + for name, implementation in vars(torch_module).items(): + original = getattr(numpy_module, name, None) + if ( + not name.startswith("_") + and callable(original) + and callable(implementation) + ): + mappings[original] = implementation + return mappings + + +class _NumpyCaptureError(Exception): + """A valid NumPy call whose tensor implementation could not be captured.""" + + +class _NumpyMutationError(NotImplementedError): + """An unsupported mutation cannot be treated as an unused pure branch.""" + + +def _failed_numpy_capture(message, *inputs): + """Dependency marker; live markers are rejected before graph export.""" + raise NotImplementedError(message) + + +def capture_numpy_call(source, func, args, kwargs): + """Log recoverable translation errors and retain their output dependencies.""" + label = f"NumPy {func.__module__}.{func.__name__} in node {source.context}" + implementation = NUMPY_OVERRIDES.get(func, _implementations().get(func)) + if func is np.result_type: + return func( + *TracedData.unwrap_traced_data(args), + **TracedData.unwrap_traced_data(kwargs), + ) + + try: + if implementation is None: + try: + native_bound = inspect.signature(func).bind(*args, **kwargs).arguments + except (TypeError, ValueError): + native_bound = kwargs + positional_out = isinstance(func, np.ufunc) and any( + value is not None for value in args[func.nin : func.nin + func.nout] + ) + if ( + native_bound.get("out") is not None + or native_bound.get("overwrite_input", False) + or (func is np.nan_to_num and not native_bound.get("copy", True)) + or positional_out + ): + raise NotImplementedError(f"{label}: untraced mutation is unsupported") + raise _NumpyCaptureError(f"{label}: no torch._numpy lowering") + try: + bound = inspect.signature(implementation).bind(*args, **kwargs) + except TypeError as exc: + raise _NumpyCaptureError(f"{label}: unsupported signature: {exc}") from exc + if ( + bound.arguments.get("out") is not None + or bound.arguments.get("overwrite_input", False) + or (func is np.nan_to_num and not bound.arguments.get("copy", True)) + ): + raise NotImplementedError( + f"{label}: mutation is unsupported by the ATen adapter" + ) + return _capture_numpy_call(source, func, args, kwargs, implementation, label) + except _NumpyCaptureError as exc: + # A failed lowering cannot establish that the native call is functional. + # Probe only this error path on copies, protecting caller storage from + # unsupported mutators such as copyto/put (which return None). Native + # NumPy errors propagate before a recoverable capture error is logged. + native_args, native_kwargs = TracedData._map_structure( + TracedData.unwrap_traced_data((args, kwargs)), + lambda x: x.copy() if isinstance(x, np.ndarray) else x, + ) + if func(*native_args, **native_kwargs) is None: + raise NotImplementedError(f"{label}: untraced mutation is unsupported") + message = str(exc) + _get_logger().error( + f"{message}. Continuing eager execution; export requires this operation " + "to be independent of the declared outputs." + ) + # A traced marker preserves dependency information until output pruning. + # Returning a plain ndarray here could freeze a live failed computation. + expected = func( + *TracedData.unwrap_traced_data(args), + **TracedData.unwrap_traced_data(kwargs), + ) + inputs = [] + TracedData._map_structure( + (args, kwargs), + lambda x: inputs.append(x.proxy) if isinstance(x, TracedData) else x, + ) + + def wrap(value): + proxy = source._context.tracer.create_proxy( + "call_function", _failed_numpy_capture, (message, *inputs), {} + ) + proxy.node.meta["leapp_capture_error"] = message + return source._new(np.asarray(value), proxy) + + return TracedData._map_structure(expected, wrap) + + +def _capture_numpy_call(source, func, args, kwargs, implementation, label): + inputs = [] + TracedData._map_structure( + (args, kwargs), lambda x: inputs.append(x) if isinstance(x, TracedData) else x + ) + + def lower(*tensors): + values = iter(tensors) + + def convert(x): + if isinstance(x, TracedData): + return tnp.asarray(next(values)) + if isinstance(x, np.dtype) or ( + isinstance(x, type) and issubclass(x, np.generic) + ): + return tnp.dtype(np.dtype(x).name) + if isinstance(x, (np.ndarray, np.generic)): + return tnp.asarray(torch.from_numpy(np.array(x, copy=True))) + return x + + converted_args, converted_kwargs = TracedData._map_structure( + (args, kwargs), convert + ) + result = implementation(*converted_args, **converted_kwargs) + return TracedData._map_structure( + result, lambda x: x.tensor if isinstance(x, tnp.ndarray) else x + ) + + try: + examples = tuple(torch.from_numpy(np.array(x.data, copy=True)) for x in inputs) + captured = make_fx(lower)(*examples) + # Composite ATen kernels (e.g. isclose) introduce temporary mutations + # during decomposition. Functionalize the decomposed graph itself. + captured = make_fx(torch.func.functionalize(captured))(*examples) + for node in captured.graph.nodes: + if ( + isinstance(node.target, torch._ops.OpOverload) + and node.target._schema.is_mutable + ): + raise _NumpyMutationError(f"{label}: mutable operation: {node.target}") + output_nodes = [] + output_node = next(node for node in captured.graph.nodes if node.op == "output") + TracedData._map_structure(output_node.args[0], output_nodes.append) + if not output_nodes or any( + not isinstance(node, torch.fx.Node) + or not isinstance(node.meta.get("val"), torch.Tensor) + for node in output_nodes + ): + raise ValueError("non-tensor output would lose runtime dependence") + for node in captured.graph.nodes: + if node.op not in ("placeholder", "get_attr", "call_function", "output"): + raise ValueError(f"unsupported graph node: {node.op}") + if node.op == "call_function" and node.target is not operator.getitem: + if not isinstance(node.target, torch._ops.OpOverload): + raise ValueError(f"non-ATen operation: {node.target}") + + except _NumpyMutationError: + raise + except Exception as exc: + raise _NumpyCaptureError(f"{label}: ATen capture failed: {exc}") from exc + + # Once the lowering is known to be functional, execute NumPy on the real + # operands so eager view/storage behavior is retained (e.g. broadcast_to). + expected = func( + *TracedData.unwrap_traced_data(args), **TracedData.unwrap_traced_data(kwargs) + ) + expected_leaves = [] + TracedData._map_structure(expected, expected_leaves.append) + if len(output_nodes) != len(expected_leaves): + raise _NumpyCaptureError(f"{label}: output structure differs from NumPy") + tracer = source._context.tracer + bindings = {} + proxies = iter(x.proxy for x in inputs) + for node in captured.graph.nodes: + if node.op == "placeholder": + bindings[node] = next(proxies) + elif node.op == "get_attr": + value = captured + for part in node.target.split("."): + value = getattr(value, part) + bindings[node] = ( + value.detach().clone() if isinstance(value, torch.Tensor) else value + ) + elif node.op == "call_function": + # TorchScript resolves overload packets, but not .default/.Tensor + # attributes. ATen's dispatcher selects the captured signature. + target = ( + node.target.overloadpacket + if isinstance(node.target, torch._ops.OpOverload) + else node.target + ) + if node.target is torch.ops.aten.lift_fresh_copy.default: + # Constants need a copy, but legacy ONNX has no lift lowering. + target = torch.clone + bindings[node] = tracer.create_proxy( + node.op, + target, + map_arg(node.args, bindings.__getitem__), + map_arg(node.kwargs, bindings.__getitem__), + ) + elif node.op == "output": + output = map_arg(node.args[0], bindings.__getitem__) + + output_leaves = [] + TracedData._map_structure(output, output_leaves.append) + output_proxies = iter(output_leaves) + return TracedData._map_structure( + expected, lambda x: source._new(np.asarray(x), next(output_proxies)) + ) diff --git a/leapp/leapp_graph/datatypes/numpy/overrides.py b/leapp/leapp_graph/datatypes/numpy/overrides.py new file mode 100644 index 00000000..d2506010 --- /dev/null +++ b/leapp/leapp_graph/datatypes/numpy/overrides.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""NumPy semantic/export exceptions to the automatic torch._numpy lowering. + +Implementations accept torch._numpy arrays and NumPy argument conventions and +return tensors or torch._numpy arrays. The shared adapter captures and validates +both overrides and upstream implementations through the same ATen path. +""" + +import numpy as np +import torch +import torch._numpy as tnp +import torch.nn.functional as F + + +def _cbrt(x, *, dtype=None): + # torch._numpy uses pow(x, 1/3), which produces NaN for negative inputs. + x = tnp.asarray(x, dtype=dtype).tensor + if not x.is_floating_point(): + x = x.to(torch.float64) + return torch.copysign(torch.abs(x).pow(1 / 3), x) + + +def _pad(array, pad_width, mode="constant", *, constant_values=0): + # Keep static widths out of tensor normalization: F.pad needs Python ints. + if mode != "constant" or isinstance(pad_width, tnp.ndarray): + raise NotImplementedError( + "pad requires constant mode and static padding widths" + ) + widths = np.broadcast_to(np.asarray(pad_width), (array.ndim, 2)) + if widths.dtype.kind not in "iu" or np.any(widths < 0): + raise NotImplementedError("pad widths must be nonnegative integers") + if not isinstance(constant_values, (int, float, bool)): + raise NotImplementedError("pad requires a scalar constant_values") + return F.pad( + array.tensor, tuple(int(x) for x in widths[::-1].flat), value=constant_values + ) + + +def _nanmean(a, axis=None, dtype=None, out=None, keepdims=False): + # Missing from torch._numpy. NumPy promotes integer inputs to float64. + a = tnp.asarray(a, dtype=dtype).tensor + if not a.is_floating_point() and not a.is_complex(): + a = a.to(torch.float64) + if axis == (): + return a + return torch.nanmean(a, dim=axis, keepdim=keepdims) + + +def _sorted(a, axis, kind, order, stable): + # Preserve LEAPP's static-K ONNX lowering. Stable/structured sorting needs + # a separate implementation; do not silently ignore those NumPy options. + if kind not in (None, "quicksort") or order is not None or stable: + raise NotImplementedError("only default, unstructured sorting is supported") + a = tnp.asarray(a).tensor + if axis is None: + a, axis = a.flatten(), -1 + return torch.topk(a, a.shape[axis], dim=axis, largest=False, sorted=True) + + +def _sort(a, axis=-1, kind=None, order=None, *, stable=None): + return _sorted(a, axis, kind, order, stable).values + + +def _argsort(a, axis=-1, kind=None, order=None, *, stable=None): + return _sorted(a, axis, kind, order, stable).indices + + +NUMPY_OVERRIDES = { + np.cbrt: _cbrt, + np.pad: _pad, + np.nanmean: _nanmean, + np.sort: _sort, + np.argsort: _argsort, +} diff --git a/leapp/leapp_graph/datatypes/numpy/traced_np_array.py b/leapp/leapp_graph/datatypes/numpy/traced_np_array.py index d09d38ea..f6a8a117 100644 --- a/leapp/leapp_graph/datatypes/numpy/traced_np_array.py +++ b/leapp/leapp_graph/datatypes/numpy/traced_np_array.py @@ -22,6 +22,7 @@ from leapp.utils.dtype import DtypeCodec, dtype_to_name, register_dtype_codec from ..proxy_view import bind_new_view, may_adopt_view, share_view from ..traced_data import TracedData +from .aten_adapter import capture_numpy_call # numpy dtype object -> common name string. Lives with the numpy node library @@ -61,210 +62,6 @@ def _torch_dtype_for(dtype): return None -# ============================================================================= -# NumPy to Torch Mappings -# ============================================================================= - -# Ufuncs are element-wise operations that numpy broadcasts automatically. -# These are intercepted via __array_ufunc__ protocol. -NUMPY_UFUNC_TO_TORCH = { - # Arithmetic operations - np.add: torch.add, - np.subtract: torch.sub, - np.multiply: torch.mul, - np.divide: torch.div, - np.true_divide: torch.div, - np.floor_divide: torch.floor_divide, - np.power: torch.pow, - np.negative: torch.neg, - np.positive: lambda x: x, # No-op - np.mod: torch.remainder, - np.remainder: torch.remainder, - np.fmod: torch.fmod, - # Absolute and sign - np.absolute: torch.abs, - np.abs: torch.abs, - np.sign: torch.sign, - # Powers and roots - np.sqrt: torch.sqrt, - np.square: torch.square, - np.exp: torch.exp, - np.exp2: lambda x: torch.pow(2, x), - np.expm1: torch.expm1, - # Logarithms - np.log: torch.log, - np.log2: torch.log2, - np.log10: torch.log10, - np.log1p: torch.log1p, - # Trigonometric functions - np.sin: torch.sin, - np.cos: torch.cos, - np.tan: torch.tan, - np.arcsin: torch.asin, - np.arccos: torch.acos, - np.arctan: torch.atan, - np.arctan2: torch.atan2, - np.hypot: torch.hypot, - # Hyperbolic functions - np.sinh: torch.sinh, - np.cosh: torch.cosh, - np.tanh: torch.tanh, - np.arcsinh: torch.asinh, - np.arccosh: torch.acosh, - np.arctanh: torch.atanh, - # Rounding - np.floor: torch.floor, - np.ceil: torch.ceil, - np.trunc: torch.trunc, - np.rint: torch.round, - # Comparison (element-wise, return boolean tensor) - np.greater: torch.gt, - np.greater_equal: torch.ge, - np.less: torch.lt, - np.less_equal: torch.le, - np.equal: torch.eq, - np.not_equal: torch.ne, - np.maximum: torch.maximum, - np.minimum: torch.minimum, - # Matrix multiplication. ``np.matmul`` is a ufunc, so it never reaches - # ``__array_function__``; ``np.dot`` is not, and is mapped there instead. - np.matmul: torch.matmul, - # Logical operations - np.logical_and: torch.logical_and, - np.logical_or: torch.logical_or, - np.logical_xor: torch.logical_xor, - np.logical_not: torch.logical_not, - # Bitwise operations - np.bitwise_and: torch.bitwise_and, - np.bitwise_or: torch.bitwise_or, - np.bitwise_xor: torch.bitwise_xor, - np.invert: torch.bitwise_not, - np.left_shift: torch.bitwise_left_shift, - np.right_shift: torch.bitwise_right_shift, - # Special values - np.isnan: torch.isnan, - np.isinf: torch.isinf, - np.isfinite: torch.isfinite, - # Clipping - np.clip: torch.clamp, -} - -# Higher-level array functions intercepted via __array_function__ protocol. -NUMPY_FUNC_TO_TORCH = { - # Reduction operations - np.sum: torch.sum, - np.prod: torch.prod, - np.mean: torch.mean, - np.std: torch.std, - np.var: torch.var, - # Use amax/amin instead of max/min because torch.max/min with dim returns (values, indices) - # while numpy just returns values. amax/amin always return just values. - np.min: torch.amin, - np.max: torch.amax, - np.argmin: torch.argmin, - np.argmax: torch.argmax, - np.cumsum: torch.cumsum, - np.cumprod: torch.cumprod, - np.all: torch.all, - np.any: torch.any, - # Array manipulation - np.concatenate: torch.cat, - np.stack: torch.stack, - np.vstack: torch.vstack, - np.hstack: torch.hstack, - # Integer sections mean "N parts" in NumPy and "chunk size" in torch.split. - # tensor_split matches both np.split and np.array_split. - np.split: torch.tensor_split, - np.array_split: torch.tensor_split, - np.squeeze: torch.squeeze, - np.expand_dims: torch.unsqueeze, - np.reshape: torch.reshape, - np.transpose: torch.permute, - np.swapaxes: torch.swapaxes, - np.moveaxis: torch.moveaxis, - np.flip: torch.flip, - np.roll: torch.roll, - np.rot90: torch.rot90, - # Sorting and searching - np.sort: torch.sort, - np.argsort: torch.argsort, - np.where: torch.where, - np.nonzero: torch.nonzero, - # Element-wise (also available as ufuncs) - np.clip: torch.clamp, - np.abs: torch.abs, - np.absolute: torch.abs, - np.sqrt: torch.sqrt, - np.square: torch.square, - np.exp: torch.exp, - np.log: torch.log, - np.sin: torch.sin, - np.cos: torch.cos, - np.tan: torch.tan, - np.tanh: torch.tanh, - np.round: torch.round, - np.around: torch.round, - # Linear algebra - np.dot: torch.matmul, # Note: torch.dot is only for 1D vectors - np.tensordot: torch.tensordot, - np.einsum: torch.einsum, - np.trace: torch.trace, - np.diagonal: torch.diagonal, - np.tril: torch.tril, - np.triu: torch.triu, - # Creation functions (when operating on TracedTensor) - np.zeros_like: torch.zeros_like, - np.ones_like: torch.ones_like, - np.full_like: torch.full_like, - np.empty_like: torch.empty_like, - # Standalone creation (less commonly needed with TracedTensor) - np.eye: torch.eye, - np.zeros: torch.zeros, - np.ones: torch.ones, - np.full: torch.full, - np.arange: torch.arange, - np.linspace: torch.linspace, -} - -# Functions that need axis -> dim conversion -AXIS_TO_DIM_FUNCTIONS = { - torch.sum, - torch.prod, - torch.mean, - torch.std, - torch.var, - torch.min, - torch.max, - torch.amin, - torch.amax, - torch.argmin, - torch.argmax, - torch.cumsum, - torch.cumprod, - torch.all, - torch.any, - torch.cat, - torch.stack, - torch.tensor_split, - torch.squeeze, - torch.unsqueeze, - torch.sort, - torch.argsort, -} - -# torch.flip / torch.roll take dims, not dim. A bare int must be a 1-tuple for flip. -AXIS_TO_DIMS_FUNCTIONS = { - torch.flip, - torch.roll, -} - -# These reject Python scalars, unlike torch.clamp and the arithmetic ops. -SCALAR_TENSOR_FUNCTIONS = { - torch.maximum, - torch.minimum, -} - - # Combined metaclass to resolve conflict between ABCMeta (from TracedData) # and np.ndarray's metaclass class _TracedNpArrayMeta(ABCMeta, type(np.ndarray)): @@ -342,6 +139,9 @@ def __array_finalize__(self, obj): def tensor(self) -> np.ndarray: """Get the underlying numpy array (for TracedData compatibility).""" array = np.ascontiguousarray(self.view(np.ndarray)) + # NumPy considers singleton axes contiguous even with negative strides. + if any(stride < 0 for stride in array.strides): + array = array.copy() return torch.from_numpy(array) @property @@ -420,264 +220,35 @@ def convert(item): return TracedData._map_structure(obj, convert) - def _convert_numpy_call_to_torch( - self, proxy_args, proxy_kwargs, torch_func, original_kwargs=None, args=None - ): - """Convert numpy args/kwargs to the matching torch call. - - Handles: - - 'axis' → 'dim' for functions in AXIS_TO_DIM_FUNCTIONS - - 'axis' → 'dims' for functions in AXIS_TO_DIMS_FUNCTIONS - - omitted / integer axis for torch.flip - - scalar operands for torch.maximum / torch.minimum - - 'keepdims' → 'keepdim' (numpy uses 's', torch doesn't) - - 'keepdims' without axis: torch requires dim when keepdim is specified - - 'axes' for transpose: None or missing → reversed dims tuple - """ - proxy_args = tuple(proxy_args) - proxy_kwargs = proxy_kwargs.copy() - if original_kwargs is None: - original_kwargs = {} - - # Get input array for dimension info - input_arr = args[0] if args and len(args) > 0 else None - - # Convert axis → dim - has_dim = False - if torch_func in AXIS_TO_DIM_FUNCTIONS and "axis" in proxy_kwargs: - proxy_kwargs["dim"] = proxy_kwargs.pop("axis") - has_dim = True - - if torch_func in AXIS_TO_DIMS_FUNCTIONS: - if "axis" in proxy_kwargs: - dims = proxy_kwargs.pop("axis") - if ( - dims is None - and torch_func is torch.flip - and hasattr(input_arr, "ndim") - ): - dims = tuple(range(input_arr.ndim)) - elif torch_func is torch.flip and isinstance(dims, int): - dims = (dims,) - if dims is not None: - proxy_kwargs["dims"] = dims - elif torch_func is torch.flip: - if args is not None and len(args) > 1: - dims = proxy_args[1] - if isinstance(dims, int): - proxy_args = (proxy_args[0], (dims,)) + proxy_args[2:] - elif hasattr(input_arr, "ndim"): - proxy_kwargs["dims"] = tuple(range(input_arr.ndim)) - - if torch_func in SCALAR_TENSOR_FUNCTIONS: - proxy_args = tuple( - torch.tensor(arg) if isinstance(arg, (int, float, bool)) else arg - for arg in proxy_args - ) - - # Convert keepdims → keepdim (numpy uses 's', torch doesn't) - # Important: torch.sum/mean only support keepdim when dim is also specified - if "keepdims" in proxy_kwargs: - keepdims_value = proxy_kwargs.pop("keepdims") - if has_dim: - # dim is specified, we can use keepdim - proxy_kwargs["keepdim"] = keepdims_value - elif ( - keepdims_value and input_arr is not None and hasattr(input_arr, "ndim") - ): - # keepdims=True but no axis specified - # torch requires dim when using keepdim, so specify all dims - proxy_kwargs["dim"] = tuple(range(input_arr.ndim)) - proxy_kwargs["keepdim"] = True - # If keepdims=False without axis, we can just omit both (default behavior) - - # np.var/std default to ddof=0 (population), torch defaults to correction=1 (sample) - if torch_func in (torch.var, torch.std): - if "ddof" in proxy_kwargs: - proxy_kwargs["correction"] = proxy_kwargs.pop("ddof") - elif "correction" not in proxy_kwargs: - proxy_kwargs["correction"] = 0 - - if torch_func is torch.tensordot and "axes" in proxy_kwargs: - proxy_kwargs["dims"] = proxy_kwargs.pop("axes") - - # Handle np.transpose with axes=None or missing - # torch.permute requires explicit dims, but numpy reverses all dims when axes is omitted - if torch_func == torch.permute: - # Check if axes was passed as positional arg (2nd argument) - axes_in_args = len(args) > 1 if args else False - - if "axes" in proxy_kwargs: - axes = proxy_kwargs.pop("axes") - if ( - axes is None - and input_arr is not None - and hasattr(input_arr, "ndim") - ): - # axes=None means reverse all dimensions - proxy_kwargs["dims"] = tuple(range(input_arr.ndim - 1, -1, -1)) - elif axes is not None: - proxy_kwargs["dims"] = axes - elif not axes_in_args and "axes" not in original_kwargs: - # axes not provided at all (neither positional nor kwarg) - numpy reverses all dims - if input_arr is not None and hasattr(input_arr, "ndim"): - proxy_kwargs["dims"] = tuple(range(input_arr.ndim - 1, -1, -1)) - # If axes_in_args is True, it will be passed as positional arg, don't add to kwargs - - return proxy_args, proxy_kwargs - # ========================================================================= - # NumPy Ufunc Interception (__array_ufunc__) + # NumPy dispatch # ========================================================================= def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): - """Execute NumPy ufuncs and record them only while the source is active.""" - traced_array = TracedNpArray.find_traced_array([inputs, kwargs]) - unwrapped_inputs = tuple(TracedData.unwrap_traced_data(inp) for inp in inputs) - unwrapped_kwargs = { - key: TracedData.unwrap_traced_data(value) for key, value in kwargs.items() - } - - if method != "__call__": - # Reduction/accumulation protocols are not export mappings yet, but - # they must remain usable after the producing node has finished. - if traced_array is None or traced_array.validate_status(inputs, kwargs): - return NotImplemented - return getattr(ufunc, method)(*unwrapped_inputs, **unwrapped_kwargs) - - result_array = ufunc(*unwrapped_inputs, **unwrapped_kwargs) - if traced_array is None: - return result_array - - # An inactive traced array is a finished boundary value. Derived results - # are new data that this node never published, so they stay native. - if not traced_array.validate_status(inputs, kwargs): - return result_array - - torch_func = NUMPY_UFUNC_TO_TORCH.get(ufunc) - if torch_func is None: - _get_logger().warning( - f"No torch equivalent for numpy ufunc {ufunc.__name__}. " - f"Operation will not be traced." - ) - return result_array - - proxy_inputs, proxy_kwargs = self._convert_numpy_call_to_torch( - tuple(self._extract_proxy(inp) for inp in inputs), - {key: self._extract_proxy(value) for key, value in kwargs.items()}, - torch_func, - original_kwargs=kwargs, - args=inputs, - ) - proxy_out = traced_array._context.tracer.create_proxy( - "call_function", torch_func, proxy_inputs, proxy_kwargs - ) - - if isinstance(result_array, np.ndarray): - return traced_array._new(result_array, proxy_out) - return result_array - - # ========================================================================= - # NumPy Function Interception (__array_function__) - # ========================================================================= - - # List of numpy functions we implement - HANDLED_FUNCTIONS = {} - - def _patch_sort(self, torch_func, traced_array, proxy_args, args, kwargs): - """Patch sort/argsort to use torch.topk for ONNX compatibility. - - torch.sort/argsort export to ONNX TopK with a dynamically computed K - (via Shape+Gather) that ONNX Runtime's shape inference cannot verify. - Using torch.topk with a concrete int k bakes K as a constant initializer - with the correct shape [1], which ONNX Runtime accepts. - - Returns the proxy output if patched, or None if not applicable. - """ - if torch_func not in (torch.sort, torch.argsort): - return None - - input_array = args[0] - dim = kwargs.get("axis", -1) - if dim is None: - dim = -1 - k = input_array.shape[dim] - - topk_kwargs = {"dim": dim, "largest": False, "sorted": True} - proxy_out = traced_array._context.tracer.create_proxy( - "call_function", torch.topk, (proxy_args[0], k), topk_kwargs - ) - # topk returns (values, indices): index 0 for sort, 1 for argsort - item_index = 0 if torch_func == torch.sort else 1 - return traced_array._context.tracer.create_proxy( - "call_function", operator.getitem, (proxy_out, item_index), {} + """Capture active calls; retain native NumPy behavior in transit.""" + source = self.find_traced_array((inputs, kwargs)) + active = source is not None and source.validate_status(inputs, kwargs) + if method != "__call__" and active: + return NotImplemented + if active: + return capture_numpy_call(source, ufunc, inputs, kwargs) + return getattr(ufunc, method)( + *TracedData.unwrap_traced_data(inputs), + **TracedData.unwrap_traced_data(kwargs), ) def __array_function__(self, func, types, args, kwargs): - """Execute NumPy functions and record them only while the source is active.""" - traced_array = TracedNpArray.find_traced_array([args, kwargs]) - unwrapped_args = tuple(TracedData.unwrap_traced_data(arg) for arg in args) - unwrapped_kwargs = { - key: TracedData.unwrap_traced_data(value) for key, value in kwargs.items() - } - result_array = func(*unwrapped_args, **unwrapped_kwargs) - - if traced_array is None: - return result_array - - # Inactive carriers use full NumPy behavior, including functions that - # have no Torch export mapping. - if not traced_array.validate_status(args, kwargs): - if self._is_equivalent_copy(func, traced_array, result_array, args): - return traced_array.preserve_port(result_array) - return result_array - - torch_func = NUMPY_FUNC_TO_TORCH.get(func) - if torch_func is None: - _get_logger().warning( - f"No torch equivalent for numpy function {func.__name__}. " - f"Operation will not be traced." - ) - return result_array - - proxy_args, proxy_kwargs = self._convert_numpy_call_to_torch( - tuple(self._extract_proxy(arg) for arg in args), - {key: self._extract_proxy(value) for key, value in kwargs.items()}, - torch_func, - original_kwargs=kwargs, - args=args, + """Use the compatibility layer plus overrides for active calls.""" + source = self.find_traced_array((args, kwargs)) + if source is not None and source.validate_status(args, kwargs): + return capture_numpy_call(source, func, args, kwargs) + result = func( + *TracedData.unwrap_traced_data(args), + **TracedData.unwrap_traced_data(kwargs), ) - - proxy_out = self._patch_sort(torch_func, traced_array, proxy_args, args, kwargs) - if proxy_out is None: - proxy_out = traced_array._context.tracer.create_proxy( - "call_function", torch_func, proxy_args, proxy_kwargs - ) - - if isinstance(result_array, (tuple, list)): - result = [] - for index, array in enumerate(result_array): - if isinstance(array, np.ndarray): - if torch_func is torch.nonzero: - # torch.nonzero returns one (N, ndim) tensor while numpy - # returns one index array per dimension, so the numpy - # element is a column of the torch result, not a row. - item_proxy = traced_array._context.tracer.create_proxy( - "call_function", torch.select, (proxy_out, 1, index), {} - ) - else: - item_proxy = traced_array._context.tracer.create_proxy( - "call_function", operator.getitem, (proxy_out, index), {} - ) - result.append(traced_array._new(array, item_proxy)) - else: - result.append(array) - return type(result_array)(result) - if isinstance(result_array, np.ndarray): - return traced_array._new(result_array, proxy_out) - if isinstance(result_array, np.generic): - return traced_array._new(np.asarray(result_array), proxy_out) - return result_array + if source is not None and self._is_equivalent_copy(func, source, result, args): + return source.preserve_port(result) + return result # ========================================================================= # Arithmetic Operators diff --git a/leapp/leapp_graph/datatypes/traced_data.py b/leapp/leapp_graph/datatypes/traced_data.py index f2fc7d3d..05d9c4d9 100644 --- a/leapp/leapp_graph/datatypes/traced_data.py +++ b/leapp/leapp_graph/datatypes/traced_data.py @@ -7,6 +7,7 @@ from abc import ABC, abstractmethod import operator +from numbers import Integral from typing import Any, Set, Optional from torch.fx.proxy import Proxy @@ -578,8 +579,93 @@ def _lower_index_key(self, proxy, key): {}, ) + def _lower_negative_slices(self, key): + """Lower static basic reverse reads before applying the remaining key. + + Preselect on the original input axes: integer indices and new axes are + applied afterwards, so they cannot shift the dimensions being reversed. + General bounds are normalized against the fixed traced input shape. + """ + elements = key if isinstance(key, tuple) else (key,) + if not any( + isinstance(item, slice) + and isinstance(item.step, Integral) + and item.step < 0 + for item in elements + ): + return None + + def basic(item): + if isinstance(item, slice): + return all( + bound is None or isinstance(bound, Integral) + for bound in (item.start, item.stop, item.step) + ) + return ( + item is None + or item is Ellipsis + or isinstance(item, Integral) + and not isinstance(item, bool) + ) + + if not all(basic(item) for item in elements): + _get_logger().fatal( + "Negative-stride slices require static basic indices; " + "runtime slice bounds and advanced indices are not supported.", + error_type=NotImplementedError, + ) + + shape = tuple(self.tensor.shape) + consumed = sum(item is not None and item is not Ellipsis for item in elements) + residual = [] + for item in elements: + if item is Ellipsis: + residual.extend([slice(None)] * (len(shape) - consumed)) + elif isinstance(item, slice): + residual.append( + slice( + *( + None if bound is None else operator.index(bound) + for bound in (item.start, item.stop, item.step) + ) + ) + ) + else: + residual.append(None if item is None else operator.index(item)) + + proxy = self.proxy + axis = 0 + for position, item in enumerate(residual): + if item is None: + continue + if isinstance(item, slice) and item.step is not None and item.step < 0: + start, stop, step = item.indices(shape[axis]) + if (start, stop, step) == (shape[axis] - 1, -1, -1): + proxy = self._context.tracer.create_proxy( + "call_function", torch.flip, (proxy, (axis,)), {} + ) + else: + indices = self._register_setitem_tensor( + torch.tensor(list(range(start, stop, step)), dtype=torch.int64), + "_slice_indices", + ) + proxy = self._context.tracer.create_proxy( + "call_function", torch.index_select, (proxy, axis, indices), {} + ) + residual[position] = slice(None) + axis += 1 + + if all(isinstance(item, slice) and item == slice(None) for item in residual): + return proxy + return self._context.tracer.create_proxy( + "call_function", operator.getitem, (proxy, tuple(residual)), {} + ) + def _create_getitem_proxy(self, key): """Create a functional getitem proxy for a supported index key.""" + negative_slices = self._lower_negative_slices(key) + if negative_slices is not None: + return negative_slices if not self._is_supported_index_key(key): return None return self._lower_index_key(self.proxy, key) diff --git a/leapp/leapp_graph/traced_node.py b/leapp/leapp_graph/traced_node.py index 03300e95..3f66e437 100644 --- a/leapp/leapp_graph/traced_node.py +++ b/leapp/leapp_graph/traced_node.py @@ -808,6 +808,16 @@ def mark_used(node): for output_node in output_nodes: mark_used(output_node) + # Capture can fail in work unrelated to the declared outputs. Only a + # failure that survives the output dependency traversal blocks export. + for node in self.graph.nodes: + if node in used_nodes and (error := node.meta.get("leapp_capture_error")): + _get_logger().fatal( + f"Cannot export node '{self.name}': a declared output depends on " + f"a failed capture. {error}", + error_type=NotImplementedError, + ) + # Remove unused nodes (both placeholders and call_function nodes) nodes_to_remove = [] for node in self.graph.nodes: diff --git a/tests/unit_tests/test_numpy_compatibility.py b/tests/unit_tests/test_numpy_compatibility.py index cdd58494..471e0739 100644 --- a/tests/unit_tests/test_numpy_compatibility.py +++ b/tests/unit_tests/test_numpy_compatibility.py @@ -4,7 +4,12 @@ # import leapp +import pytest from leapp.leapp import _MANAGER as annotate +from leapp.leapp_graph.datatypes.numpy.patching import NumpyPatchBackend +from leapp.leapp_graph.datatypes.numpy import aten_adapter +from leapp.leapp_graph.datatypes.torch.patching import TorchPatchBackend +from leapp.leapp_graph.traced_node import TracedTensorNode import numpy as np import torch import unittest @@ -231,3 +236,74 @@ def test_np_asarray_returns_array_when_patching_disabled(self): annotate.output_tensors("test_node", {"tensor": traced}, export_with="jit") leapp.stop() + + +def test_conversion_patches_during_capture(): + numpy_patch, torch_patch = NumpyPatchBackend(), TorchPatchBackend() + try: + numpy_patch.install() + torch_patch.install() + fn = lambda x: np.linalg.norm(x, axis=-1) + x = np.arange(6, dtype=np.float32).reshape(2, 3) + ctx = TracedTensorNode(name="numpy_patches", node_index=0) + traced = ctx.create_input(x.copy(), name="x") + ctx.compile_trace({"y": fn(traced)}) + model = ctx.m + finally: + torch_patch.uninstall() + numpy_patch.uninstall() + torch.testing.assert_close( + model(torch.from_numpy(x + 2)), torch.from_numpy(fn(x + 2)) + ) + + +def test_public_lifecycle_export_and_replay(tmp_path): + x = np.arange(1, 7, dtype=np.float32).reshape(2, 3) + leapp.start(name="numpy_aten", save_path=str(tmp_path)) + try: + traced = leapp.annotate.input_tensors("norm", {"x": x}) + np.spacing(traced) # A failed, unused branch must not prevent export. + result = np.linalg.norm(traced, axis=-1) + leapp.annotate.output_tensors("norm", {"y": result}, export_with="jit") + finally: + leapp.stop() + leapp.compile_graph(visualize=False) + model = torch.jit.load(str(next(tmp_path.rglob("norm.pt")))) + torch.testing.assert_close( + model(torch.from_numpy(x + 2)), torch.from_numpy(np.linalg.norm(x + 2, axis=-1)) + ) + + +@pytest.mark.parametrize("case", ["dtype_only", "bad_values", "cancelled"]) +def test_numpy_numerical_differences_checked_at_final_validation( + case, monkeypatch, tmp_path +): + original = aten_adapter._implementations()[np.sin] + + def different_sin(x): + return original(x).astype("float64") + (0 if case == "dtype_only" else 0.5) + + monkeypatch.setitem(aten_adapter.NUMPY_OVERRIDES, np.sin, different_sin) + sample = np.array([1.0, 2.0, 3.0], dtype=np.float32) + leapp.start(name="numpy_final_validation", save_path=str(tmp_path), max_cached_io=2) + try: + for raw in (sample, sample + 1): + x = leapp.annotate.input_tensors("sine", {"x": raw}) + y = np.sin(x) + if case == "cancelled": + y = y - y + x + y = y.astype(np.float32) + leapp.annotate.output_tensors("sine", {"y": y}, export_with="jit") + finally: + leapp.stop() + if case == "bad_values": + with pytest.raises(Exception, match="Model validation failed"): + leapp.compile_graph(visualize=False) + else: + leapp.compile_graph(visualize=False) + model = torch.jit.load(str(next(tmp_path.rglob("sine.pt")))) + changed = sample + 2 + expected = changed if case == "cancelled" else np.sin(changed) + torch.testing.assert_close( + model(torch.from_numpy(changed)), torch.from_numpy(expected) + ) diff --git a/tests/unit_tests/test_numpy_slicing.py b/tests/unit_tests/test_numpy_slicing.py new file mode 100644 index 00000000..f60e7750 --- /dev/null +++ b/tests/unit_tests/test_numpy_slicing.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reverse reads retain NumPy slice semantics in exported tensor graphs.""" + +import numpy as np +import onnxruntime as ort +import pytest +import torch + +from leapp.leapp_graph.datatypes import TracedNpArray +from leapp.leapp_graph.traced_node import TracedTensorNode + + +CASES = [ + pytest.param(slice(None, None, -1), id="full_reverse"), + pytest.param(slice(None, None, -2), id="negative_step"), + pytest.param(slice(3, 0, -2), id="bounded_step"), + pytest.param(slice(-2, None, -1), id="negative_start"), + pytest.param(slice(None, -1, -1), id="explicit_minus_one_stop_is_empty"), + pytest.param(slice(100, -100, -3), id="clipped_bounds"), + pytest.param(slice(-100, 100, -1), id="empty_clipped_bounds"), + pytest.param(slice(np.int64(3), None, np.int64(-2)), id="numpy_integer_bounds"), + pytest.param( + (slice(None, None, -1), slice(1, 5, 2), slice(None, None, -2)), + id="multiple_axes", + ), + pytest.param( + (1, slice(None, None, -2), slice(None, None, -1)), id="integer_before_slices" + ), + pytest.param((slice(None, None, -1), 2, slice(None)), id="integer_after_slice"), + pytest.param((Ellipsis, slice(None, None, -1)), id="ellipsis"), + pytest.param((None, Ellipsis, slice(5, 0, -2), None), id="new_axes"), + pytest.param((1, None, slice(None, None, -1), Ellipsis, None), id="mixed_basic"), +] + + +def capture(key, x): + ctx = TracedTensorNode(name="reverse_slice", node_index=0) + traced = ctx.create_input(x.copy(), name="x") + result = traced[key] + assert isinstance(result, TracedNpArray) + ctx.compile_trace({"y": result}) + return ctx.m + + +@pytest.mark.parametrize("key", CASES) +@pytest.mark.parametrize("backend", ["fx", "jit", "pt2"]) +def test_reverse_slice_changed_input_replay(key, backend): + x = np.arange(120, dtype=np.float32).reshape(4, 5, 6) + model = capture(key, x) + if backend == "jit": + model = torch.jit.script(model) + elif backend == "pt2": + model = torch.export.export( + model, (torch.from_numpy(x),), strict=False + ).module() + for sample in (x, x * -2 + 7, np.roll(x, 1, axis=-1)): + expected = torch.from_numpy(np.array(sample[key], copy=True)) + torch.testing.assert_close(model(torch.from_numpy(sample)), expected) + + +@pytest.mark.parametrize("key", CASES[::3]) +def test_reverse_slice_onnx_replay(key): + x = np.arange(120, dtype=np.float32).reshape(4, 5, 6) + model = capture(key, x) + exported = torch.onnx.export(model, (torch.from_numpy(x),), dynamo=True) + session = ort.InferenceSession( + exported.model_proto.SerializeToString(), providers=["CPUExecutionProvider"] + ) + for sample in (x, x * -2 + 7): + actual = session.run(None, {session.get_inputs()[0].name: sample})[0] + expected = sample[key] + assert actual.dtype == expected.dtype + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("size", [0, 1, 7]) +@pytest.mark.parametrize("key", [slice(None, None, -1), slice(None, None, -3)]) +def test_empty_and_singleton_dimensions(size, key): + x = np.arange(size, dtype=np.int64) + model = capture(key, x) + torch.testing.assert_close( + model(torch.from_numpy(x + 3)), + torch.from_numpy(np.array((x + 3)[key], copy=True)), + ) + + +def test_lowering_and_chained_numpy_operations(): + x = np.arange(30, dtype=np.float32).reshape(5, 6) + full = capture((slice(None, None, -1), Ellipsis), x) + assert any(node.target is torch.flip for node in full.graph.nodes) + partial = capture((Ellipsis, slice(4, 0, -2)), x) + assert any(node.target is torch.index_select for node in partial.graph.nodes) + + ctx = TracedTensorNode(name="reverse_chain", node_index=0) + traced = ctx.create_input(x.copy(), name="x") + result = np.sin(traced[::-1][:, ::-2]) + 1 + ctx.compile_trace({"y": result}) + sample = x * -2 + 7 + torch.testing.assert_close( + ctx.m(torch.from_numpy(sample)), + torch.from_numpy(np.sin(sample[::-1][:, ::-2]) + 1), + ) + + +def test_eager_and_inactive_reverse_views(): + ctx = TracedTensorNode(name="reverse_view", node_index=0) + x = ctx.create_input(np.arange(8, dtype=np.float32), name="x") + y = x[::-2] + assert np.shares_memory(x.data, y.data) + assert y.strides == x.data[::-2].strides + ctx.compile_trace({"y": y}) + inactive = x[::-1] + assert type(inactive) is np.ndarray + np.testing.assert_array_equal(inactive, x.data[::-1]) + + +def test_unsupported_combinations_fail_during_capture(): + ctx = TracedTensorNode(name="reverse_unsupported", node_index=0) + x = ctx.create_input(np.arange(12, dtype=np.float32).reshape(3, 4), name="x") + bound = ctx.create_input(np.array(2, dtype=np.int64), name="bound") + for key in ((slice(None, None, -1), [1, 3]), (slice(bound, None, -1), slice(None))): + before = tuple(ctx.graph.nodes) + with pytest.raises(NotImplementedError, match="static basic indices"): + x[key] + assert tuple(ctx.graph.nodes) == before + + +def test_zero_step_preserves_numpy_error(): + ctx = TracedTensorNode(name="zero_step", node_index=0) + x = ctx.create_input(np.arange(4, dtype=np.float32), name="x") + with pytest.raises(ValueError, match="slice step cannot be zero"): + x[::0] diff --git a/tests/unit_tests/test_traced_np_array.py b/tests/unit_tests/test_traced_np_array.py index fef32093..ae58dbfc 100644 --- a/tests/unit_tests/test_traced_np_array.py +++ b/tests/unit_tests/test_traced_np_array.py @@ -16,8 +16,10 @@ import onnxruntime as ort import torch +from leapp.utils.logging import _get_logger from leapp.leapp_graph.traced_node import TracedTensorNode from leapp.leapp_graph.datatypes import TracedNpArray +from leapp.leapp_graph.datatypes.numpy import aten_adapter from leapp.leapp_graph.datatypes.patching import TracingPatcher from tests.unit_tests.export_format_validation import ( @@ -25,224 +27,44 @@ verify_onnx_dynamo, ) -Array = np.ndarray | TracedNpArray - warnings.filterwarnings("ignore", message=".*legacy TorchScript-based ONNX export.*") warnings.filterwarnings("ignore", message=".*on empty non-base types in `__init__`.*") warnings.filterwarnings("ignore", message=".*feature will be removed.*") warnings.filterwarnings("ignore", message=".*Casting complex values to real.*") -class NpArrayArithmeticFunctions: - """NumPy array arithmetic functions for testing.""" - - @staticmethod - def scalar_arithmetic(a: Array) -> Array: - """Basic arithmetic operations.""" - b = a + 1 - c = b - 2 - d = c * 3 - e = d / 4 - f = e**5 - return f - - @staticmethod - def scalar_arithmetic_reverse(a: Array) -> Array: - """Reverse arithmetic where scalar comes first.""" - b = 1 + a - c = 2 - b - d = 3 * c - e = 4 / d - f = 5**e - return f - - @staticmethod - def array_arithmetic(a: Array) -> Array: - """Array-array arithmetic operations.""" - b = np.array([4.0, 5.0, 6.0]) - c = np.array([7.0, 8.0, 9.0]) - d = a + 2 * b - 0.5 * c - return d - - @staticmethod - def array_arithmetic_reuse(a: Array) -> Array: - """Array self arithmetic operations.""" - b = np.array([4.0, 5.0, 6.0]) - a = a + b - a = a + 2 - a = a - 1 - a = a * 3 - a = a / 4 - a = a**5 - return a - - @staticmethod - def matmul_operator(a: Array) -> Array: - """Matrix multiplication operator on array.""" - b = np.array([4.0, 5.0, 6.0]) - c = a @ b - return c - - @staticmethod - def matmul_function(a: Array) -> Array: - """np.matmul operation.""" - b = np.array([4.0, 5.0, 6.0]) - c = np.matmul(a, b) - return c - - @staticmethod - def sum_method(a: Array) -> Array: - """Sum operation via method.""" - b = a.sum() - return b - - @staticmethod - def sum_function(a: Array) -> Array: - """Sum operation via np.sum.""" - b = np.sum(a) - return b - - @staticmethod - def mean_method(a: Array) -> Array: - """Mean operation via method.""" - b = a.mean() - return b - - @staticmethod - def mean_function(a: Array) -> Array: - """Mean operation via np.mean.""" - b = np.mean(a) - return b - - @staticmethod - def max_method(a: Array) -> Array: - """Max operation via method.""" - b = a.max() - return b - - @staticmethod - def max_function(a: Array) -> Array: - """Max operation via np.max.""" - b = np.max(a) - return b - - @staticmethod - def min_method(a: Array) -> Array: - """Min operation via method.""" - b = a.min() - return b - - @staticmethod - def min_function(a: Array) -> Array: - """Min operation via np.min.""" - b = np.min(a) - return b - - @staticmethod - def argmax_method(a: Array) -> Array: - """Argmax operation via method.""" - b = a.argmax() - return b - - @staticmethod - def argmax_function(a: Array) -> Array: - """Argmax operation via np.argmax.""" - b = np.argmax(a) - return b - - @staticmethod - def argmin_method(a: Array) -> Array: - """Argmin operation via method.""" - b = a.argmin() - return b - - @staticmethod - def argmin_function(a: Array) -> Array: - """Argmin operation via np.argmin.""" - b = np.argmin(a) - return b - - @staticmethod - def reshape_method(a: Array) -> Array: - """Reshape operation via method.""" - b = a.reshape(1, 3) - c = b * 2 - d = c.reshape(3) - return d - - @staticmethod - def reshape_function(a: Array) -> Array: - """Reshape operation via np.reshape.""" - b = np.reshape(a, (1, 3)) - c = b * 2 - d = np.reshape(c, (3,)) - return d - - @staticmethod - def transpose_method(a: Array) -> Array: - """array.transpose() method.""" - b = a.reshape(3, 1) - c = b.transpose(1, 0) - d = c * 2 - return d - - @staticmethod - def transpose_function(a: Array) -> Array: - """np.transpose operation.""" - b = a.reshape(3, 1) - c = np.transpose(b, (1, 0)) - d = c * 2 - return d - - @staticmethod - def squeeze_method(a: Array) -> Array: - """array.squeeze() method.""" - b = a.reshape(1, 3, 1) - c = b.squeeze() - d = c * 2 - return d - - @staticmethod - def squeeze_function(a: Array) -> Array: - """np.squeeze operation.""" - b = a.reshape(1, 3, 1) - c = np.squeeze(b) - d = c * 2 - return d - - @staticmethod - def broadcasting_operation(a: Array) -> Array: - """Broadcasting operation.""" - b = a.reshape(3, 1) - c = np.array([[4.0, 5.0]]) - d = b + c - return d - - @staticmethod - def indexing_single(a: Array) -> Array: - """Single element indexing.""" - b = a[1] - return b - - @staticmethod - def indexing_slice(a: Array) -> Array: - """Slice indexing.""" - b = a[0:2] - c = b * 2 - return c - - @staticmethod - def indexing_negative(a: Array) -> Array: - """Negative indexing.""" - b = a[-1] - return b - - @staticmethod - def indexing_step(a: Array) -> Array: - """Step indexing.""" - b = a[::2] - c = b * 2 - return c +def _assert_numpy_output(actual, expected, atol=1e-5): + if isinstance(expected, (tuple, list)): + assert isinstance(actual, (tuple, list)) + assert len(actual) == len(expected) + for value, reference in zip(actual, expected): + _assert_numpy_output(value, reference, atol) + return + torch.testing.assert_close( + actual, + torch.from_numpy(np.array(expected, copy=True)), + rtol=1e-7, + atol=atol, + equal_nan=True, + ) + + +def _capture_numpy(fn, *inputs): + expected = fn(*(x.copy() for x in inputs)) + ctx = TracedTensorNode(name="numpy_operation", node_index=0) + result = fn( + *(ctx.create_input(x.copy(), name=f"x{i}") for i, x in enumerate(inputs)) + ) + outputs = result if isinstance(result, (tuple, list)) else (result,) + assert all(isinstance(x, TracedNpArray) for x in outputs) + references = expected if isinstance(expected, (tuple, list)) else (expected,) + assert len(outputs) == len(references) + for output, reference in zip(outputs, references): + _assert_numpy_output( + torch.from_numpy(np.array(output.data, copy=True)), reference + ) + ctx.compile_trace({f"y{i}": x for i, x in enumerate(outputs)}) + return ctx.m class TestTracedNpArray(unittest.TestCase): @@ -250,34 +72,13 @@ class TestTracedNpArray(unittest.TestCase): # ==================== Basic Property Tests ==================== - def test_shape(self): - """Test shape property.""" + def test_array_properties(self): ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.random.randn(3, 4, 5), name="x") + x = ctx.create_input(np.arange(60, dtype=np.int32).reshape(3, 4, 5), name="x") self.assertEqual(x.shape, (3, 4, 5)) - - def test_ndim(self): - """Test ndim property.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.random.randn(3, 4, 5), name="x") self.assertEqual(x.ndim, 3) - - def test_dtype(self): - """Test dtype property.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1, 2, 3], dtype=np.int32), name="x") self.assertEqual(x.dtype, np.int32) - - def test_len(self): - """Test len() function.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.random.randn(10, 5), name="x") - self.assertEqual(len(x), 10) - - def test_size(self): - """Test size property.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.random.randn(3, 4, 5), name="x") + self.assertEqual(len(x), 3) self.assertEqual(x.size, 60) def test_T_property(self): @@ -290,62 +91,6 @@ def test_T_property(self): # ==================== Arithmetic Operation Tests ==================== - def test_add(self): - """Test addition.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - y = x + 1 - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([2.0, 3.0, 4.0])) - - def test_radd(self): - """Test reverse addition.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - y = 1 + x - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([2.0, 3.0, 4.0])) - - def test_sub(self): - """Test subtraction.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - y = x - 1 - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([0.0, 1.0, 2.0])) - - def test_rsub(self): - """Test reverse subtraction.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - y = 10 - x - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([9.0, 8.0, 7.0])) - - def test_mul(self): - """Test multiplication.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - y = x * 2 - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([2.0, 4.0, 6.0])) - - def test_rmul(self): - """Test reverse multiplication.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - y = 2 * x - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([2.0, 4.0, 6.0])) - - def test_truediv(self): - """Test division.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([2.0, 4.0, 6.0]), name="x") - y = x / 2 - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([1.0, 2.0, 3.0])) - def test_rtruediv(self): """Test reverse division.""" ctx = TracedTensorNode(name="test", node_index=0) @@ -354,22 +99,6 @@ def test_rtruediv(self): self.assertIsInstance(y, TracedNpArray) np.testing.assert_array_almost_equal(y, np.array([8.0, 4.0, 2.0])) - def test_pow(self): - """Test power.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - y = x**2 - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([1.0, 4.0, 9.0])) - - def test_neg(self): - """Test negation.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, -2.0, 3.0]), name="x") - y = -x - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([-1.0, 2.0, -3.0])) - def test_abs(self): """Test absolute value.""" ctx = TracedTensorNode(name="test", node_index=0) @@ -412,111 +141,8 @@ def test_comparison_eq(self): self.assertIsInstance(y, TracedNpArray) np.testing.assert_array_equal(y, np.array([False, True, False])) - # ==================== NumPy Ufunc Tests ==================== - - def test_np_sin(self): - """Test np.sin ufunc.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([0.0, np.pi / 2, np.pi]), name="x") - y = np.sin(x) - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([0.0, 1.0, 0.0]), decimal=5) - - def test_np_cos(self): - """Test np.cos ufunc.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([0.0, np.pi / 2, np.pi]), name="x") - y = np.cos(x) - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([1.0, 0.0, -1.0]), decimal=5) - - def test_np_exp(self): - """Test np.exp ufunc.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([0.0, 1.0, 2.0]), name="x") - y = np.exp(x) - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.exp(np.array([0.0, 1.0, 2.0]))) - - def test_np_log(self): - """Test np.log ufunc.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, np.e, np.e**2]), name="x") - y = np.log(x) - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([0.0, 1.0, 2.0])) - - def test_np_sqrt(self): - """Test np.sqrt ufunc.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 4.0, 9.0]), name="x") - y = np.sqrt(x) - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([1.0, 2.0, 3.0])) - - def test_np_abs(self): - """Test np.abs ufunc.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([-1.0, 2.0, -3.0]), name="x") - y = np.abs(x) - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([1.0, 2.0, 3.0])) - # ==================== NumPy Function Tests ==================== - def test_np_sum(self): - """Test np.sum function.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - y = np.sum(x) - # Result might be scalar or 0-d array - self.assertAlmostEqual(float(y), 6.0) - - def test_np_mean(self): - """Test np.mean function.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - y = np.mean(x) - self.assertAlmostEqual(float(y), 2.0) - - def test_np_max(self): - """Test np.max function.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 3.0, 2.0]), name="x") - y = np.max(x) - self.assertAlmostEqual(float(y), 3.0) - - def test_np_min(self): - """Test np.min function.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 3.0, 2.0]), name="x") - y = np.min(x) - self.assertAlmostEqual(float(y), 1.0) - - def test_np_reshape(self): - """Test np.reshape function.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]), name="x") - y = np.reshape(x, (2, 3)) - self.assertIsInstance(y, TracedNpArray) - self.assertEqual(y.shape, (2, 3)) - - def test_np_transpose(self): - """Test np.transpose function.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.random.randn(3, 4), name="x") - y = np.transpose(x) - self.assertIsInstance(y, TracedNpArray) - self.assertEqual(y.shape, (4, 3)) - - def test_np_squeeze(self): - """Test np.squeeze function.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.random.randn(1, 3, 1), name="x") - y = np.squeeze(x) - self.assertIsInstance(y, TracedNpArray) - self.assertEqual(y.shape, (3,)) - def test_np_concatenate(self): """Test np.concatenate function.""" ctx = TracedTensorNode(name="test", node_index=0) @@ -552,31 +178,6 @@ def test_np_where(self): self.assertIsInstance(y, TracedNpArray) np.testing.assert_array_almost_equal(y, np.array([1.0, 0.0, 3.0])) - def test_np_sort(self): - """Test np.sort function.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([3.0, 1.0, 2.0]), name="x") - y = np.sort(x) - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, np.array([1.0, 2.0, 3.0])) - - def test_np_sort_axis(self): - """Test np.sort with axis parameter.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([[3.0, 1.0], [2.0, 4.0]]), name="x") - y = np.sort(x, axis=1) - self.assertIsInstance(y, TracedNpArray) - expected = np.array([[1.0, 3.0], [2.0, 4.0]]) - np.testing.assert_array_almost_equal(y, expected) - - def test_np_argsort(self): - """Test np.argsort function.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([3.0, 1.0, 2.0]), name="x") - y = np.argsort(x) - self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_equal(y, np.array([1, 2, 0])) - def test_np_nonzero(self): """Test np.nonzero function.""" ctx = TracedTensorNode(name="test", node_index=0) @@ -705,7 +306,7 @@ def test_method_copy(self): x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") y = x.copy() self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, x) + np.testing.assert_array_almost_equal(y.data, x.data) def test_method_astype(self): """Test .astype() method.""" @@ -714,7 +315,7 @@ def test_method_astype(self): y = x.astype(np.float32) self.assertIsInstance(y, TracedNpArray) self.assertEqual(y.dtype, np.float32) - np.testing.assert_array_almost_equal(y, x) + np.testing.assert_array_almost_equal(y.data, x.data) def test_method_astype_int(self): """Test .astype() to integer type.""" @@ -732,7 +333,7 @@ def test_method_view(self): # View as same dtype should return TracedNpArray y = x.view(np.float64) self.assertIsInstance(y, TracedNpArray) - np.testing.assert_array_almost_equal(y, x) + np.testing.assert_array_almost_equal(y.data, x.data) def test_method_copy_compiled(self): """Test .copy() method compiles correctly.""" @@ -761,19 +362,7 @@ def test_graph_basic_operations(self): ctx.compile_trace({"output": z}) self.assertIsNotNone(ctx.m) self.assertIsNotNone(ctx.graph) - - def test_graph_has_nodes(self): - """Test that operations are recorded in the graph.""" - ctx = TracedTensorNode(name="test", node_index=0) - x = ctx.create_input(np.array([1.0, 2.0, 3.0]), name="x") - y = x + 1 - z = y * 2 - - ctx.compile_trace({"output": z}) - - # Check graph has nodes - nodes = list(ctx.graph.nodes) - self.assertGreater(len(nodes), 0) + self.assertGreater(len(list(ctx.graph.nodes)), 0) def test_graph_multiple_inputs(self): """Test graph with multiple inputs.""" @@ -806,6 +395,7 @@ def test_operations_after_compile_return_native_arrays(self): x.copy(), x.astype(np.float32), np.array(x), + np.cbrt(x), ] # These alias the boundary value rather than producing new data. aliases = [np.asarray(x), np.asanyarray(x)] @@ -819,6 +409,8 @@ def test_operations_after_compile_return_native_arrays(self): for result in aliases: self.assertIs(result, x) + np.testing.assert_array_equal(derived[-1], np.cbrt(x.data)) + # The boundary value itself is untouched and still carries its node. self.assertIsInstance(x, TracedNpArray) self.assertIs(x.context_obj, ctx) @@ -952,81 +544,6 @@ def test_complex_computation(self): np.testing.assert_array_almost_equal(w, np.ones(3), decimal=5) -class TestArithmeticFunctions(unittest.TestCase): - """Test all arithmetic functions can be traced and produce correct results.""" - - def setUp(self): - # Use [2.0, 3.0, 4.0] to avoid zeros in reverse arithmetic: - # b = 1 + a = [3, 4, 5], c = 2 - b = [-1, -2, -3] (no zeros) - self.test_data = np.array([2.0, 3.0, 4.0]) - - def _test_function(self, func, name): - """Helper to test a function produces correct results.""" - # Get expected result from regular numpy - expected = func(self.test_data.copy()) - - # Create traced version - ctx = TracedTensorNode(name=f"test_{name}", node_index=0) - traced_input = ctx.create_input(self.test_data.copy(), name="input") - traced_result = func(traced_input) - - # Compare results - if isinstance(traced_result, TracedNpArray): - actual = traced_result - else: - actual = np.array(traced_result) - - if isinstance(expected, np.ndarray): - np.testing.assert_array_almost_equal(actual, expected, decimal=5) - else: - self.assertAlmostEqual(float(actual), float(expected), places=5) - - def test_scalar_arithmetic(self): - """Test scalar arithmetic function.""" - self._test_function( - NpArrayArithmeticFunctions.scalar_arithmetic, "scalar_arithmetic" - ) - - def test_scalar_arithmetic_reverse(self): - """Test reverse scalar arithmetic function.""" - self._test_function( - NpArrayArithmeticFunctions.scalar_arithmetic_reverse, - "scalar_arithmetic_reverse", - ) - - def test_array_arithmetic(self): - """Test array arithmetic function.""" - self._test_function( - NpArrayArithmeticFunctions.array_arithmetic, "array_arithmetic" - ) - - def test_sum_method(self): - """Test sum method.""" - self._test_function(NpArrayArithmeticFunctions.sum_method, "sum_method") - - def test_sum_function(self): - """Test sum function.""" - self._test_function(NpArrayArithmeticFunctions.sum_function, "sum_function") - - def test_mean_method(self): - """Test mean method.""" - self._test_function(NpArrayArithmeticFunctions.mean_method, "mean_method") - - def test_mean_function(self): - """Test mean function.""" - self._test_function(NpArrayArithmeticFunctions.mean_function, "mean_function") - - def test_reshape_method(self): - """Test reshape method.""" - self._test_function(NpArrayArithmeticFunctions.reshape_method, "reshape_method") - - def test_reshape_function(self): - """Test reshape function.""" - self._test_function( - NpArrayArithmeticFunctions.reshape_function, "reshape_function" - ) - - # Suppress TracerWarning from setitem tests where torch.tensor/torch.as_tensor constants # are baked into the trace (e.g. "torch.as_tensor results are registered as constants in the trace"). # This is expected behavior for constant index/value assignments like x[0] = 10.0. @@ -1043,105 +560,158 @@ class TestCompiledGraphExecution(unittest.TestCase): def _test_operation( self, - op_name: str, + op_name, np_func, - input_data: np.ndarray, - expected_output: np.ndarray, - atol: float = 1e-5, + input_data, + expected_output, + atol=1e-5, + *, + backends=("fx", "jit", "onnx", "pt2"), + replay_inputs=None, ): - """Helper to test a single operation through all export formats. - - Args: - op_name: Name of the operation (for error messages) - np_func: Function that takes TracedNpArray and returns result - input_data: Input numpy array - expected_output: Expected numpy output - atol: Absolute tolerance for comparisons - """ - # Create traced version - ctx = TracedTensorNode(name=f"test_{op_name}", node_index=0) - traced_input = ctx.create_input(input_data.copy(), name="input") - traced_result = np_func(traced_input) - - self.assertIsInstance( - traced_result, - TracedNpArray, - f"{op_name}: result should be TracedNpArray, got {type(traced_result)}", + """Capture once; check eager and changed-input replay without dtype/shape coercion.""" + graph_module = _capture_numpy(np_func, input_data) + input_tensor = torch.from_numpy(input_data.copy()) + samples = (input_data,) + tuple( + (input_data + 2,) if replay_inputs is None else replay_inputs ) - - # Compile the trace - ctx.compile_trace({"output": traced_result}) - graph_module = ctx.m - - # Convert input to torch tensor for execution - input_tensor = torch.from_numpy(input_data.copy()).float() - - # Test 1: FX GraphModule execution - try: - graph_output = graph_module(input_tensor) - graph_output_np = graph_output.detach().numpy() - np.testing.assert_allclose( - graph_output_np.flatten(), - expected_output.flatten(), - atol=atol, - err_msg=f"{op_name}: FX GraphModule output doesn't match expected", - ) - except Exception as e: - self.fail(f"{op_name}: FX GraphModule execution failed: {e}") - - # Test 2: TorchScript export and execution - try: - scripted = torch.jit.script(graph_module) - script_output = scripted(input_tensor) - script_output_np = script_output.detach().numpy() - np.testing.assert_allclose( - script_output_np.flatten(), - expected_output.flatten(), - atol=atol, - err_msg=f"{op_name}: TorchScript output doesn't match expected", - ) - except Exception as e: - self.fail(f"{op_name}: TorchScript export/execution failed: {e}") - - # Test 3: ONNX export and execution - try: - with tempfile.TemporaryDirectory() as tmpdir: - onnx_path = pathlib.Path(tmpdir) / f"{op_name}.onnx" - torch.onnx.export( - graph_module, - (input_tensor,), - onnx_path, - dynamo=False, - export_params=True, - opset_version=17, - do_constant_folding=True, - input_names=["input"], - output_names=["output"], + for backend in backends: + with self.subTest(operation=op_name, backend=backend): + model = graph_module + if backend == "jit": + model = torch.jit.script(graph_module) + elif backend == "jit-trace": + model = torch.jit.trace(graph_module, (input_tensor.clone(),)) + elif backend == "pt2": + model = torch.export.export( + graph_module, (input_tensor.clone(),), strict=False + ).module() + elif backend in ("onnx", "onnx-dynamo"): + if backend == "onnx-dynamo": + exported = torch.onnx.export( + graph_module, (input_tensor.clone(),), dynamo=True + ) + session = ort.InferenceSession( + exported.model_proto.SerializeToString(), + providers=["CPUExecutionProvider"], + ) + else: + with tempfile.TemporaryDirectory() as tmpdir: + path = pathlib.Path(tmpdir) / f"{op_name}.onnx" + torch.onnx.export( + graph_module, + (input_tensor.clone(),), + path, + dynamo=False, + opset_version=17, + input_names=["input"], + output_names=["output"], + ) + session = ort.InferenceSession( + str(path), providers=["CPUExecutionProvider"] + ) + for index, sample in enumerate(samples): + with self.subTest(sample=index): + reference = ( + expected_output if index == 0 else np_func(sample.copy()) + ) + if backend in ("onnx", "onnx-dynamo"): + actual = torch.from_numpy( + session.run( + None, {session.get_inputs()[0].name: sample.copy()} + )[0] + ) + else: + actual = model(torch.from_numpy(sample.copy())) + _assert_numpy_output(actual, reference, atol) + + def test_numpy_function_replay(self): + x = np.array( + [[3.0, 1.0, 2.0], [1.0, 4.0, 1.0], [2.0, 1.0, 5.0]], dtype=np.float32 + ) + corner = x.copy() + corner[0, 0] += 4 + operations = { + "inverse": np.linalg.inv, + "det": np.linalg.det, + "slogdet": lambda x: tuple(np.linalg.slogdet(x)), + "cholesky": np.linalg.cholesky, + "solve": lambda x: np.linalg.solve(x, np.ones(3, dtype=np.float32)), + "matrix_power": lambda x: np.linalg.matrix_power(x, 3), + "svd_values": lambda x: np.linalg.svd(x, compute_uv=False), + "eigvalsh": np.linalg.eigvalsh, + "pinv": np.linalg.pinv, + "norm": lambda x: x / np.linalg.norm(x, axis=-1, keepdims=True), + "cross": lambda x: np.cross(x, np.array([1.0, 2.0, 3.0], dtype=np.float32)), + "isclose_where": lambda x: np.where(np.isclose(x, 3.0), x, -x), + "sinc": np.sinc, + "degrees": np.degrees, + "radians": np.radians, + "logaddexp": lambda x: np.logaddexp(x, x + 1), + "heaviside": lambda x: np.heaviside(x - 4, 0.5), + "copysign": lambda x: np.copysign(x, -1.0), + "gradient": lambda x: tuple(np.gradient(x)), + "repeat": lambda x: np.repeat(x, 2, axis=1), + "tile": lambda x: np.tile(x, (2, 1)), + "take": lambda x: np.take(x, [0, 2], axis=1), + "take_along_axis": lambda x: np.take_along_axis( + x, np.array([[2, 0], [0, 1], [1, 2]]), axis=1 + ), + "broadcast_to": lambda x: np.broadcast_to(x, (2, 3, 3)), + "median": lambda x: np.median(x, axis=1), + "ptp": lambda x: np.ptp(x, axis=1), + "unique": np.unique, + "diff": lambda x: np.diff(x, axis=-1), + "quantile": lambda x: np.quantile(x, 0.25, axis=1), + "cbrt": np.cbrt, + "pad": lambda x: np.pad(x, ((1, 0), (0, 2)), constant_values=-2), + "nanmean": lambda x: np.nanmean(x, axis=-1, keepdims=True), + } + for name, fn in operations.items(): + backends = ("fx", "jit", "jit-trace", "pt2") + if name in ("norm", "cross", "isclose_where"): + backends += ("onnx-dynamo",) + with self.subTest(operation=name): + self._test_operation( + name, + fn, + x, + fn(x), + backends=backends, + # Cholesky requires positive-definite inputs. + replay_inputs=(x + 2, corner) + if name == "cholesky" + else (x + 2, corner, -x), ) - session = ort.InferenceSession(str(onnx_path)) - onnx_output = session.run( - None, {"input": input_data.astype(np.float32)} - )[0] - - np.testing.assert_allclose( - onnx_output.flatten(), - expected_output.flatten(), - atol=atol, - err_msg=f"{op_name}: ONNX output doesn't match expected", - ) - except Exception as e: - self.fail(f"{op_name}: ONNX export/execution failed: {e}") + def test_integer_mean(self): + x = np.arange(6, dtype=np.int32).reshape(2, 3) + for fn in (np.mean, np.nanmean): + operation = lambda x: fn(x, axis=-1, keepdims=True) + self._test_operation( + fn.__name__, + operation, + x, + operation(x), + backends=("fx", "jit", "pt2"), + replay_inputs=(x + 3,), + ) - # Test 4: ExportedProgram export and execution - verify_exported_program( - self, - graph_module, - input_tensor, - torch.from_numpy(expected_output).float(), - test_name=op_name, - atol=atol, - ) + def test_method_and_arithmetic_chains(self): + x = np.array([2.0, 3.0, 4.0]) + operations = { + "scalar_arithmetic": lambda x: (((x + 1) - 2) * 3 / 4) ** 5, + "reverse_arithmetic": lambda x: 5 ** (4 / (3 * (2 - (1 + x)))), + "array_arithmetic": lambda x: ( + x + 2 * np.array([4.0, 5.0, 6.0]) - 0.5 * np.array([7.0, 8.0, 9.0]) + ), + "sum_method": lambda x: x.sum(), + "mean_method": lambda x: x.mean(), + "reshape_method": lambda x: (x.reshape(1, 3) * 2).reshape(3), + "reshape_function": lambda x: np.reshape(np.reshape(x, (1, 3)) * 2, (3,)), + } + for name, fn in operations.items(): + self._test_operation(name, fn, x, fn(x)) def test_add_scalar(self): """Test: x + 1""" @@ -1185,7 +755,7 @@ def test_mul_np_float64(self): ) def test_div_norm(self): - """Test: x / np.linalg.norm(x) — unmapped norm returns a numpy scalar.""" + """Test: normalization retains its runtime-dependent denominator.""" input_data = np.array([3.0, 4.0], dtype=np.float32) expected = input_data / np.linalg.norm(input_data) self._test_operation( @@ -1240,15 +810,22 @@ def test_rdiv(self): def test_np_sin(self): """Test: np.sin(x)""" - input_data = np.array([0.0, 0.5, 1.0]) + input_data = np.array([0.0, 0.5, 1.0, np.pi / 2, np.pi]) expected = np.sin(input_data) self._test_operation("np_sin", lambda x: np.sin(x), input_data, expected) def test_np_cos(self): """Test: np.cos(x)""" - input_data = np.array([0.0, 0.5, 1.0]) - expected = np.cos(input_data) - self._test_operation("np_cos", lambda x: np.cos(x), input_data, expected) + for dtype in (np.float32, np.float64): + with self.subTest(dtype=dtype): + input_data = np.array([0.0, 0.5, 1.0, np.pi / 2, np.pi], dtype=dtype) + # ONNX Runtime 1.23.2's CPU Cos kernel only supports float32. + backends = ("fx", "jit", "pt2") + if dtype is np.float32: + backends += ("onnx",) + self._test_operation( + "np_cos", np.cos, input_data, np.cos(input_data), backends=backends + ) def test_np_exp(self): """Test: np.exp(x)""" @@ -1258,7 +835,7 @@ def test_np_exp(self): def test_np_log(self): """Test: np.log(x)""" - input_data = np.array([1.0, 2.0, 3.0]) + input_data = np.array([1.0, 2.0, 3.0, np.e, np.e**2]) expected = np.log(input_data) self._test_operation("np_log", lambda x: np.log(x), input_data, expected) @@ -1578,25 +1155,26 @@ def test_np_sort(self): self._test_operation("np_sort", lambda x: np.sort(x), input_data, expected) def test_np_sort_axis(self): - """Test: np.sort(x, axis=1)""" - input_data = np.array([[3.0, 1.0, 2.0], [6.0, 4.0, 5.0]]) - expected = np.sort(input_data, axis=1) - self._test_operation( - "np_sort_axis", lambda x: np.sort(x, axis=1), input_data, expected - ) + for axis in (None, 1): + x = np.array([[3.0, 1.0, 2.0], [6.0, 4.0, 5.0]], dtype=np.float32) + fn = lambda x: np.sort(x, axis=axis) + self._test_operation( + f"sort_{axis}", fn, x, fn(x), replay_inputs=(x + 2, -x) + ) def test_np_argsort(self): - """Test: np.argsort(x)""" - input_data = np.array([3.0, 1.0, 4.0, 1.0, 5.0]) - expected = np.argsort(input_data).astype( - np.float32 - ) # Convert to float for comparison - # Note: argsort returns indices as int64, but we compare as float for consistency + # Retain the default 1-D duplicate-value case and its integer output. + x = np.array([3.0, 1.0, 4.0, 1.0, 5.0], dtype=np.float32) + self._test_operation("argsort", np.argsort, x, np.argsort(x)) + + x = np.array( + [[3.0, 1.0, 2.0], [1.0, 4.0, 1.0], [2.0, 1.0, 5.0]], dtype=np.float32 + ) + fn = lambda x: np.argsort(x, axis=0) + corner = x.copy() + corner[0, 0] += 4 self._test_operation( - "np_argsort", - lambda x: np.argsort(x).astype(np.float32), - input_data, - expected, + "argsort_axis", fn, x, fn(x), replay_inputs=(x + 2, corner, -x) ) def test_np_where_compiled(self): @@ -1609,112 +1187,30 @@ def test_np_where_compiled(self): def test_np_nonzero_rows(self): """Test: np.nonzero(x)[0] — numpy returns one index array per dimension.""" - input_data = np.array([[0.0, 0.0, 1.0], [1.0, 0.0, 1.0]]) - expected = np.nonzero(input_data)[0].astype(np.float32) + # The legacy ONNX Runtime NonZero kernel requires float32 here. + input_data = np.array([[0.0, 0.0, 1.0], [1.0, 0.0, 1.0]], dtype=np.float32) + expected = np.nonzero(input_data)[0] self._test_operation( "np_nonzero_rows", - lambda x: np.nonzero(x)[0].astype(np.float32), + lambda x: np.nonzero(x)[0], input_data, expected, ) def test_np_nonzero_cols(self): """Test: np.nonzero(x)[1] — the second element indexes the last dimension.""" - input_data = np.array([[0.0, 0.0, 1.0], [1.0, 0.0, 1.0]]) - expected = np.nonzero(input_data)[1].astype(np.float32) + # The legacy ONNX Runtime NonZero kernel requires float32 here. + input_data = np.array([[0.0, 0.0, 1.0], [1.0, 0.0, 1.0]], dtype=np.float32) + expected = np.nonzero(input_data)[1] self._test_operation( "np_nonzero_cols", - lambda x: np.nonzero(x)[1].astype(np.float32), + lambda x: np.nonzero(x)[1], input_data, expected, ) # ==================== Setitem Operations ==================== - def _test_setitem( - self, - test_name: str, - input_data: np.ndarray, - setitem_func, - expected_output: np.ndarray, - atol: float = 1e-5, - ): - """Helper to test setitem operations through all export formats. - - Args: - test_name: Name of the test (for error messages and ONNX filename) - input_data: Input numpy array - setitem_func: Function that takes TracedNpArray and does setitem + returns result - expected_output: Expected numpy output - atol: Absolute tolerance for comparisons - """ - # Create traced version - ctx = TracedTensorNode(name=f"test_{test_name}", node_index=0) - traced_input = ctx.create_input(input_data.copy(), name="input") - traced_result = setitem_func(traced_input) - - # Compile the trace - ctx.compile_trace({"output": traced_result}) - graph_module = ctx.m - - # Convert input to torch tensor for execution - torch_input = torch.from_numpy(input_data.copy()).float() - torch_expected = torch.from_numpy(expected_output).float() - - # Test 1: FX GraphModule execution - try: - output = graph_module(torch_input) - self.assertTrue( - torch.allclose(output, torch_expected, atol=atol), - f"{test_name}: FX GraphModule execution failed - output doesn't match", - ) - except Exception as e: - self.fail(f"{test_name}: FX GraphModule execution failed: {e}") - - # Test 2: TorchScript export - try: - scripted = torch.jit.script(graph_module) - output_ts = scripted(torch_input) - self.assertTrue( - torch.allclose(output_ts, torch_expected, atol=atol), - f"{test_name}: TorchScript output doesn't match", - ) - except Exception as e: - self.fail(f"{test_name}: TorchScript export/execution failed: {e}") - - # Test 3: ONNX export - try: - with tempfile.TemporaryDirectory() as tmpdir: - onnx_path = pathlib.Path(tmpdir) / f"{test_name}.onnx" - torch.onnx.export( - graph_module, - (torch_input,), - onnx_path, - dynamo=False, - export_params=True, - opset_version=17, - input_names=["input"], - output_names=["output"], - ) - session = ort.InferenceSession(str(onnx_path)) - output_onnx = session.run(None, {"input": torch_input.numpy()})[0] - self.assertTrue( - np.allclose(output_onnx, expected_output, atol=atol), - f"{test_name}: ONNX output doesn't match", - ) - except Exception as e: - self.fail(f"{test_name}: ONNX export/execution failed: {e}") - - # Test 4: ExportedProgram export - verify_exported_program( - self, - graph_module, - torch_input, - torch_expected, - test_name=test_name, - atol=atol, - ) - def test_setitem_single_index(self): """Test: x[0] = 10""" input_data = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) @@ -1724,7 +1220,7 @@ def setitem_func(x): return x * 2 expected = np.array([20.0, 4.0, 6.0, 8.0, 10.0]) - self._test_setitem("setitem_single", input_data, setitem_func, expected) + self._test_operation("setitem_single", setitem_func, input_data, expected) def test_setitem_slice(self): """Test: x[1:3] = [10, 20]""" @@ -1735,7 +1231,7 @@ def setitem_func(x): return x * 2 expected = np.array([2.0, 20.0, 40.0, 8.0, 10.0]) - self._test_setitem("setitem_slice", input_data, setitem_func, expected) + self._test_operation("setitem_slice", setitem_func, input_data, expected) def test_setitem_full_slice(self): """Test: x[:] = constant""" @@ -1746,7 +1242,7 @@ def setitem_func(x): return x * 2 expected = np.array([20.0, 40.0, 60.0]) - self._test_setitem("setitem_full", input_data, setitem_func, expected) + self._test_operation("setitem_full", setitem_func, input_data, expected) def test_setitem_step_slice(self): """Test: x[::2] = [10, 30, 50] with step""" @@ -1757,7 +1253,7 @@ def setitem_func(x): return x * 2 expected = np.array([20.0, 4.0, 60.0, 8.0, 100.0]) - self._test_setitem("setitem_step", input_data, setitem_func, expected) + self._test_operation("setitem_step", setitem_func, input_data, expected) def _validate_dynamic_setitem(self, graph_module, inputs, expected, test_name): """Validate export-first backends for traced NumPy index inputs.""" @@ -1849,5 +1345,219 @@ def test_setitem_traced_index_plain_numpy_source(self): ) +@pytest.mark.parametrize("equal_nan", [False, True]) +def test_isclose_special_values_and_broadcasting(equal_nan): + fn = lambda x, y: np.isclose(x, y, rtol=1e-4, atol=1e-6, equal_nan=equal_nan) + x = np.array([[0.0, 1.00009, np.nan, np.inf, -np.inf]], dtype=np.float64) + y = np.array([0.0, 1.0, np.nan, np.inf, np.inf], dtype=np.float64) + model = _capture_numpy(fn, x, y) + changed = np.array([[1e-5, 1.0002, 3.0, -np.inf, -np.inf]], dtype=np.float64) + for a, b in ((x, y), (changed, y), (changed, changed[0])): + torch.testing.assert_close( + model(torch.from_numpy(a), torch.from_numpy(b)), torch.from_numpy(fn(a, b)) + ) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_cbrt_signed_zero_and_special_values(dtype): + x = np.array([-np.inf, -8.0, -0.0, 0.0, 27.0, np.inf, np.nan], dtype=dtype) + model = _capture_numpy(np.cbrt, x) + for sample in (x, -x): + actual = model(torch.from_numpy(sample)) + expected = torch.from_numpy(np.cbrt(sample)) + torch.testing.assert_close(actual, expected, equal_nan=True) + torch.testing.assert_close( + torch.signbit(actual[2:4]), torch.signbit(expected[2:4]) + ) + + +@pytest.mark.parametrize("axis", [None, -1, (0, 1), ()]) +@pytest.mark.filterwarnings("ignore:Mean of empty slice:RuntimeWarning") +def test_nanmean_nan_slices_and_explicit_dtype(axis): + fn = lambda x: np.nanmean(x, axis=axis, dtype=np.float64, keepdims=True) + x = np.array([[np.nan, np.nan], [1.0, np.nan]], dtype=np.float32) + model = _capture_numpy(fn, x) + changed = np.array([[2.0, 3.0], [np.nan, 4.0]], dtype=np.float32) + for sample in (x, changed, changed + 2): + torch.testing.assert_close( + model(torch.from_numpy(sample)), + torch.from_numpy(fn(sample)), + equal_nan=True, + ) + + +def test_structured_outputs_and_constant_snapshot(): + constant = np.array([1.0, 2.0, 3.0], dtype=np.float32) + fn = lambda x: (np.cross(x, constant), *np.gradient(x)) + x = np.arange(9, dtype=np.float32).reshape(3, 3) + model = _capture_numpy(fn, x) + expected = fn(x + 2) + constant[:] = 100 + for actual, reference in zip(model(torch.from_numpy(x + 2)), expected): + torch.testing.assert_close(actual, torch.from_numpy(reference)) + + +@pytest.mark.parametrize("backend", ["fx", "jit", "pt2"]) +@pytest.mark.parametrize( + "fn, message", + [ + (np.spacing, "no torch._numpy lowering"), + (lambda x: np.allclose(x, x), "ATen capture failed"), + (lambda x: np.sort(x, stable=True), "sorting"), + (lambda x: np.argsort(x, kind="mergesort"), "sorting"), + (lambda x: np.pad(x, 1, mode="reflect"), "pad"), + (lambda x: np.pad(x, 1, constant_values=(1, 2)), "pad"), + (lambda x: np.cbrt(x, order="C"), "unsupported signature"), + ], +) +def test_failed_capture_on_unused_branch(fn, message, backend, monkeypatch): + errors = [] + monkeypatch.setattr(_get_logger(), "error", errors.append) + ctx = TracedTensorNode(name="unused_failure", node_index=0) + sample = np.array([3.0, 1.0, 2.0], dtype=np.float32) + x = ctx.create_input(sample.copy(), name="x") + unused = fn(x) + np.testing.assert_equal(unused.data, fn(sample)) + # Consumers of a failed branch can also be discarded with that branch. + unused + 1 + assert any(message in error and "unused_failure" in error for error in errors) + ctx.compile_trace({"y": np.sin(x)}) + assert not any("leapp_capture_error" in node.meta for node in ctx.graph.nodes) + model = ctx.m + if backend == "jit": + model = torch.jit.script(model) + elif backend == "pt2": + model = torch.export.export( + model, (torch.from_numpy(sample),), strict=False + ).module() + for changed in (sample, sample + 2, -sample): + _assert_numpy_output(model(torch.from_numpy(changed)), np.sin(changed)) + + +@pytest.mark.parametrize("dependent", [False, True]) +@pytest.mark.parametrize("fn", [np.spacing, lambda x: np.pad(x, 1, mode="reflect")]) +def test_failed_capture_in_output_tree_is_rejected(fn, dependent, monkeypatch): + errors = [] + monkeypatch.setattr(_get_logger(), "error", errors.append) + ctx = TracedTensorNode(name="required_failure", node_index=0) + x = ctx.create_input(np.array([3.0, 1.0, 2.0], dtype=np.float32), name="x") + failed = fn(x) + output = np.sin(failed) if dependent else failed + assert errors + with pytest.raises( + NotImplementedError, match="declared output depends on a failed capture" + ): + ctx.compile_trace({"y": output}) + + +def test_native_numpy_errors_still_propagate(monkeypatch): + errors = [] + monkeypatch.setattr(_get_logger(), "error", errors.append) + ctx = TracedTensorNode(name="invalid_numpy", node_index=0) + x = ctx.create_input(np.ones(3, dtype=np.float32), name="x") + before = tuple(ctx.graph.nodes) + with pytest.raises(ValueError): + np.reshape(x, (2, 2)) + assert not errors + assert tuple(ctx.graph.nodes) == before + + +@pytest.mark.parametrize("positional", [False, True]) +def test_reject_out_before_mutating_storage(positional): + ctx = TracedTensorNode(name="out", node_index=0) + x = ctx.create_input(np.array([2.0, 1.0, 3.0]), name="x") + out = np.zeros(2) + with pytest.raises(NotImplementedError, match="mutation"): + if positional: + np.take(x, [0, 2], None, out) + else: + np.take(x, [0, 2], out=out) + np.testing.assert_array_equal(out, np.zeros(2)) + + +@pytest.mark.parametrize("positional", [False, True]) +def test_missing_ufunc_rejects_out_before_mutating_storage(positional): + ctx = TracedTensorNode(name="missing_out", node_index=0) + x = ctx.create_input(np.array([2.0, 1.0, 3.0]), name="x") + out = np.zeros(3) + with pytest.raises(NotImplementedError, match="mutation"): + if positional: + np.spacing(x, out) + else: + np.spacing(x, out=out) + np.testing.assert_array_equal(out, np.zeros(3)) + + +def test_reject_overwrite_input_before_changing_storage(): + ctx = TracedTensorNode(name="overwrite", node_index=0) + original = np.array([3.0, 1.0, 2.0], dtype=np.float32) + x = ctx.create_input(original, name="x") + with pytest.raises(NotImplementedError, match="mutation"): + np.nanmedian(x, overwrite_input=True) + np.testing.assert_array_equal(original, [3.0, 1.0, 2.0]) + + +def test_reject_mutating_call_before_changing_input(): + ctx = TracedTensorNode(name="mutation", node_index=0) + original = np.array([np.nan, 1.0], dtype=np.float32) + x = ctx.create_input(original, name="x") + with pytest.raises(NotImplementedError, match="mutation"): + np.nan_to_num(x, copy=False) + np.testing.assert_array_equal(original, [np.nan, 1.0]) + + +def test_successful_capture_executes_native_numpy_once(monkeypatch): + calls = [] + + def native_sin(x): + calls.append(x) + return np.sin(x) + + monkeypatch.setitem( + aten_adapter.NUMPY_OVERRIDES, + native_sin, + aten_adapter._implementations()[np.sin], + ) + sample = np.array([1.0, 2.0], dtype=np.float32) + ctx = TracedTensorNode(name="single_native_call", node_index=0) + x = ctx.create_input(sample, name="x") + y = aten_adapter.capture_numpy_call(x, native_sin, (x,), {}) + assert len(calls) == 1 + assert np.shares_memory(calls[0], sample) + ctx.compile_trace({"y": y}) + _assert_numpy_output(ctx.m(torch.from_numpy(sample + 1)), np.sin(sample + 1)) + + +@pytest.mark.parametrize("implementation", [lambda x: 1, lambda x: (x, x)]) +def test_invalid_capture_outputs_retain_failure_dependency(implementation, monkeypatch): + monkeypatch.setitem(aten_adapter.NUMPY_OVERRIDES, np.sin, implementation) + ctx = TracedTensorNode(name="invalid_structure", node_index=0) + x = ctx.create_input(np.array([1.0, 2.0], dtype=np.float32), name="x") + y = np.sin(x) + np.testing.assert_array_equal(y.data, np.sin(x.data)) + with pytest.raises(NotImplementedError, match="declared output depends"): + ctx.compile_trace({"y": y}) + + +@pytest.mark.parametrize( + "fn", [lambda x: np.copyto(x, 0), lambda x: np.fill_diagonal(x, 0)] +) +def test_reject_none_returning_mutator_before_changing_storage(fn): + sample = np.ones((2, 2), dtype=np.float32) + ctx = TracedTensorNode(name="native_mutator", node_index=0) + x = ctx.create_input(sample, name="x") + with pytest.raises(NotImplementedError, match="mutat"): + fn(x) + np.testing.assert_array_equal(sample, np.ones((2, 2), dtype=np.float32)) + + +def test_eager_view_keeps_numpy_storage(): + ctx = TracedTensorNode(name="view", node_index=0) + x = ctx.create_input(np.arange(3, dtype=np.float32), name="x") + y = np.broadcast_to(x, (2, 3)) + assert np.shares_memory(x.data, y.data) + assert not y.flags.writeable + + if __name__ == "__main__": unittest.main() From 1b92b4187055a0f671c722624923a4ee76ebc7c1 Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Wed, 21 Oct 2026 12:22:28 -0700 Subject: [PATCH 19/21] test: organize bundle inspection tests Move the bundle regression target under tests/unit_tests and remove the redundant C++ fixture-loading check. Keep the inspection-only build instructions aligned with the test layout. Signed-off-by: Marc Alwan --- packages/leapp-runtime-rust/Cargo.toml | 4 +++ packages/leapp-runtime-rust/README.md | 2 ++ .../tests/{ => unit_tests}/bundle.rs | 27 +------------------ 3 files changed, 7 insertions(+), 26 deletions(-) rename packages/leapp-runtime-rust/tests/{ => unit_tests}/bundle.rs (95%) diff --git a/packages/leapp-runtime-rust/Cargo.toml b/packages/leapp-runtime-rust/Cargo.toml index aaa45100..9859579a 100644 --- a/packages/leapp-runtime-rust/Cargo.toml +++ b/packages/leapp-runtime-rust/Cargo.toml @@ -27,6 +27,10 @@ safetensors = "=0.6.2" [dev-dependencies] tempfile = "3" +[[test]] +name = "bundle" +path = "tests/unit_tests/bundle.rs" + [[test]] name = "binding_lifecycle" path = "bindings/common/tests/lifecycle.rs" diff --git a/packages/leapp-runtime-rust/README.md b/packages/leapp-runtime-rust/README.md index 2e4941e0..904e77a9 100644 --- a/packages/leapp-runtime-rust/README.md +++ b/packages/leapp-runtime-rust/README.md @@ -90,6 +90,8 @@ the complete workspace also requires a C++20 compiler and Python 3.10 or newer. See [Python bindings](bindings/python/README.md) and [C++ bindings](bindings/cxx/README.md) for installation and binding tests. +Core bundle tests live in `tests/unit_tests/`, with shared fixtures in +`tests/fixtures/`. Language-binding tests stay under their respective bindings. All inspection tests create opaque artifact bytes; no model generation is needed. ## Shared lifecycle diff --git a/packages/leapp-runtime-rust/tests/bundle.rs b/packages/leapp-runtime-rust/tests/unit_tests/bundle.rs similarity index 95% rename from packages/leapp-runtime-rust/tests/bundle.rs rename to packages/leapp-runtime-rust/tests/unit_tests/bundle.rs index 9e14dbdc..79412c6d 100644 --- a/packages/leapp-runtime-rust/tests/bundle.rs +++ b/packages/leapp-runtime-rust/tests/unit_tests/bundle.rs @@ -6,7 +6,7 @@ use safetensors::{tensor::TensorView, Dtype}; use std::{fs, path::PathBuf}; use tempfile::TempDir; -const YAML: &str = include_str!("fixtures/pipeline.yaml"); +const YAML: &str = include_str!("../fixtures/pipeline.yaml"); struct Fixture { directory: TempDir, @@ -508,31 +508,6 @@ fn yaml_and_artifacts_use_the_same_normalized_path() { assert_eq!(Bundle::open(&path).unwrap_err().code(), ErrorCode::NotFound); } -#[test] -fn cpp_bundle_fixtures_load() { - let f = Fixture::new(); - fs::create_dir(f.directory.path().join("artifacts")).unwrap(); - fs::write( - f.directory.path().join("artifacts/policy.pt2"), - b"model bytes\n", - ) - .unwrap(); - fs::write( - f.directory.path().join("artifacts/preprocess.onnx"), - b"model bytes\n", - ) - .unwrap(); - fs::write(&f.path, include_str!("fixtures/policy.yaml")).unwrap(); - let bundle = Bundle::open(&f.path).unwrap(); - assert_eq!(bundle.model_names(), ["policy", "preprocess"]); - assert_eq!(bundle.pipeline_inputs()[0].name(), "preprocess/observation"); - let text = - include_str!("fixtures/feedback.yaml").replace(" initial_values: state.safetensors\n", ""); - fs::write(&f.path, text).unwrap(); - let bundle = Bundle::open(&f.path).unwrap(); - assert_eq!(bundle.input_info("policy", "state").unwrap().shape, [2]); -} - #[test] fn uses_yaml12_scalar_resolution() { let f = Fixture::new(); From 47bb35bee5d3ba0f084ba5f968210b49b1377bc1 Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Wed, 21 Oct 2026 13:31:55 -0700 Subject: [PATCH 20/21] feat: expose the core Session API Export Session, construction options, result/error types, and backend-independent DLPack ownership. Compile the common binding wrapper against the real core API. Execution returns Unsupported in this contract revision and is supplied by the runtime integration. Signed-off-by: Marc Alwan --- packages/leapp-runtime-rust/Cargo.lock | 7 + packages/leapp-runtime-rust/Cargo.toml | 8 +- packages/leapp-runtime-rust/README.md | 68 +++- .../bindings/common/Cargo.toml | 11 + .../bindings/common/README.md | 65 ++-- .../bindings/common/src/lib.rs | 4 +- .../bindings/common/src/runtime.rs | 188 ---------- .../bindings/common/src/session.rs | 112 ++++++ .../bindings/common/tests/lifecycle.rs | 244 +++++-------- packages/leapp-runtime-rust/src/dlpack.rs | 218 ++++++++++++ .../leapp-runtime-rust/src/dlpack/README.md | 18 + packages/leapp-runtime-rust/src/dlpack/cpu.rs | 161 +++++++++ .../src/dlpack/generate_bindings.py | 93 +++++ .../src/dlpack/include/dlpack/dlpack.h | 332 ++++++++++++++++++ packages/leapp-runtime-rust/src/dlpack/sys.rs | 75 ++++ packages/leapp-runtime-rust/src/engines.rs | 25 ++ packages/leapp-runtime-rust/src/error.rs | 70 ++++ packages/leapp-runtime-rust/src/lib.rs | 30 +- packages/leapp-runtime-rust/src/session.rs | 106 ++++++ .../tests/unit_tests/dlpack.rs | 9 + .../tests/unit_tests/dlpack/abi.rs | 65 ++++ .../tests/unit_tests/dlpack/cpu.rs | 67 ++++ .../tests/unit_tests/dlpack/ownership.rs | 117 ++++++ 23 files changed, 1680 insertions(+), 413 deletions(-) create mode 100644 packages/leapp-runtime-rust/bindings/common/Cargo.toml delete mode 100644 packages/leapp-runtime-rust/bindings/common/src/runtime.rs create mode 100644 packages/leapp-runtime-rust/bindings/common/src/session.rs create mode 100644 packages/leapp-runtime-rust/src/dlpack.rs create mode 100644 packages/leapp-runtime-rust/src/dlpack/README.md create mode 100644 packages/leapp-runtime-rust/src/dlpack/cpu.rs create mode 100644 packages/leapp-runtime-rust/src/dlpack/generate_bindings.py create mode 100644 packages/leapp-runtime-rust/src/dlpack/include/dlpack/dlpack.h create mode 100644 packages/leapp-runtime-rust/src/dlpack/sys.rs create mode 100644 packages/leapp-runtime-rust/src/engines.rs create mode 100644 packages/leapp-runtime-rust/src/session.rs create mode 100644 packages/leapp-runtime-rust/tests/unit_tests/dlpack.rs create mode 100644 packages/leapp-runtime-rust/tests/unit_tests/dlpack/abi.rs create mode 100644 packages/leapp-runtime-rust/tests/unit_tests/dlpack/cpu.rs create mode 100644 packages/leapp-runtime-rust/tests/unit_tests/dlpack/ownership.rs diff --git a/packages/leapp-runtime-rust/Cargo.lock b/packages/leapp-runtime-rust/Cargo.lock index fd5be50c..6653df41 100644 --- a/packages/leapp-runtime-rust/Cargo.lock +++ b/packages/leapp-runtime-rust/Cargo.lock @@ -288,6 +288,13 @@ dependencies = [ "tempfile", ] +[[package]] +name = "leapp-runtime-bindings" +version = "0.1.0" +dependencies = [ + "leapp-runtime", +] + [[package]] name = "leapp-runtime-cxx" version = "0.1.0" diff --git a/packages/leapp-runtime-rust/Cargo.toml b/packages/leapp-runtime-rust/Cargo.toml index 9859579a..5209b5ec 100644 --- a/packages/leapp-runtime-rust/Cargo.toml +++ b/packages/leapp-runtime-rust/Cargo.toml @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 [workspace] -members = ["bindings/python", "bindings/cxx"] +members = ["bindings/common", "bindings/python", "bindings/cxx"] default-members = ["."] resolver = "2" @@ -12,7 +12,7 @@ version = "0.1.0" edition = "2021" license = "Apache-2.0" publish = false -description = "Backend-independent LEAPP bundle loading and inspection" +description = "LEAPP bundle inspection and core Session API" [dependencies] anyhow = "1" @@ -31,6 +31,10 @@ tempfile = "3" name = "bundle" path = "tests/unit_tests/bundle.rs" +[[test]] +name = "dlpack" +path = "tests/unit_tests/dlpack.rs" + [[test]] name = "binding_lifecycle" path = "bindings/common/tests/lifecycle.rs" diff --git a/packages/leapp-runtime-rust/README.md b/packages/leapp-runtime-rust/README.md index 904e77a9..52355cc7 100644 --- a/packages/leapp-runtime-rust/README.md +++ b/packages/leapp-runtime-rust/README.md @@ -3,12 +3,13 @@ SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All SPDX-License-Identifier: Apache-2.0 --> -# LEAPP Rust bundle API and shared lifecycle sources +# LEAPP Rust bundle and Session API This package loads and validates exported LEAPP bundles and exposes their public -inspection API in Rust, Python, and C++. It also provides shared runtime/session -lifecycle sources for ownership, error handling, and shutdown. Bundle loading -does not initialize an inference backend. +inspection API in Rust, Python, and C++. It also exports the core Rust Session API and compiles its shared binding +wrapper. This revision defines the execution contract: `Session::open()` returns +`ErrorCode::Unsupported` until the runtime integration supplies execution. +Bundle loading does not initialize an inference backend. The inspection and full-runtime Python distributions share the `leapp_runtime` import name and should be installed in separate environments. @@ -19,17 +20,19 @@ import name and should be installed in separate environments. | --- | --- | | Rust core | `Bundle`, YAML schema parsing, graph and artifact validation, SHA-256 checks, safetensors initial-state validation | | Public inspection types | `PortRef`, `PortSpec`, `DataType`, `Metadata`, `MetadataValue`, `BundleError`, `ErrorCode` | -| Common lifecycle sources | `RuntimeHandle`, `SessionHandle`, error-source preservation, and cleanup after poisoned locks; five focused regression tests | +| Core Session API | `Session`, `SessionOptions`, `OutputPlacement`, `RunResult`, device/memory types, and `RuntimeError` | +| Tensor API | Backend-independent DLPack descriptors and ownership types used by Session inputs and results | +| Common lifecycle crate | `SessionHandle`, independent closure, error-source preservation, and cleanup after poisoned locks; three focused regression tests | | Python | PyO3 bundle bindings, detached metadata and port descriptors, typed stubs and structured exceptions | | C++ | C++20 bundle facade, typed metadata, status and data-type headers, private cxx transport, installable CMake target | | Tests | Rust bundle regression suite and independent Python/C++ inspection and lifetime tests | -Concrete engine runtimes, sessions, results, tensor/DLPack code, graph planners, -and backend adapters are supplied by the full runtime packages. -Python and C++ share the core bundle implementation. The `bindings/common` -lifecycle source depends on concrete execution types, so it is not a separate -compiled crate in this inspection-only package. See the -[common lifecycle sources](bindings/common/README.md) for integration and testing. +Execution, feedback state, graph planning, native storage, and backend adapters +are supplied by the runtime integration. The Rust Session contract and the +`bindings/common` crate compile here against the same exported API types. +Python and C++ currently expose bundle inspection; their executable Session +bindings arrive with runtime integration. See the +[common lifecycle crate](bindings/common/README.md) for integration and testing. This package has no Triton, CUDA, ONNX Runtime, PyTorch, or NumPy dependency. The C++ binding currently supports Unix platforms (Linux and macOS), using native path bytes. @@ -84,7 +87,7 @@ cargo clippy --locked --workspace --all-targets -- -D warnings cargo fmt --all --check ``` -The default workspace member is the core crate, including the five common +The default workspace member is the core crate, including the three common lifecycle tests. These tests require no engine or language-binding SDK. Building the complete workspace also requires a C++20 compiler and Python 3.10 or newer. See @@ -94,10 +97,41 @@ Core bundle tests live in `tests/unit_tests/`, with shared fixtures in `tests/fixtures/`. Language-binding tests stay under their respective bindings. All inspection tests create opaque artifact bytes; no model generation is needed. +## Core Session contract + +`Session::open(&bundle, options)` accepts immutable construction options: + +```rust +use leapp_runtime::{ExecutionDevice, OutputPlacement, Session, SessionOptions}; + +let options = SessionOptions { + device: ExecutionDevice::Cpu, + output_placement: OutputPlacement::default(), +}; +# let bundle = leapp_runtime::Bundle::open("export/policy.yaml")?; +let session = Session::open(&bundle, options); +// This API revision returns a RuntimeError with ErrorCode::Unsupported. +# Ok::<(), leapp_runtime::BundleError>(()) +``` + +The public lifecycle is `open`, `info`, `is_open`, `new_session`, `run`, `reset`, +and `close`. `run` accepts named `dlpack::BorrowedTensor` inputs and returns a +`RunResult` retaining `SharedTensor` outputs. There is no separate public Runtime +owner, no per-call RunOptions, and no engine field in SessionOptions. Output +placement is configured at construction and inherited by every new session. + +The execution contract requires `new_session()` to share loaded models and +immutable options while starting from captured initial state, not current +feedback. `reset()` restores that initial state. Closing one session preserves +its peers and retained results. Execution is deliberately unavailable in this +revision; the runtime integration implements these operations without changing +the public signatures. There is no runtime selector or optional backend feature. + ## Shared lifecycle -`bindings/common/src/runtime.rs` preserves runtime error sources and completes -shutdown cleanup even when locks are poisoned. The lifecycle regression target -compiles this implementation against a small test-only runtime API, without -inference SDKs. See the [session reset contract](bindings/common/README.md#session-reset) -for state ownership and reset behavior. +`bindings/common` compiles its `SessionHandle` wrapper against the exported core +Session API. It provides serialization, error-source preservation, and independent +closure. Its existing three lifecycle tests use a small test-only session to +exercise waiting and poisoned-lock cleanup before native execution is available. +The wrapper itself is built against the real core crate by workspace checks. +See the [session reset contract](bindings/common/README.md#session-reset). diff --git a/packages/leapp-runtime-rust/bindings/common/Cargo.toml b/packages/leapp-runtime-rust/bindings/common/Cargo.toml new file mode 100644 index 00000000..7ee6bc40 --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/common/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "leapp-runtime-bindings" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +publish = false +# Lifecycle regressions run in the root test target. +autotests = false + +[dependencies] +leapp-runtime = { path = "../.." } diff --git a/packages/leapp-runtime-rust/bindings/common/README.md b/packages/leapp-runtime-rust/bindings/common/README.md index 07e6466c..dc9bf093 100644 --- a/packages/leapp-runtime-rust/bindings/common/README.md +++ b/packages/leapp-runtime-rust/bindings/common/README.md @@ -3,25 +3,29 @@ SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All SPDX-License-Identifier: Apache-2.0 --> -# Shared lifecycle sources - -`src/lib.rs` and `src/runtime.rs` provide session ownership, locking, and explicit -closure. Wrapped runtime errors expose their source chain, and close finishes -releasing every session and the runtime even when locks are poisoned. -Close reports the lock error after cleanup; subsequent -close calls succeed. Normal run/reset operations still fail on poisoned locks. - -This source depends on the concrete `Runtime`, `Session`, tensor, options, and -error types in `leapp_runtime`. Those execution types are not part of this -inspection-only package, so this directory is not a separate Cargo crate here. -The Python and C++ inspection bindings use the core bundle crate directly. -Runtime packages compile these sources with their execution types and features. - -The root Cargo test target includes `runtime.rs` against a small runtime API -defined only in the test binary. Five tests cover session ownership, close -invalidation and retained bundle info, waiting for an active run, error-source -preservation, and cleanup after poisoned locks. Execution behavior is tested in -the full runtime integration. +# Shared lifecycle crate + +`src/lib.rs` and `src/session.rs` provide session ownership, locking, and +explicit closure. Wrapped runtime errors expose their source chain, and close +releases that session even when its lock is poisoned. Other sessions keep their +engine references. +Close reports the lock error after cleanup; subsequent close calls succeed. +Normal run/reset operations still fail on poisoned locks. + +This crate compiles against the exported `Session`, `SessionOptions`, DLPack +input/result types, and errors in `leapp_runtime`. Execution is not implemented +in this API revision: opening a SessionHandle preserves the core Unsupported +error. The runtime integration supplies execution using these same signatures. +The Python and C++ bindings currently expose bundle inspection only. + +`SessionHandle::open()` accepts construction-time `SessionOptions`. `run()` +accepts named borrowed inputs without per-call options. `new_session()` creates +an independent session sharing the loaded engine and immutable options. + +The root Cargo test target includes `session.rs` against a small session API +defined only in the test binary. Three tests cover waiting for an active run, +error-source preservation, and cleanup after poisoned locks. Execution behavior +is covered by the full runtime package's core execution tests. ```sh cargo test --locked --test binding_lifecycle @@ -30,17 +34,18 @@ cargo test --locked --test binding_lifecycle ## Session reset `SessionHandle::reset()` delegates to the execution session. State belongs to -each session, so reset restores its feedback and default inputs to their captured -initial values. Those values are loaded when the runtime opens: exported -initial tensors are used when present, otherwise zeros are used for statically -shaped state. Dynamic state requires an exported initial tensor with a concrete -shape. +each session. `new_session()` shares its engine and construction options, but +starts from the captured initial state rather than current feedback. Reset +restores its feedback and default inputs to their captured initial values. +Those values are loaded when the session opens: +exported initial tensors are used when present, otherwise zeros are used for +statically shaped state. Dynamic state requires an exported initial tensor with +a concrete shape. Reset serializes with runs on the same session. It preserves other sessions and -retained results, and it does not reread files, reload models, reset -backend-internal mutable state, or flush shared allocation pools. The shared -wrapper rejects reset after explicit runtime closure. +retained results, and it does not reread files, reload models, or reset +backend-internal mutable state. The shared wrapper rejects reset after explicit +session closure. -This inspection-only package provides the reset wrapper as source; it does not -export an executable runtime/session API. Tensor-state restoration requires the -execution engine and is outside the lifecycle tests here. +The crate is compiled against the real public Session API. Tensor-state +restoration requires execution and is outside the lifecycle tests here. diff --git a/packages/leapp-runtime-rust/bindings/common/src/lib.rs b/packages/leapp-runtime-rust/bindings/common/src/lib.rs index 4eaec5f5..45f3a82a 100644 --- a/packages/leapp-runtime-rust/bindings/common/src/lib.rs +++ b/packages/leapp-runtime-rust/bindings/common/src/lib.rs @@ -3,5 +3,5 @@ //! Shared synchronous lifecycle for language bindings. #![forbid(unsafe_code)] -mod runtime; -pub use runtime::{Error, RuntimeHandle, SessionHandle}; +mod session; +pub use session::{Error, SessionHandle}; diff --git a/packages/leapp-runtime-rust/bindings/common/src/runtime.rs b/packages/leapp-runtime-rust/bindings/common/src/runtime.rs deleted file mode 100644 index 3a28982d..00000000 --- a/packages/leapp-runtime-rust/bindings/common/src/runtime.rs +++ /dev/null @@ -1,188 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use leapp_runtime::{ - Bundle, ErrorCode, RunOptions, RunResult, Runtime, RuntimeError, RuntimeOptions, Session, - TensorView, -}; -use std::{ - fmt, - sync::{Arc, Mutex, RwLock, Weak}, -}; - -/// Binding lifecycle failures preserve native execution diagnostics. -/// Wrapped runtime failures remain accessible through [`std::error::Error::source`]. -#[derive(Debug)] -pub enum Error { - Closed, - Lock, - Runtime(RuntimeError), -} - -impl Error { - /// Category shared with the core runtime. - pub fn code(&self) -> ErrorCode { - match self { - Self::Closed => ErrorCode::InvalidArgument, - Self::Lock => ErrorCode::Internal, - Self::Runtime(e) => e.code(), - } - } - - /// Operation that failed. - pub fn phase(&self) -> &str { - match self { - Self::Runtime(e) => e.phase(), - _ => "execute", - } - } -} - -impl fmt::Display for Error { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Closed => f.write_str("runtime is closed"), - Self::Lock => f.write_str("runtime lock poisoned"), - Self::Runtime(e) => e.fmt(f), - } - } -} - -impl std::error::Error for Error { - fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { - match self { - Self::Runtime(error) => Some(error), - Self::Closed | Self::Lock => None, - } - } -} - -type SessionSlot = Mutex>; - -struct State { - bundle: Bundle, - // Session operations hold a read guard; close waits with the write guard. - runtime: RwLock>, - sessions: Mutex>>, -} - -/// Shared runtime with explicit closure of all associated sessions. -#[derive(Clone)] -pub struct RuntimeHandle(Arc); - -/// Independent session that retains its runtime and serializes operations. -pub struct SessionHandle { - state: Arc, - slot: Arc, -} - -impl RuntimeHandle { - /// Load models using the existing core runtime. - pub fn open(bundle: &Bundle, options: RuntimeOptions) -> Result { - let runtime = Runtime::open(bundle, options).map_err(Error::Runtime)?; - Ok(Self(Arc::new(State { - bundle: bundle.clone(), - runtime: RwLock::new(Some(runtime)), - sessions: Mutex::new(Vec::new()), - }))) - } - - /// Snapshot remains available after close. - pub fn info(&self) -> &Bundle { - &self.0.bundle - } - - /// Check whether explicit close has released the runtime. - pub fn is_open(&self) -> Result { - Ok(self.0.runtime.read().map_err(|_| Error::Lock)?.is_some()) - } - - /// Create independently initialized feedback state. - pub fn create_session(&self) -> Result { - let guard = self.0.runtime.read().map_err(|_| Error::Lock)?; - let runtime = guard.as_ref().ok_or(Error::Closed)?; - let session = runtime.create_session().map_err(Error::Runtime)?; - let slot = Arc::new(Mutex::new(Some(session))); - let mut sessions = self.0.sessions.lock().map_err(|_| Error::Lock)?; - sessions.retain(|s| s.strong_count() != 0); - sessions.push(Arc::downgrade(&slot)); - Ok(SessionHandle { - state: self.0.clone(), - slot, - }) - } - - /// Wait for active native calls and close all sessions. - /// Poisoned locks are recovered only for cleanup, then reported as `Error::Lock`. - /// Once cleanup completes, repeated calls succeed. - pub fn close(&self) -> Result<(), Error> { - let mut poisoned = false; - let mut runtime = self.0.runtime.write().unwrap_or_else(|error| { - poisoned = true; - error.into_inner() - }); - let mut sessions = self.0.sessions.lock().unwrap_or_else(|error| { - poisoned = true; - error.into_inner() - }); - // The runtime write guard excludes active calls. Recover poisoned state - // only to drop it, and visit every session before reporting the failure. - for weak in sessions.drain(..) { - if let Some(session) = weak.upgrade() { - let mut slot = session.lock().unwrap_or_else(|error| { - poisoned = true; - error.into_inner() - }); - slot.take(); - session.clear_poison(); - } - } - runtime.take(); - // Empty, closed state is valid again. Clear poison while holding the guards. - self.0.sessions.clear_poison(); - self.0.runtime.clear_poison(); - if poisoned { - Err(Error::Lock) - } else { - Ok(()) - } - } -} - -impl SessionHandle { - /// Borrow inputs until synchronous execution releases them. - pub fn run( - &self, - inputs: &[(&str, TensorView<'_>)], - options: &RunOptions, - ) -> Result { - let guard = self.state.runtime.read().map_err(|_| Error::Lock)?; - if guard.is_none() { - return Err(Error::Closed); - } - self.slot - .lock() - .map_err(|_| Error::Lock)? - .as_mut() - .ok_or(Error::Closed)? - .run_with_options(inputs, options) - .map_err(Error::Runtime) - } - - /// Restore captured initial feedback and default inputs for this session. - /// Retained results are unaffected; files and models are not reloaded. - /// Backend-internal mutable state is not reset. - pub fn reset(&self) -> Result<(), Error> { - let guard = self.state.runtime.read().map_err(|_| Error::Lock)?; - if guard.is_none() { - return Err(Error::Closed); - } - self.slot - .lock() - .map_err(|_| Error::Lock)? - .as_mut() - .ok_or(Error::Closed)? - .reset() - .map_err(Error::Runtime) - } -} diff --git a/packages/leapp-runtime-rust/bindings/common/src/session.rs b/packages/leapp-runtime-rust/bindings/common/src/session.rs new file mode 100644 index 00000000..4ae13f0c --- /dev/null +++ b/packages/leapp-runtime-rust/bindings/common/src/session.rs @@ -0,0 +1,112 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use leapp_runtime::dlpack::BorrowedTensor; +use leapp_runtime::{Bundle, ErrorCode, RunResult, RuntimeError, Session, SessionOptions}; +use std::{fmt, sync::Mutex}; + +/// Binding lifecycle failures preserve native execution diagnostics. +/// Wrapped runtime failures remain accessible through [`std::error::Error::source`]. +#[derive(Debug)] +pub enum Error { + Lock, + Runtime(RuntimeError), +} + +impl Error { + /// Category shared with the core runtime. + pub fn code(&self) -> ErrorCode { + match self { + Self::Lock => ErrorCode::Internal, + Self::Runtime(e) => e.code(), + } + } + + /// Operation that failed. + pub fn phase(&self) -> &str { + match self { + Self::Runtime(e) => e.phase(), + _ => "execute", + } + } +} + +impl fmt::Display for Error { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Lock => f.write_str("session lock poisoned"), + Self::Runtime(e) => e.fmt(f), + } + } +} + +impl std::error::Error for Error { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Runtime(error) => Some(error), + Self::Lock => None, + } + } +} + +/// Serializes one session's operations and independent closure. +pub struct SessionHandle { + bundle: Bundle, + session: Mutex, +} +impl SessionHandle { + pub fn open(bundle: &Bundle, options: SessionOptions) -> Result { + let session = Session::open(bundle, options).map_err(Error::Runtime)?; + Ok(Self { + bundle: bundle.clone(), + session: Mutex::new(session), + }) + } + pub fn info(&self) -> &Bundle { + &self.bundle + } + pub fn is_open(&self) -> Result { + Ok(self.session.lock().map_err(|_| Error::Lock)?.is_open()) + } + pub fn new_session(&self) -> Result { + let session = self + .session + .lock() + .map_err(|_| Error::Lock)? + .new_session() + .map_err(Error::Runtime)?; + Ok(Self { + bundle: self.bundle.clone(), + session: Mutex::new(session), + }) + } + /// Wait for this session's current operation. Recover poison only for cleanup. + pub fn close(&self) -> Result<(), Error> { + let mut poisoned = false; + let mut session = self.session.lock().unwrap_or_else(|error| { + poisoned = true; + error.into_inner() + }); + session.close(); + self.session.clear_poison(); + if poisoned { + Err(Error::Lock) + } else { + Ok(()) + } + } + pub fn run(&self, inputs: &[(&str, BorrowedTensor<'_>)]) -> Result { + self.session + .lock() + .map_err(|_| Error::Lock)? + .run(inputs) + .map_err(Error::Runtime) + } + pub fn reset(&self) -> Result<(), Error> { + self.session + .lock() + .map_err(|_| Error::Lock)? + .reset() + .map_err(Error::Runtime) + } +} diff --git a/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs b/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs index 28cffcbb..a610050c 100644 --- a/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs +++ b/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs @@ -1,11 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -// Compile the shared lifecycle source against a test-only runtime API. -// The real dependency still supplies Bundle and ErrorCode; no production shim. +// Compile the binding lifecycle against a test-only session. extern crate leapp_runtime as bundle_api; extern crate self as leapp_runtime; - pub use bundle_api::{Bundle, ErrorCode}; use std::{ fmt, @@ -15,125 +13,102 @@ use std::{ }; mod lifecycle { - include!("../src/runtime.rs"); - + include!("../src/session.rs"); + pub fn set_wait( + session: &SessionHandle, + entered: std::sync::mpsc::Sender<()>, + resume: std::sync::mpsc::Receiver<()>, + ) { + session.session.lock().unwrap().wait = Some((entered, resume)); + } #[test] - fn close_recovers_poisoned_locks_and_releases_every_session() { - use std::panic::{catch_unwind, AssertUnwindSafe}; - - for poisoned in ["session", "registry", "runtime"] { - let (runtime, owner) = crate::runtime(); - let sessions: Vec<_> = (0..3).map(|_| runtime.create_session().unwrap()).collect(); - assert!(catch_unwind(AssertUnwindSafe(|| match poisoned { - "session" => { - // Cleanup must visit sessions on both sides of the poisoned one. - let _guard = sessions[1].slot.lock().unwrap(); - panic!("simulate session panic"); - } - "registry" => { - let _guard = runtime.0.sessions.lock().unwrap(); - panic!("simulate registry panic"); - } - "runtime" => { - let _guard = runtime.0.runtime.write().unwrap(); - panic!("simulate runtime panic"); - } - _ => unreachable!(), - })) - .is_err()); - - assert!(matches!(runtime.close(), Err(Error::Lock)), "{poisoned}"); - assert!(!runtime.is_open().unwrap()); - assert!(owner.upgrade().is_none()); - assert!(runtime.0.sessions.lock().unwrap().is_empty()); - for session in &sessions { - assert!(session.slot.lock().unwrap().is_none()); - assert!(matches!(session.reset(), Err(Error::Closed))); - assert!(matches!( - session.run(&[], &RunOptions::default()), - Err(Error::Closed) - )); - } - assert!(matches!(runtime.create_session(), Err(Error::Closed))); - runtime.close().unwrap(); - } + fn close_recovers_poison_without_closing_peers() { + let (session, owner) = crate::session(); + let peer = session.new_session().unwrap(); + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _guard = session.session.lock().unwrap(); + panic!("simulate session panic"); + })); + assert!(matches!(session.reset(), Err(Error::Lock))); + assert!(matches!(session.close(), Err(Error::Lock))); + assert!(!session.is_open().unwrap()); + assert!(session.info().model_names().is_empty()); + session.close().unwrap(); + peer.run(&[]).unwrap(); + peer.close().unwrap(); + assert!(owner.upgrade().is_none()); } } -use lifecycle::{Error, RuntimeHandle}; - -pub type RuntimeOptions = Arc<()>; -pub type TensorView<'a> = &'a [u8]; +use lifecycle::{Error, SessionHandle}; +pub mod dlpack { + pub type BorrowedTensor<'a> = &'a [u8]; +} pub type RunResult = (); - -#[derive(Default)] -pub struct RunOptions(Option<(mpsc::Sender<()>, mpsc::Receiver<()>)>); - -pub struct Runtime(Arc<()>); - -pub struct Session(Weak<()>); - +pub struct SessionOptions(Arc<()>); +pub struct Session { + owner: Option>, + wait: Option<(mpsc::Sender<()>, mpsc::Receiver<()>)>, +} #[derive(Debug)] pub struct RuntimeError(std::io::Error); - impl RuntimeError { pub fn code(&self) -> ErrorCode { - ErrorCode::Internal + ErrorCode::InvalidArgument } - pub fn phase(&self) -> &str { "execute" } } - impl fmt::Display for RuntimeError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "test runtime error: {}", self.0) + self.0.fmt(f) } } - impl std::error::Error for RuntimeError { fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { Some(&self.0) } } - -impl Runtime { - pub fn open(_: &Bundle, owner: RuntimeOptions) -> Result { - Ok(Self(owner)) +impl Session { + pub fn open(_: &Bundle, options: SessionOptions) -> Result { + Ok(Self { + owner: Some(options.0), + wait: None, + }) } - - pub fn create_session(&self) -> Result { - Ok(Session(Arc::downgrade(&self.0))) + pub fn is_open(&self) -> bool { + self.owner.is_some() } -} - -impl Session { - pub fn run_with_options( - &mut self, - _: &[(&str, TensorView<'_>)], - options: &RunOptions, - ) -> Result { - assert!( - self.0.upgrade().is_some(), - "runtime dropped during session lifetime" - ); - if let Some((entered, resume)) = &options.0 { - entered.send(()).unwrap(); - resume.recv_timeout(Duration::from_secs(5)).unwrap(); + pub fn new_session(&self) -> Result { + self.reset_check()?; + Ok(Self { + owner: self.owner.clone(), + wait: None, + }) + } + pub fn close(&mut self) { + self.owner.take(); + } + fn reset_check(&self) -> Result<(), RuntimeError> { + if self.is_open() { + Ok(()) + } else { + Err(RuntimeError(std::io::Error::other("session is closed"))) } - Ok(()) } - pub fn reset(&mut self) -> Result<(), RuntimeError> { - assert!( - self.0.upgrade().is_some(), - "runtime dropped during session lifetime" - ); + self.reset_check() + } + pub fn run(&mut self, _: &[(&str, dlpack::BorrowedTensor<'_>)]) -> Result<(), RuntimeError> { + self.reset_check()?; + if let Some((entered, resume)) = &self.wait { + entered.send(()).unwrap(); + resume.recv_timeout(Duration::from_secs(5)).unwrap(); + } Ok(()) } } - -fn runtime() -> (RuntimeHandle, Weak<()>) { +fn session() -> (SessionHandle, Weak<()>) { let directory = tempfile::tempdir().unwrap(); let path = directory.path().join("empty.yaml"); std::fs::write( @@ -144,92 +119,40 @@ fn runtime() -> (RuntimeHandle, Weak<()>) { let owner = Arc::new(()); let weak = Arc::downgrade(&owner); ( - RuntimeHandle::open(&Bundle::open(path).unwrap(), owner).unwrap(), + SessionHandle::open(&Bundle::open(path).unwrap(), SessionOptions(owner)).unwrap(), weak, ) } - #[test] -fn runtime_errors_preserve_the_source_chain_and_diagnostics() { - use std::error::Error as _; - - let runtime = RuntimeError(std::io::Error::other("test backend error")); - let message = runtime.to_string(); - let error = Error::Runtime(runtime); - let source = error - .source() +fn errors_preserve_the_source_chain_and_diagnostics() { + let error = Error::Runtime(RuntimeError(std::io::Error::other("backend diagnostic"))); + assert!(error.to_string().contains("backend diagnostic")); + assert_eq!(error.phase(), "execute"); + assert_eq!(error.code(), ErrorCode::InvalidArgument); + assert!(std::error::Error::source(&error) .unwrap() - .downcast_ref::() - .unwrap(); - assert_eq!(error.code(), source.code()); - assert_eq!(error.phase(), source.phase()); - assert_eq!(error.to_string(), message); - assert_eq!(source.source().unwrap().to_string(), "test backend error"); - assert!(Error::Closed.source().is_none()); - assert!(Error::Lock.source().is_none()); -} - -#[test] -fn sessions_keep_runtime_alive_after_the_last_runtime_handle_is_dropped() { - let (runtime, owner) = runtime(); - let session = runtime.create_session().unwrap(); - drop(runtime); - assert!(owner.upgrade().is_some()); - session.run(&[], &RunOptions::default()).unwrap(); - session.reset().unwrap(); - drop(session); - assert!(owner.upgrade().is_none()); -} - -#[test] -fn close_invalidates_all_handles_and_retains_bundle_info() { - let (runtime, owner) = runtime(); - let clone = runtime.clone(); - let path = runtime.info().yaml_path().to_owned(); - let sessions = [ - runtime.create_session().unwrap(), - runtime.create_session().unwrap(), - ]; - clone.close().unwrap(); - assert!(!runtime.is_open().unwrap()); - assert!(owner.upgrade().is_none()); - assert_eq!(runtime.info().yaml_path(), path); - assert!(matches!(runtime.create_session(), Err(Error::Closed))); - for session in sessions { - assert!(matches!( - session.run(&[], &RunOptions::default()), - Err(Error::Closed) - )); - let error = session.reset().unwrap_err(); - assert!(matches!(error, Error::Closed)); - assert_eq!(error.code(), ErrorCode::InvalidArgument); - assert_eq!(error.phase(), "execute"); - } - runtime.close().unwrap(); + .source() + .is_some()); } - #[test] -fn close_waits_until_an_active_run_finishes() { - let (runtime, owner) = runtime(); - let session = Arc::new(runtime.create_session().unwrap()); +fn close_waits_until_this_sessions_run_finishes() { + let (session, owner) = session(); let (entered, running) = mpsc::channel(); let (resume, paused) = mpsc::channel(); + lifecycle::set_wait(&session, entered, paused); + let session = Arc::new(session); let run = { let session = session.clone(); - thread::spawn(move || { - session - .run(&[], &RunOptions(Some((entered, paused)))) - .unwrap() - }) + thread::spawn(move || session.run(&[]).unwrap()) }; running.recv_timeout(Duration::from_secs(5)).unwrap(); let (started, closing) = mpsc::channel(); let (finished, closed) = mpsc::channel(); let close = { - let runtime = runtime.clone(); + let session = session.clone(); thread::spawn(move || { started.send(()).unwrap(); - runtime.close().unwrap(); + session.close().unwrap(); finished.send(()).unwrap(); }) }; @@ -243,5 +166,4 @@ fn close_waits_until_an_active_run_finishes() { run.join().unwrap(); close.join().unwrap(); assert!(owner.upgrade().is_none()); - assert!(matches!(session.reset(), Err(Error::Closed))); } diff --git a/packages/leapp-runtime-rust/src/dlpack.rs b/packages/leapp-runtime-rust/src/dlpack.rs new file mode 100644 index 00000000..2d5de980 --- /dev/null +++ b/packages/leapp-runtime-rust/src/dlpack.rs @@ -0,0 +1,218 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Backend-independent DLPack ABI, lifetime guards, and borrowed CPU-slice access. +#![deny(unsafe_op_in_unsafe_fn)] + +#[path = "dlpack/cpu.rs"] +mod cpu; +#[path = "dlpack/sys.rs"] +pub mod sys; +pub use cpu::CpuElement; +use std::{ + marker::PhantomData, + mem::ManuallyDrop, + ptr::NonNull, + sync::{Arc, Mutex}, +}; +pub use sys::*; + +/// An exclusively owned producer lease. Foreign owners remain on their importing thread. +pub struct ManagedTensor(ManagedTensorPointer); +enum ManagedTensorPointer { + Legacy(NonNull), + Versioned(NonNull), +} + +impl ManagedTensor { + /// Adopt a producer lease, including responsibility for its deleter. + /// # Safety + /// `ptr` must be a live, exclusively owned DLPack managed descriptor. Its metadata + /// and allocation must remain valid until its deleter is called. The producer's + /// deleter must not unwind. No other owner may release this lease. + pub unsafe fn from_legacy(ptr: *mut DLManagedTensor) -> Option { + NonNull::new(ptr).map(|p| Self(ManagedTensorPointer::Legacy(p))) + } + + /// Adopt a versioned lease. The same contract as `from_legacy` applies. + /// # Safety + /// The descriptor must satisfy the DLPack versioned ABI, including the stable + /// version/context/deleter prefix when its major version is unsupported. + pub unsafe fn from_versioned(ptr: *mut DLManagedTensorVersioned) -> Option { + NonNull::new(ptr).map(|p| Self(ManagedTensorPointer::Versioned(p))) + } + + /// Access metadata without dereferencing the allocation or shape/stride pointers. + pub fn descriptor(&self) -> Result<&DLTensor, &'static str> { + // SAFETY: adoption guarantees a live descriptor; inspect version first. + unsafe { + match &self.0 { + ManagedTensorPointer::Legacy(p) => Ok(&p.as_ref().dl_tensor), + ManagedTensorPointer::Versioned(p) => { + if (*p.as_ptr()).version.major != DLPACK_MAJOR_VERSION { + return Err("unsupported DLPack major version"); + } + Ok(&p.as_ref().dl_tensor) + } + } + } + } + + /// Return the raw lease, transferring responsibility for invoking its deleter. + pub fn into_raw(self) -> (*mut std::ffi::c_void, bool) { + match &ManuallyDrop::new(self).0 { + ManagedTensorPointer::Legacy(p) => (p.as_ptr().cast(), false), + ManagedTensorPointer::Versioned(p) => (p.as_ptr().cast(), true), + } + } +} + +impl Drop for ManagedTensor { + fn drop(&mut self) { + // SAFETY: this guard exclusively owns the producer lease, even on rejection. + unsafe { + match &self.0 { + ManagedTensorPointer::Legacy(p) => { + if let Some(f) = (*p.as_ptr()).deleter { + f(p.as_ptr()); + } + } + ManagedTensorPointer::Versioned(p) => { + if let Some(f) = (*p.as_ptr()).deleter { + f(p.as_ptr()); + } + } + } + } + } +} + +/// A descriptor borrowed together with its metadata and ready, immutable storage. +#[derive(Clone, Copy)] +pub struct BorrowedTensor<'a> { + descriptor: DLTensor, + lifetime: PhantomData<&'a DLTensor>, +} +impl<'a> BorrowedTensor<'a> { + /// # Safety + /// All descriptor pointers must be valid for the declared sizes and lifetime. + /// Storage must be ready on every execution stream, readable on the declared + /// device, and unmodified for this borrow, including concurrent consumers. + pub unsafe fn new(descriptor: &'a DLTensor) -> Self { + Self { + descriptor: *descriptor, + lifetime: PhantomData, + } + } + pub fn descriptor(&self) -> &DLTensor { + &self.descriptor + } +} +// SAFETY: constructors require or retain immutable, ready storage for the borrow; +// it never transfers or invokes the producer's deleter. +unsafe impl Send for BorrowedTensor<'_> {} +unsafe impl Sync for BorrowedTensor<'_> {} + +/// Shared ownership for engine-produced leases with thread-independent cleanup. +#[derive(Clone)] +pub struct SharedTensor(Arc); +/// Synchronously copy the tensor's contiguous bytes into caller-owned CPU storage. +pub type CopyToCpu = dyn Fn(&mut [u8]) -> Result<(), String> + Send + Sync; + +struct ThreadSafeTensorOwner { + owner: ManagedTensor, + exposed: Mutex, + copy_to_cpu: Option>, + on_expose: Option>, +} +// SAFETY: SharedTensor::new requires thread-independent metadata lifetime and cleanup. +unsafe impl Send for ThreadSafeTensorOwner {} +unsafe impl Sync for ThreadSafeTensorOwner {} +impl SharedTensor { + /// # Safety + /// The producer must support metadata access and final deletion on any thread. + /// The descriptor must use a supported ABI and remain unchanged. Its allocation + /// must permit writable exports. Storage must be initialized and ready for + /// reads, with no writes except through this owner's exposed descriptors or + /// exports. Internal unsafe users must not race with CPU copies. This is not + /// automatically true for foreign producers; imported leases stay `ManagedTensor`. + /// `on_expose`, when present, runs before exposing pointers through descriptor + /// access or export. It must be idempotent and must not modify the descriptor. + /// `copy_to_cpu`, when supplied, must copy exactly the descriptor's contiguous + /// bytes, complete all transfer work before returning, and retain its source. + /// It must not expose the allocation or re-enter this owner. Without it, only + /// host-accessible storage supports safe copies. + pub unsafe fn new( + owner: ManagedTensor, + on_expose: Option>, + copy_to_cpu: Option>, + ) -> Result { + owner.descriptor()?; + Ok(Self(Arc::new(ThreadSafeTensorOwner { + owner, + exposed: Mutex::new(false), + copy_to_cpu, + on_expose, + }))) + } + + /// Expose the descriptor, notifying the producer before releasing any pointers. + pub fn descriptor(&self) -> &DLTensor { + // Serialize the first exposure against safe CPU reads on every clone. + *self.0.exposed.lock().unwrap_or_else(|e| e.into_inner()) = true; + if let Some(notify) = &self.0.on_expose { + notify(); + } + self.0.owner.descriptor().expect("validated ABI") + } + + /// Borrow storage for a bounded operation without exposing it externally. + /// # Safety + /// Storage must be ready and unmodified for the borrow. No pointers may escape + /// it, and all work using them must finish before the borrow ends. + pub unsafe fn borrow(&self) -> BorrowedTensor<'_> { + // SAFETY: the caller supplies readiness, immutability, and completion; + // this owner retains the descriptor, metadata, and allocation. + unsafe { BorrowedTensor::new(self.0.owner.descriptor().expect("validated ABI")) } + } + + /// Produce a fresh lease retaining this allocation. The consumer owns this lease. + /// Flags deliberately do not claim a private copy: multiple exports may alias. + pub fn export(&self, versioned: bool) -> ManagedTensor { + let descriptor = *self.descriptor(); + let context = Box::into_raw(Box::new(self.clone())).cast(); + if versioned { + let p = Box::into_raw(Box::new(DLManagedTensorVersioned { + version: DLPackVersion { major: 1, minor: 0 }, + manager_ctx: context, + deleter: Some(delete_versioned), + flags: 0, + dl_tensor: descriptor, + })); + // SAFETY: this new descriptor owns the context retaining its metadata/storage. + unsafe { ManagedTensor::from_versioned(p).unwrap() } + } else { + let p = Box::into_raw(Box::new(DLManagedTensor { + dl_tensor: descriptor, + manager_ctx: context, + deleter: Some(delete_legacy), + })); + // SAFETY: as above, using the legacy layout. + unsafe { ManagedTensor::from_legacy(p).unwrap() } + } + } +} +unsafe extern "C" fn delete_legacy(p: *mut DLManagedTensor) { + // SAFETY: export allocated both boxes; this callback is called exactly once. + unsafe { + let p = Box::from_raw(p); + drop(Box::from_raw(p.manager_ctx.cast::())); + } +} +unsafe extern "C" fn delete_versioned(p: *mut DLManagedTensorVersioned) { + // SAFETY: export allocated both boxes; this callback is called exactly once. + unsafe { + let p = Box::from_raw(p); + drop(Box::from_raw(p.manager_ctx.cast::())); + } +} diff --git a/packages/leapp-runtime-rust/src/dlpack/README.md b/packages/leapp-runtime-rust/src/dlpack/README.md new file mode 100644 index 00000000..d81275c8 --- /dev/null +++ b/packages/leapp-runtime-rust/src/dlpack/README.md @@ -0,0 +1,18 @@ + + +# DLPack API types + +This module supplies the backend-independent tensor descriptors and ownership +contracts used by `Session::run` and `RunResult`. It has no native SDK dependency. + +`BorrowedTensor` retains a descriptor borrow for synchronous execution. +`ManagedTensor` owns a producer lease and invokes its deleter exactly once. +`SharedTensor` retains a thread-safe allocation and can export managed leases. +The CPU helpers support checked typed borrows and copies of host-accessible data. + +The vendored DLPack header, generated Rust ABI, and generator are kept together. +Existing tests cover ABI layout, CPU views, and lease ownership. Native allocation, +CUDA producer handoff, and Python/C++ protocol integration arrive with execution. diff --git a/packages/leapp-runtime-rust/src/dlpack/cpu.rs b/packages/leapp-runtime-rust/src/dlpack/cpu.rs new file mode 100644 index 00000000..f2660eff --- /dev/null +++ b/packages/leapp-runtime-rust/src/dlpack/cpu.rs @@ -0,0 +1,161 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use std::mem::{align_of, size_of}; + +mod private { + pub trait Sealed {} +} + +/// Rust numeric primitives with a DLPack dtype and no invalid bit patterns. +/// Sealed to prevent incorrect dtype declarations. Half floats and booleans use +/// the raw descriptor API. +pub trait CpuElement: private::Sealed + Copy + Default + Send + Sync { + const DTYPE: DLDataType; +} + +macro_rules! elements { + ($($ty:ty => $code:expr),* $(,)?) => {$ ( + impl private::Sealed for $ty {} + impl CpuElement for $ty { + const DTYPE: DLDataType = DLDataType { + code: $code as u8, bits: (size_of::() * 8) as u8, lanes: 1, + }; + } + )*}; +} +elements!( + u8 => kDLUInt, u16 => kDLUInt, u32 => kDLUInt, u64 => kDLUInt, + i8 => kDLInt, i16 => kDLInt, i32 => kDLInt, i64 => kDLInt, + f32 => kDLFloat, f64 => kDLFloat, +); + +fn element_count(shape: &[i64]) -> Result { + shape.iter().try_fold(1usize, |count, &dimension| { + count + .checked_mul(usize::try_from(dimension).map_err(|_| "negative dimension")?) + .ok_or("tensor size overflow") + }) +} + +impl<'a> BorrowedTensor<'a> { + /// Borrow contiguous CPU storage without copying. The values and shape remain + /// immutably borrowed for the tensor's lifetime. An empty shape is a scalar. + pub fn cpu(values: &'a [T], shape: &'a [i64]) -> Result { + let ndim = i32::try_from(shape.len()).map_err(|_| "too many dimensions")?; + if element_count(shape)? != values.len() { + return Err("shape does not match slice length"); + } + Ok(Self { + descriptor: DLTensor { + data: values.as_ptr().cast_mut().cast(), + device: DLDevice { + device_type: kDLCPU as i32, + device_id: 0, + }, + ndim, + dtype: T::DTYPE, + shape: shape.as_ptr().cast_mut(), + strides: std::ptr::null_mut(), + byte_offset: 0, + }, + lifetime: PhantomData, + }) + } + + /// Read a contiguous CPU tensor with the requested primitive type. + /// The borrow already guarantees ready, immutable storage and valid pointers. + pub fn as_cpu_slice(self) -> Result<&'a [T], &'static str> { + let t = self.descriptor(); + if t.device.device_type != kDLCPU as i32 && t.device.device_type != kDLCUDAHost as i32 { + return Err("tensor is not in CPU memory"); + } + let (data, count) = self.typed_layout::()?; + if count == 0 { + return Ok(&[]); + } + // SAFETY: this borrow guarantees initialized, immutable CPU storage for 'a; + // typed_layout checks dtype, contiguity, alignment, and length. + Ok(unsafe { std::slice::from_raw_parts(data, count) }) + } + + fn typed_layout(self) -> Result<(*const T, usize), &'static str> { + let t = self.descriptor(); + if (t.dtype.code, t.dtype.bits, t.dtype.lanes) != (T::DTYPE.code, T::DTYPE.bits, 1) { + return Err("tensor dtype does not match slice element type"); + } + let ndim = usize::try_from(t.ndim).map_err(|_| "negative dimension count")?; + if ndim > isize::MAX as usize / size_of::() || (ndim != 0 && t.shape.is_null()) { + return Err("invalid shape descriptor"); + } + // SAFETY: the tensor borrow guarantees valid, immutable metadata. + let shape = if ndim == 0 { + &[] + } else { + unsafe { std::slice::from_raw_parts(t.shape, ndim) } + }; + let count = element_count(shape)?; + let bytes = count + .checked_mul(size_of::()) + .filter(|&n| n <= isize::MAX as usize) + .ok_or("tensor byte size overflow")?; + if count == 0 { + return Ok((NonNull::::dangling().as_ptr(), 0)); + } + if !t.strides.is_null() { + let mut expected = 1i64; + for i in (0..ndim).rev() { + // SAFETY: the tensor borrow guarantees ndim valid stride entries. + if shape[i] > 1 && unsafe { *t.strides.add(i) } != expected { + return Err("tensor is not C-contiguous"); + } + expected = expected.checked_mul(shape[i]).ok_or("stride overflow")?; + } + } + let offset = usize::try_from(t.byte_offset).map_err(|_| "byte offset overflow")?; + let address = (t.data as usize) + .checked_add(offset) + .ok_or("tensor address overflow")?; + address + .checked_add(bytes) + .ok_or("tensor address overflow")?; + if t.data.is_null() || address % align_of::() != 0 { + return Err("null or unaligned tensor storage"); + } + Ok((t.data.cast::().wrapping_add(offset).cast(), count)) + } +} + +impl SharedTensor { + /// Copy a contiguous output into an owned Rust vector, including device + /// storage when the engine supplies a synchronous transfer callback. + /// + /// Rejects storage previously exposed through `descriptor()` or `export()` + /// on any clone: a foreign consumer could still be writing it. Copies and + /// first exposure are synchronized; copying does not prevent pool reuse. + pub fn to_cpu_vec(&self) -> Result, String> { + let exposed = self.0.exposed.lock().unwrap_or_else(|e| e.into_inner()); + if *exposed { + return Err("tensor storage has been exposed to a writable consumer".into()); + } + // SAFETY: construction guarantees ready storage; the lock excludes + // first exposure until the copy completes, and no pointers escape. + let borrowed = unsafe { BorrowedTensor::new(self.0.owner.descriptor()?) }; + let values = if let Some(copy) = &self.0.copy_to_cpu { + let (_, count) = borrowed.typed_layout::()?; + let mut values = vec![T::default(); count]; + // SAFETY: CpuElement is sealed to padding-free primitives accepting + // every bit pattern. The callback cannot retain this mutable borrow. + let bytes = unsafe { + std::slice::from_raw_parts_mut(values.as_mut_ptr().cast(), count * size_of::()) + }; + copy(bytes)?; + values + } else { + borrowed.as_cpu_slice::()?.to_vec() + }; + drop(exposed); + Ok(values) + } +} diff --git a/packages/leapp-runtime-rust/src/dlpack/generate_bindings.py b/packages/leapp-runtime-rust/src/dlpack/generate_bindings.py new file mode 100644 index 00000000..aa6a1520 --- /dev/null +++ b/packages/leapp-runtime-rust/src/dlpack/generate_bindings.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Generate the pinned DLPack ABI with clang; normal builds need no generator.""" + +import json +from pathlib import Path +import re +import subprocess + +root = Path(__file__).resolve().parent +header = root / "include/dlpack/dlpack.h" +source = header.read_text() +ast = json.loads( + subprocess.check_output( + ["clang", "-x", "c", "-Xclang", "-ast-dump=json", "-fsyntax-only", str(header)] + ) +) +records = {n["id"]: n for n in ast["inner"] if n["kind"] == "RecordDecl"} +types = { + "int32_t": "i32", + "int64_t": "i64", + "uint8_t": "u8", + "uint16_t": "u16", + "uint32_t": "u32", + "uint64_t": "u64", + "void": "std::ffi::c_void", +} +selected = { + "DLPackVersion", + "DLDevice", + "DLDataType", + "DLTensor", + "DLManagedTensor", + "DLManagedTensorVersioned", +} + + +def rust_type(value): + if value.endswith(" *"): + return "*mut " + rust_type(value[:-2]) + value = value.removeprefix("struct ") + return types.get(value, value) + + +out = [ + "// Generated from upstream DLPack v1.0 by generate_bindings.py. Do not edit.", + "#![allow(non_camel_case_types, non_upper_case_globals, non_snake_case)]", + "pub type DLDeviceType = i32;", +] +for suffix in ["MAJOR_VERSION", "MINOR_VERSION"]: + name = f"DLPACK_{suffix}" + value = re.search(rf"#define {name} (\d+)", source)[1] + out.append(f"pub const {name}: u32 = {value};") +for suffix in ["READ_ONLY", "IS_COPIED"]: + name = f"DLPACK_FLAG_BITMASK_{suffix}" + shift = re.search(rf"#define {name} \(1UL << (\d+)UL\)", source)[1] + out.append(f"pub const {name}: u64 = {1 << int(shift)};") + +for node in ast["inner"]: + if node["kind"] == "EnumDecl": + for item in node.get("inner", []): + if item["kind"] != "EnumConstantDecl": + continue + value_node = item["inner"][0] + while "value" not in value_node: + value_node = value_node["inner"][0] + out.append(f'pub const {item["name"]}: u32 = {value_node["value"]};') + if node["kind"] not in {"TypedefDecl", "RecordDecl"}: + continue + if node.get("name") not in selected: + continue + # In v1.0 the versioned descriptor is a named struct without a C typedef. + if node["kind"] == "RecordDecl" and node["name"] != "DLManagedTensorVersioned": + continue + name = node["name"] + if node["kind"] == "RecordDecl": + record = node + else: + declaration = next(n for n in node["inner"] if n["kind"] == "ElaboratedType") + record = records[declaration["ownedTagDecl"]["id"]] + out.extend(["#[repr(C)]", "#[derive(Clone, Copy, Debug)]", f"pub struct {name} {{"]) + for field in record["inner"]: + if field["kind"] != "FieldDecl": + continue + value = ( + f'Option' + if field["name"] == "deleter" + else rust_type(field["type"]["qualType"]) + ) + out.append(f' pub {field["name"]}: {value},') + out.append("}") +(root / "sys.rs").write_text("\n".join(out) + "\n") +subprocess.run(["rustfmt", str(root / "sys.rs")], check=True) diff --git a/packages/leapp-runtime-rust/src/dlpack/include/dlpack/dlpack.h b/packages/leapp-runtime-rust/src/dlpack/include/dlpack/dlpack.h new file mode 100644 index 00000000..bcb77949 --- /dev/null +++ b/packages/leapp-runtime-rust/src/dlpack/include/dlpack/dlpack.h @@ -0,0 +1,332 @@ +/*! + * Copyright (c) 2017 by Contributors + * \file dlpack.h + * \brief The common header of DLPack. + */ +#ifndef DLPACK_DLPACK_H_ +#define DLPACK_DLPACK_H_ + +/** + * \brief Compatibility with C++ + */ +#ifdef __cplusplus +#define DLPACK_EXTERN_C extern "C" +#else +#define DLPACK_EXTERN_C +#endif + +/*! \brief The current major version of dlpack */ +#define DLPACK_MAJOR_VERSION 1 + +/*! \brief The current minor version of dlpack */ +#define DLPACK_MINOR_VERSION 0 + +/*! \brief DLPACK_DLL prefix for windows */ +#ifdef _WIN32 +#ifdef DLPACK_EXPORTS +#define DLPACK_DLL __declspec(dllexport) +#else +#define DLPACK_DLL __declspec(dllimport) +#endif +#else +#define DLPACK_DLL +#endif + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/*! + * \brief The DLPack version. + * + * A change in major version indicates that we have changed the + * data layout of the ABI - DLManagedTensorVersioned. + * + * A change in minor version indicates that we have added new + * code, such as a new device type, but the ABI is kept the same. + * + * If an obtained DLPack tensor has a major version that disagrees + * with the version number specified in this header file + * (i.e. major != DLPACK_MAJOR_VERSION), the consumer must call the deleter + * (and it is safe to do so). It is not safe to access any other fields + * as the memory layout will have changed. + * + * In the case of a minor version mismatch, the tensor can be safely used as + * long as the consumer knows how to interpret all fields. Minor version + * updates indicate the addition of enumeration values. + */ +typedef struct { + /*! \brief DLPack major version. */ + uint32_t major; + /*! \brief DLPack minor version. */ + uint32_t minor; +} DLPackVersion; + +/*! + * \brief The device type in DLDevice. + */ +#ifdef __cplusplus +typedef enum : int32_t { +#else +typedef enum { +#endif + /*! \brief CPU device */ + kDLCPU = 1, + /*! \brief CUDA GPU device */ + kDLCUDA = 2, + /*! + * \brief Pinned CUDA CPU memory by cudaMallocHost + */ + kDLCUDAHost = 3, + /*! \brief OpenCL devices. */ + kDLOpenCL = 4, + /*! \brief Vulkan buffer for next generation graphics. */ + kDLVulkan = 7, + /*! \brief Metal for Apple GPU. */ + kDLMetal = 8, + /*! \brief Verilog simulator buffer */ + kDLVPI = 9, + /*! \brief ROCm GPUs for AMD GPUs */ + kDLROCM = 10, + /*! + * \brief Pinned ROCm CPU memory allocated by hipMallocHost + */ + kDLROCMHost = 11, + /*! + * \brief Reserved extension device type, + * used for quickly test extension device + * The semantics can differ depending on the implementation. + */ + kDLExtDev = 12, + /*! + * \brief CUDA managed/unified memory allocated by cudaMallocManaged + */ + kDLCUDAManaged = 13, + /*! + * \brief Unified shared memory allocated on a oneAPI non-partititioned + * device. Call to oneAPI runtime is required to determine the device + * type, the USM allocation type and the sycl context it is bound to. + * + */ + kDLOneAPI = 14, + /*! \brief GPU support for next generation WebGPU standard. */ + kDLWebGPU = 15, + /*! \brief Qualcomm Hexagon DSP */ + kDLHexagon = 16, + /*! \brief Microsoft MAIA devices */ + kDLMAIA = 17, +} DLDeviceType; + +/*! + * \brief A Device for Tensor and operator. + */ +typedef struct { + /*! \brief The device type used in the device. */ + DLDeviceType device_type; + /*! + * \brief The device index. + * For vanilla CPU memory, pinned memory, or managed memory, this is set to 0. + */ + int32_t device_id; +} DLDevice; + +/*! + * \brief The type code options DLDataType. + */ +typedef enum { + /*! \brief signed integer */ + kDLInt = 0U, + /*! \brief unsigned integer */ + kDLUInt = 1U, + /*! \brief IEEE floating point */ + kDLFloat = 2U, + /*! + * \brief Opaque handle type, reserved for testing purposes. + * Frameworks need to agree on the handle data type for the exchange to be well-defined. + */ + kDLOpaqueHandle = 3U, + /*! \brief bfloat16 */ + kDLBfloat = 4U, + /*! + * \brief complex number + * (C/C++/Python layout: compact struct per complex number) + */ + kDLComplex = 5U, + /*! \brief boolean */ + kDLBool = 6U, +} DLDataTypeCode; + +/*! + * \brief The data type the tensor can hold. The data type is assumed to follow the + * native endian-ness. An explicit error message should be raised when attempting to + * export an array with non-native endianness + * + * Examples + * - float: type_code = 2, bits = 32, lanes = 1 + * - float4(vectorized 4 float): type_code = 2, bits = 32, lanes = 4 + * - int8: type_code = 0, bits = 8, lanes = 1 + * - std::complex: type_code = 5, bits = 64, lanes = 1 + * - bool: type_code = 6, bits = 8, lanes = 1 (as per common array library convention, the underlying storage size of bool is 8 bits) + */ +typedef struct { + /*! + * \brief Type code of base types. + * We keep it uint8_t instead of DLDataTypeCode for minimal memory + * footprint, but the value should be one of DLDataTypeCode enum values. + * */ + uint8_t code; + /*! + * \brief Number of bits, common choices are 8, 16, 32. + */ + uint8_t bits; + /*! \brief Number of lanes in the type, used for vector types. */ + uint16_t lanes; +} DLDataType; + +/*! + * \brief Plain C Tensor object, does not manage memory. + */ +typedef struct { + /*! + * \brief The data pointer points to the allocated data. This will be CUDA + * device pointer or cl_mem handle in OpenCL. It may be opaque on some device + * types. This pointer is always aligned to 256 bytes as in CUDA. The + * `byte_offset` field should be used to point to the beginning of the data. + * + * Note that as of Nov 2021, multiply libraries (CuPy, PyTorch, TensorFlow, + * TVM, perhaps others) do not adhere to this 256 byte aligment requirement + * on CPU/CUDA/ROCm, and always use `byte_offset=0`. This must be fixed + * (after which this note will be updated); at the moment it is recommended + * to not rely on the data pointer being correctly aligned. + * + * For given DLTensor, the size of memory required to store the contents of + * data is calculated as follows: + * + * \code{.c} + * static inline size_t GetDataSize(const DLTensor* t) { + * size_t size = 1; + * for (tvm_index_t i = 0; i < t->ndim; ++i) { + * size *= t->shape[i]; + * } + * size *= (t->dtype.bits * t->dtype.lanes + 7) / 8; + * return size; + * } + * \endcode + * + * Note that if the tensor is of size zero, then the data pointer should be + * set to `NULL`. + */ + void* data; + /*! \brief The device of the tensor */ + DLDevice device; + /*! \brief Number of dimensions */ + int32_t ndim; + /*! \brief The data type of the pointer*/ + DLDataType dtype; + /*! \brief The shape of the tensor */ + int64_t* shape; + /*! + * \brief strides of the tensor (in number of elements, not bytes) + * can be NULL, indicating tensor is compact and row-majored. + */ + int64_t* strides; + /*! \brief The offset in bytes to the beginning pointer to data */ + uint64_t byte_offset; +} DLTensor; + +/*! + * \brief C Tensor object, manage memory of DLTensor. This data structure is + * intended to facilitate the borrowing of DLTensor by another framework. It is + * not meant to transfer the tensor. When the borrowing framework doesn't need + * the tensor, it should call the deleter to notify the host that the resource + * is no longer needed. + * + * \note This data structure is used as Legacy DLManagedTensor + * in DLPack exchange and is deprecated after DLPack v0.8 + * Use DLManagedTensorVersioned instead. + * This data structure may get renamed or deleted in future versions. + * + * \sa DLManagedTensorVersioned + */ +typedef struct DLManagedTensor { + /*! \brief DLTensor which is being memory managed */ + DLTensor dl_tensor; + /*! \brief the context of the original host framework of DLManagedTensor in + * which DLManagedTensor is used in the framework. It can also be NULL. + */ + void * manager_ctx; + /*! + * \brief Destructor - this should be called + * to destruct the manager_ctx which backs the DLManagedTensor. It can be + * NULL if there is no way for the caller to provide a reasonable destructor. + * The destructor deletes the argument self as well. + */ + void (*deleter)(struct DLManagedTensor * self); +} DLManagedTensor; + +// bit masks used in in the DLManagedTensorVersioned + +/*! \brief bit mask to indicate that the tensor is read only. */ +#define DLPACK_FLAG_BITMASK_READ_ONLY (1UL << 0UL) + +/*! + * \brief bit mask to indicate that the tensor is a copy made by the producer. + * + * If set, the tensor is considered solely owned throughout its lifetime by the + * consumer, until the producer-provided deleter is invoked. + */ +#define DLPACK_FLAG_BITMASK_IS_COPIED (1UL << 1UL) + +/*! + * \brief A versioned and managed C Tensor object, manage memory of DLTensor. + * + * This data structure is intended to facilitate the borrowing of DLTensor by + * another framework. It is not meant to transfer the tensor. When the borrowing + * framework doesn't need the tensor, it should call the deleter to notify the + * host that the resource is no longer needed. + * + * \note This is the current standard DLPack exchange data structure. + */ +struct DLManagedTensorVersioned { + /*! + * \brief The API and ABI version of the current managed Tensor + */ + DLPackVersion version; + /*! + * \brief the context of the original host framework. + * + * Stores DLManagedTensorVersioned is used in the + * framework. It can also be NULL. + */ + void *manager_ctx; + /*! + * \brief Destructor. + * + * This should be called to destruct manager_ctx which holds the DLManagedTensorVersioned. + * It can be NULL if there is no way for the caller to provide a reasonable + * destructor. The destructor deletes the argument self as well. + */ + void (*deleter)(struct DLManagedTensorVersioned *self); + /*! + * \brief Additional bitmask flags information about the tensor. + * + * By default the flags should be set to 0. + * + * \note Future ABI changes should keep everything until this field + * stable, to ensure that deleter can be correctly called. + * + * \sa DLPACK_FLAG_BITMASK_READ_ONLY + * \sa DLPACK_FLAG_BITMASK_IS_COPIED + */ + uint64_t flags; + /*! \brief DLTensor which is being memory managed */ + DLTensor dl_tensor; +}; + +#ifdef __cplusplus +} // DLPACK_EXTERN_C +#endif +#endif // DLPACK_DLPACK_H_ diff --git a/packages/leapp-runtime-rust/src/dlpack/sys.rs b/packages/leapp-runtime-rust/src/dlpack/sys.rs new file mode 100644 index 00000000..8b2204aa --- /dev/null +++ b/packages/leapp-runtime-rust/src/dlpack/sys.rs @@ -0,0 +1,75 @@ +// Generated from upstream DLPack v1.0 by generate_bindings.py. Do not edit. +#![allow(non_camel_case_types, non_upper_case_globals, non_snake_case)] +pub type DLDeviceType = i32; +pub const DLPACK_MAJOR_VERSION: u32 = 1; +pub const DLPACK_MINOR_VERSION: u32 = 0; +pub const DLPACK_FLAG_BITMASK_READ_ONLY: u64 = 1; +pub const DLPACK_FLAG_BITMASK_IS_COPIED: u64 = 2; +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct DLPackVersion { + pub major: u32, + pub minor: u32, +} +pub const kDLCPU: u32 = 1; +pub const kDLCUDA: u32 = 2; +pub const kDLCUDAHost: u32 = 3; +pub const kDLOpenCL: u32 = 4; +pub const kDLVulkan: u32 = 7; +pub const kDLMetal: u32 = 8; +pub const kDLVPI: u32 = 9; +pub const kDLROCM: u32 = 10; +pub const kDLROCMHost: u32 = 11; +pub const kDLExtDev: u32 = 12; +pub const kDLCUDAManaged: u32 = 13; +pub const kDLOneAPI: u32 = 14; +pub const kDLWebGPU: u32 = 15; +pub const kDLHexagon: u32 = 16; +pub const kDLMAIA: u32 = 17; +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct DLDevice { + pub device_type: DLDeviceType, + pub device_id: i32, +} +pub const kDLInt: u32 = 0; +pub const kDLUInt: u32 = 1; +pub const kDLFloat: u32 = 2; +pub const kDLOpaqueHandle: u32 = 3; +pub const kDLBfloat: u32 = 4; +pub const kDLComplex: u32 = 5; +pub const kDLBool: u32 = 6; +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct DLDataType { + pub code: u8, + pub bits: u8, + pub lanes: u16, +} +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct DLTensor { + pub data: *mut std::ffi::c_void, + pub device: DLDevice, + pub ndim: i32, + pub dtype: DLDataType, + pub shape: *mut i64, + pub strides: *mut i64, + pub byte_offset: u64, +} +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct DLManagedTensor { + pub dl_tensor: DLTensor, + pub manager_ctx: *mut std::ffi::c_void, + pub deleter: Option, +} +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct DLManagedTensorVersioned { + pub version: DLPackVersion, + pub manager_ctx: *mut std::ffi::c_void, + pub deleter: Option, + pub flags: u64, + pub dl_tensor: DLTensor, +} diff --git a/packages/leapp-runtime-rust/src/engines.rs b/packages/leapp-runtime-rust/src/engines.rs new file mode 100644 index 00000000..98ad7461 --- /dev/null +++ b/packages/leapp-runtime-rust/src/engines.rs @@ -0,0 +1,25 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Backend-independent device and memory types, re-exported at the crate root. + +/// Device used by model instances, independently of input and output storage. +#[derive(Debug, Clone, Copy, Default)] +pub enum ExecutionDevice { + #[default] + Cpu, + Cuda { + device_id: u32, + }, +} + +/// Storage location, independent of a model's execution device. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum MemoryLocation { + #[default] + Cpu, + PinnedCpu, + Cuda { + device_id: u32, + }, +} diff --git a/packages/leapp-runtime-rust/src/error.rs b/packages/leapp-runtime-rust/src/error.rs index fb332775..1c49a606 100644 --- a/packages/leapp-runtime-rust/src/error.rs +++ b/packages/leapp-runtime-rust/src/error.rs @@ -61,6 +61,17 @@ impl std::error::Error for BundleError { } } +#[derive(Debug)] +pub(crate) struct Unsupported(pub String); + +impl fmt::Display for Unsupported { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +impl std::error::Error for Unsupported {} + pub(crate) fn in_phase( path: &Path, phase: &'static str, @@ -77,14 +88,73 @@ pub(crate) fn in_phase( }) } +#[derive(Debug)] +pub(crate) struct Internal(pub String); + +impl fmt::Display for Internal { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +impl std::error::Error for Internal {} + fn classify_error(source: &anyhow::Error) -> ErrorCode { + if let Some(error) = source.downcast_ref::() { + return error.code(); + } + if source.is::() { + return ErrorCode::Internal; + } if let Some(error) = source.downcast_ref::() { if error.kind() == io::ErrorKind::NotFound { ErrorCode::NotFound } else { ErrorCode::Io } + } else if source.is::() { + ErrorCode::Unsupported } else { ErrorCode::InvalidArgument } } + +/// Execution failure with its phase and underlying diagnostic. +#[derive(Debug)] +pub struct RuntimeError { + code: ErrorCode, + phase: &'static str, + source: anyhow::Error, +} + +impl RuntimeError { + pub(crate) fn new(phase: &'static str, source: anyhow::Error) -> Self { + Self { + code: classify_error(&source), + phase, + source, + } + } + + /// Category suitable for programmatic error handling. + pub fn code(&self) -> ErrorCode { + self.code + } + + /// Operation that failed, such as `convert`, `load`, `state`, or `execute`. + pub fn phase(&self) -> &'static str { + self.phase + } +} + +impl fmt::Display for RuntimeError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}: {:#}", self.phase, self.source) + } +} + +impl std::error::Error for RuntimeError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(self.source.as_ref()) + } +} diff --git a/packages/leapp-runtime-rust/src/lib.rs b/packages/leapp-runtime-rust/src/lib.rs index 0abd9730..4030402e 100644 --- a/packages/leapp-runtime-rust/src/lib.rs +++ b/packages/leapp-runtime-rust/src/lib.rs @@ -1,27 +1,31 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Load and inspect an immutable LEAPP bundle without initializing a backend. +//! Bundle inspection and the public Session API. +//! +//! Session execution is supplied by the runtime integration; in this API revision, +//! [`Session::open`] returns [`ErrorCode::Unsupported`]. //! -//! ```no_run -//! use leapp_runtime::Bundle; -//! let bundle = Bundle::open("export/policy.yaml")?; -//! for input in bundle.pipeline_inputs() { -//! if let Some(spec) = bundle.input_info(&input.node, &input.port) { -//! println!("{}: {:?}", input.name(), spec.shape); -//! } -//! } -//! # Ok::<(), leapp_runtime::BundleError>(()) -//! ``` //! Implementation schema types are intentionally private: //! //! ```compile_fail //! use leapp_runtime::{BundleDescription, NodeSpec, Pipeline}; //! ``` -#![forbid(unsafe_code)] +//! +//! Engine implementations are an internal compile-time choice: +//! +//! ```compile_fail +//! use leapp_runtime::engines; +//! ``` +#![deny(unsafe_op_in_unsafe_fn)] mod bundle; +pub mod dlpack; +mod engines; mod error; +mod session; pub use bundle::{Bundle, DataType, Metadata, MetadataValue, PortRef, PortSpec}; -pub use error::{BundleError, ErrorCode}; +pub use engines::{ExecutionDevice, MemoryLocation}; +pub use error::{BundleError, ErrorCode, RuntimeError}; +pub use session::{OutputPlacement, RunResult, Session, SessionOptions}; diff --git a/packages/leapp-runtime-rust/src/session.rs b/packages/leapp-runtime-rust/src/session.rs new file mode 100644 index 00000000..900c4fc3 --- /dev/null +++ b/packages/leapp-runtime-rust/src/session.rs @@ -0,0 +1,106 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Public Session API. Execution is supplied by the runtime integration. + +use crate::{ + dlpack::{BorrowedTensor, SharedTensor}, + Bundle, ExecutionDevice, MemoryLocation, RuntimeError, +}; +use indexmap::IndexMap; +use std::collections::HashMap; + +/// Exact allocation locations keyed by logical output names. +#[derive(Debug, Clone, Default)] +pub struct OutputPlacement { + pub default_location: MemoryLocation, + pub overrides: HashMap, +} +/// Immutable construction options, inherited by every new session. +#[derive(Debug, Clone, Default)] +pub struct SessionOptions { + pub device: ExecutionDevice, + pub output_placement: OutputPlacement, +} + +/// Managed DLPack public outputs in bundle declaration order, independent of the session. +/// Repeated output declarations share one named entry. Cloning shares storage. +#[derive(Clone)] +pub struct RunResult { + outputs: IndexMap, +} + +impl RunResult { + /// Look up an output by its qualified LEAPP `node/port` name. + pub fn get(&self, name: &str) -> Option<&SharedTensor> { + self.outputs.get(name) + } + + /// Iterate through public names and tensors in declaration order. + pub fn iter(&self) -> impl Iterator { + self.outputs + .iter() + .map(|(name, tensor)| (name.as_str(), tensor)) + } +} + +/// Independent feedback state sharing loaded models with sessions from `new_session`. +/// Separate sessions can execute concurrently. Model-internal state remains shared. +/// +/// This API revision defines the contract but does not implement execution. +/// [`Session::open`] returns [`crate::ErrorCode::Unsupported`]. +pub struct Session { + bundle: Bundle, +} + +impl Session { + /// Load models and capture initial values on the execution device. + /// Dynamic state requires an exported initial value; other absent values use zeros. + /// Execution is unavailable in this API-only revision. + pub fn open(_bundle: &Bundle, _options: SessionOptions) -> Result { + Err(execution_unavailable("load")) + } + + /// Immutable bundle snapshot, also available after close. + pub fn info(&self) -> &Bundle { + &self.bundle + } + + /// Whether this session retains its loaded execution engine. + pub fn is_open(&self) -> bool { + false + } + + /// Create an independently initialized session without reloading models or files. + /// Configuration and captured initial values are inherited; current feedback is not. + pub fn new_session(&self) -> Result { + Err(execution_unavailable("load")) + } + + /// Release this session's state and engine reference. Other sessions and results survive. + pub fn close(&mut self) {} + + /// Restore captured state without rereading files or resetting model-internal state. + pub fn reset(&mut self) -> Result<(), RuntimeError> { + Err(execution_unavailable("state")) + } + + /// Borrow every public input exactly once until synchronous execution completes. + /// Feedback advances only after every output and isolation copy succeeds. + pub fn run( + &mut self, + _inputs: &[(&str, BorrowedTensor<'_>)], + ) -> Result { + Err(execution_unavailable("execute")) + } +} + +fn execution_unavailable(phase: &'static str) -> RuntimeError { + RuntimeError::new( + phase, + crate::error::Unsupported( + "Session execution is not implemented in this API revision".into(), + ) + .into(), + ) +} diff --git a/packages/leapp-runtime-rust/tests/unit_tests/dlpack.rs b/packages/leapp-runtime-rust/tests/unit_tests/dlpack.rs new file mode 100644 index 00000000..a055e74d --- /dev/null +++ b/packages/leapp-runtime-rust/tests/unit_tests/dlpack.rs @@ -0,0 +1,9 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#[path = "dlpack/abi.rs"] +mod abi; +#[path = "dlpack/cpu.rs"] +mod cpu; +#[path = "dlpack/ownership.rs"] +mod ownership; diff --git a/packages/leapp-runtime-rust/tests/unit_tests/dlpack/abi.rs b/packages/leapp-runtime-rust/tests/unit_tests/dlpack/abi.rs new file mode 100644 index 00000000..d412bf5d --- /dev/null +++ b/packages/leapp-runtime-rust/tests/unit_tests/dlpack/abi.rs @@ -0,0 +1,65 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +use leapp_runtime::dlpack::*; +use std::{ + mem::{align_of, offset_of, size_of}, + process::Command, +}; + +#[test] +fn generated_layout_matches_the_pinned_c_header() { + let dir = std::env::temp_dir().join(format!("leapp-dlpack-abi-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let c = dir.join("abi.c"); + let exe = dir.join("abi"); + std::fs::write(&c,r#" +#include +#include +#include +#define PRINT(T) printf("%zu %zu ", sizeof(T), _Alignof(T)) +int main(void) { + PRINT(DLDevice); PRINT(DLDataType); PRINT(DLTensor); PRINT(DLManagedTensor); PRINT(struct DLManagedTensorVersioned); + printf("%zu %zu %zu %zu %zu %zu", offsetof(DLTensor, shape), offsetof(DLTensor, byte_offset), + offsetof(DLManagedTensor, deleter), offsetof(struct DLManagedTensorVersioned, deleter), + offsetof(struct DLManagedTensorVersioned, flags), offsetof(struct DLManagedTensorVersioned, dl_tensor)); +} +"#).unwrap(); + assert!(Command::new("cc") + .args(["-std=c11", "-I"]) + .arg(concat!(env!("CARGO_MANIFEST_DIR"), "/src/dlpack/include")) + .arg(&c) + .arg("-o") + .arg(&exe) + .status() + .unwrap() + .success()); + let output = Command::new(&exe).output().unwrap(); + assert!(output.status.success()); + let actual: Vec = String::from_utf8(output.stdout) + .unwrap() + .split_whitespace() + .map(|s| s.parse().unwrap()) + .collect(); + assert_eq!( + actual, + [ + size_of::(), + align_of::(), + size_of::(), + align_of::(), + size_of::(), + align_of::(), + size_of::(), + align_of::(), + size_of::(), + align_of::(), + offset_of!(DLTensor, shape), + offset_of!(DLTensor, byte_offset), + offset_of!(DLManagedTensor, deleter), + offset_of!(DLManagedTensorVersioned, deleter), + offset_of!(DLManagedTensorVersioned, flags), + offset_of!(DLManagedTensorVersioned, dl_tensor) + ] + ); + std::fs::remove_dir_all(dir).unwrap(); +} diff --git a/packages/leapp-runtime-rust/tests/unit_tests/dlpack/cpu.rs b/packages/leapp-runtime-rust/tests/unit_tests/dlpack/cpu.rs new file mode 100644 index 00000000..ca7084b0 --- /dev/null +++ b/packages/leapp-runtime-rust/tests/unit_tests/dlpack/cpu.rs @@ -0,0 +1,67 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use leapp_runtime::dlpack::*; + +#[test] +fn cpu_slices_infer_dtype_and_validate_shape_without_copying() { + let values = [1f32, 2.]; + let input = BorrowedTensor::cpu(&values, &[1, 2]).unwrap(); + let dtype = input.descriptor().dtype; + assert_eq!((dtype.code, dtype.bits, dtype.lanes), (2, 32, 1)); + assert_eq!( + input.as_cpu_slice::().unwrap().as_ptr(), + values.as_ptr() + ); + for (values, shape) in [(&[7i32][..], &[][..]), (&[][..], &[2, 0][..])] { + assert_eq!( + BorrowedTensor::cpu(values, shape) + .unwrap() + .as_cpu_slice::() + .unwrap(), + values + ); + } + for shape in [&[3][..], &[-1][..], &[i64::MAX, 8][..], &[0, -1][..]] { + assert!(BorrowedTensor::cpu(&[1i32, 2], shape).is_err()); + } +} + +#[test] +fn cpu_reads_check_layout_and_apply_byte_offsets() { + let values = [1f32, 2., 3., 4.]; + let input = BorrowedTensor::cpu(&values, &[4]).unwrap(); + let mut descriptor = *input.descriptor(); + let shape = [2]; + let strides = [1]; + descriptor.shape = shape.as_ptr().cast_mut(); + descriptor.strides = strides.as_ptr().cast_mut(); + descriptor.byte_offset = 8; + // SAFETY: descriptors below refer to these live, immutable arrays; error + // cases are rejected before reading storage through incompatible metadata. + let borrowed = unsafe { BorrowedTensor::new(&descriptor) }; + assert_eq!(borrowed.as_cpu_slice::().unwrap(), &[3., 4.]); + assert!(borrowed.as_cpu_slice::().is_err()); + let strided = [2]; + for scenario in 0..4 { + let mut bad = descriptor; + match scenario { + 0 => bad.device.device_type = kDLCUDA as i32, + 1 => { + bad.strides = strided.as_ptr().cast_mut(); + bad.byte_offset = 0; + } + 2 => { + // Offset a naturally aligned f32 pointer by one byte. + bad.data = values.as_ptr().cast_mut().cast(); + bad.byte_offset = 1; + } + _ => { + bad.dtype.lanes = 2; + } + } + assert!(unsafe { BorrowedTensor::new(&bad) } + .as_cpu_slice::() + .is_err()); + } +} diff --git a/packages/leapp-runtime-rust/tests/unit_tests/dlpack/ownership.rs b/packages/leapp-runtime-rust/tests/unit_tests/dlpack/ownership.rs new file mode 100644 index 00000000..f192b648 --- /dev/null +++ b/packages/leapp-runtime-rust/tests/unit_tests/dlpack/ownership.rs @@ -0,0 +1,117 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +use leapp_runtime::dlpack::*; +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, +}; + +struct Context { + dropped: Arc, + values: [f32; 2], + shape: [i64; 1], +} +unsafe extern "C" fn delete(p: *mut DLManagedTensorVersioned) { + unsafe { + let p = Box::from_raw(p); + let context = Box::from_raw(p.manager_ctx.cast::()); + context.dropped.fetch_add(1, Ordering::SeqCst); + } +} +fn producer() -> (ManagedTensor, Arc) { + let dropped = Arc::new(AtomicUsize::new(0)); + let context = Box::new(Context { + dropped: dropped.clone(), + values: [2., 3.], + shape: [2], + }); + let dl_tensor = *BorrowedTensor::cpu(&context.values, &context.shape) + .unwrap() + .descriptor(); + let p = Box::into_raw(Box::new(DLManagedTensorVersioned { + version: DLPackVersion { major: 1, minor: 0 }, + manager_ctx: Box::into_raw(context).cast(), + deleter: Some(delete), + flags: 0, + dl_tensor, + })); + ( + unsafe { ManagedTensor::from_versioned(p).unwrap() }, + dropped, + ) +} +#[test] +fn leases_preserve_ownership_and_exposure_across_threads() { + for versioned in [false, true] { + let (owner, dropped) = producer(); + // SAFETY: the fixture owns ready CPU storage and can be deleted on any thread. + let shared = unsafe { SharedTensor::new(owner, None, None).unwrap() }; + let saved = shared.to_cpu_vec::().unwrap(); + assert_eq!(saved, [2., 3.]); + let descriptor = shared.clone().descriptor().data; + // SAFETY: this CPU fixture has no outstanding readers or foreign consumers. + unsafe { *descriptor.cast::() = 99. }; + assert_eq!(saved, [2., 3.]); + assert!(shared.to_cpu_vec::().is_err()); + let released = dropped.clone(); + std::thread::spawn(move || { + let sibling = shared.export(!versioned); + let (pointer, actual) = shared.export(versioned).into_raw(); + assert_eq!(actual, versioned); + drop(shared); + drop(sibling); + assert_eq!(released.load(Ordering::SeqCst), 0); + // Adopt the transferred lease exactly once, on the consumer thread. + let consumer = unsafe { + if versioned { + ManagedTensor::from_versioned(pointer.cast()) + } else { + ManagedTensor::from_legacy(pointer.cast()) + } + } + .unwrap(); + let tensor = consumer.descriptor().unwrap(); + assert_eq!(unsafe { *tensor.shape }, 2); + assert_eq!( + unsafe { std::slice::from_raw_parts(tensor.data.cast::(), 2) }, + [99., 3.] + ); + }) + .join() + .unwrap(); + assert_eq!(dropped.load(Ordering::SeqCst), 1); + } +} +#[test] +fn transfer_failures_preserve_the_owner_and_allow_retry() { + let (owner, dropped) = producer(); + let attempts = Arc::new(AtomicUsize::new(0)); + let copies = attempts.clone(); + // SAFETY: the fixture is ready; this synchronous callback copies its values + // without retaining the destination or exposing any source pointers. + let shared = unsafe { + SharedTensor::new( + owner, + None, + Some(Box::new(move |destination| { + if copies.fetch_add(1, Ordering::SeqCst) == 0 { + return Err("simulated device transfer failure".into()); + } + destination[..4].copy_from_slice(&2f32.to_ne_bytes()); + destination[4..].copy_from_slice(&3f32.to_ne_bytes()); + Ok(()) + })), + ) + .unwrap() + }; + assert_eq!( + shared.to_cpu_vec::().unwrap_err(), + "simulated device transfer failure" + ); + assert_eq!(shared.to_cpu_vec::().unwrap(), [2., 3.]); + drop(shared.export(true)); + assert!(shared.to_cpu_vec::().unwrap_err().contains("exposed")); + assert_eq!(attempts.load(Ordering::SeqCst), 2); + drop(shared); + assert_eq!(dropped.load(Ordering::SeqCst), 1); +} From 5fb5cdd08c415383603daf0f2aa5d6cf9a8a4acc Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Wed, 28 Oct 2026 17:03:53 -0700 Subject: [PATCH 21/21] refactor: align runtime API names and session options Expose the public API as lrt across Rust, Python, and C++. Accept borrowed input maps and define session-owned output-cache retention with bundle-derived byte sizing. Update bindings, documentation, and API checks. Signed-off-by: Marc Alwan --- packages/leapp-runtime-rust/Cargo.toml | 3 + packages/leapp-runtime-rust/README.md | 63 ++++++++++++++----- .../bindings/common/README.md | 8 ++- .../bindings/common/src/session.rs | 8 +-- .../bindings/common/tests/lifecycle.rs | 14 +++-- .../bindings/cxx/CMakeLists.txt | 6 +- .../leapp-runtime-rust/bindings/cxx/README.md | 8 ++- .../bindings/cxx/include/leapp/bundle.h | 4 +- .../bindings/cxx/include/leapp/data_type.h | 4 +- .../bindings/cxx/include/leapp/metadata.h | 4 +- .../bindings/cxx/include/leapp/status.h | 4 +- .../bindings/cxx/src/bridge.rs | 17 ++++- .../bindings/cxx/src/bundle.cc | 4 +- .../bindings/cxx/src/bundle.rs | 2 +- .../bindings/cxx/src/dtype.rs | 2 +- .../bindings/cxx/src/lib.rs | 2 +- .../bindings/cxx/tests/bundle_test.cc | 20 +++--- .../cxx/tests/consumer/CMakeLists.txt | 2 +- .../bindings/cxx/tests/consumer/main.cc | 6 +- .../bindings/python/README.md | 9 +-- .../bindings/python/pyproject.toml | 2 +- .../python/{leapp_runtime => lrt}/__init__.py | 0 .../python/{leapp_runtime => lrt}/_native.pyi | 0 .../python/{leapp_runtime => lrt}/py.typed | 0 .../bindings/python/src/bundle.rs | 26 +++++--- .../bindings/python/src/lib.rs | 6 +- .../bindings/python/tests/test_bundle.py | 31 +++++---- packages/leapp-runtime-rust/src/lib.rs | 6 +- packages/leapp-runtime-rust/src/session.rs | 50 ++++++++++++++- .../tests/unit_tests/bundle.rs | 8 ++- .../tests/unit_tests/dlpack/abi.rs | 2 +- .../tests/unit_tests/dlpack/cpu.rs | 2 +- .../tests/unit_tests/dlpack/ownership.rs | 2 +- 33 files changed, 224 insertions(+), 101 deletions(-) rename packages/leapp-runtime-rust/bindings/python/python/{leapp_runtime => lrt}/__init__.py (100%) rename packages/leapp-runtime-rust/bindings/python/python/{leapp_runtime => lrt}/_native.pyi (100%) rename packages/leapp-runtime-rust/bindings/python/python/{leapp_runtime => lrt}/py.typed (100%) diff --git a/packages/leapp-runtime-rust/Cargo.toml b/packages/leapp-runtime-rust/Cargo.toml index 5209b5ec..31b2af07 100644 --- a/packages/leapp-runtime-rust/Cargo.toml +++ b/packages/leapp-runtime-rust/Cargo.toml @@ -14,6 +14,9 @@ license = "Apache-2.0" publish = false description = "LEAPP bundle inspection and core Session API" +[lib] +name = "lrt" + [dependencies] anyhow = "1" indexmap = { version = "2", features = ["serde"] } diff --git a/packages/leapp-runtime-rust/README.md b/packages/leapp-runtime-rust/README.md index 52355cc7..ddb0213a 100644 --- a/packages/leapp-runtime-rust/README.md +++ b/packages/leapp-runtime-rust/README.md @@ -11,8 +11,10 @@ wrapper. This revision defines the execution contract: `Session::open()` returns `ErrorCode::Unsupported` until the runtime integration supplies execution. Bundle loading does not initialize an inference backend. -The inspection and full-runtime Python distributions share the `leapp_runtime` -import name and should be installed in separate environments. +The public API uses `lrt` in all three languages: the Rust crate, Python import, +and C++ namespace. Cargo and Python distribution names remain `leapp-runtime`. +The inspection and full-runtime Python distributions should be installed in +separate environments. ## Included code @@ -20,7 +22,7 @@ import name and should be installed in separate environments. | --- | --- | | Rust core | `Bundle`, YAML schema parsing, graph and artifact validation, SHA-256 checks, safetensors initial-state validation | | Public inspection types | `PortRef`, `PortSpec`, `DataType`, `Metadata`, `MetadataValue`, `BundleError`, `ErrorCode` | -| Core Session API | `Session`, `SessionOptions`, `OutputPlacement`, `RunResult`, device/memory types, and `RuntimeError` | +| Core Session API | `Session`, `SessionOptions`, `OutputPoolOptions`, `OutputPlacement`, `RunResult`, device/memory types, and `RuntimeError` | | Tensor API | Backend-independent DLPack descriptors and ownership types used by Session inputs and results | | Common lifecycle crate | `SessionHandle`, independent closure, error-source preservation, and cleanup after poisoned locks; three focused regression tests | | Python | PyO3 bundle bindings, detached metadata and port descriptors, typed stubs and structured exceptions | @@ -40,14 +42,14 @@ path bytes. ## Rust ```rust -use leapp_runtime::Bundle; +use lrt::Bundle; let bundle = Bundle::open("export/policy.yaml")?; for output in bundle.pipeline_outputs() { let spec = bundle.output_info(&output.node, &output.port).unwrap(); println!("{}: {:?} {:?}", output.name(), spec.dtype, spec.shape); } -# Ok::<(), leapp_runtime::BundleError>(()) +# Ok::<(), lrt::BundleError>(()) ``` `model_io(node)` returns `Option<(Vec<&str>, Vec<&str>)>`: input names followed @@ -102,31 +104,62 @@ All inspection tests create opaque artifact bytes; no model generation is needed `Session::open(&bundle, options)` accepts immutable construction options: ```rust -use leapp_runtime::{ExecutionDevice, OutputPlacement, Session, SessionOptions}; +use lrt::{ExecutionDevice, Session, SessionOptions}; let options = SessionOptions { device: ExecutionDevice::Cpu, - output_placement: OutputPlacement::default(), + ..Default::default() }; -# let bundle = leapp_runtime::Bundle::open("export/policy.yaml")?; +# let bundle = lrt::Bundle::open("export/policy.yaml")?; let session = Session::open(&bundle, options); // This API revision returns a RuntimeError with ErrorCode::Unsupported. -# Ok::<(), leapp_runtime::BundleError>(()) +# Ok::<(), lrt::BundleError>(()) ``` The public lifecycle is `open`, `info`, `is_open`, `new_session`, `run`, `reset`, -and `close`. `run` accepts named `dlpack::BorrowedTensor` inputs and returns a -`RunResult` retaining `SharedTensor` outputs. There is no separate public Runtime -owner, no per-call RunOptions, and no engine field in SessionOptions. Output -placement is configured at construction and inherited by every new session. +and `close`. `run` borrows a `HashMap<&str, dlpack::BorrowedTensor<'_>>` and returns +a `RunResult` retaining `SharedTensor` outputs: + +```rust +use lrt::{dlpack::BorrowedTensor, RunResult, RuntimeError, Session}; +use std::collections::HashMap; + +fn step(session: &mut Session) -> Result { + let values = [1_f32, 2., 3.]; + let inputs = HashMap::from([ + ("model/input", BorrowedTensor::cpu(&values, &[3]).unwrap()), + ]); + session.run(&inputs) +} +``` + +Input names and tensors are borrowed. Matching is by name, independent of map +iteration order; missing and unknown inputs are rejected by execution. Inserting +a duplicate key replaces the previous value before the runtime sees the map. +Results preserve the bundle's output declaration order. + +There is no separate public Runtime owner, no per-call RunOptions, and no engine +field in SessionOptions. Output placement is configured at construction and +inherited by every new session. The execution contract requires `new_session()` to share loaded models and immutable options while starting from captured initial state, not current -feedback. `reset()` restores that initial state. Closing one session preserves -its peers and retained results. Execution is deliberately unavailable in this +feedback. Each session owns its output cache; peers start with an empty cache +and inherit the retention policy. `reset()` restores initial state. Closing one +session preserves its peers and retained results. Execution is unavailable in this revision; the runtime integration implements these operations without changing the public signatures. There is no runtime selector or optional backend feature. +`SessionOptions::output_pool` contains `OutputPoolOptions { max_entries }`, +defaulting to 64 cached buffers per session. Zero disables caching without +preventing allocations. The execution implementation derives the byte budget +from the bundle's tensor shapes, dtypes, and placement, growing it for actual +dynamic output sizes. There is no fixed byte cap or public `max_bytes` option; +large tensors remain eligible for reuse. This policy does not cap model memory, +fixed workspaces, or retained results. Retained or externally exposed outputs +must not be reused, and result pooling must preserve native workspace addresses +and cached metadata. Pool allocation and sizing arrive with runtime integration. + ## Shared lifecycle `bindings/common` compiles its `SessionHandle` wrapper against the exported core diff --git a/packages/leapp-runtime-rust/bindings/common/README.md b/packages/leapp-runtime-rust/bindings/common/README.md index dc9bf093..7ff14928 100644 --- a/packages/leapp-runtime-rust/bindings/common/README.md +++ b/packages/leapp-runtime-rust/bindings/common/README.md @@ -13,14 +13,16 @@ Close reports the lock error after cleanup; subsequent close calls succeed. Normal run/reset operations still fail on poisoned locks. This crate compiles against the exported `Session`, `SessionOptions`, DLPack -input/result types, and errors in `leapp_runtime`. Execution is not implemented +input/result types, and errors in `lrt`. Execution is not implemented in this API revision: opening a SessionHandle preserves the core Unsupported error. The runtime integration supplies execution using these same signatures. The Python and C++ bindings currently expose bundle inspection only. `SessionHandle::open()` accepts construction-time `SessionOptions`. `run()` -accepts named borrowed inputs without per-call options. `new_session()` creates -an independent session sharing the loaded engine and immutable options. +borrows a `HashMap<&str, BorrowedTensor<'_>>` without per-call options and passes +it directly to the core. `new_session()` creates an independent session sharing +the loaded engine and immutable options, with an empty private output cache +using the inherited `OutputPoolOptions` retention policy. The root Cargo test target includes `session.rs` against a small session API defined only in the test binary. Three tests cover waiting for an active run, diff --git a/packages/leapp-runtime-rust/bindings/common/src/session.rs b/packages/leapp-runtime-rust/bindings/common/src/session.rs index 4ae13f0c..3565c59c 100644 --- a/packages/leapp-runtime-rust/bindings/common/src/session.rs +++ b/packages/leapp-runtime-rust/bindings/common/src/session.rs @@ -1,9 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use leapp_runtime::dlpack::BorrowedTensor; -use leapp_runtime::{Bundle, ErrorCode, RunResult, RuntimeError, Session, SessionOptions}; -use std::{fmt, sync::Mutex}; +use lrt::dlpack::BorrowedTensor; +use lrt::{Bundle, ErrorCode, RunResult, RuntimeError, Session, SessionOptions}; +use std::{collections::HashMap, fmt, sync::Mutex}; /// Binding lifecycle failures preserve native execution diagnostics. /// Wrapped runtime failures remain accessible through [`std::error::Error::source`]. @@ -95,7 +95,7 @@ impl SessionHandle { Ok(()) } } - pub fn run(&self, inputs: &[(&str, BorrowedTensor<'_>)]) -> Result { + pub fn run(&self, inputs: &HashMap<&str, BorrowedTensor<'_>>) -> Result { self.session .lock() .map_err(|_| Error::Lock)? diff --git a/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs b/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs index a610050c..9b53b5b6 100644 --- a/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs +++ b/packages/leapp-runtime-rust/bindings/common/tests/lifecycle.rs @@ -2,10 +2,11 @@ // SPDX-License-Identifier: Apache-2.0 // Compile the binding lifecycle against a test-only session. -extern crate leapp_runtime as bundle_api; -extern crate self as leapp_runtime; +extern crate lrt as bundle_api; +extern crate self as lrt; pub use bundle_api::{Bundle, ErrorCode}; use std::{ + collections::HashMap, fmt, sync::{mpsc, Arc, Weak}, thread, @@ -34,7 +35,7 @@ mod lifecycle { assert!(!session.is_open().unwrap()); assert!(session.info().model_names().is_empty()); session.close().unwrap(); - peer.run(&[]).unwrap(); + peer.run(&std::collections::HashMap::new()).unwrap(); peer.close().unwrap(); assert!(owner.upgrade().is_none()); } @@ -99,7 +100,10 @@ impl Session { pub fn reset(&mut self) -> Result<(), RuntimeError> { self.reset_check() } - pub fn run(&mut self, _: &[(&str, dlpack::BorrowedTensor<'_>)]) -> Result<(), RuntimeError> { + pub fn run( + &mut self, + _: &HashMap<&str, dlpack::BorrowedTensor<'_>>, + ) -> Result<(), RuntimeError> { self.reset_check()?; if let Some((entered, resume)) = &self.wait { entered.send(()).unwrap(); @@ -143,7 +147,7 @@ fn close_waits_until_this_sessions_run_finishes() { let session = Arc::new(session); let run = { let session = session.clone(); - thread::spawn(move || session.run(&[]).unwrap()) + thread::spawn(move || session.run(&HashMap::new()).unwrap()) }; running.recv_timeout(Duration::from_secs(5)).unwrap(); let (started, closing) = mpsc::channel(); diff --git a/packages/leapp-runtime-rust/bindings/cxx/CMakeLists.txt b/packages/leapp-runtime-rust/bindings/cxx/CMakeLists.txt index 44afe34a..5ec78294 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/CMakeLists.txt +++ b/packages/leapp-runtime-rust/bindings/cxx/CMakeLists.txt @@ -22,7 +22,7 @@ add_custom_target(leapp_runtime_rust VERBATIM) add_library(leapp_runtime SHARED src/bundle.cc) -add_library(leapp::runtime ALIAS leapp_runtime) +add_library(lrt::runtime ALIAS leapp_runtime) add_dependencies(leapp_runtime leapp_runtime_rust) set_target_properties(leapp_runtime PROPERTIES EXPORT_NAME runtime VERSION ${PROJECT_VERSION} SOVERSION 0) @@ -46,7 +46,7 @@ target_compile_options(leapp_runtime PRIVATE -Wall -Wextra -Wpedantic -Werror -f install(TARGETS leapp_runtime EXPORT leapp_runtimeTargets LIBRARY DESTINATION "${CMAKE_INSTALL_LIBDIR}") install(DIRECTORY include/leapp DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}") -install(EXPORT leapp_runtimeTargets NAMESPACE leapp:: +install(EXPORT leapp_runtimeTargets NAMESPACE lrt:: DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/leapp_runtime") configure_package_config_file(cmake/leapp_runtimeConfig.cmake.in @@ -62,7 +62,7 @@ install(FILES "${CMAKE_CURRENT_BINARY_DIR}/leapp_runtimeConfig.cmake" if(BUILD_TESTING) add_executable(leapp_bundle_test tests/bundle_test.cc) - target_link_libraries(leapp_bundle_test PRIVATE leapp::runtime Threads::Threads) + target_link_libraries(leapp_bundle_test PRIVATE lrt::runtime Threads::Threads) target_compile_options(leapp_bundle_test PRIVATE -Wall -Wextra -Wpedantic -Werror -fno-exceptions) add_test(NAME leapp_bundle COMMAND leapp_bundle_test diff --git a/packages/leapp-runtime-rust/bindings/cxx/README.md b/packages/leapp-runtime-rust/bindings/cxx/README.md index 2234466f..61d28897 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/README.md +++ b/packages/leapp-runtime-rust/bindings/cxx/README.md @@ -18,13 +18,15 @@ cmake --build target/consumer --parallel ``` Consumers use `find_package(leapp_runtime CONFIG REQUIRED)` and link -`leapp::runtime`. Include `leapp/bundle.h` for inspection or `leapp/data_type.h` +`lrt::runtime`. Public types live in namespace `lrt`; the CMake package and +installed header paths retain their LEAPP names. +Include `leapp/bundle.h` for inspection or `leapp/data_type.h` for scalar type declarations. Custom YAML tags are unsupported; `Bundle::Open` returns an invalid-argument status in the `parse` phase. ```cpp -leapp::Bundle bundle; -auto status = leapp::Bundle::Open("export/policy.yaml", &bundle); +lrt::Bundle bundle; +auto status = lrt::Bundle::Open("export/policy.yaml", &bundle); if (!status.ok()) { // status.code(), status.phase(), and status.message() describe the failure. return 1; diff --git a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/bundle.h b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/bundle.h index 19d8c4fd..83b8bb8f 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/bundle.h +++ b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/bundle.h @@ -16,7 +16,7 @@ #include "leapp/metadata.h" #include "leapp/status.h" -namespace leapp { +namespace lrt { struct PortRef { std::string node; @@ -86,6 +86,6 @@ class Bundle { std::shared_ptr data_; }; -} // namespace leapp +} // namespace lrt #endif // LEAPP_BUNDLE_H_ diff --git a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/data_type.h b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/data_type.h index 65930339..0a975878 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/data_type.h +++ b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/data_type.h @@ -6,7 +6,7 @@ #include -namespace leapp { +namespace lrt { /// @brief Fixed-width scalar types; half and Boolean values use packed bytes. enum class DataType : uint8_t { @@ -25,6 +25,6 @@ enum class DataType : uint8_t { kFloat64 }; -} // namespace leapp +} // namespace lrt #endif // LEAPP_DATA_TYPE_H_ diff --git a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/metadata.h b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/metadata.h index 2b45b4c0..63f9391d 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/metadata.h +++ b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/metadata.h @@ -10,7 +10,7 @@ #include #include -namespace leapp { +namespace lrt { /// @brief Detached typed metadata, preserving ordered mappings and non-string /// keys. @@ -34,6 +34,6 @@ struct Metadata::Entry { Metadata value; }; -} // namespace leapp +} // namespace lrt #endif // LEAPP_METADATA_H_ diff --git a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/status.h b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/status.h index 42565841..faeb613c 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/include/leapp/status.h +++ b/packages/leapp-runtime-rust/bindings/cxx/include/leapp/status.h @@ -7,7 +7,7 @@ #include #include -namespace leapp { +namespace lrt { /// @brief Failure categories shared with the Rust runtime. enum class ErrorCode { @@ -41,6 +41,6 @@ class [[nodiscard]] Status { std::string message_; }; -} // namespace leapp +} // namespace lrt #endif // LEAPP_STATUS_H_ diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs b/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs index 307de97f..da273890 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs +++ b/packages/leapp-runtime-rust/bindings/cxx/src/bridge.rs @@ -3,7 +3,7 @@ use crate::bundle::*; -#[cxx::bridge(namespace = "leapp::detail")] +#[cxx::bridge(namespace = "lrt::detail")] pub mod ffi { struct StatusData { code: i32, @@ -85,13 +85,26 @@ pub mod ffi { extern "Rust" { type BundleHandle; + fn bundle_new() -> Box; + fn bundle_open(path: &[u8], output: &mut BundleHandle) -> StatusData; + fn bundle_path(bundle: &BundleHandle) -> Vec; + fn bundle_models(bundle: &BundleHandle) -> Vec; + fn bundle_model_io(bundle: &BundleHandle, node: &[u8]) -> ModelIoData; + fn bundle_ports(bundle: &BundleHandle, inputs: bool) -> Vec; - fn bundle_port(bundle: &BundleHandle, node: &[u8], port: &[u8], input: bool) -> PortData; + + fn bundle_port( + bundle: &BundleHandle, + node: &[u8], + port: &[u8], + input: bool, + ) -> PortData; + fn bundle_metadata( bundle: &BundleHandle, field: MetadataField, diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc index 8dd650ca..6667535c 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc +++ b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.cc @@ -16,7 +16,7 @@ #include "leapp-runtime-cxx/src/bridge.rs.h" -namespace leapp { +namespace lrt { namespace { namespace wire = detail; @@ -311,4 +311,4 @@ const PortSpec* Bundle::output_info(std::string_view node, return data_ ? data_->Port(node, port, false) : nullptr; } -} // namespace leapp +} // namespace lrt diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs index 5bd9685b..016ac2eb 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs +++ b/packages/leapp-runtime-rust/bindings/cxx/src/bundle.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use crate::{dtype::wire_dtype, ffi}; -use leapp_runtime::{Bundle, Metadata, MetadataValue}; +use lrt::{Bundle, Metadata, MetadataValue}; use std::{ffi::OsStr, os::unix::ffi::OsStrExt, path::Path}; pub struct BundleHandle(pub Option); diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/dtype.rs b/packages/leapp-runtime-rust/bindings/cxx/src/dtype.rs index bf9f111c..cceb5e36 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/src/dtype.rs +++ b/packages/leapp-runtime-rust/bindings/cxx/src/dtype.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use crate::ffi; -use leapp_runtime::DataType; +use lrt::DataType; pub fn wire_dtype(value: DataType) -> ffi::Dtype { use DataType as D; diff --git a/packages/leapp-runtime-rust/bindings/cxx/src/lib.rs b/packages/leapp-runtime-rust/bindings/cxx/src/lib.rs index 046953cb..51a830ab 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/src/lib.rs +++ b/packages/leapp-runtime-rust/bindings/cxx/src/lib.rs @@ -9,7 +9,7 @@ mod bundle; mod dtype; use bridge::ffi; -use leapp_runtime::ErrorCode; +use lrt::ErrorCode; fn success() -> ffi::StatusData { ffi::StatusData { diff --git a/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc b/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc index b4e7e7b7..c7a8c4aa 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc +++ b/packages/leapp-runtime-rust/bindings/cxx/tests/bundle_test.cc @@ -38,19 +38,19 @@ int main(int argc, char** argv) { "sequence_keys: {[x, y]: coordinates}\n" "unsigned: 18446744073709551615\n"; - leapp::Bundle empty; + lrt::Bundle empty; Check(!empty.is_valid() && empty.model_names().empty(), "empty handle"); Check(empty.input_info("missing", "port") == nullptr, "empty descriptor"); Check(empty.model_io("missing") == nullptr, "empty model I/O"); - Check(!leapp::Bundle::Open(path, nullptr).ok(), "null destination"); - auto status = leapp::Bundle::Open(directory / "absent.yaml", &empty); - Check(status.code() == leapp::ErrorCode::kNotFound, "missing YAML category"); + Check(!lrt::Bundle::Open(path, nullptr).ok(), "null destination"); + auto status = lrt::Bundle::Open(directory / "absent.yaml", &empty); + Check(status.code() == lrt::ErrorCode::kNotFound, "missing YAML category"); Check(status.phase() == "open" && !status.message().empty(), "error diagnostics"); Check(!empty.is_valid(), "failure preserves empty destination"); - leapp::Bundle bundle; - Check(leapp::Bundle::Open(path, &bundle).ok(), "load opaque artifacts"); + lrt::Bundle bundle; + Check(lrt::Bundle::Open(path, &bundle).ok(), "load opaque artifacts"); Check(bundle.yaml_path() == path, "absolute path"); Check(bundle.model_names() == std::vector({"z_observation", "a_policy"}), @@ -65,7 +65,7 @@ int main(int argc, char** argv) { Check(bundle.pipeline_outputs()[0].first == "a_policy/next_state", "output order"); const auto* port = bundle.output_info("a_policy", "actions"); - Check(port != nullptr && port->dtype == leapp::DataType::kFloat32, + Check(port != nullptr && port->dtype == lrt::DataType::kFloat32, "output dtype"); Check(port->shape == std::vector({3}), "output shape"); Check(bundle.input_info("absent", "port") == nullptr, "absent descriptor"); @@ -76,15 +76,15 @@ int main(int argc, char** argv) { Check(sensor && sensor->Find("axes"), "nested mapping metadata"); const auto* mapping = bundle.metadata().Find("sequence_keys"); Check(mapping != nullptr, "mapping metadata"); - const auto& entries = std::get(mapping->value); - Check(std::get(entries[0].key.value).size() == 2, + const auto& entries = std::get(mapping->value); + Check(std::get(entries[0].key.value).size() == 2, "non-string mapping key"); const auto* unsigned_value = bundle.metadata().Find("unsigned"); Check(unsigned_value && std::get(unsigned_value->value) == UINT64_MAX, "unsigned metadata"); - Check(!leapp::Bundle::Open(path, &bundle).ok(), + Check(!lrt::Bundle::Open(path, &bundle).ok(), "reject occupied destination"); Check(bundle.output_info("a_policy", "actions") == port, "failed open preserves snapshot"); diff --git a/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/CMakeLists.txt b/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/CMakeLists.txt index 15f51a8a..bfb3397e 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/CMakeLists.txt +++ b/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/CMakeLists.txt @@ -4,4 +4,4 @@ project(leapp_bundle_consumer LANGUAGES CXX) find_package(leapp_runtime CONFIG REQUIRED) add_executable(consumer main.cc) -target_link_libraries(consumer PRIVATE leapp::runtime) +target_link_libraries(consumer PRIVATE lrt::runtime) diff --git a/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/main.cc b/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/main.cc index ab42c01f..e59bb08d 100644 --- a/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/main.cc +++ b/packages/leapp-runtime-rust/bindings/cxx/tests/consumer/main.cc @@ -4,11 +4,11 @@ #include "leapp/bundle.h" int main() { - leapp::Bundle bundle; - const auto status = leapp::Bundle::Open("", &bundle); + lrt::Bundle bundle; + const auto status = lrt::Bundle::Open("", &bundle); return !status.ok() && !bundle.is_valid() && bundle.model_io("absent") == nullptr && - status.code() == leapp::ErrorCode::kInvalidArgument + status.code() == lrt::ErrorCode::kInvalidArgument ? 0 : 1; } diff --git a/packages/leapp-runtime-rust/bindings/python/README.md b/packages/leapp-runtime-rust/bindings/python/README.md index fd08da6f..268fc2e2 100644 --- a/packages/leapp-runtime-rust/bindings/python/README.md +++ b/packages/leapp-runtime-rust/bindings/python/README.md @@ -12,9 +12,9 @@ python -m pytest bindings/python/tests ``` ```python -from leapp_runtime import Bundle, BundleError +import lrt -bundle = Bundle.open("export/policy.yaml") +bundle = lrt.Bundle.open("export/policy.yaml") for name, spec in bundle.pipeline_outputs().items(): print(name, spec.dtype, spec.shape) for model in bundle.model_names(): @@ -34,6 +34,7 @@ with unhashable or colliding Python keys become `MetadataMapping` so no entries are lost. Custom YAML tags are unsupported and raise `BundleError` during parsing. `BundleError` exposes `code` (an `ErrorCode`), `phase`, `path`, and `message`. -Only bundle inspection is exported. This distribution uses the same -`leapp_runtime` package name as the full runtimes, so use a separate environment. +Only bundle inspection is exported. Install the `leapp-runtime` distribution +and import `lrt`. The full runtimes use the same import name, so use a separate +environment. No inference SDK or NumPy is required. diff --git a/packages/leapp-runtime-rust/bindings/python/pyproject.toml b/packages/leapp-runtime-rust/bindings/python/pyproject.toml index 02c2e9c5..d6b58bf1 100644 --- a/packages/leapp-runtime-rust/bindings/python/pyproject.toml +++ b/packages/leapp-runtime-rust/bindings/python/pyproject.toml @@ -15,4 +15,4 @@ test = ["pytest>=8"] [tool.maturin] python-source = "python" -module-name = "leapp_runtime._native" +module-name = "lrt._native" diff --git a/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/__init__.py b/packages/leapp-runtime-rust/bindings/python/python/lrt/__init__.py similarity index 100% rename from packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/__init__.py rename to packages/leapp-runtime-rust/bindings/python/python/lrt/__init__.py diff --git a/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi b/packages/leapp-runtime-rust/bindings/python/python/lrt/_native.pyi similarity index 100% rename from packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/_native.pyi rename to packages/leapp-runtime-rust/bindings/python/python/lrt/_native.pyi diff --git a/packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/py.typed b/packages/leapp-runtime-rust/bindings/python/python/lrt/py.typed similarity index 100% rename from packages/leapp-runtime-rust/bindings/python/python/leapp_runtime/py.typed rename to packages/leapp-runtime-rust/bindings/python/python/lrt/py.typed diff --git a/packages/leapp-runtime-rust/bindings/python/src/bundle.rs b/packages/leapp-runtime-rust/bindings/python/src/bundle.rs index e21391b2..cd82dc45 100644 --- a/packages/leapp-runtime-rust/bindings/python/src/bundle.rs +++ b/packages/leapp-runtime-rust/bindings/python/src/bundle.rs @@ -1,14 +1,14 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use leapp_runtime::{Bundle, Metadata, MetadataValue}; +use lrt::{Bundle, Metadata, MetadataValue}; use pyo3::{ prelude::*, types::{PyDict, PyList}, }; use std::path::PathBuf; -#[pyclass(name = "PortRef", frozen, get_all, module = "leapp_runtime")] +#[pyclass(name = "PortRef", frozen, get_all, module = "lrt")] pub struct PortRef { node: String, port: String, @@ -25,7 +25,7 @@ impl PortRef { } } -#[pyclass(frozen, get_all, module = "leapp_runtime")] +#[pyclass(frozen, get_all, module = "lrt")] pub struct PortSpec { name: String, dtype: String, @@ -35,7 +35,7 @@ pub struct PortSpec { } /// Lossless fallback for YAML mappings containing keys Python cannot hash. -#[pyclass(frozen, get_all, module = "leapp_runtime")] +#[pyclass(frozen, get_all, module = "lrt")] pub struct MetadataMapping { entries: Py, } @@ -101,7 +101,7 @@ pub fn metadata(py: Python<'_>, data: &Metadata) -> PyResult> { Ok(result.into_any().unbind()) } -fn spec(py: Python<'_>, v: Option<&leapp_runtime::PortSpec>) -> PyResult> { +fn spec(py: Python<'_>, v: Option<&lrt::PortSpec>) -> PyResult> { v.map(|v| { Ok(PortSpec { name: v.name.clone(), @@ -114,7 +114,7 @@ fn spec(py: Python<'_>, v: Option<&leapp_runtime::PortSpec>) -> PyResult, node: &str, port: &str) -> PyResult> { + fn input_info( + &self, + py: Python<'_>, + node: &str, + port: &str, + ) -> PyResult> { spec(py, self.inner.input_info(node, port)) } - fn output_info(&self, py: Python<'_>, node: &str, port: &str) -> PyResult> { + fn output_info( + &self, + py: Python<'_>, + node: &str, + port: &str, + ) -> PyResult> { spec(py, self.inner.output_info(node, port)) } } diff --git a/packages/leapp-runtime-rust/bindings/python/src/lib.rs b/packages/leapp-runtime-rust/bindings/python/src/lib.rs index 6c609e24..70165923 100644 --- a/packages/leapp-runtime-rust/bindings/python/src/lib.rs +++ b/packages/leapp-runtime-rust/bindings/python/src/lib.rs @@ -5,16 +5,16 @@ mod bundle; -use leapp_runtime::ErrorCode; +use lrt::ErrorCode; use pyo3::{create_exception, exceptions::PyException, prelude::*}; -create_exception!(leapp_runtime, BundleError, PyException); +create_exception!(lrt, BundleError, PyException); fn diagnostic(py: Python<'_>, error: PyErr, code: ErrorCode, phase: &str) -> PyErr { let value = error.value(py); let result = (|| -> PyResult<()> { let code = py - .import("leapp_runtime")? + .import("lrt")? .getattr("ErrorCode")? .call1((code as i32,))?; value.setattr("code", code)?; diff --git a/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py b/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py index 7780faba..a8f80806 100644 --- a/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py +++ b/packages/leapp-runtime-rust/bindings/python/tests/test_bundle.py @@ -6,12 +6,17 @@ from concurrent.futures import ThreadPoolExecutor from pathlib import Path -import leapp_runtime as lr +import lrt import pytest YAML = (Path(__file__).parents[3] / "tests/fixtures/pipeline.yaml").read_text() +def test_public_types_use_lrt_module(): + for name in lrt.__all__: + assert getattr(lrt, name).__module__ == "lrt" + + @pytest.fixture def bundle_path(tmp_path): (tmp_path / "model.bin").write_bytes(b"abc") @@ -21,7 +26,7 @@ def bundle_path(tmp_path): def test_ordered_inspection_and_detached_values(bundle_path): - bundle = lr.Bundle.open(bundle_path) + bundle = lrt.Bundle.open(bundle_path) assert bundle.yaml_path() == bundle_path assert bundle.model_names() == ["z_observation", "a_policy"] assert list(bundle.pipeline_inputs()) == ["z_observation/joint_pos"] @@ -58,7 +63,7 @@ def test_ordered_inspection_and_detached_values(bundle_path): def test_model_io_returns_detached_internal_port_names(bundle_path): - bundle = lr.Bundle.open(bundle_path) + bundle = lrt.Bundle.open(bundle_path) expected = (["features", "state"], ["actions", "next_state"]) inputs, outputs = bundle.model_io(model_name="a_policy") assert (inputs, outputs) == expected @@ -74,7 +79,7 @@ def test_custom_backend_and_unresolved_input(bundle_path): "inputs: {z_observation: [joint_pos]}", "inputs: {absent: [port]}" ) ) - bundle = lr.Bundle.open(bundle_path) + bundle = lrt.Bundle.open(bundle_path) assert bundle.model_parameters("a_policy")["backend"] == "custom" assert bundle.pipeline_inputs() == {"absent/port": None} @@ -89,11 +94,11 @@ def test_lossless_metadata(bundle_path): unsigned: 18446744073709551615 """ ) - metadata = lr.Bundle.open(bundle_path).metadata() + metadata = lrt.Bundle.open(bundle_path).metadata() assert metadata["sensor"] == {"axes": ["x", "y"], "enabled": True} - assert isinstance(metadata["sequence_keys"], lr.MetadataMapping) + assert isinstance(metadata["sequence_keys"], lrt.MetadataMapping) assert metadata["sequence_keys"].entries == [(["x", "y"], "coordinates")] - assert isinstance(metadata["colliding_keys"], lr.MetadataMapping) + assert isinstance(metadata["colliding_keys"], lrt.MetadataMapping) assert metadata["colliding_keys"].entries == [(True, "boolean"), (1, "integer")] assert metadata["unsigned"] == 18446744073709551615 @@ -101,9 +106,9 @@ def test_lossless_metadata(bundle_path): @pytest.mark.parametrize( "failure,code,phase", [ - ("missing", lr.ErrorCode.NOT_FOUND, "open"), - ("parse", lr.ErrorCode.INVALID_ARGUMENT, "parse"), - ("artifact", lr.ErrorCode.INVALID_ARGUMENT, "artifact"), + ("missing", lrt.ErrorCode.NOT_FOUND, "open"), + ("parse", lrt.ErrorCode.INVALID_ARGUMENT, "parse"), + ("artifact", lrt.ErrorCode.INVALID_ARGUMENT, "artifact"), ], ) def test_diagnostics(bundle_path, failure, code, phase): @@ -113,8 +118,8 @@ def test_diagnostics(bundle_path, failure, code, phase): bundle_path.write_text("models: [") else: (bundle_path.parent / "model.bin").write_bytes(b"wrong hash") - with pytest.raises(lr.BundleError) as caught: - lr.Bundle.open(bundle_path) + with pytest.raises(lrt.BundleError) as caught: + lrt.Bundle.open(bundle_path) assert caught.value.code is code assert caught.value.phase == phase assert caught.value.path == bundle_path @@ -122,7 +127,7 @@ def test_diagnostics(bundle_path, failure, code, phase): def test_snapshot_and_concurrent_reads(bundle_path): - bundle = lr.Bundle.open(bundle_path) + bundle = lrt.Bundle.open(bundle_path) bundle_path.unlink() with ThreadPoolExecutor(max_workers=4) as pool: results = list( diff --git a/packages/leapp-runtime-rust/src/lib.rs b/packages/leapp-runtime-rust/src/lib.rs index 4030402e..728e16f4 100644 --- a/packages/leapp-runtime-rust/src/lib.rs +++ b/packages/leapp-runtime-rust/src/lib.rs @@ -9,13 +9,13 @@ //! Implementation schema types are intentionally private: //! //! ```compile_fail -//! use leapp_runtime::{BundleDescription, NodeSpec, Pipeline}; +//! use lrt::{BundleDescription, NodeSpec, Pipeline}; //! ``` //! //! Engine implementations are an internal compile-time choice: //! //! ```compile_fail -//! use leapp_runtime::engines; +//! use lrt::engines; //! ``` #![deny(unsafe_op_in_unsafe_fn)] @@ -28,4 +28,4 @@ mod session; pub use bundle::{Bundle, DataType, Metadata, MetadataValue, PortRef, PortSpec}; pub use engines::{ExecutionDevice, MemoryLocation}; pub use error::{BundleError, ErrorCode, RuntimeError}; -pub use session::{OutputPlacement, RunResult, Session, SessionOptions}; +pub use session::{OutputPlacement, OutputPoolOptions, RunResult, Session, SessionOptions}; diff --git a/packages/leapp-runtime-rust/src/session.rs b/packages/leapp-runtime-rust/src/session.rs index 900c4fc3..540fd654 100644 --- a/packages/leapp-runtime-rust/src/session.rs +++ b/packages/leapp-runtime-rust/src/session.rs @@ -16,11 +16,41 @@ pub struct OutputPlacement { pub default_location: MemoryLocation, pub overrides: HashMap, } + +/// Retention policy for each session's reusable output buffers. +/// +/// The runtime derives the byte budget from bundle tensor shapes, dtypes, and +/// placement, growing it for actual dynamic output sizes. There is no fixed +/// byte limit that excludes large tensors from reuse. This policy does not cap +/// memory held by models, fixed workspaces, or retained results. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct OutputPoolOptions { + /// Maximum number of cached buffers. Zero disables caching, not allocation. + pub max_entries: usize, +} + +impl Default for OutputPoolOptions { + fn default() -> Self { + Self { max_entries: 64 } + } +} + /// Immutable construction options, inherited by every new session. +/// +/// ``` +/// use lrt::{OutputPoolOptions, SessionOptions}; +/// +/// let options = SessionOptions { +/// output_pool: OutputPoolOptions { max_entries: 0 }, +/// ..Default::default() +/// }; +/// ``` #[derive(Debug, Clone, Default)] pub struct SessionOptions { pub device: ExecutionDevice, pub output_placement: OutputPlacement, + /// Each session has its own cache; peers inherit this policy, not its buffers. + pub output_pool: OutputPoolOptions, } /// Managed DLPack public outputs in bundle declaration order, independent of the session. @@ -73,6 +103,7 @@ impl Session { /// Create an independently initialized session without reloading models or files. /// Configuration and captured initial values are inherited; current feedback is not. + /// The peer starts with an empty output cache using the same retention policy. pub fn new_session(&self) -> Result { Err(execution_unavailable("load")) } @@ -86,10 +117,27 @@ impl Session { } /// Borrow every public input exactly once until synchronous execution completes. + /// Names and tensors are borrowed from the map; matching is by name, independent + /// of iteration order. Missing or unknown inputs are errors. Inserting a duplicate + /// key replaces its earlier value before the runtime receives the map. /// Feedback advances only after every output and isolation copy succeeds. + /// + /// ``` + /// use lrt::{dlpack::BorrowedTensor, RunResult, RuntimeError, Session}; + /// use std::collections::HashMap; + /// + /// fn step(session: &mut Session) -> Result { + /// let values = [1_f32, 2., 3.]; + /// let name = String::from("model/input"); + /// let inputs = HashMap::from([ + /// (name.as_str(), BorrowedTensor::cpu(&values, &[3]).unwrap()), + /// ]); + /// session.run(&inputs) + /// } + /// ``` pub fn run( &mut self, - _inputs: &[(&str, BorrowedTensor<'_>)], + _inputs: &HashMap<&str, BorrowedTensor<'_>>, ) -> Result { Err(execution_unavailable("execute")) } diff --git a/packages/leapp-runtime-rust/tests/unit_tests/bundle.rs b/packages/leapp-runtime-rust/tests/unit_tests/bundle.rs index 79412c6d..d34f22f7 100644 --- a/packages/leapp-runtime-rust/tests/unit_tests/bundle.rs +++ b/packages/leapp-runtime-rust/tests/unit_tests/bundle.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use leapp_runtime::{Bundle, DataType, ErrorCode, MetadataValue as Value, PortRef}; +use lrt::{Bundle, DataType, ErrorCode, MetadataValue as Value, PortRef}; use safetensors::{tensor::TensorView, Dtype}; use std::{fs, path::PathBuf}; use tempfile::TempDir; @@ -229,7 +229,8 @@ fn collections_cannot_be_null_or_scalar() { #[test] fn accepts_empty_graph_and_unspecified_version() { let mut f = Fixture::new(); - f.value = value("models: {}\npipeline: {inputs: {}, outputs: {}}\nsystem information: {}\n"); + f.value = + value("models: {}\npipeline: {inputs: {}, outputs: {}}\nsystem information: {}\n"); let b = f.open(); assert!(b.model_names().is_empty()); assert!(b.pipeline_inputs().is_empty()); @@ -617,7 +618,8 @@ fn artifacts_fall_back_to_the_working_directory() { let relative_path = artifact_path.strip_prefix(&working_directory).unwrap(); let mut f = Fixture::new(); for model in f.value["models"].as_mapping_mut().unwrap().values_mut() { - model["parameters"]["model_path"] = Value::String(relative_path.to_str().unwrap().into()); + model["parameters"]["model_path"] = + Value::String(relative_path.to_str().unwrap().into()); } f.open(); diff --git a/packages/leapp-runtime-rust/tests/unit_tests/dlpack/abi.rs b/packages/leapp-runtime-rust/tests/unit_tests/dlpack/abi.rs index d412bf5d..1f255cd5 100644 --- a/packages/leapp-runtime-rust/tests/unit_tests/dlpack/abi.rs +++ b/packages/leapp-runtime-rust/tests/unit_tests/dlpack/abi.rs @@ -1,6 +1,6 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use leapp_runtime::dlpack::*; +use lrt::dlpack::*; use std::{ mem::{align_of, offset_of, size_of}, process::Command, diff --git a/packages/leapp-runtime-rust/tests/unit_tests/dlpack/cpu.rs b/packages/leapp-runtime-rust/tests/unit_tests/dlpack/cpu.rs index ca7084b0..963c8d1f 100644 --- a/packages/leapp-runtime-rust/tests/unit_tests/dlpack/cpu.rs +++ b/packages/leapp-runtime-rust/tests/unit_tests/dlpack/cpu.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use leapp_runtime::dlpack::*; +use lrt::dlpack::*; #[test] fn cpu_slices_infer_dtype_and_validate_shape_without_copying() { diff --git a/packages/leapp-runtime-rust/tests/unit_tests/dlpack/ownership.rs b/packages/leapp-runtime-rust/tests/unit_tests/dlpack/ownership.rs index f192b648..761932f6 100644 --- a/packages/leapp-runtime-rust/tests/unit_tests/dlpack/ownership.rs +++ b/packages/leapp-runtime-rust/tests/unit_tests/dlpack/ownership.rs @@ -1,6 +1,6 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use leapp_runtime::dlpack::*; +use lrt::dlpack::*; use std::sync::{ atomic::{AtomicUsize, Ordering}, Arc,