From 35ab8a79cccdeaf8a392f8a558f7df8b56bf7672 Mon Sep 17 00:00:00 2001 From: Pradeep Garigipati Date: Fri, 7 Aug 2026 07:17:17 +0530 Subject: [PATCH 1/2] feat(nvtx): wire viewers and timeline client --- Cargo.lock | 2 + crates/open/src/viewer.rs | 23 ++- crates/open/src/wrapper.rs | 148 ++++++++++++------ examples/simulator/server/Cargo.toml | 1 + examples/simulator/server/src/main.rs | 9 +- examples/simulator/ui-bindings/Cargo.toml | 1 + examples/simulator/ui-bindings/src/lib.rs | 7 +- ui/packages/@quent/client/src/api.ts | 60 ++++++- ui/packages/@quent/client/src/index.ts | 11 ++ ui/packages/@quent/client/src/nvtx.test.ts | 103 ++++++++++++ ui/packages/@quent/client/src/nvtx.ts | 67 ++++++++ .../@quent/client/src/nvtxCanonical.ts | 39 +++++ ui/packages/@quent/utils/src/types/index.ts | 17 ++ 13 files changed, 438 insertions(+), 50 deletions(-) create mode 100644 ui/packages/@quent/client/src/nvtx.test.ts create mode 100644 ui/packages/@quent/client/src/nvtx.ts create mode 100644 ui/packages/@quent/client/src/nvtxCanonical.ts diff --git a/Cargo.lock b/Cargo.lock index 19427e926..2739944e8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3184,6 +3184,7 @@ version = "0.1.0" dependencies = [ "axum", "clap", + "nvtx-server", "quent-io", "quent-query-engine-server", "quent-simulator-analyzer", @@ -3207,6 +3208,7 @@ dependencies = [ name = "quent-simulator-ui-bindings" version = "0.1.0" dependencies = [ + "nvtx-ui", "quent-query-engine-ui", "quent-simulator-ui", "quent-ui", diff --git a/crates/open/src/viewer.rs b/crates/open/src/viewer.rs index 5473d94a6..d9f2c49e1 100644 --- a/crates/open/src/viewer.rs +++ b/crates/open/src/viewer.rs @@ -90,7 +90,7 @@ async fn build_one(group: ViewerGroup) -> Result { println!("building: {label}"); let crate_dir = build_dir(&spec)?; - wrapper::generate(&spec, &crate_dir, wrapper::IO_PACKAGE)?; + wrapper::generate(&spec, &crate_dir, wrapper::IO_PACKAGE, true)?; let bin = match cargo_build(&crate_dir).await { // The pinned quent revision predates the `quent-exporter` → `quent-io` // rename (cargo found no `quent-io` package there, failing resolution @@ -102,7 +102,26 @@ async fn build_one(group: ViewerGroup) -> Result { wrapper::IO_PACKAGE, wrapper::LEGACY_IO_PACKAGE ); - wrapper::generate(&spec, &crate_dir, wrapper::LEGACY_IO_PACKAGE)?; + wrapper::generate(&spec, &crate_dir, wrapper::LEGACY_IO_PACKAGE, true)?; + cargo_build(&crate_dir).await? + } + Err(error) if missing_package(&error, "nvtx-server") => { + println!( + "note: pinned quent has no `nvtx-server` package; retrying without NVTX routes" + ); + wrapper::generate(&spec, &crate_dir, wrapper::IO_PACKAGE, false)?; + cargo_build(&crate_dir).await? + } + Err(error) + if missing_package(&error, wrapper::IO_PACKAGE) + || missing_package(&error, wrapper::LEGACY_IO_PACKAGE) => + { + let io_package = if missing_package(&error, wrapper::LEGACY_IO_PACKAGE) { + wrapper::LEGACY_IO_PACKAGE + } else { + wrapper::IO_PACKAGE + }; + wrapper::generate(&spec, &crate_dir, io_package, false)?; cargo_build(&crate_dir).await? } result => result?, diff --git a/crates/open/src/wrapper.rs b/crates/open/src/wrapper.rs index 443cc750c..aa8e82eb0 100644 --- a/crates/open/src/wrapper.rs +++ b/crates/open/src/wrapper.rs @@ -35,10 +35,21 @@ pub const ADDR_ENV: &str = "QUENT_OPEN_ADDR"; /// Write the wrapper crate (`Cargo.toml` + `src/main.rs`) into `crate_dir`. /// `io_package` is the name of quent's I/O crate at the pinned revision /// ([`IO_PACKAGE`], or [`LEGACY_IO_PACKAGE`] for revisions predating the rename). -pub fn generate(spec: &ViewerSpec, crate_dir: &Path, io_package: &str) -> Result<()> { +pub fn generate( + spec: &ViewerSpec, + crate_dir: &Path, + io_package: &str, + with_nvtx_routes: bool, +) -> Result<()> { std::fs::create_dir_all(crate_dir.join("src"))?; - std::fs::write(crate_dir.join("Cargo.toml"), cargo_toml(spec, io_package))?; - std::fs::write(crate_dir.join("src/main.rs"), main_rs(spec))?; + std::fs::write( + crate_dir.join("Cargo.toml"), + cargo_toml(spec, io_package, with_nvtx_routes), + )?; + std::fs::write( + crate_dir.join("src/main.rs"), + main_rs(spec, with_nvtx_routes), + )?; Ok(()) } @@ -55,10 +66,10 @@ fn git_dep(url: String, rev: &str, features: &[&str]) -> Dependency { /// Wrapper `Cargo.toml`, built with `cargo-manifest`: pin quent crates to /// `quent.{remote,commit}` and the analyzer to `analyzer.{remote,commit}`; the /// empty `[workspace]` keeps the generated crate out of any parent workspace. -fn cargo_toml(spec: &ViewerSpec, io_package: &str) -> String { +fn cargo_toml(spec: &ViewerSpec, io_package: &str, with_nvtx_routes: bool) -> String { let quent = spec.quent.cargo_url(); let q_rev = spec.quent.commit.as_str(); - let dependencies = BTreeMap::from([ + let mut dependencies = BTreeMap::from([ ( "quent-query-engine-server".to_string(), git_dep(quent.clone(), q_rev, &["ui"]), @@ -70,7 +81,7 @@ fn cargo_toml(spec: &ViewerSpec, io_package: &str) -> String { ( // All formats enabled so the analyzer can detect the artifact's format at runtime. io_package.to_string(), - git_dep(quent, q_rev, &["ndjson", "msgpack", "postcard"]), + git_dep(quent.clone(), q_rev, &["ndjson", "msgpack", "postcard"]), ), ( spec.analyzer_package.clone(), @@ -92,6 +103,9 @@ fn cargo_toml(spec: &ViewerSpec, io_package: &str) -> String { ), ("uuid".to_string(), Dependency::Simple("1".to_string())), ]); + if with_nvtx_routes { + dependencies.insert("nvtx-server".to_string(), git_dep(quent, q_rev, &[])); + } let mut package = Package::new(WRAPPER_PACKAGE.to_string(), "0.0.0".to_string()); package.edition = Some(MaybeInherited::Local(Edition::E2024)); @@ -115,43 +129,85 @@ fn cargo_toml(spec: &ViewerSpec, io_package: &str) -> String { /// Wrapper `src/main.rs`: wire `::Viewer`'s analyzer/importer into /// `analyzer_service_router` and serve it. Root (`/` subdirs) and /// bind address come from env so one built binary serves any artifacts. -fn main_rs(spec: &ViewerSpec) -> String { +fn main_rs(spec: &ViewerSpec, with_nvtx_routes: bool) -> String { let analyzer_crate = format_ident!("{}", spec.analyzer_crate()); let (root_env, addr_env) = (ROOT_ENV, ADDR_ENV); - let tokens = quote! { - use std::net::SocketAddr; - use std::path::PathBuf; - - use quent_query_engine_analyzer::ui::QuentViewer; - use quent_query_engine_server::analyzer_cache::index_query_engines; - use quent_query_engine_server::analyzer_service_router; - use #analyzer_crate::Viewer; - - type Analyzer = ::Analyzer; - - #[tokio::main] - async fn main() -> Result<(), Box> { - let root = PathBuf::from(std::env::var(#root_env)?); - let addr: SocketAddr = std::env::var(#addr_env)?.parse()?; - - let import_root = root.clone(); - let importer = move |id: uuid::Uuid| { - Ok(::import_events( - &import_root.join(id.to_string()), - )?) - }; - let lister_root = root.clone(); - let lister = move || index_query_engines(&lister_root); - - let router = analyzer_service_router::( - Box::new(importer), - Box::new(lister), - None, - )?; - - let listener = tokio::net::TcpListener::bind(addr).await?; - axum::serve(listener, router.into_make_service()).await?; - Ok(()) + let tokens = if with_nvtx_routes { + quote! { + use std::net::SocketAddr; + use std::path::PathBuf; + + use quent_query_engine_analyzer::ui::QuentViewer; + use quent_query_engine_server::analyzer_cache::index_query_engines; + use quent_query_engine_server::analyzer_service_router_with_routes; + use nvtx_server::{import_context_events, routes as nvtx_routes}; + use #analyzer_crate::Viewer; + + type Analyzer = ::Analyzer; + + #[tokio::main] + async fn main() -> Result<(), Box> { + let root = PathBuf::from(std::env::var(#root_env)?); + let addr: SocketAddr = std::env::var(#addr_env)?.parse()?; + + let import_root = root.clone(); + let importer = move |id: uuid::Uuid| { + Ok(::import_events( + &import_root.join(id.to_string()), + &import_root.join(id.to_string()), + )?) + }; + let lister_root = root.clone(); + let lister = move || index_query_engines(&lister_root); + let nvtx_root = root.clone(); + let nvtx_importer = move |id: uuid::Uuid| import_context_events(&nvtx_root, id); + + let router = analyzer_service_router_with_routes::( + Box::new(importer), + Box::new(lister), + None, + nvtx_routes(Box::new(nvtx_importer)), + )?; + + let listener = tokio::net::TcpListener::bind(addr).await?; + axum::serve(listener, router.into_make_service()).await?; + Ok(()) + } + } + } else { + quote! { + use std::net::SocketAddr; + use std::path::PathBuf; + + use quent_query_engine_analyzer::ui::QuentViewer; + use quent_query_engine_server::analyzer_cache::index_query_engines; + use quent_query_engine_server::analyzer_service_router; + use #analyzer_crate::Viewer; + + type Analyzer = ::Analyzer; + + #[tokio::main] + async fn main() -> Result<(), Box> { + let root = PathBuf::from(std::env::var(#root_env)?); + let addr: SocketAddr = std::env::var(#addr_env)?.parse()?; + + let import_root = root.clone(); + let importer = move |id: uuid::Uuid| { + Ok(::import_events(&import_root.join(id.to_string()))?) + }; + let lister_root = root.clone(); + let lister = move || index_query_engines(&lister_root); + + let router = analyzer_service_router::( + Box::new(importer), + Box::new(lister), + None, + )?; + + let listener = tokio::net::TcpListener::bind(addr).await?; + axum::serve(listener, router.into_make_service()).await?; + Ok(()) + } } }; let file = syn::parse2(tokens).expect("generated wrapper main.rs is valid Rust"); @@ -182,13 +238,14 @@ mod tests { #[test] fn cargo_toml_pins_quent_and_analyzer() { - let manifest: toml::Value = toml::from_str(&cargo_toml(&spec(), IO_PACKAGE)).unwrap(); + let manifest: toml::Value = toml::from_str(&cargo_toml(&spec(), IO_PACKAGE, true)).unwrap(); assert!(manifest.get("workspace").is_some(), "standalone workspace"); let deps = &manifest["dependencies"]; let server = &deps["quent-query-engine-server"]; assert_eq!(server["git"].as_str().unwrap(), "https://example.com/quent"); assert_eq!(server["rev"].as_str().unwrap(), "quentcommit"); assert_eq!(server["features"][0].as_str().unwrap(), "ui"); + assert_eq!(deps["nvtx-server"]["rev"].as_str().unwrap(), "quentcommit"); // The exporter enables all formats so the analyzer detects the artifact's format at runtime. let exporter_features = deps["quent-io"]["features"].as_array().unwrap(); for format in ["ndjson", "msgpack", "postcard"] { @@ -207,10 +264,11 @@ mod tests { // Artifacts pinned to quent revisions predating the `quent-exporter` → // `quent-io` rename depend on the legacy package instead, same features. let manifest: toml::Value = - toml::from_str(&cargo_toml(&spec(), LEGACY_IO_PACKAGE)).unwrap(); + toml::from_str(&cargo_toml(&spec(), LEGACY_IO_PACKAGE, false)).unwrap(); let deps = &manifest["dependencies"]; assert!(deps.get("quent-io").is_none()); let exporter = &deps["quent-exporter"]; + assert!(deps.get("nvtx-server").is_none()); assert_eq!( exporter["git"].as_str().unwrap(), "https://example.com/quent" @@ -224,9 +282,11 @@ mod tests { #[test] fn main_rs_wires_the_viewer() { - let main = main_rs(&spec()); + let main = main_rs(&spec(), true); assert!(main.contains("use quent_simulator_analyzer::Viewer;")); assert!(main.contains("import_events")); + assert!(main.contains("import_context_events")); + assert!(main.contains("analyzer_service_router_with_routes")); assert!(main.contains("QUENT_OPEN_ADDR")); // bind address is configurable } } diff --git a/examples/simulator/server/Cargo.toml b/examples/simulator/server/Cargo.toml index fa3692f18..9ba8582a3 100644 --- a/examples/simulator/server/Cargo.toml +++ b/examples/simulator/server/Cargo.toml @@ -10,6 +10,7 @@ swagger = ["quent-query-engine-server/swagger"] [dependencies] axum = { version = "0.8.7" } clap = { version = "4.5", features = ["derive", "env"] } +nvtx-server = { path = "../../../integrations/nvtx/server" } quent-io = { path = "../../../crates/io" } quent-query-engine-server = { path = "../../../domains/query_engine/server" } quent-simulator-analyzer = { path = "../analyzer" } diff --git a/examples/simulator/server/src/main.rs b/examples/simulator/server/src/main.rs index b367d5eff..a6acf61d2 100644 --- a/examples/simulator/server/src/main.rs +++ b/examples/simulator/server/src/main.rs @@ -4,10 +4,11 @@ use std::{net::ToSocketAddrs, path::PathBuf}; use clap::Parser; +use nvtx_server::{import_context_events, routes as nvtx_routes}; use quent_io::ExporterOptions; use quent_io::filesystem::{self, Format}; use quent_query_engine_server::{ - analyzer_cache::index_query_engines, analyzer_service_router, collector_service, + analyzer_cache::index_query_engines, analyzer_service_router_with_routes, collector_service, initialize_tracing, }; use quent_simulator_analyzer::SimulatorUiAnalyzer; @@ -88,6 +89,7 @@ async fn main() -> Result<(), Box> { let importer_output_dir = output_dir.clone(); let lister_output_dir = output_dir.clone(); + let nvtx_output_dir = output_dir.clone(); let format = match exporter.as_str() { "ndjson" => Format::Ndjson, @@ -131,10 +133,13 @@ async fn main() -> Result<(), Box> { let analyzer = async { axum::serve( TcpListener::bind(analyzer_addr).await?, - analyzer_service_router::( + analyzer_service_router_with_routes::( Box::new(importer), Box::new(lister), cors_address, + nvtx_routes(Box::new(move |context_id| { + import_context_events(&nvtx_output_dir, context_id) + })), )? .into_make_service(), ) diff --git a/examples/simulator/ui-bindings/Cargo.toml b/examples/simulator/ui-bindings/Cargo.toml index 36b3c30bd..4632e7f8b 100644 --- a/examples/simulator/ui-bindings/Cargo.toml +++ b/examples/simulator/ui-bindings/Cargo.toml @@ -5,6 +5,7 @@ edition.workspace = true publish.workspace = true [dependencies] +nvtx-ui = { path = "../../../integrations/nvtx/ui" } quent-query-engine-ui = { path = "../../../domains/query_engine/ui" } quent-simulator-ui = { path = "../ui" } quent-ui = { path = "../../../crates/ui" } diff --git a/examples/simulator/ui-bindings/src/lib.rs b/examples/simulator/ui-bindings/src/lib.rs index 600cacb14..f2d77af29 100644 --- a/examples/simulator/ui-bindings/src/lib.rs +++ b/examples/simulator/ui-bindings/src/lib.rs @@ -5,8 +5,9 @@ use std::path::Path; +use nvtx_ui::{NvtxCatalog, NvtxViewportRequest, NvtxViewportResponse}; use quent_query_engine_ui::DataFlowTimelineBinned; -use quent_query_engine_ui::{OperatorFilter, QueryBundle, QueryFilter}; +use quent_query_engine_ui::{EngineContexts, OperatorFilter, QueryBundle, QueryFilter}; use quent_simulator_ui::EntityRef; use quent_ui::entities::{request::EntityListRequest, response::EntityListResponse}; use quent_ui::timeline::{ @@ -34,6 +35,10 @@ pub fn generate(output_dir: &Path) -> Result<(), Box> { ::export_all(&cfg)?; as TS>::export_all(&cfg)?; ::export_all(&cfg)?; + ::export_all(&cfg)?; + ::export_all(&cfg)?; + ::export_all(&cfg)?; + ::export_all(&cfg)?; as TS>::export_all(&cfg)?; ::export_all(&cfg)?; diff --git a/ui/packages/@quent/client/src/api.ts b/ui/packages/@quent/client/src/api.ts index 9c224225a..284ca46e8 100644 --- a/ui/packages/@quent/client/src/api.ts +++ b/ui/packages/@quent/client/src/api.ts @@ -3,6 +3,7 @@ import { parseJsonWithBigInt } from '@quent/utils'; import { getApiBaseUrl } from './config'; +import { canonicalizeNvtxRequest } from './nvtxCanonical'; import type { QueryBundle, QueryGroup, @@ -20,10 +21,14 @@ import type { TimelineConfig, EntityListRequest, EntityListResponse, + EngineContexts, + NvtxCatalog, + NvtxViewportRequest, + NvtxViewportResponse, } from '@quent/utils'; interface ApiFetchOptions { - params?: Record; + params?: Record; fetchOptions?: RequestInit; } @@ -86,6 +91,59 @@ export async function fetchListEngines(): Promise { return apiFetch('/engines', { params: { with_metadata: true } }); } +export async function fetchEngineContexts(engineId: string): Promise { + return apiFetch(`/engines/${engineId}/contexts`); +} + +/** Fetch stable NVTX metadata, resolving a 404 to optional absence. */ +export async function fetchNvtxCatalog( + contextId: string, + queryStartUnixNs: bigint +): Promise { + const response = await apiFetchResponse(`/nvtx/contexts/${contextId}/catalog`, { + params: { query_start: queryStartUnixNs }, + }); + if (response.status === 404) return null; + if (!response.ok) { + throw new Error(`API Error: ${response.status} ${response.statusText}`); + } + return parseJsonWithBigInt(await response.text()); +} + +export async function fetchNvtxViewport( + contextId: string, + queryStartUnixNs: bigint, + request: NvtxViewportRequest +): Promise { + const canonical = canonicalizeNvtxRequest(request); + const response = await apiFetchResponse(`/nvtx/contexts/${contextId}/viewport`, { + params: { query_start: queryStartUnixNs }, + fetchOptions: { + method: 'POST', + body: JSON.stringify(canonical), + }, + }); + if (!response.ok) { + throw new Error(`API Error: ${response.status} ${response.statusText}`); + } + return normalizeNvtxViewport(parseJsonWithBigInt(await response.text())); +} + +function asBigInt(value: bigint | number): bigint { + return typeof value === 'bigint' ? value : BigInt(value); +} + +function normalizeNvtxViewport(viewport: NvtxViewportResponse): NvtxViewportResponse { + return { + ...viewport, + statistics: viewport.statistics.map(statistics => ({ + ...statistics, + count: asBigInt(statistics.count), + observed_count: asBigInt(statistics.observed_count), + })), + }; +} + export async function fetchListCoordinators(engineId: string): Promise { return apiFetch(`/engines/${engineId}/query-groups`); } diff --git a/ui/packages/@quent/client/src/index.ts b/ui/packages/@quent/client/src/index.ts index 0a6085cb3..2790a70e9 100644 --- a/ui/packages/@quent/client/src/index.ts +++ b/ui/packages/@quent/client/src/index.ts @@ -15,6 +15,9 @@ export { fetchBulkTimelines, fetchDataFlow, fetchEntityList, + fetchEngineContexts, + fetchNvtxCatalog, + fetchNvtxViewport, } from './api'; // queryOptions factories @@ -26,6 +29,13 @@ export { singleTimelineQueryOptions } from './timeline'; export { bulkTimelineQueryOptions } from './bulkTimelines'; export { dataFlowQueryOptions } from './dataFlow'; export { entityListInfiniteQueryOptions, entityListQueryOptions } from './entityList'; +export { + canonicalizeNvtxRequest, + canonicalizeNvtxSelections, + engineContextsQueryOptions, + nvtxCatalogQueryOptions, + nvtxViewportQueryOptions, +} from './nvtx'; // Hooks export { useQueryBundle } from './queryBundle'; @@ -35,3 +45,4 @@ export { useQueries } from './queries'; export { useTimeline } from './timeline'; export { useDataFlow } from './dataFlow'; export { useEntityList, useInfiniteEntityList } from './entityList'; +export { useEngineContexts, useNvtxCatalog, useNvtxViewport } from './nvtx'; diff --git a/ui/packages/@quent/client/src/nvtx.test.ts b/ui/packages/@quent/client/src/nvtx.test.ts new file mode 100644 index 000000000..74c78b8e3 --- /dev/null +++ b/ui/packages/@quent/client/src/nvtx.test.ts @@ -0,0 +1,103 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { afterEach, describe, expect, it, vi } from 'vitest'; +import type { NvtxViewportRequest } from '@quent/utils'; +import { fetchNvtxCatalog, fetchNvtxViewport } from './api'; +import { canonicalizeNvtxRequest, nvtxCatalogQueryOptions, nvtxViewportQueryOptions } from './nvtx'; + +function stubFetch(response: Response) { + const fetchMock = vi.fn().mockResolvedValue(response); + vi.stubGlobal('fetch', fetchMock); + return fetchMock; +} + +const QUERY_START_UNIX_NS = 1_800_000_000_000_000_001n; + +const request: NvtxViewportRequest = { + viewport: { start: -0.25, end: 1.5 }, + selections: [ + { domain_id: '9', category_ids: [7, 3, 7], include_uncategorized: false }, + { domain_id: '2', category_ids: [], include_uncategorized: true }, + ], +}; + +describe('NVTX client', () => { + afterEach(() => vi.unstubAllGlobals()); + + it('treats catalog 404 as optional absence and fetches catalogs once', async () => { + const fetchMock = stubFetch(new Response(null, { status: 404, statusText: 'Not Found' })); + await expect(fetchNvtxCatalog('context-1', QUERY_START_UNIX_NS)).resolves.toBeNull(); + expect(fetchMock.mock.calls[0]?.[0]).toContain( + `/api/nvtx/contexts/context-1/catalog?query_start=${QUERY_START_UNIX_NS}` + ); + expect(nvtxCatalogQueryOptions('context-1', QUERY_START_UNIX_NS).staleTime).toBe(Infinity); + }); + + it('propagates non-404 catalog failures', async () => { + stubFetch(new Response(null, { status: 500, statusText: 'Internal Server Error' })); + await expect(fetchNvtxCatalog('context-1', QUERY_START_UNIX_NS)).rejects.toThrow( + 'API Error: 500 Internal Server Error' + ); + }); + + it('uses the same canonical selector order for body and query key', async () => { + const fetchMock = stubFetch( + new Response('{"viewport":{"start":-0.25,"end":1.5},"domains":[],"statistics":[]}', { + status: 200, + }) + ); + const canonical = canonicalizeNvtxRequest(request); + expect(canonical.selections.map(selection => selection.domain_id)).toEqual(['2', '9']); + expect(canonical.selections[1].category_ids).toEqual([3, 7]); + + await fetchNvtxViewport('context-1', QUERY_START_UNIX_NS, request); + const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]; + expect(url).toContain( + `/api/nvtx/contexts/context-1/viewport?query_start=${QUERY_START_UNIX_NS}` + ); + expect(init.body).toBe( + '{"viewport":{"start":-0.25,"end":1.5},"selections":[{"domain_id":"2","category_ids":[],"include_uncategorized":true},{"domain_id":"9","category_ids":[3,7],"include_uncategorized":false}]}' + ); + + const options = nvtxViewportQueryOptions('context-1', QUERY_START_UNIX_NS, request); + expect(options.queryKey).toEqual([ + 'nvtxViewport', + 'context-1', + QUERY_START_UNIX_NS.toString(10), + -0.25, + 1.5, + [ + ['2', [], true], + ['9', [3, 7], false], + ], + ]); + expect(options.placeholderData).toBeTypeOf('function'); + }); + + it('preserves relative seconds and decimal-string identifiers from the catalog', async () => { + stubFetch( + new Response( + '{"trace_start":-0.5,"trace_end":2,"domains":[{"domain_id":"3","name":"d","color":"#000000","threads":[],"categories":[],"has_uncategorized":false}],"anomalies":{"orphan_range_ends":"0","orphan_range_pops":"0","orphan_resource_destroys":"0","reused_range_ids":"0","reused_resource_handles":"0","total":"0","is_faithful":true}}', + { status: 200 } + ) + ); + const catalog = await fetchNvtxCatalog('context-1', QUERY_START_UNIX_NS); + expect(catalog?.trace_start).toBe(-0.5); + expect(catalog?.domains[0].domain_id).toBe('3'); + expect(catalog?.anomalies.total).toBe('0'); + }); + + it('normalizes even safe u64 statistic counts to bigint', async () => { + stubFetch( + new Response( + '{"viewport":{"start":-0.25,"end":1.5},"domains":[],"statistics":[{"message":"work","domain_id":"3","domain_name":"d","category_id":null,"category_name":null,"count":1,"observed_count":1,"total_duration":1.25,"avg_duration":1.25,"min_duration":1.25,"max_duration":1.25,"saturated":false}]}', + { status: 200 } + ) + ); + const viewport = await fetchNvtxViewport('context-1', QUERY_START_UNIX_NS, request); + expect(viewport.statistics[0]?.count).toBe(1n); + expect(viewport.statistics[0]?.observed_count).toBe(1n); + expect(viewport.statistics[0]?.total_duration).toBe(1.25); + }); +}); diff --git a/ui/packages/@quent/client/src/nvtx.ts b/ui/packages/@quent/client/src/nvtx.ts new file mode 100644 index 000000000..b530dcb9c --- /dev/null +++ b/ui/packages/@quent/client/src/nvtx.ts @@ -0,0 +1,67 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { keepPreviousData, queryOptions, useQuery } from '@tanstack/react-query'; +import type { NvtxViewportRequest } from '@quent/utils'; +import { fetchEngineContexts, fetchNvtxCatalog, fetchNvtxViewport } from './api'; +import { DEFAULT_STALE_TIME } from './constants'; +import { canonicalizeNvtxRequest } from './nvtxCanonical'; +export { canonicalizeNvtxRequest, canonicalizeNvtxSelections } from './nvtxCanonical'; + +export const engineContextsQueryOptions = (engineId: string) => + queryOptions({ + queryKey: ['engineContexts', engineId], + queryFn: () => fetchEngineContexts(engineId), + staleTime: DEFAULT_STALE_TIME, + }); + +export const nvtxCatalogQueryOptions = (contextId: string, queryStartUnixNs: bigint) => { + const queryStartKey = queryStartUnixNs.toString(10); + return queryOptions({ + queryKey: ['nvtxCatalog', contextId, queryStartKey], + queryFn: () => fetchNvtxCatalog(contextId, queryStartUnixNs), + staleTime: Infinity, + }); +}; + +export const nvtxViewportQueryOptions = ( + contextId: string, + queryStartUnixNs: bigint, + request: NvtxViewportRequest, + options?: { enabled?: boolean; staleTime?: number } +) => { + const canonical = canonicalizeNvtxRequest(request); + const queryStartKey = queryStartUnixNs.toString(10); + const selectionKey = canonical.selections.map(selection => [ + selection.domain_id, + selection.category_ids, + selection.include_uncategorized, + ]); + return queryOptions({ + queryKey: [ + 'nvtxViewport', + contextId, + queryStartKey, + canonical.viewport.start, + canonical.viewport.end, + selectionKey, + ], + queryFn: () => fetchNvtxViewport(contextId, queryStartUnixNs, canonical), + enabled: options?.enabled ?? true, + staleTime: options?.staleTime ?? DEFAULT_STALE_TIME, + placeholderData: keepPreviousData, + }); +}; + +export const useEngineContexts = (engineId: string) => + useQuery(engineContextsQueryOptions(engineId)); + +export const useNvtxCatalog = (contextId: string, queryStartUnixNs: bigint) => + useQuery(nvtxCatalogQueryOptions(contextId, queryStartUnixNs)); + +export const useNvtxViewport = ( + contextId: string, + queryStartUnixNs: bigint, + request: NvtxViewportRequest, + options?: { enabled?: boolean; staleTime?: number } +) => useQuery(nvtxViewportQueryOptions(contextId, queryStartUnixNs, request, options)); diff --git a/ui/packages/@quent/client/src/nvtxCanonical.ts b/ui/packages/@quent/client/src/nvtxCanonical.ts new file mode 100644 index 000000000..e29304679 --- /dev/null +++ b/ui/packages/@quent/client/src/nvtxCanonical.ts @@ -0,0 +1,39 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import type { NvtxDomainSelection, NvtxViewportRequest } from '@quent/utils'; + +export function canonicalizeNvtxSelections( + selections: readonly NvtxDomainSelection[] +): NvtxDomainSelection[] { + const seen = new Set(); + const canonical = selections.map(selection => { + const domainKey = BigInt(selection.domain_id).toString(10); + if (seen.has(domainKey)) throw new Error(`duplicate NVTX domain ${domainKey}`); + seen.add(domainKey); + const category_ids = [...new Set(selection.category_ids)].sort((left, right) => left - right); + if (category_ids.length === 0 && !selection.include_uncategorized) { + throw new Error(`NVTX domain ${domainKey} selects no categories`); + } + return { ...selection, domain_id: domainKey, category_ids }; + }); + canonical.sort((left, right) => + BigInt(left.domain_id) < BigInt(right.domain_id) + ? -1 + : BigInt(left.domain_id) > BigInt(right.domain_id) + ? 1 + : 0 + ); + return canonical; +} + +export function canonicalizeNvtxRequest(request: NvtxViewportRequest): NvtxViewportRequest { + const { start, end } = request.viewport; + if (!Number.isFinite(start) || !Number.isFinite(end) || start > end) { + throw new Error('NVTX viewport bounds must be finite and ordered'); + } + return { + viewport: { start, end }, + selections: canonicalizeNvtxSelections(request.selections), + }; +} diff --git a/ui/packages/@quent/utils/src/types/index.ts b/ui/packages/@quent/utils/src/types/index.ts index 6cd8e70fc..2be404ef2 100644 --- a/ui/packages/@quent/utils/src/types/index.ts +++ b/ui/packages/@quent/utils/src/types/index.ts @@ -15,6 +15,7 @@ export type { DataFlowTimelineBinned } from '../../../../../generated/ts-binding export type { DimensionKeyDecl } from '../../../../../generated/ts-bindings/DimensionKeyDecl'; export type { Edge } from '../../../../../generated/ts-bindings/Edge'; export type { Engine } from '../../../../../generated/ts-bindings/Engine'; +export type { EngineContexts } from '../../../../../generated/ts-bindings/EngineContexts'; export type { EngineImplementationAttributes } from '../../../../../generated/ts-bindings/EngineImplementationAttributes'; export type { EntityFilter } from '../../../../../generated/ts-bindings/EntityFilter'; export type { EntityListEntry } from '../../../../../generated/ts-bindings/EntityListEntry'; @@ -69,3 +70,19 @@ export type { TimelineRequest } from '../../../../../generated/ts-bindings/Timel export type { TimeWindow } from '../../../../../generated/ts-bindings/TimeWindow'; export type { DynamicValue } from '../../../../../generated/ts-bindings/DynamicValue'; export type { Worker } from '../../../../../generated/ts-bindings/Worker'; +export type { NvtxCatalog } from '../../../../../generated/ts-bindings/NvtxCatalog'; +export type { NvtxCatalogAnomalies } from '../../../../../generated/ts-bindings/NvtxCatalogAnomalies'; +export type { NvtxCatalogCategory } from '../../../../../generated/ts-bindings/NvtxCatalogCategory'; +export type { NvtxCatalogDomain } from '../../../../../generated/ts-bindings/NvtxCatalogDomain'; +export type { NvtxCatalogThread } from '../../../../../generated/ts-bindings/NvtxCatalogThread'; +export type { NvtxDomainLaneGroup } from '../../../../../generated/ts-bindings/NvtxDomainLaneGroup'; +export type { NvtxDomainSelection } from '../../../../../generated/ts-bindings/NvtxDomainSelection'; +export type { NvtxLane } from '../../../../../generated/ts-bindings/NvtxLane'; +export type { NvtxLaneIdentity } from '../../../../../generated/ts-bindings/NvtxLaneIdentity'; +export type { NvtxMarkItem } from '../../../../../generated/ts-bindings/NvtxMarkItem'; +export type { NvtxRangeItem } from '../../../../../generated/ts-bindings/NvtxRangeItem'; +export type { NvtxRangeKind } from '../../../../../generated/ts-bindings/NvtxRangeKind'; +export type { NvtxRangeStatistics } from '../../../../../generated/ts-bindings/NvtxRangeStatistics'; +export type { NvtxViewportRequest } from '../../../../../generated/ts-bindings/NvtxViewportRequest'; +export type { NvtxViewportResponse } from '../../../../../generated/ts-bindings/NvtxViewportResponse'; +export type { NvtxViewportWindow } from '../../../../../generated/ts-bindings/NvtxViewportWindow'; From e41737cb76c3c4a2358d35f0122f7db0f2fa1d90 Mon Sep 17 00:00:00 2001 From: Pradeep Garigipati Date: Thu, 13 Aug 2026 21:02:51 +0530 Subject: [PATCH 2/2] fix(nvtx): reconcile viewer and client integration --- crates/open/src/viewer.rs | 66 ++++--- crates/open/src/wrapper.rs | 202 ++++++++++++--------- ui/packages/@quent/client/src/nvtx.test.ts | 25 ++- ui/packages/@quent/client/src/nvtx.ts | 23 ++- 4 files changed, 183 insertions(+), 133 deletions(-) diff --git a/crates/open/src/viewer.rs b/crates/open/src/viewer.rs index d9f2c49e1..0a7f13c2b 100644 --- a/crates/open/src/viewer.rs +++ b/crates/open/src/viewer.rs @@ -90,41 +90,39 @@ async fn build_one(group: ViewerGroup) -> Result { println!("building: {label}"); let crate_dir = build_dir(&spec)?; - wrapper::generate(&spec, &crate_dir, wrapper::IO_PACKAGE, true)?; - let bin = match cargo_build(&crate_dir).await { - // The pinned quent revision predates the `quent-exporter` → `quent-io` - // rename (cargo found no `quent-io` package there, failing resolution - // before anything compiles); regenerate the wrapper against the legacy - // package name and build again. - Err(error) if missing_package(&error, wrapper::IO_PACKAGE) => { - println!( - "note: pinned quent has no `{}` package; retrying with `{}`", - wrapper::IO_PACKAGE, - wrapper::LEGACY_IO_PACKAGE - ); - wrapper::generate(&spec, &crate_dir, wrapper::LEGACY_IO_PACKAGE, true)?; - cargo_build(&crate_dir).await? - } - Err(error) if missing_package(&error, "nvtx-server") => { - println!( - "note: pinned quent has no `nvtx-server` package; retrying without NVTX routes" - ); - wrapper::generate(&spec, &crate_dir, wrapper::IO_PACKAGE, false)?; - cargo_build(&crate_dir).await? - } - Err(error) - if missing_package(&error, wrapper::IO_PACKAGE) - || missing_package(&error, wrapper::LEGACY_IO_PACKAGE) => - { - let io_package = if missing_package(&error, wrapper::LEGACY_IO_PACKAGE) { - wrapper::LEGACY_IO_PACKAGE - } else { - wrapper::IO_PACKAGE - }; - wrapper::generate(&spec, &crate_dir, io_package, false)?; - cargo_build(&crate_dir).await? + let mut io_package = wrapper::IO_PACKAGE; + let mut nvtx_routes = wrapper::NvtxRoutes::Enabled; + let bin = loop { + wrapper::generate(&spec, &crate_dir, io_package, nvtx_routes)?; + match cargo_build(&crate_dir).await { + Ok(bin) => break bin, + Err(error) + if nvtx_routes == wrapper::NvtxRoutes::Enabled + && missing_package(&error, wrapper::NVTX_SERVER_PACKAGE) => + { + println!( + "note: pinned quent has no `{}` package; retrying without NVTX routes", + wrapper::NVTX_SERVER_PACKAGE + ); + nvtx_routes = wrapper::NvtxRoutes::Disabled; + } + // The pinned quent revision predates both the `quent-exporter` → + // `quent-io` rename and the NVTX routes. Switch both capabilities + // before retrying, regardless of which missing package Cargo reports first. + Err(error) + if io_package == wrapper::IO_PACKAGE + && missing_package(&error, wrapper::IO_PACKAGE) => + { + println!( + "note: pinned quent has no `{}` package; retrying with `{}` and without NVTX routes", + wrapper::IO_PACKAGE, + wrapper::LEGACY_IO_PACKAGE + ); + io_package = wrapper::LEGACY_IO_PACKAGE; + nvtx_routes = wrapper::NvtxRoutes::Disabled; + } + Err(error) => return Err(error), } - result => result?, }; Ok(BuiltViewer { bin, diff --git a/crates/open/src/wrapper.rs b/crates/open/src/wrapper.rs index aa8e82eb0..d7428b4b2 100644 --- a/crates/open/src/wrapper.rs +++ b/crates/open/src/wrapper.rs @@ -25,6 +25,21 @@ pub const IO_PACKAGE: &str = "quent-io"; /// Cargo package of the I/O crate before its rename to [`IO_PACKAGE`], for /// artifacts pinned to quent revisions that predate the rename. pub const LEGACY_IO_PACKAGE: &str = "quent-exporter"; +/// Cargo package that provides the optional NVTX HTTP routes. +pub const NVTX_SERVER_PACKAGE: &str = "nvtx-server"; + +/// Whether a generated wrapper targets a quent revision with NVTX routes. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum NvtxRoutes { + Enabled, + Disabled, +} + +impl NvtxRoutes { + fn is_enabled(self) -> bool { + self == Self::Enabled + } +} /// Wrapper env var for the output root: a directory of `/` /// context directories. @@ -39,17 +54,14 @@ pub fn generate( spec: &ViewerSpec, crate_dir: &Path, io_package: &str, - with_nvtx_routes: bool, + nvtx_routes: NvtxRoutes, ) -> Result<()> { std::fs::create_dir_all(crate_dir.join("src"))?; std::fs::write( crate_dir.join("Cargo.toml"), - cargo_toml(spec, io_package, with_nvtx_routes), - )?; - std::fs::write( - crate_dir.join("src/main.rs"), - main_rs(spec, with_nvtx_routes), + cargo_toml(spec, io_package, nvtx_routes), )?; + std::fs::write(crate_dir.join("src/main.rs"), main_rs(spec, nvtx_routes))?; Ok(()) } @@ -66,10 +78,16 @@ fn git_dep(url: String, rev: &str, features: &[&str]) -> Dependency { /// Wrapper `Cargo.toml`, built with `cargo-manifest`: pin quent crates to /// `quent.{remote,commit}` and the analyzer to `analyzer.{remote,commit}`; the /// empty `[workspace]` keeps the generated crate out of any parent workspace. -fn cargo_toml(spec: &ViewerSpec, io_package: &str, with_nvtx_routes: bool) -> String { +fn cargo_toml(spec: &ViewerSpec, io_package: &str, nvtx_routes: NvtxRoutes) -> String { let quent = spec.quent.cargo_url(); let q_rev = spec.quent.commit.as_str(); - let mut dependencies = BTreeMap::from([ + let nvtx_dependency = nvtx_routes.is_enabled().then(|| { + ( + NVTX_SERVER_PACKAGE.to_string(), + git_dep(quent.clone(), q_rev, &[]), + ) + }); + let dependencies: BTreeMap = BTreeMap::from([ ( "quent-query-engine-server".to_string(), git_dep(quent.clone(), q_rev, &["ui"]), @@ -81,7 +99,7 @@ fn cargo_toml(spec: &ViewerSpec, io_package: &str, with_nvtx_routes: bool) -> St ( // All formats enabled so the analyzer can detect the artifact's format at runtime. io_package.to_string(), - git_dep(quent.clone(), q_rev, &["ndjson", "msgpack", "postcard"]), + git_dep(quent, q_rev, &["ndjson", "msgpack", "postcard"]), ), ( spec.analyzer_package.clone(), @@ -102,10 +120,10 @@ fn cargo_toml(spec: &ViewerSpec, io_package: &str, with_nvtx_routes: bool) -> St }), ), ("uuid".to_string(), Dependency::Simple("1".to_string())), - ]); - if with_nvtx_routes { - dependencies.insert("nvtx-server".to_string(), git_dep(quent, q_rev, &[])); - } + ]) + .into_iter() + .chain(nvtx_dependency) + .collect(); let mut package = Package::new(WRAPPER_PACKAGE.to_string(), "0.0.0".to_string()); package.edition = Some(MaybeInherited::Local(Edition::E2024)); @@ -129,85 +147,67 @@ fn cargo_toml(spec: &ViewerSpec, io_package: &str, with_nvtx_routes: bool) -> St /// Wrapper `src/main.rs`: wire `::Viewer`'s analyzer/importer into /// `analyzer_service_router` and serve it. Root (`/` subdirs) and /// bind address come from env so one built binary serves any artifacts. -fn main_rs(spec: &ViewerSpec, with_nvtx_routes: bool) -> String { +fn main_rs(spec: &ViewerSpec, nvtx_routes: NvtxRoutes) -> String { let analyzer_crate = format_ident!("{}", spec.analyzer_crate()); let (root_env, addr_env) = (ROOT_ENV, ADDR_ENV); - let tokens = if with_nvtx_routes { - quote! { - use std::net::SocketAddr; - use std::path::PathBuf; - - use quent_query_engine_analyzer::ui::QuentViewer; - use quent_query_engine_server::analyzer_cache::index_query_engines; - use quent_query_engine_server::analyzer_service_router_with_routes; - use nvtx_server::{import_context_events, routes as nvtx_routes}; - use #analyzer_crate::Viewer; - - type Analyzer = ::Analyzer; - - #[tokio::main] - async fn main() -> Result<(), Box> { - let root = PathBuf::from(std::env::var(#root_env)?); - let addr: SocketAddr = std::env::var(#addr_env)?.parse()?; - - let import_root = root.clone(); - let importer = move |id: uuid::Uuid| { - Ok(::import_events( - &import_root.join(id.to_string()), - &import_root.join(id.to_string()), - )?) - }; - let lister_root = root.clone(); - let lister = move || index_query_engines(&lister_root); + let (route_imports, route_setup, router_call) = match nvtx_routes { + NvtxRoutes::Enabled => ( + quote! { + use quent_query_engine_server::analyzer_service_router_with_routes; + use nvtx_server::{import_context_events, routes as nvtx_routes}; + }, + quote! { let nvtx_root = root.clone(); let nvtx_importer = move |id: uuid::Uuid| import_context_events(&nvtx_root, id); - - let router = analyzer_service_router_with_routes::( + }, + quote! { + analyzer_service_router_with_routes::( Box::new(importer), Box::new(lister), None, nvtx_routes(Box::new(nvtx_importer)), - )?; - - let listener = tokio::net::TcpListener::bind(addr).await?; - axum::serve(listener, router.into_make_service()).await?; - Ok(()) - } - } - } else { - quote! { - use std::net::SocketAddr; - use std::path::PathBuf; + ) + }, + ), + NvtxRoutes::Disabled => ( + quote! { + use quent_query_engine_server::analyzer_service_router; + }, + quote! {}, + quote! { + analyzer_service_router::(Box::new(importer), Box::new(lister), None) + }, + ), + }; + let tokens = quote! { + use std::net::SocketAddr; + use std::path::PathBuf; - use quent_query_engine_analyzer::ui::QuentViewer; - use quent_query_engine_server::analyzer_cache::index_query_engines; - use quent_query_engine_server::analyzer_service_router; - use #analyzer_crate::Viewer; + use quent_query_engine_analyzer::ui::QuentViewer; + use quent_query_engine_server::analyzer_cache::index_query_engines; + #route_imports + use #analyzer_crate::Viewer; - type Analyzer = ::Analyzer; + type Analyzer = ::Analyzer; - #[tokio::main] - async fn main() -> Result<(), Box> { - let root = PathBuf::from(std::env::var(#root_env)?); - let addr: SocketAddr = std::env::var(#addr_env)?.parse()?; + #[tokio::main] + async fn main() -> Result<(), Box> { + let root = PathBuf::from(std::env::var(#root_env)?); + let addr: SocketAddr = std::env::var(#addr_env)?.parse()?; - let import_root = root.clone(); - let importer = move |id: uuid::Uuid| { - Ok(::import_events(&import_root.join(id.to_string()))?) - }; - let lister_root = root.clone(); - let lister = move || index_query_engines(&lister_root); + let import_root = root.clone(); + let importer = move |id: uuid::Uuid| { + Ok(::import_events(&import_root.join(id.to_string()))?) + }; + let lister_root = root.clone(); + let lister = move || index_query_engines(&lister_root); + #route_setup - let router = analyzer_service_router::( - Box::new(importer), - Box::new(lister), - None, - )?; + let router = #router_call?; - let listener = tokio::net::TcpListener::bind(addr).await?; - axum::serve(listener, router.into_make_service()).await?; - Ok(()) - } + let listener = tokio::net::TcpListener::bind(addr).await?; + axum::serve(listener, router.into_make_service()).await?; + Ok(()) } }; let file = syn::parse2(tokens).expect("generated wrapper main.rs is valid Rust"); @@ -238,7 +238,8 @@ mod tests { #[test] fn cargo_toml_pins_quent_and_analyzer() { - let manifest: toml::Value = toml::from_str(&cargo_toml(&spec(), IO_PACKAGE, true)).unwrap(); + let manifest: toml::Value = + toml::from_str(&cargo_toml(&spec(), IO_PACKAGE, NvtxRoutes::Enabled)).unwrap(); assert!(manifest.get("workspace").is_some(), "standalone workspace"); let deps = &manifest["dependencies"]; let server = &deps["quent-query-engine-server"]; @@ -262,9 +263,13 @@ mod tests { #[test] fn cargo_toml_supports_the_legacy_io_package() { // Artifacts pinned to quent revisions predating the `quent-exporter` → - // `quent-io` rename depend on the legacy package instead, same features. - let manifest: toml::Value = - toml::from_str(&cargo_toml(&spec(), LEGACY_IO_PACKAGE, false)).unwrap(); + // `quent-io` rename also predate the NVTX server package. + let manifest: toml::Value = toml::from_str(&cargo_toml( + &spec(), + LEGACY_IO_PACKAGE, + NvtxRoutes::Disabled, + )) + .unwrap(); let deps = &manifest["dependencies"]; assert!(deps.get("quent-io").is_none()); let exporter = &deps["quent-exporter"]; @@ -281,12 +286,39 @@ mod tests { } #[test] - fn main_rs_wires_the_viewer() { - let main = main_rs(&spec(), true); + fn cargo_toml_can_disable_nvtx_with_the_current_io_package() { + let manifest: toml::Value = + toml::from_str(&cargo_toml(&spec(), IO_PACKAGE, NvtxRoutes::Disabled)).unwrap(); + assert!(manifest["dependencies"].get(NVTX_SERVER_PACKAGE).is_none()); + } + + #[test] + fn main_rs_wires_the_nvtx_viewer() { + let main = main_rs(&spec(), NvtxRoutes::Enabled); assert!(main.contains("use quent_simulator_analyzer::Viewer;")); - assert!(main.contains("import_events")); assert!(main.contains("import_context_events")); assert!(main.contains("analyzer_service_router_with_routes")); assert!(main.contains("QUENT_OPEN_ADDR")); // bind address is configurable } + + #[test] + fn main_rs_without_nvtx_uses_the_legacy_router() { + let main = main_rs(&spec(), NvtxRoutes::Disabled); + assert!(main.contains("use quent_query_engine_server::analyzer_service_router;")); + assert!(!main.contains("nvtx_server")); + assert!(!main.contains("analyzer_service_router_with_routes")); + } + + #[test] + fn main_rs_imports_each_context_once_in_both_modes() { + for nvtx_routes in [NvtxRoutes::Enabled, NvtxRoutes::Disabled] { + let main = main_rs(&spec(), nvtx_routes); + assert_eq!( + main.matches("::import_events") + .count(), + 1 + ); + assert_eq!(main.matches("&import_root.join(id.to_string())").count(), 1); + } + } } diff --git a/ui/packages/@quent/client/src/nvtx.test.ts b/ui/packages/@quent/client/src/nvtx.test.ts index 74c78b8e3..b244452aa 100644 --- a/ui/packages/@quent/client/src/nvtx.test.ts +++ b/ui/packages/@quent/client/src/nvtx.test.ts @@ -4,7 +4,12 @@ import { afterEach, describe, expect, it, vi } from 'vitest'; import type { NvtxViewportRequest } from '@quent/utils'; import { fetchNvtxCatalog, fetchNvtxViewport } from './api'; -import { canonicalizeNvtxRequest, nvtxCatalogQueryOptions, nvtxViewportQueryOptions } from './nvtx'; +import { + canonicalizeNvtxRequest, + nvtxCatalogQueryOptions, + nvtxCatalogStaleTime, + nvtxViewportQueryOptions, +} from './nvtx'; function stubFetch(response: Response) { const fetchMock = vi.fn().mockResolvedValue(response); @@ -25,13 +30,17 @@ const request: NvtxViewportRequest = { describe('NVTX client', () => { afterEach(() => vi.unstubAllGlobals()); - it('treats catalog 404 as optional absence and fetches catalogs once', async () => { + it('treats catalog 404 as optional absence and keeps absent streams retryable', async () => { const fetchMock = stubFetch(new Response(null, { status: 404, statusText: 'Not Found' })); await expect(fetchNvtxCatalog('context-1', QUERY_START_UNIX_NS)).resolves.toBeNull(); expect(fetchMock.mock.calls[0]?.[0]).toContain( `/api/nvtx/contexts/context-1/catalog?query_start=${QUERY_START_UNIX_NS}` ); - expect(nvtxCatalogQueryOptions('context-1', QUERY_START_UNIX_NS).staleTime).toBe(Infinity); + expect(nvtxCatalogQueryOptions('context-1', QUERY_START_UNIX_NS).staleTime).toBeTypeOf( + 'function' + ); + expect(nvtxCatalogStaleTime(null)).toBe(0); + expect(nvtxCatalogStaleTime(undefined)).toBe(Infinity); }); it('propagates non-404 catalog failures', async () => { @@ -41,7 +50,7 @@ describe('NVTX client', () => { ); }); - it('uses the same canonical selector order for body and query key', async () => { + it('canonicalizes at the fetch boundary without validating during render', async () => { const fetchMock = stubFetch( new Response('{"viewport":{"start":-0.25,"end":1.5},"domains":[],"statistics":[]}', { status: 200, @@ -68,11 +77,17 @@ describe('NVTX client', () => { -0.25, 1.5, [ + ['9', [7, 3, 7], false], ['2', [], true], - ['9', [3, 7], false], ], ]); expect(options.placeholderData).toBeTypeOf('function'); + + const invalid = { ...request, viewport: { start: 2, end: 1 } }; + expect(() => nvtxViewportQueryOptions('context-1', QUERY_START_UNIX_NS, invalid)).not.toThrow(); + await expect(fetchNvtxViewport('context-1', QUERY_START_UNIX_NS, invalid)).rejects.toThrow( + 'NVTX viewport bounds must be finite and ordered' + ); }); it('preserves relative seconds and decimal-string identifiers from the catalog', async () => { diff --git a/ui/packages/@quent/client/src/nvtx.ts b/ui/packages/@quent/client/src/nvtx.ts index b530dcb9c..60efe4266 100644 --- a/ui/packages/@quent/client/src/nvtx.ts +++ b/ui/packages/@quent/client/src/nvtx.ts @@ -2,10 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 import { keepPreviousData, queryOptions, useQuery } from '@tanstack/react-query'; -import type { NvtxViewportRequest } from '@quent/utils'; +import type { NvtxCatalog, NvtxViewportRequest } from '@quent/utils'; import { fetchEngineContexts, fetchNvtxCatalog, fetchNvtxViewport } from './api'; import { DEFAULT_STALE_TIME } from './constants'; -import { canonicalizeNvtxRequest } from './nvtxCanonical'; export { canonicalizeNvtxRequest, canonicalizeNvtxSelections } from './nvtxCanonical'; export const engineContextsQueryOptions = (engineId: string) => @@ -15,12 +14,17 @@ export const engineContextsQueryOptions = (engineId: string) => staleTime: DEFAULT_STALE_TIME, }); +export const nvtxCatalogStaleTime = (catalog: NvtxCatalog | null | undefined) => + catalog === null ? 0 : Infinity; + export const nvtxCatalogQueryOptions = (contextId: string, queryStartUnixNs: bigint) => { const queryStartKey = queryStartUnixNs.toString(10); return queryOptions({ queryKey: ['nvtxCatalog', contextId, queryStartKey], queryFn: () => fetchNvtxCatalog(contextId, queryStartUnixNs), - staleTime: Infinity, + // A present catalog is immutable, but the server deliberately leaves an absent + // stream retryable so telemetry that appears later can be discovered on remount. + staleTime: query => nvtxCatalogStaleTime(query.state.data), }); }; @@ -30,11 +34,10 @@ export const nvtxViewportQueryOptions = ( request: NvtxViewportRequest, options?: { enabled?: boolean; staleTime?: number } ) => { - const canonical = canonicalizeNvtxRequest(request); const queryStartKey = queryStartUnixNs.toString(10); - const selectionKey = canonical.selections.map(selection => [ + const selectionKey = request.selections.map(selection => [ selection.domain_id, - selection.category_ids, + [...selection.category_ids], selection.include_uncategorized, ]); return queryOptions({ @@ -42,11 +45,13 @@ export const nvtxViewportQueryOptions = ( 'nvtxViewport', contextId, queryStartKey, - canonical.viewport.start, - canonical.viewport.end, + request.viewport.start, + request.viewport.end, selectionKey, ], - queryFn: () => fetchNvtxViewport(contextId, queryStartUnixNs, canonical), + // Validation and canonicalization belong to the fetch boundary, where a bad + // request rejects the query rather than throwing during React render. + queryFn: () => fetchNvtxViewport(contextId, queryStartUnixNs, request), enabled: options?.enabled ?? true, staleTime: options?.staleTime ?? DEFAULT_STALE_TIME, placeholderData: keepPreviousData,