Skip to content

Commit b3cb3bc

Browse files
calebschoeppcalebschoepp
committed
Move kv and llm instrumentation up a level of abstraction
Signed-off-by: Caleb Schoepp <[email protected]>
1 parent 3e62d2e commit b3cb3bc

File tree

14 files changed

+43
-58
lines changed

14 files changed

+43
-58
lines changed

‎Cargo.lock

Lines changed: 0 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/factor-key-value/src/host.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use spin_locked_app::MetadataKey;
55
use spin_world::v2::key_value;
66
use std::{collections::HashSet, sync::Arc};
77
use table::Table;
8+
use tracing::{instrument, Level};
89

910
pub const KEY_VALUE_STORES_KEY: MetadataKey<Vec<String>> = MetadataKey::new("key_value_stores");
1011

@@ -79,6 +80,7 @@ impl key_value::Host for KeyValueDispatch {}
7980

8081
#[async_trait]
8182
impl key_value::HostStore for KeyValueDispatch {
83+
#[instrument(name = "spin_key_value.open", skip(self), err(level = Level::INFO), fields(otel.kind = "client", kv.backend=self.manager.summary(&name).unwrap_or("unknown".to_string())))]
8284
async fn open(&mut self, name: String) -> Result<Result<Resource<key_value::Store>, Error>> {
8385
Ok(async {
8486
if self.allowed_stores.contains(&name) {
@@ -94,6 +96,7 @@ impl key_value::HostStore for KeyValueDispatch {
9496
.await)
9597
}
9698

99+
#[instrument(name = "spin_key_value.get", skip(self, store), err(level = Level::INFO), fields(otel.kind = "client"))]
97100
async fn get(
98101
&mut self,
99102
store: Resource<key_value::Store>,
@@ -103,6 +106,7 @@ impl key_value::HostStore for KeyValueDispatch {
103106
Ok(store.get(&key).await)
104107
}
105108

109+
#[instrument(name = "spin_key_value.set", skip(self, store, value), err(level = Level::INFO), fields(otel.kind = "client"))]
106110
async fn set(
107111
&mut self,
108112
store: Resource<key_value::Store>,
@@ -113,6 +117,7 @@ impl key_value::HostStore for KeyValueDispatch {
113117
Ok(store.set(&key, &value).await)
114118
}
115119

120+
#[instrument(name = "spin_key_value.delete", skip(self, store), err(level = Level::INFO), fields(otel.kind = "client"))]
116121
async fn delete(
117122
&mut self,
118123
store: Resource<key_value::Store>,
@@ -122,6 +127,7 @@ impl key_value::HostStore for KeyValueDispatch {
122127
Ok(store.delete(&key).await)
123128
}
124129

130+
#[instrument(name = "spin_key_value.exists", skip(self, store), err(level = Level::INFO), fields(otel.kind = "client"))]
125131
async fn exists(
126132
&mut self,
127133
store: Resource<key_value::Store>,
@@ -131,6 +137,7 @@ impl key_value::HostStore for KeyValueDispatch {
131137
Ok(store.exists(&key).await)
132138
}
133139

140+
#[instrument(name = "spin_key_value.get_keys", skip(self, store), err(level = Level::INFO), fields(otel.kind = "client"))]
134141
async fn get_keys(
135142
&mut self,
136143
store: Resource<key_value::Store>,

‎crates/factor-llm/src/host.rs

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
use async_trait::async_trait;
22
use spin_world::v1::llm::{self as v1};
33
use spin_world::v2::llm::{self as v2};
4+
use tracing::field::Empty;
5+
use tracing::{instrument, Level};
46

57
use crate::InstanceState;
68

79
#[async_trait]
810
impl v2::Host for InstanceState {
11+
#[instrument(name = "spin_llm.infer", skip(self, prompt), err(level = Level::INFO), fields(otel.kind = "client", llm.backend = Empty))]
912
async fn infer(
1013
&mut self,
1114
model: v2::InferencingModel,
@@ -15,9 +18,9 @@ impl v2::Host for InstanceState {
1518
if !self.allowed_models.contains(&model) {
1619
return Err(access_denied_error(&model));
1720
}
18-
self.engine
19-
.lock()
20-
.await
21+
let mut engine = self.engine.lock().await;
22+
tracing::Span::current().record("llm.backend", engine.summary());
23+
engine
2124
.infer(
2225
model,
2326
prompt,
@@ -33,15 +36,18 @@ impl v2::Host for InstanceState {
3336
.await
3437
}
3538

39+
#[instrument(name = "spin_llm.generate_embeddings", skip(self, data), err(level = Level::INFO), fields(otel.kind = "client", llm.backend = Empty))]
3640
async fn generate_embeddings(
3741
&mut self,
38-
m: v1::EmbeddingModel,
42+
model: v1::EmbeddingModel,
3943
data: Vec<String>,
4044
) -> Result<v2::EmbeddingsResult, v2::Error> {
41-
if !self.allowed_models.contains(&m) {
42-
return Err(access_denied_error(&m));
45+
if !self.allowed_models.contains(&model) {
46+
return Err(access_denied_error(&model));
4347
}
44-
self.engine.lock().await.generate_embeddings(m, data).await
48+
let mut engine = self.engine.lock().await;
49+
tracing::Span::current().record("llm.backend", engine.summary());
50+
engine.generate_embeddings(model, data).await
4551
}
4652

4753
fn convert_error(&mut self, error: v2::Error) -> anyhow::Result<v2::Error> {

‎crates/factor-llm/src/lib.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,13 @@ pub trait LlmEngine: Send + Sync {
127127
model: v2::EmbeddingModel,
128128
data: Vec<String>,
129129
) -> Result<v2::EmbeddingsResult, v2::Error>;
130+
131+
/// A human-readable summary of the given engine's configuration
132+
///
133+
/// Example: "local model"
134+
fn summary(&self) -> Option<String> {
135+
None
136+
}
130137
}
131138

132139
/// A creator for an LLM engine.

‎crates/factor-llm/src/spin.rs

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,10 @@ mod local {
3434
) -> Result<v2::EmbeddingsResult, v2::Error> {
3535
self.generate_embeddings(model, data).await
3636
}
37+
38+
fn summary(&self) -> Option<String> {
39+
Some("local model".to_string())
40+
}
3741
}
3842
}
3943

@@ -78,6 +82,10 @@ impl LlmEngine for RemoteHttpLlmEngine {
7882
) -> Result<v2::EmbeddingsResult, v2::Error> {
7983
self.generate_embeddings(model, data).await
8084
}
85+
86+
fn summary(&self) -> Option<String> {
87+
Some(format!("model at {}", self.url()))
88+
}
8189
}
8290

8391
pub fn runtime_config_from_toml(
@@ -161,5 +169,9 @@ mod noop {
161169
"Local LLM operations are not supported in this version of Spin.".into(),
162170
))
163171
}
172+
173+
fn summary(&self) -> Option<String> {
174+
Some("noop model".to_owned())
175+
}
164176
}
165177
}

‎crates/key-value-azure/Cargo.toml

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ serde = { version = "1.0", features = ["derive", "rc"] }
1717
spin-core = { path = "../core" }
1818
spin-factor-key-value = { path = "../factor-key-value" }
1919
tokio = "1"
20-
tracing = { workspace = true }
2120
url = "2"
2221

2322
[lints]

‎crates/key-value-azure/src/store.rs

Lines changed: 0 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ use futures::StreamExt;
99
use serde::{Deserialize, Serialize};
1010
use spin_core::async_trait;
1111
use spin_factor_key_value::{log_error, Error, Store, StoreManager};
12-
use tracing::{instrument, Level};
1312

1413
pub struct KeyValueAzureCosmos {
1514
client: CollectionClient,
@@ -119,20 +118,11 @@ struct AzureCosmosStore {
119118

120119
#[async_trait]
121120
impl Store for AzureCosmosStore {
122-
#[instrument(name = "spin_key_value_azure.get", skip(self), err(level = Level::INFO), fields(
123-
otel.kind = "client"
124-
))]
125121
async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, Error> {
126122
let pair = self.get_pair(key).await?;
127123
Ok(pair.map(|p| p.value))
128124
}
129125

130-
#[instrument(
131-
name = "spin_key_value_azure.set",
132-
skip(self, value),
133-
err(level = Level::INFO),
134-
fields(otel.kind = "client")
135-
)]
136126
async fn set(&self, key: &str, value: &[u8]) -> Result<(), Error> {
137127
let pair = Pair {
138128
id: key.to_string(),
@@ -146,9 +136,6 @@ impl Store for AzureCosmosStore {
146136
Ok(())
147137
}
148138

149-
#[instrument(name = "spin_key_value_azure.delete", skip(self), err(level = Level::INFO), fields(
150-
otel.kind = "client"
151-
))]
152139
async fn delete(&self, key: &str) -> Result<(), Error> {
153140
if self.exists(key).await? {
154141
let document_client = self.client.document_client(key, &key).map_err(log_error)?;
@@ -157,19 +144,10 @@ impl Store for AzureCosmosStore {
157144
Ok(())
158145
}
159146

160-
#[instrument(name = "spin_key_value_azure.exists", skip(self), err(level = Level::INFO), fields(
161-
otel.kind = "client"
162-
))]
163147
async fn exists(&self, key: &str) -> Result<bool, Error> {
164148
Ok(self.get_pair(key).await?.is_some())
165149
}
166150

167-
#[instrument(
168-
name = "spin_key_value_azure.get_keys",
169-
skip(self),
170-
err(level = Level::INFO),
171-
fields(otel.kind = "client")
172-
)]
173151
async fn get_keys(&self) -> Result<Vec<String>, Error> {
174152
self.get_keys().await
175153
}

‎crates/key-value-redis/Cargo.toml

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ spin-core = { path = "../core" }
1212
spin-factor-key-value = { path = "../factor-key-value" }
1313
spin-world = { path = "../world" }
1414
tokio = "1"
15-
tracing = { workspace = true }
1615
url = "2"
1716

1817
[lints]

‎crates/key-value-redis/src/store.rs

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ use spin_core::async_trait;
44
use spin_factor_key_value::{log_error, Error, Store, StoreManager};
55
use std::sync::Arc;
66
use tokio::sync::{Mutex, OnceCell};
7-
use tracing::{instrument, Level};
87
use url::Url;
98

109
pub struct KeyValueRedis {
@@ -25,7 +24,6 @@ impl KeyValueRedis {
2524

2625
#[async_trait]
2726
impl StoreManager for KeyValueRedis {
28-
#[instrument(name = "spin_key_value_redis.get_store", skip(self), err(level = Level::INFO), fields(otel.kind = "client"))]
2927
async fn get(&self, _name: &str) -> Result<Arc<dyn Store>, Error> {
3028
let connection = self
3129
.connection
@@ -60,13 +58,11 @@ struct RedisStore {
6058

6159
#[async_trait]
6260
impl Store for RedisStore {
63-
#[instrument(name = "spin_key_value_redis.get", skip(self), err(level = Level::INFO), fields(otel.kind = "client"))]
6461
async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, Error> {
6562
let mut conn = self.connection.lock().await;
6663
conn.get(key).await.map_err(log_error)
6764
}
6865

69-
#[instrument(name = "spin_key_value_redis.set", skip(self, value), err(level = Level::INFO), fields(otel.kind = "client"))]
7066
async fn set(&self, key: &str, value: &[u8]) -> Result<(), Error> {
7167
self.connection
7268
.lock()
@@ -76,7 +72,6 @@ impl Store for RedisStore {
7672
.map_err(log_error)
7773
}
7874

79-
#[instrument(name = "spin_key_value_redis.delete", skip(self), err(level = Level::INFO), fields(otel.kind = "client"))]
8075
async fn delete(&self, key: &str) -> Result<(), Error> {
8176
self.connection
8277
.lock()
@@ -86,7 +81,6 @@ impl Store for RedisStore {
8681
.map_err(log_error)
8782
}
8883

89-
#[instrument(name = "spin_key_value_redis.exists", skip(self), err(level = Level::INFO), fields(otel.kind = "client"))]
9084
async fn exists(&self, key: &str) -> Result<bool, Error> {
9185
self.connection
9286
.lock()
@@ -96,7 +90,6 @@ impl Store for RedisStore {
9690
.map_err(log_error)
9791
}
9892

99-
#[instrument(name = "spin_key_value_redis.get_keys", skip(self), err(level = Level::INFO), fields(otel.kind = "client"))]
10093
async fn get_keys(&self) -> Result<Vec<String>, Error> {
10194
self.connection
10295
.lock()

‎crates/key-value-spin/Cargo.toml

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ spin-core = { path = "../core" }
1313
spin-factor-key-value = { path = "../factor-key-value" }
1414
spin-world = { path = "../world" }
1515
tokio = { version = "1", features = ["rt-multi-thread"] }
16-
tracing = { workspace = true }
1716

1817
[lints]
1918
workspace = true

0 commit comments

Comments
 (0)