Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .cargo/config.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,3 +10,4 @@ rustflags = ["-C", "target-cpu=x86-64-v3", "-C", "force-frame-pointers=yes"]
# Force frame pointers to make performance and memory profiling much easier
[target.aarch64-unknown-linux-gnu]
rustflags = ["-C", "target-cpu=neoverse-n1", "-C", "force-frame-pointers=yes"]

4 changes: 3 additions & 1 deletion .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ jobs:
steps:
- uses: actions/checkout@v4
- name: Install Rust components
run: rustup component add rustfmt clippy
run: rustup component add rustfmt clippy && rustup target add wasm32-unknown-unknown
- uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1
with:
cache-bin: false
Expand All @@ -120,6 +120,8 @@ jobs:
run: cargo clippy --workspace --all-targets --locked -- -D warnings
- name: cargo clippy (prefill-router)
run: cargo clippy -p switchyard-server --all-targets --features prefill-router --locked -- -D warnings
- name: cargo clippy (libsy, wasm32)
run: cargo clippy -p switchyard-libsy --target wasm32-unknown-unknown --locked -- -D warnings
- name: cargo test

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a compile/lint check, not a Worker or browser runtime test. I think that's fine if we call this compile support for Rust wasm hosts. If we're claiming runtime support, we should add a small smoke test.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Compile support is the claim. The PR body's first line says "compile for wasm32-unknown-unknown", and the CI step is a clippy check on that target. Our Worker runs it at runtime on our side; a smoke test in this repo can come in a follow-up if the team wants to claim runtime support.

run: cargo test --workspace --locked
- name: cargo test (prefill-router)
Expand Down
3 changes: 3 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

17 changes: 15 additions & 2 deletions crates/libsy/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -33,12 +33,25 @@ rand.workspace = true
regex.workspace = true
switchyard-protocol.workspace = true
thiserror.workspace = true
tokio.workspace = true
tokio-stream = "0.1"
tokio-stream = { version = "0.1", default-features = false }
tracing.workspace = true
tracing-opentelemetry.workspace = true
uuid.workspace = true

[target.'cfg(not(all(target_arch = "wasm32", target_os = "unknown")))'.dependencies]
tokio.workspace = true

[target.'cfg(all(target_arch = "wasm32", target_os = "unknown"))'.dependencies]
# Workers/browser targets have no OS threads or epoll: only tokio's
# runtime-agnostic pieces (channels, mutexes, macros) are linked.
tokio = { version = "1", default-features = false, features = ["sync", "macros"] }
wasm-bindgen-futures = "0.4"
web-time = "1"
# `rand` 0.10 pulls getrandom 0.4; route it through the JS crypto API.
# (`jsonschema` already does the same for the 0.3 line that `ahash` uses.)
getrandom = { version = "0.4", default-features = false, features = ["wasm_js"] }
uuid = { workspace = true, features = ["js"] }

[dev-dependencies]
# SDK + in-memory exporter to assert what the observability layer records.
opentelemetry_sdk = { version = "0.32", features = ["metrics", "testing", "trace"] }
Expand Down
3 changes: 2 additions & 1 deletion crates/libsy/src/algorithms/advisor_gate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,8 @@
//! holds the only mutable state.

use std::sync::Arc;
use std::time::Instant;

use crate::rt::Instant;

