Use normal tokio tasks for core functions.

crates/ai00-core/src/lib.rs

@@ -442,15 +442,14 @@ async fn load_runtime(
     Ok(runtime)
 }
 
-#[tokio::main]
 pub async fn model_route(receiver: Receiver<ThreadRequest>) -> Result<()> {
     let env: Arc<RwLock<Environment>> = Default::default();
     let queue: Arc<Mutex<Vec<GenerateContext>>> = Default::default();
 
     let sender = {
         let (sender, receiver) = flume::unbounded();
         let env = env.clone();
-        tokio::task::spawn_blocking(move || crate::run::run(receiver, env));
+        tokio::task::spawn(crate::run::run(receiver, env));
         sender
     };
 

crates/ai00-core/src/run.rs

@@ -366,10 +366,7 @@ impl Environment {
                 match runtime.queue(context).await.expect("queue task error") {
                     SlotResult::Success(batch) => log::info!("queued task at slot {batch}"),
                     SlotResult::Fault(batch) => log::info!("swapped task at slot {batch}"),
-                    SlotResult::Failure(context) => {
-                        log::info!("failed to queue task");
-                        queue.push(*context);
-                    }
+                    SlotResult::Failure(context) => queue.push(*context),
                     SlotResult::Error(reason) => log::warn!("queue task failed: {}", reason),
                 }
             }
@@ -973,7 +970,6 @@ impl Runtime {
     }
 }
 
