Add typed tensor artifact outputs to V2 - #524
Conversation
Greptile SummaryThe PR adds typed tensor-artifact schemas and a generation-aware model-output sink to runtime V2.
Confidence Score: 5/5The PR appears safe to merge because no blocking failure remains. No blocking failure remains. Important Files Changed
Sequence DiagramsequenceDiagram
participant CLI
participant Runner as Session Runner
participant Model
participant Sink as Tensor Artifact Sink
CLI->>Sink: Construct with output directory
Runner->>Sink: open(session description)
Sink->>Sink: Publish incomplete manifest
loop Model steps
Model-->>Runner: Complete channel batch
Runner->>Sink: write(generation, results)
Sink->>Sink: Validate and buffer artifacts
end
Runner->>Sink: "close(commit=true)"
Sink->>Sink: Stage arrays, backups, and complete manifest
Sink->>Sink: Publish arrays and manifest
Reviews (4): Last reviewed commit: "Address tensor artifact review feedback" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| f"dimensions from {tuple(first.shape)} to {tuple(tensor.shape)}." | ||
| ) | ||
|
|
||
| def close(self) -> None: |
There was a problem hiding this comment.
Is it intentional that a failed run writes its artifacts? run_session's finally (session_runner.py:224) closes every sink regardless of how the run ended, and close() here persists whatever chunks it managed to buffer before re-raising the failure.
For a streaming model that emits one chunk per step, a run that dies at step 6 of 10 leaves an actions.npy with the right name, the right dtype, and a shape that looks perfectly reasonable, nothing in the file says how many chunks were supposed to be in it. The metrics sink and the MP4 writer don't really have this problem, since a short JSON or a short video is obviously short, but a tensor isn't self-describing.
Would it make sense to either skip the write when the run failed, or record somewhere the consumer can see what the file actually contains?
There was a problem hiding this comment.
This was not intentional, and the lifecycle now distinguishes commit from cleanup. ModelOutputSink.close(commit=False) is used when model execution fails, so the tensor sink discards buffered chunks. It writes an incomplete manifest during open() and changes it to complete: true only after successful array publication. An end-to-end regression emits one chunk, fails on the next step, and verifies that no partial .npy file is committed.
Evaluated and implemented by GPT-5.6 Sol via Codex.
| if not outputs: | ||
| return | ||
|
|
||
| self._output_dir.mkdir(parents=True, exist_ok=True) |
There was a problem hiding this comment.
The mkdir happens in close(), so if --tensor-artifact-dir points at something that's already a file, or a directory the process can't write to, the run fails after the whole rollout has finished. open() on line 37 already receives the session and could mkdir(parents=True, exist_ok=True) there instead.
Asking because cli.py:80 already fails fast on the other half of this flag, it rejects the directory before window startup when the application declares no artifacts, so the intent seems to be to catch bad input early.
There was a problem hiding this comment.
Agreed and moved this validation to open(). The sink now creates the directory and atomically writes its initial incomplete manifest before generation, so a path that is already a file or cannot be written fails at startup. The invalid-file destination is covered by a focused regression.
Evaluated and implemented by GPT-5.6 Sol via Codex.
| name: str | ||
| """Stable name used to route and persist the artifact.""" | ||
|
|
||
| dimension_names: tuple[str, ...] |
There was a problem hiding this comment.
dimension_names is doing one functional job here, giving the expected rank so TensorArtifactOutput can check it. The names themselves never reach the output: np.save with allow_pickle=False stores dtype and shape only, so someone loading actions.npy sees (64, 16) and has to read the integration source to find out which axis is which.
Since the API goes to the trouble of declaring them, would it be worth writing them out too? Like a sidecar JSON, or .npz with the names attached? Otherwise they're documentation that only exists on the producing side.
There was a problem hiding this comment.
Agreed. I kept the .npy format for compatibility and added a tensor_artifacts.json sidecar manifest. It preserves each declared artifact dimension name and concatenation axis, and records emitted state, path, dtype, and shape. The manifest also acts as the completeness marker. CLI, sink, and documentation tests verify the consumer-visible metadata.
Evaluated and implemented by GPT-5.6 Sol via Codex.
jarcherNV
left a comment
jarcherNV
left a comment
There was a problem hiding this comment.
Overall LGTM. I left a few comments but I don't think any are blocking. Up to you if you think they are worth addressing or not.
|
Review follow-up: current head 0f0eb2a addresses both Greptile findings and all three Jesse Archer suggestions. Successful runs remove omitted declared artifact files and publish a complete schema manifest; failed model execution discards buffered arrays and leaves an incomplete marker; destination validation happens before generation; and dimension metadata is consumer-visible. Publication now stages backups and restores overwritten or removed prior artifacts if any array or final-manifest publication step fails. After rebasing on current main, the complete V2 CPU suite passes with 276 tests, and pre-commit passes including Ruff and ty. Evaluated and implemented by GPT-5.6 Sol via Codex. |
Signed-off-by: Jonathan McCaffrey <jmccaffrey@nvidia.com>
Signed-off-by: Jonathan McCaffrey <jmccaffrey@nvidia.com>
5a6a109 to
b1be7c9
Compare
b1be7c9 to
0f0eb2a
Compare
2 participants
Summary
ModelOutputSinkpath that receives complete model-step channel batches independently of UI presentation and the legacy metrics sinkContract details
--tensor-artifact-dirfails before application/window startup when an application declares no artifact schemasopen, before generation startstensor_artifacts.jsonrecords completeness plus each declared artifact’s emitted state, dimension names, concatenation axis, path, dtype, and shapecomplete: truemanifest; failed execution discards buffered arrays and leaves the manifest incompleteValidation
python -m pytest flashdreams/test_v2 -m ci_cpu -q— 276 passedpre-commit run --files <changed files>— passed, including Ruff andtyStack
This is the model-neutral prerequisite for #520. It contains no LingBot-specific code, tests, documentation, or assets. PR #520 will be rebased onto this branch and reduced to the LingBot model plus its V2 adapter.