feat(nvtx): wire viewers and timeline client

crates/open/src/viewer.rs

@@ -90,22 +90,39 @@ async fn build_one(group: ViewerGroup) -> Result<BuiltViewer> {
     println!("building: {label}");
 
     let crate_dir = build_dir(&spec)?;
-    wrapper::generate(&spec, &crate_dir, wrapper::IO_PACKAGE)?;
-    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)?;
-            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,

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-uuid>/`
 /// context directories.
@@ -35,10 +50,18 @@ 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,
+    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))?;
-    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, nvtx_routes),
+    )?;
+    std::fs::write(crate_dir.join("src/main.rs"), main_rs(spec, nvtx_routes))?;
     Ok(())
 }
 
@@ -55,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) -> 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 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<String, Dependency> = BTreeMap::from([
         (
             "quent-query-engine-server".to_string(),
             git_dep(quent.clone(), q_rev, &["ui"]),
@@ -91,7 +120,10 @@ fn cargo_toml(spec: &ViewerSpec, io_package: &str) -> String {
             }),
         ),
         ("uuid".to_string(), Dependency::Simple("1".to_string())),
-    ]);
+    ])
+    .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));
@@ -115,16 +147,45 @@ fn cargo_toml(spec: &ViewerSpec, io_package: &str) -> String {
 /// Wrapper `src/main.rs`: wire `<analyzer>::Viewer`'s analyzer/importer into
 /// `analyzer_service_router` and serve it. Root (`<context-uuid>/` 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, nvtx_routes: NvtxRoutes) -> String {
     let analyzer_crate = format_ident!("{}", spec.analyzer_crate());
     let (root_env, addr_env) = (ROOT_ENV, ADDR_ENV);
+    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);
+            },
+            quote! {
+                analyzer_service_router_with_routes::<Analyzer>(
+                    Box::new(importer),
+                    Box::new(lister),
+                    None,
+                    nvtx_routes(Box::new(nvtx_importer)),
+                )
+            },
+        ),
+        NvtxRoutes::Disabled => (
+            quote! {
+                use quent_query_engine_server::analyzer_service_router;
+            },
+            quote! {},
+            quote! {
+                analyzer_service_router::<Analyzer>(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;
+        #route_imports
         use #analyzer_crate::Viewer;
 
         type Analyzer = <Viewer as QuentViewer>::Analyzer;
@@ -136,18 +197,13 @@ fn main_rs(spec: &ViewerSpec) -> String {
 
             let import_root = root.clone();
             let importer = move |id: uuid::Uuid| {
-                Ok(<Viewer as QuentViewer>::import_events(
-                    &import_root.join(id.to_string()),
-                )?)
+                Ok(<Viewer as QuentViewer>::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::<Analyzer>(
-                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?;
@@ -182,13 +238,15 @@ 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, NvtxRoutes::Enabled)).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"] {
@@ -205,12 +263,17 @@ 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)).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"];
+        assert!(deps.get("nvtx-server").is_none());
         assert_eq!(
             exporter["git"].as_str().unwrap(),
             "https://example.com/quent"
@@ -223,10 +286,39 @@ mod tests {
     }
 
     #[test]
-    fn main_rs_wires_the_viewer() {
-        let main = main_rs(&spec());
+    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("<Viewer as QuentViewer>::import_events")
+                    .count(),
+                1
+            );
+            assert_eq!(main.matches("&import_root.join(id.to_string())").count(), 1);
+        }
+    }
 }

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" }

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<dyn std::error::Error>> {
 
     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<dyn std::error::Error>> {
     let analyzer = async {
         axum::serve(
             TcpListener::bind(analyzer_addr).await?,
-            analyzer_service_router::<SimulatorUiAnalyzer>(
+            analyzer_service_router_with_routes::<SimulatorUiAnalyzer>(
                 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(),
         )

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" }

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<dyn std::error::Error>> {
     <BulkTimelinesResponse as TS>::export_all(&cfg)?;
     <CategoricalTimelineRequest<QueryFilter> as TS>::export_all(&cfg)?;
     <DataFlowTimelineBinned as TS>::export_all(&cfg)?;
+    <EngineContexts as TS>::export_all(&cfg)?;
+    <NvtxCatalog as TS>::export_all(&cfg)?;
+    <NvtxViewportRequest as TS>::export_all(&cfg)?;
+    <NvtxViewportResponse as TS>::export_all(&cfg)?;
 
     <EntityListRequest<QueryFilter, OperatorFilter> as TS>::export_all(&cfg)?;
     <EntityListResponse as TS>::export_all(&cfg)?;