Skip to content
Merged
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
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.

1 change: 1 addition & 0 deletions crates/adaptive/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ workspace = true

[dependencies]
nemo-relay.workspace = true
log = { version = "0.4", features = ["kv"] }
uuid = { workspace = true, features = ["v4", "v7", "serde"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
Expand Down
23 changes: 18 additions & 5 deletions crates/adaptive/src/drain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -179,8 +179,12 @@ async fn store_run(
backend: &Arc<dyn StorageBackendDyn + Send + Sync>,
completed_run: &RunRecord,
) -> bool {
if let Err(error) = backend.store_run_dyn(completed_run).await {
eprintln!("nemo-relay-adaptive drain: store_run failed: {error}");
if backend.store_run_dyn(completed_run).await.is_err() {
log::warn!(
target: "nemo_relay.runtime",
event = "adaptive_run_storage_failed";
"Adaptive runtime could not persist a completed run"
);
return false;
}
true
Expand All @@ -193,11 +197,16 @@ async fn run_learners(
hot_cache: &Arc<RwLock<HotCache>>,
) {
for learner in learners {
if let Err(error) = learner
if learner
.process_run(completed_run, backend.as_ref(), hot_cache)
.await
.is_err()
{
eprintln!("nemo-relay-adaptive drain: learner failed: {error}");
log::warn!(
target: "nemo_relay.runtime",
event = "adaptive_learner_failed";
"Adaptive runtime learner failed while processing a completed run"
);
}
}
}
Expand All @@ -213,7 +222,11 @@ async fn refresh_hot_cache_plan(
guard.plan = plan;
}
}
Err(error) => eprintln!("nemo-relay-adaptive drain: load_plan failed: {error}"),
Err(_) => log::warn!(
target: "nemo_relay.runtime",
event = "adaptive_plan_refresh_failed";
"Adaptive runtime could not refresh the cached plan"
),
}
}

Expand Down
46 changes: 40 additions & 6 deletions crates/adaptive/src/redis.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,12 +55,46 @@ impl RedisBackend {
/// Returns [`AdaptiveError::Storage`] if the client cannot be created or the
/// connection cannot be established.
pub async fn new(url: &str, key_prefix: impl Into<String>) -> Result<Self> {
let client = redis::Client::open(url)
.map_err(|e| AdaptiveError::Storage(format!("redis client: {e}")))?;
let conn = client
.get_connection_manager()
.await
.map_err(|e| AdaptiveError::Storage(format!("redis connection: {e}")))?;
log::info!(
target: "nemo_relay.plugin",
event = "plugin_resource_access_pending",
plugin_kind = "adaptive",
resource_kind = "redis",
permission = "connect";
"Plugin Redis connectivity validation started"
);
let client = redis::Client::open(url).map_err(|e| {
log::warn!(
target: "nemo_relay.plugin",
event = "plugin_resource_access_failed",
plugin_kind = "adaptive",
resource_kind = "redis",
permission = "connect",
reason = "client_configuration";
"Plugin resource access validation failed"
);
AdaptiveError::Storage(format!("redis client: {e}"))
})?;
let conn = client.get_connection_manager().await.map_err(|e| {
log::warn!(
target: "nemo_relay.plugin",
event = "plugin_resource_access_failed",
plugin_kind = "adaptive",
resource_kind = "redis",
permission = "connect",
reason = "connection_failed";
"Plugin resource access validation failed"
);
AdaptiveError::Storage(format!("redis connection: {e}"))
})?;
log::info!(
target: "nemo_relay.plugin",
event = "plugin_resource_connected",
plugin_kind = "adaptive",
resource_kind = "redis",
permission = "connect";
"Plugin Redis connectivity established"
);
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Ok(Self {
client,
conn,
Expand Down
11 changes: 10 additions & 1 deletion crates/adaptive/src/runtime/backend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,16 @@ pub async fn build_backend(
backend: &BackendSpec,
) -> Result<Arc<dyn StorageBackendDyn + Send + Sync>> {
match backend.kind.as_str() {
"in_memory" => Ok(Arc::new(InMemoryBackend::new())),
"in_memory" => {
log::info!(
target: "nemo_relay.plugin",
event = "plugin_resource_validation_completed",
plugin_kind = "adaptive",
resource_count = 0;
"Plugin resource validation completed"
);
Ok(Arc::new(InMemoryBackend::new()))
}
#[cfg(feature = "redis-backend")]
"redis" => {
let url = backend
Expand Down
14 changes: 11 additions & 3 deletions crates/adaptive/src/runtime/features.rs
Original file line number Diff line number Diff line change
Expand Up @@ -392,10 +392,14 @@ impl AdaptiveRuntime {

if self.config.acg.is_some()
&& let Some(backend) = self.backend.as_ref()
&& let Err(error) =
&& let Err(_) =
load_persisted_acg_state(&agent_id, backend.as_ref(), &self.hot_cache).await
{
eprintln!("nemo-relay-adaptive: acg hot cache seeding failed: {error}");
log::warn!(
target: "nemo_relay.runtime",
event = "adaptive_acg_cache_seed_failed";
"Adaptive runtime could not seed the ACG hot cache"
);
}

let mut pending = self.pending_features(&agent_id);
Expand Down Expand Up @@ -431,7 +435,11 @@ impl AdaptiveRuntime {
guard.plan = plan;
}
}
Err(error) => eprintln!("nemo-relay-adaptive: hot cache seeding failed: {error}"),
Err(_) => log::warn!(
target: "nemo_relay.runtime",
event = "adaptive_hot_cache_seed_failed";
"Adaptive runtime could not seed the hot cache"
),
}
}

Expand Down
14 changes: 13 additions & 1 deletion crates/adaptive/tests/integration/redis_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@

#![cfg(feature = "redis-backend")]

use std::sync::{Arc, RwLock};
use std::sync::{Arc, Once, RwLock};

use chrono::Utc;
use nemo_relay::codec::request::{AnnotatedLlmRequest, Message, MessageContent};
Expand All @@ -29,6 +29,7 @@ use nemo_relay_adaptive::types::plan::ExecutionPlan;
use nemo_relay_adaptive::types::records::{CallKind, CallRecord, RunRecord};

const REDIS_TEST_ENV: &str = "NEMO_RELAY_RUN_REDIS_TESTS";
static TEST_LOGGING: Once = Once::new();

// ---------------------------------------------------------------------------
// Helpers
Expand All @@ -41,9 +42,19 @@ fn env_value_is_truthy(value: Option<&str>) -> bool {
)
}

fn enable_operational_logs() {
TEST_LOGGING.call_once(|| {
let runtime =
nemo_relay::logging::init_logging(&nemo_relay::logging::LoggingConfig::default())
.expect("test logging should initialize");
Box::leak(Box::new(runtime));
});
}

/// Attempt to connect to a local Redis instance. Returns `None` (skip) if
/// Redis tests were not explicitly enabled or Redis is unavailable.
async fn get_test_redis() -> Option<RedisBackend> {
enable_operational_logs();
let redis_test_env =
std::env::var_os(REDIS_TEST_ENV).map(|value| value.to_string_lossy().into_owned());
if !env_value_is_truthy(redis_test_env.as_deref()) {
Expand All @@ -65,6 +76,7 @@ async fn get_test_redis() -> Option<RedisBackend> {
}

async fn get_test_redis_with_prefix() -> Option<(RedisBackend, String)> {
enable_operational_logs();
let prefix = format!("test:{}:", Uuid::now_v7());
match RedisBackend::new("redis://127.0.0.1/", prefix.clone()).await {
Ok(backend) => Some((backend, prefix)),
Expand Down
13 changes: 12 additions & 1 deletion crates/adaptive/tests/integration/runtime_integration_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
//! Integration tests for runtime integration in the NeMo Relay adaptive crate.

use std::pin::Pin;
use std::sync::{Arc, Mutex as StdMutex, RwLock};
use std::sync::{Arc, Mutex as StdMutex, Once, RwLock};

use chrono::Utc;
use nemo_relay::api::event::Event;
Expand Down Expand Up @@ -48,12 +48,23 @@ use tokio_stream::StreamExt;
use uuid::Uuid;

static TEST_MUTEX: Mutex<()> = Mutex::const_new(());
static TEST_LOGGING: Once = Once::new();

fn enable_operational_logs() {
TEST_LOGGING.call_once(|| {
let runtime =
nemo_relay::logging::init_logging(&nemo_relay::logging::LoggingConfig::default())
.expect("test logging should initialize");
Box::leak(Box::new(runtime));
});
}

fn short_hash(value: &str) -> &str {
value.get(..16).unwrap_or(value)
}

fn reset_global() {
enable_operational_logs();
let _ = clear_plugin_configuration();
let _ = deregister_plugin("test.header_plugin");
let _ = deregister_plugin("test.failing_plugin");
Expand Down
1 change: 1 addition & 0 deletions crates/cli/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ hyper = "1"
hyper-util = { version = "0.1", features = ["tokio"] }
dialoguer = { version = "0.11", default-features = false, features = ["password"] }
jsonschema = { version = "0.46.6", default-features = false }
log = { version = "0.4", features = ["kv"] }
percent-encoding = "2"
reqwest = { version = "0.12", default-features = false, features = ["charset", "http2", "json", "rustls-tls-native-roots", "stream"] }
regex = "1"
Expand Down
73 changes: 71 additions & 2 deletions crates/cli/src/bootstrap/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,11 +91,48 @@ impl GatewaySpec {
/// Return the compatible gateway already bound at this endpoint, or start the existing server
/// under a per-user lock and wait for authenticated readiness.
pub(crate) fn acquire(&self) -> Result<GatewayEndpoint, String> {
acquire_gateway(self)
match acquire_gateway(self) {
Ok(endpoint) => Ok(endpoint),
Err(error) => {
log::error!(
target: "nemo_relay.bootstrap",
event = "gateway_acquisition_failed",
bind = self.bind.to_string().as_str();
"Gateway acquisition failed"
);
Err(error)
}
}
}

pub(crate) fn recover(&self, expected_instance: &str) -> Result<GatewayEndpoint, String> {
recover_gateway(self, expected_instance)
log::warn!(
target: "nemo_relay.bootstrap",
event = "gateway_recovery_started",
instance_id = expected_instance;
"Gateway recovery started"
);
match recover_gateway(self, expected_instance) {
Ok(endpoint) => {
log::info!(
target: "nemo_relay.bootstrap",
event = "gateway_recovered",
previous_instance_id = expected_instance,
instance_id = endpoint.instance_id.as_str();
"Gateway recovery completed"
);
Ok(endpoint)
}
Err(error) => {
log::error!(
target: "nemo_relay.bootstrap",
event = "gateway_recovery_failed",
instance_id = expected_instance;
"Gateway recovery failed"
);
Err(error)
}
}
}

pub(crate) fn healthy_instance(&self, url: &str) -> Option<String> {
Expand Down Expand Up @@ -235,6 +272,13 @@ fn compatible_endpoint(
instance_id: Option<String>,
) -> Result<GatewayEndpoint, String> {
let instance_id = instance_id.ok_or_else(|| foreign_listener_error(&url))?;
log::info!(
target: "nemo_relay.bootstrap",
event = "gateway_reused",
instance_id = instance_id.as_str(),
address = address.to_string().as_str();
"An existing gateway was reused"
);
Ok(GatewayEndpoint {
address,
url,
Expand All @@ -255,6 +299,12 @@ fn incompatible_relay_error(url: &str) -> String {
}

fn start_gateway(spec: &GatewaySpec, state: &Path) -> Result<GatewayEndpoint, String> {
log::info!(
target: "nemo_relay.bootstrap",
event = "gateway_spawn_started",
bind = spec.bind.to_string().as_str();
"Gateway process spawn started"
);
let relay = relay_binary()?;
let ready_path = state.join(format!(
"gateway-{}-{}.ready.json",
Expand Down Expand Up @@ -302,7 +352,14 @@ fn start_gateway(spec: &GatewaySpec, state: &Path) -> Result<GatewayEndpoint, St
.map_err(|error| format!("failed to spawn nemo-relay gateway: {error}"))?;
let mut child = ArmedChild::new(child);
let deadline = Instant::now() + BOOTSTRAP_START_TIMEOUT;
let mut attempts = 0_u64;
log::warn!(
target: "nemo_relay.bootstrap",
event = "gateway_readiness_pending";
"Waiting for the spawned gateway to become ready"
);
while Instant::now() < deadline {
attempts += 1;
if let Some(endpoint) = read_ready_file(&ready_path)?
&& (endpoint.address == spec.bind
|| (spec.bind.port() == 0 && endpoint.address.ip() == spec.bind.ip()))
Expand All @@ -314,6 +371,13 @@ fn start_gateway(spec: &GatewaySpec, state: &Path) -> Result<GatewayEndpoint, St
return Err(error);
}
let _ = fs::remove_file(&ready_path);
log::info!(
target: "nemo_relay.bootstrap",
event = "gateway_ready",
instance_id = endpoint.instance_id.as_str(),
attempt_count = attempts;
"Spawned gateway became ready"
);
return Ok(endpoint);
}
match child.try_wait() {
Expand Down Expand Up @@ -372,6 +436,11 @@ fn hand_off_to_reaper_with<T>(
}
return Err(format!("failed to start gateway reaper thread: {error}"));
}
log::info!(
target: "nemo_relay.bootstrap",
event = "gateway_reaper_handoff";
"Gateway process ownership was handed to the reaper"
);
Ok(())
}

Expand Down
Loading
Loading