use switchyard_protocol::{
Category, ContentBlock, InstructionBlock, LlmRequest, Message, ModelId, OutputParams, Request,
Expand Down
31 changes: 28 additions & 3 deletions crates/libsy/src/algorithms/fall_through.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,14 +19,19 @@

use std::{
collections::HashMap,
sync::{Arc, Once, Weak},
time::{Duration, Instant},
sync::{Arc, Once},
time::Duration,
};

#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
use std::sync::Weak;

use async_trait::async_trait;
use parking_lot::Mutex;
use tokio::sync::Mutex as AsyncMutex;

use crate::rt::Instant;

use crate::core::algorithm::{self, Algorithm, Driver};
use crate::core::classifier::{Classifier, Score};
use crate::core::processor::{Event, Processor};
Expand Down Expand Up @@ -58,6 +63,8 @@ pub struct FallThrough<S = ()> {
classifiers: Vec<Arc<dyn Classifier<S>>>,
session_states: Option<Arc<SessionStates<S>>>,
cleanup_started: Once,
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
next_cleanup: Mutex<Instant>,
}

impl FallThrough<()> {
Expand All @@ -69,6 +76,8 @@ impl FallThrough<()> {
classifiers: Vec::new(),
session_states: None,
cleanup_started: Once::new(),
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
next_cleanup: Mutex::new(Instant::now() + SESSION_CLEANUP_INTERVAL),
}
}
}
Expand All @@ -85,6 +94,8 @@ where
classifiers: Vec::new(),
session_states: Some(Arc::new(Mutex::new(HashMap::new()))),
cleanup_started: Once::new(),
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
next_cleanup: Mutex::new(Instant::now() + SESSION_CLEANUP_INTERVAL),
}
}

Expand Down Expand Up @@ -128,7 +139,12 @@ where
};
let states = Arc::downgrade(states);
self.cleanup_started.call_once(move || {
// Timer-driven background cleanup needs a Tokio runtime; wasm hosts
// sweep expired sessions inline in `session_state` instead.
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
drop(tokio::spawn(cleanup_inactive_sessions(states)));
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
drop(states);
});
}

