Skip to content
Merged
Show file tree
Hide file tree
Changes from 4 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