-#[tokio::main]
 pub async fn run(receiver: Receiver<()>, env: Arc<RwLock<Environment>>) {
     {
         // this task constantly runs, cleaning up state cache

crates/ai00-server/src/api/model.rs

@@ -1,5 +1,3 @@
-use std::time::Duration;
-
 use ai00_core::{
     run::{InitState, StateId},
     ReloadRequest, RuntimeInfo, SaveRequest, ThreadRequest,
@@ -10,7 +8,7 @@ use serde::Serialize;
 use web_rwkv::runtime::model::ModelInfo;
 
 use super::*;
-use crate::{build_path, types::ThreadState};
+use crate::{build_path, types::ThreadState, SLEEP};
 
 #[derive(Debug, Clone, Serialize)]
 pub struct InfoResponse {
@@ -33,7 +31,7 @@ pub async fn info(depot: &mut Depot) -> Json<InfoResponse> {
         model,
         states,
         ..
-    } = request_info(sender.to_owned(), Duration::from_millis(500)).await;
+    } = request_info(sender.to_owned(), SLEEP).await;
     let states = states
         .into_iter()
         .map(|(id, InitState { name, .. })| InitStateInfo { id, name })
@@ -50,7 +48,7 @@ pub async fn info(depot: &mut Depot) -> Json<InfoResponse> {
 pub async fn state(depot: &mut Depot, res: &mut Response) {
     let ThreadState { sender, .. } = depot.obtain::<ThreadState>().unwrap();
     let (info_sender, info_receiver) = flume::unbounded();
-    let task = request_info_stream(sender.to_owned(), info_sender, Duration::from_millis(500));
+    let task = request_info_stream(sender.to_owned(), info_sender, SLEEP);
     tokio::task::spawn(task);
 
     let stream = info_receiver.into_stream().map(

crates/ai00-server/src/api/oai/chat.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, sync::Arc, time::Duration};
+use std::{collections::HashMap, sync::Arc};
 
 use ai00_core::{
     run::StateId, sampler::Sampler, FinishReason, GenerateRequest, ThreadRequest, Token,
@@ -15,6 +15,7 @@ use super::*;
 use crate::{
     api::request_info,
     types::{Array, ThreadState},
+    SLEEP,
 };
 
 #[derive(Debug, Default, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, ToSchema)]
@@ -184,7 +185,7 @@ struct PartialChatResponse {
 
 async fn respond_one(depot: &mut Depot, request: ChatRequest, res: &mut Response) {
     let ThreadState { sender, .. } = depot.obtain::<ThreadState>().unwrap();
-    let info = request_info(sender.clone(), Duration::from_millis(500)).await;
+    let info = request_info(sender.clone(), SLEEP).await;
     let model_name = info.reload.model_path.to_string_lossy().into_owned();
 
     let (token_sender, token_receiver) = flume::unbounded();
@@ -233,7 +234,7 @@ async fn respond_one(depot: &mut Depot, request: ChatRequest, res: &mut Response
 
 async fn respond_stream(depot: &mut Depot, request: ChatRequest, res: &mut Response) {
     let ThreadState { sender, .. } = depot.obtain::<ThreadState>().unwrap();
-    let info = request_info(sender.clone(), Duration::from_millis(500)).await;
+    let info = request_info(sender.clone(), SLEEP).await;
     let model_name = info.reload.model_path.to_string_lossy().into_owned();
 
     let (token_sender, token_receiver) = flume::unbounded();

crates/ai00-server/src/api/oai/completion.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, sync::Arc, time::Duration};
+use std::{collections::HashMap, sync::Arc};
 
 use ai00_core::{
     run::StateId, FinishReason, GenerateRequest, ThreadRequest, Token, TokenCounter, MAX_TOKENS,
@@ -16,6 +16,7 @@ use super::*;
 use crate::{
     api::request_info,
     types::{Array, ThreadState},
+    SLEEP,
 };
 
 #[derive(Debug, Deserialize, ToSchema)]
@@ -131,7 +132,7 @@ pub struct PartialCompletionResponse {
 
 async fn respond_one(depot: &mut Depot, request: CompletionRequest, res: &mut Response) {
     let ThreadState { sender, .. } = depot.obtain::<ThreadState>().unwrap();
-    let info = request_info(sender.clone(), Duration::from_millis(500)).await;
+    let info = request_info(sender.clone(), SLEEP).await;
     let model_name = info.reload.model_path.to_string_lossy().into_owned();
 
     let (token_sender, token_receiver) = flume::unbounded();
@@ -177,7 +178,7 @@ async fn respond_one(depot: &mut Depot, request: CompletionRequest, res: &mut Re
 
 async fn respond_stream(depot: &mut Depot, request: CompletionRequest, res: &mut Response) {
     let ThreadState { sender, .. } = depot.obtain::<ThreadState>().unwrap();
-    let info = request_info(sender.clone(), Duration::from_millis(500)).await;
+    let info = request_info(sender.clone(), SLEEP).await;
     let model_name = info.reload.model_path.to_string_lossy().into_owned();
 
     let (token_sender, token_receiver) = flume::unbounded();

crates/ai00-server/src/api/oai/embedding.rs

@@ -1,5 +1,3 @@
-use std::time::Duration;
-
 use ai00_core::{GenerateRequest, ThreadRequest, Token, TokenCounter};
 use futures_util::StreamExt;
 use salvo::{
@@ -11,6 +9,7 @@ use serde::{Deserialize, Serialize};
 use crate::{
     api::request_info,
     types::{Array, ThreadState},
+    SLEEP,
 };
 
 #[derive(Debug, Default, Clone, Deserialize, ToSchema, ToParameters)]
@@ -57,7 +56,7 @@ pub async fn embeddings(
 ) -> Json<EmbeddingResponse> {
     let request = req.to_owned(); // req.parse_json::<EmbeddingRequest>().await.unwrap();
     let ThreadState { sender, .. } = depot.obtain::<ThreadState>().unwrap();
-    let info = request_info(sender.clone(), Duration::from_millis(500)).await;
+    let info = request_info(sender.clone(), SLEEP).await;
     let model_name = info.reload.model_path.to_string_lossy().into_owned();
 
     let (token_sender, token_receiver) = flume::unbounded();

crates/ai00-server/src/api/oai/info.rs

@@ -1,12 +1,10 @@
-use std::time::Duration;
-
 use salvo::{
     oapi::{ToResponse, ToSchema},
     prelude::*,
 };
 use serde::Serialize;
 
-use crate::{api::request_info, types::ThreadState};
+use crate::{api::request_info, types::ThreadState, SLEEP};
 
 #[derive(Debug, Serialize, ToSchema)]
 struct ModelChoice {
@@ -23,7 +21,7 @@ pub struct ModelResponse {
 #[endpoint]
 pub async fn models(depot: &mut Depot) -> Json<ModelResponse> {
     let ThreadState { sender, .. } = depot.obtain::<ThreadState>().unwrap();
-    let info = request_info(sender.to_owned(), Duration::from_millis(500)).await;
+    let info = request_info(sender.to_owned(), SLEEP).await;
     let model_name = info
         .reload
         .model_path

crates/ai00-server/src/main.rs

@@ -2,6 +2,7 @@ use std::{
     io::Cursor,
     net::{IpAddr, Ipv4Addr, SocketAddr},
     path::{Path, PathBuf},
+    time::Duration,
 };
 
 use ai00_core::{model_route, ThreadRequest};
@@ -30,6 +31,8 @@ mod api;
 mod config;
 mod types;
 
+const SLEEP: Duration = Duration::from_millis(500);
+
 pub fn build_path(path: impl AsRef<Path>, name: impl AsRef<Path>) -> Result<PathBuf> {
     let permitted = path.as_ref();
     let name = name.as_ref();
@@ -116,7 +119,7 @@ async fn main() {
     log::info!("{}\tversion: {}", bin_name, version);
 
     let (sender, receiver) = flume::unbounded::<ThreadRequest>();
-    tokio::task::spawn_blocking(move || model_route(receiver));
+    tokio::task::spawn(model_route(receiver));
 
     let (listen, config) = {
         let path = args