Expand Down Expand Up @@ -184,8 +200,16 @@ where
fn session_state(&self, request: &Request) -> Option<Arc<AsyncMutex<S>>> {
let states = self.session_states.as_ref()?;
let session_id = session_id(request)?;
let mut states = states.lock();
let now = Instant::now();
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
{
let mut next_cleanup = self.next_cleanup.lock();
if now >= *next_cleanup {
remove_inactive_sessions(states, now, SESSION_STATE_TTL);
*next_cleanup = now + SESSION_CLEANUP_INTERVAL;
}
}
let mut states = states.lock();
let session = states.entry(session_id).or_insert_with(|| SessionState {
state: Arc::new(AsyncMutex::new(S::default())),
last_accessed: now,
Expand Down Expand Up @@ -246,6 +270,7 @@ where
}
}

#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
async fn cleanup_inactive_sessions<S>(states: Weak<SessionStates<S>>)
where
S: Send + 'static,
Expand Down
12 changes: 5 additions & 7 deletions crates/libsy/src/core/algorithm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,7 @@
//! The [`Algorithm`] trait and its [`Driver`] — the orchestration contract every
//! algorithm implements and the offload channel it uses for routing-time model calls.

use std::{
collections::HashMap, future::Future, panic::AssertUnwindSafe, pin::Pin, sync::Arc,
time::Instant,
};
use std::{collections::HashMap, future::Future, panic::AssertUnwindSafe, pin::Pin, sync::Arc};

use async_trait::async_trait;
use futures::{FutureExt, Stream, StreamExt};
Expand All @@ -26,6 +23,7 @@ use tracing::Instrument;
/// [`switchyard_protocol::LlmResponseStream`] or the terminal aggregate.
use switchyard_protocol::{Category, ModelId, Request, Response};

use crate::rt::Instant;
use crate::{DriverError, LibsyError, Result, observability};

/// A boxed, `Send` stream of [`Step`]s — the output of
Expand Down Expand Up @@ -419,7 +417,7 @@ fn panic_message(payload: &(dyn std::any::Any + Send)) -> String {
}

/// Abort guard
struct AbortOnDrop(tokio::task::AbortHandle);
struct AbortOnDrop(crate::rt::AbortHandle);

impl Drop for AbortOnDrop {
fn drop(&mut self) {
Expand Down Expand Up @@ -526,7 +524,7 @@ pub trait Algorithm: Send + Sync + 'static {
fn run_stream(self: Arc<Self>, request: Request, models: Arc<RuntimeModels>) -> StepStream {
let (driver, step_rx) = Driver::new(self.name(), models);
let span = observability::run_span(self.name(), &request);
let handle = tokio::spawn(
let handle = crate::rt::spawn_abortable(
async move {
let algorithm = self.name().to_string();
// Catch a panicking algorithm so the run still publishes a terminal step.
Expand All @@ -548,7 +546,7 @@ pub trait Algorithm: Send + Sync + 'static {
.instrument(span),
);
// Dropping the stream aborts the algorithm task when its consumer goes away.
let abort_guard = AbortOnDrop(handle.abort_handle());
let abort_guard = AbortOnDrop(handle);
Box::pin(ReceiverStream::new(step_rx).map(move |step| {
// link abort guard to stream
let _keep_alive = &abort_guard;
Expand Down
1 change: 1 addition & 0 deletions crates/libsy/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ pub use core::processor::{Event, Processor};
pub use core::state::{State, StateValue};

mod error;
mod rt;
pub use error::{DriverError, LibsyError, Result};

mod algorithms;
Expand Down
4 changes: 3 additions & 1 deletion crates/libsy/src/observability.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,9 @@
//! negligible next to a model call.

use std::future::Future;
use std::time::{Duration, Instant};
use std::time::Duration;

use crate::rt::Instant;

use opentelemetry::metrics::Meter;
use opentelemetry::{KeyValue, global};
Expand Down
41 changes: 41 additions & 0 deletions crates/libsy/src/rt.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Runtime abstraction over task spawning and monotonic clocks, so the crate
//! runs on native Tokio hosts and on single-threaded wasm32 hosts (browsers,
//! Cloudflare Workers) alike.

/// Monotonic instant. On wasm32 `std::time::Instant::now()` aborts, so a
/// JS-clock-backed drop-in replacement is used there.
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
pub(crate) use std::time::Instant;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we scope the JS branches to all(target_arch = "wasm32", target_os = "unknown")? As written, wasm32-wasip1 selects this path too. I built and ran a WASI consumer and it panicked when it reached js-sys.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 26fe647: every JS branch (rt.rs, fall_through.rs, the two Cargo target tables) is now all(target_arch = "wasm32", target_os = "unknown"). wasm32-wasip1 takes the native path; today that fails at compile time in tokio (Only features sync,macros,io-util,rt,time are supported on wasm, because the workspace tokio enables more), instead of panicking in js-sys at runtime. Making WASI build is out of scope here.

#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
pub(crate) use web_time::Instant;

/// Handle that aborts the task returned by [`spawn_abortable`].
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
pub(crate) use futures::future::AbortHandle;
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
pub(crate) use tokio::task::AbortHandle;

/// Spawns a future on the host runtime and returns a handle that aborts it.
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
pub(crate) fn spawn_abortable<F>(future: F) -> AbortHandle
where
F: std::future::Future<Output = ()> + Send + 'static,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we leave the native path on Tokio's own abort handle? This adds a futures::Abortable wrapper and another shared allocation to every native run even though tokio::spawn(future).abort_handle() already gives us cancellation. The wrapper should only be needed on wasm.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 26fe647: native spawn_abortable is tokio::spawn(future).abort_handle() again, and rt::AbortHandle is tokio::task::AbortHandle on native, futures::future::AbortHandle only on wasm. AbortOnDrop holds rt::AbortHandle. No Abortable wrapper on native.

{
tokio::spawn(future).abort_handle()
}

/// Spawns a future on the JS microtask queue and returns a handle that aborts it.
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
pub(crate) fn spawn_abortable<F>(future: F) -> AbortHandle
where
F: std::future::Future<Output = ()> + 'static,
{
let (handle, registration) = AbortHandle::new_pair();
wasm_bindgen_futures::spawn_local(async move {
let _ = futures::future::Abortable::new(future, registration).await;
});
handle
}
Loading