feat: switch to wasi-async-runtime from tokio

Author: iambenkay

Cargo.lock

@@ -17,7 +17,7 @@ aws-smithy-wasm = { version = "0.1.4", default-features = false }
 aws-sdk-bedrockruntime = { version = "1.56.0", default-features = false }
 aws-smithy-types = { version = "1.3.1" }
 
-tokio = { version = "1.43.0", features = ["rt", "sync", "time"] }
+wasi-async-runtime = "0.1.2"
 
 # To infer mime types of downloaded images before passing to bedrock
 infer = { version = "0.19.0", default-features = false }

llm/bedrock/Cargo.toml

@@ -0,0 +1,16 @@
+use std::future::Future;
+
+pub fn get_async_runtime() -> AsyncRuntime {
+    AsyncRuntime
+}
+
+pub struct AsyncRuntime;
+
+impl AsyncRuntime {
+    pub fn block_on<F>(self, f: F) -> F::Output
+    where
+        F: Future,
+    {
+        wasi_async_runtime::block_on(|_| async { f.await })
+    }
+}

llm/bedrock/src/async_utils.rs

@@ -1,4 +1,5 @@
 use crate::{
+    async_utils,
     conversions::{self, from_converse_sdk_error, from_converse_stream_sdk_error, BedrockInput},
     stream::BedrockChatStream,
 };
@@ -17,6 +18,7 @@ use golem_llm::{
     config::{get_config_key, get_config_key_or_none},
     golem::llm::llm,
 };
+use golem_rust::bindings::wasi::clocks::monotonic_clock;
 use log::trace;
 
 #[derive(Debug)]
@@ -30,14 +32,14 @@ impl Bedrock {
 
         let wasi_http = WasiHttpClientBuilder::new().build();
 
-        let runtime = get_async_runtime();
+        let runtime = async_utils::get_async_runtime();
 
         runtime.block_on(async {
             let sdk_config = aws_config::defaults(BehaviorVersion::latest())
                 .region(environment.aws_region())
                 .http_client(wasi_http)
                 .credentials_provider(environment.aws_credentials())
-                .sleep_impl(TokioSleep)
+                .sleep_impl(WasiSleep)
                 .load()
                 .await;
             let client = bedrock::Client::new(&sdk_config);
@@ -53,7 +55,7 @@ impl Bedrock {
     ) -> llm::ChatEvent {
         let bedrock_input = BedrockInput::from(messages, config, tool_results);
 
-        let runtime = get_async_runtime();
+        let runtime = async_utils::get_async_runtime();
 
         match bedrock_input {
             Err(err) => llm::ChatEvent::Error(err),
@@ -97,7 +99,7 @@ impl Bedrock {
         match bedrock_input {
             Err(err) => BedrockChatStream::failed(err),
             Ok(input) => {
-                let runtime = get_async_runtime();
+                let runtime = async_utils::get_async_runtime();
                 trace!("Sending request to AWS Bedrock: {input:?}");
                 runtime.block_on(async {
                     let model_id = input.model_id.clone();
@@ -176,19 +178,13 @@ impl BedrockEnvironment {
     }
 }
 
-pub fn get_async_runtime() -> tokio::runtime::Runtime {
-    tokio::runtime::Builder::new_current_thread()
-        .enable_time()
-        .build()
-        .unwrap()
-}
-
 #[derive(Debug, Clone)]
-struct TokioSleep;
-impl AsyncSleep for TokioSleep {
+struct WasiSleep;
+impl AsyncSleep for WasiSleep {
     fn sleep(&self, duration: std::time::Duration) -> Sleep {
         Sleep::new(Box::pin(async move {
-            tokio::time::sleep(duration).await;
+            let nanos = duration.as_nanos() as u64;
+            monotonic_clock::subscribe_duration(nanos).block();
         }))
     }
 }

llm/bedrock/src/client.rs

@@ -6,8 +6,10 @@ use golem_llm::{
     golem::llm::llm::{self, ChatEvent, ChatStream, Config, Guest, Message, ToolCall, ToolResult},
     LOGGING_STATE,
 };
+use golem_rust::bindings::wasi::clocks::monotonic_clock;
 use stream::BedrockChatStream;
 
+mod async_utils;
 mod client;
 mod conversions;
 mod stream;
@@ -116,7 +118,8 @@ impl ExtendedGuest for BedrockComponent {
     }
 
     fn subscribe(_stream: &Self::ChatStream) -> golem_rust::wasm_rpc::Pollable {
-        unimplemented!()
+        // this function will never get called in bedrock implementation because of `golem-llm/nopoll` feature flag
+        monotonic_clock::subscribe_duration(0)
     }
 }
 

llm/bedrock/src/lib.rs

@@ -6,7 +6,7 @@ use golem_llm::golem::llm::llm;
 use std::cell::{RefCell, RefMut};
 
 use crate::{
-    client::get_async_runtime,
+    async_utils,
     conversions::{converse_stream_output_to_stream_event, custom_error, merge_metadata},
 };
 
@@ -53,7 +53,7 @@ impl BedrockChatStream {
     }
     fn get_single_event(&self) -> Option<llm::StreamEvent> {
         if let Some(stream) = self.stream_mut().as_mut() {
-            let runtime = get_async_runtime();
+            let runtime = async_utils::get_async_runtime();
 
             runtime.block_on(async move {
                 let token = stream.recv().await;

llm/bedrock/src/stream.rs

@@ -38,8 +38,8 @@ path = "wit-generated"
 
 [package.metadata.component.target.dependencies]
 "golem:llm" = { path = "wit-generated/deps/golem-llm" }
-"wasi:clocks" = { path = "wit-generated/deps/clocks" }
 "wasi:io" = { path = "wit-generated/deps/io" }
+"wasi:clocks" = { path = "wit-generated/deps/clocks" }
 "golem:rpc" = { path = "wit-generated/deps/golem-rpc" }
 "test:helper-client" = { path = "wit-generated/deps/test_helper-client" }
 "test:llm-exports" = { path = "wit-generated/deps/test_llm-exports" }