diff --git a/Cargo.lock b/Cargo.lock index 3ac48ba77..70a56e4bd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,56 @@ version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + [[package]] name = "anyhow" version = "1.0.102" @@ -199,6 +249,52 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + [[package]] name = "concurrent-queue" version = "2.5.0" @@ -265,6 +361,15 @@ version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-queue" version = "0.3.12" @@ -326,6 +431,12 @@ dependencies = [ "zeroize", ] +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" + [[package]] name = "digest" version = "0.10.7" @@ -959,6 +1070,12 @@ dependencies = [ "spider-tdl", ] +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + [[package]] name = "itertools" version = "0.14.0" @@ -1048,6 +1165,12 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "libyaml-rs" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e126dda6f34391ab7b444f9922055facc83c07a910da3eb16f1e4d9c45dc777" + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -1155,6 +1278,12 @@ dependencies = [ "zeroize", ] +[[package]] +name = "num-conv" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" + [[package]] name = "num-integer" version = "0.1.46" @@ -1191,6 +1320,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + [[package]] name = "papergrid" version = "0.17.0" @@ -1324,6 +1459,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -1890,6 +2031,8 @@ dependencies = [ "anyhow", "async-channel", "async-trait", + "bincode", + "clap", "const_format", "dashmap", "rand 0.9.4", @@ -1899,14 +2042,18 @@ dependencies = [ "serial_test", "spider-core", "spider-derive", + "spider-proto-rust", "spider-tdl", + "spider-utils", "sqlx", "subtle", "tabled", "thiserror", "tokio", "tokio-util", + "tonic", "tracing", + "yaml_serde", ] [[package]] @@ -1951,6 +2098,14 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "spider-utils" +version = "0.1.0" +dependencies = [ + "tracing-appender", + "tracing-subscriber", +] + [[package]] name = "spin" version = "0.9.8" @@ -2175,6 +2330,12 @@ dependencies = [ "unicode-properties", ] +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "strum" version = "0.28.0" @@ -2202,6 +2363,12 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" +[[package]] +name = "symlink" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7973cce6668464ea31f176d85b13c7ab3bba2cb3b77a2ed26abd7801688010a" + [[package]] name = "syn" version = "1.0.109" @@ -2362,6 +2529,36 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "time" +version = "0.3.51" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85c17d80feb7334b40c484e45ed1a5273dfd8bfda537c3be2e74a06a6686f327" +dependencies = [ + "deranged", + "num-conv", + "powerfmt", + "serde_core", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" + +[[package]] +name = "time-macros" +version = "0.2.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dcef1a61bdb119096e153208ec5cbec23944ce8bca13be5c7f60c634f7403935" +dependencies = [ + "num-conv", + "time-core", +] + [[package]] name = "tinystr" version = "0.8.3" @@ -2541,6 +2738,19 @@ dependencies = [ "tracing-core", ] +[[package]] +name = "tracing-appender" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "050686193eb999b4bb3bc2acfa891a13da00f79734704c4b8b4ef1a10b368a3c" +dependencies = [ + "crossbeam-channel", + "symlink", + "thiserror", + "time", + "tracing-subscriber", +] + [[package]] name = "tracing-attributes" version = "0.1.31" @@ -2659,6 +2869,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "valuable" version = "0.1.1" @@ -3036,6 +3252,19 @@ version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" +[[package]] +name = "yaml_serde" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c7c1b1a6a7c8a6b2741a6c21a4f8918e51899b111cfa08d1288202656e3975" +dependencies = [ + "indexmap 2.14.0", + "itoa", + "libyaml-rs", + "ryu", + "serde", +] + [[package]] name = "yoke" version = "0.8.3" diff --git a/Cargo.toml b/Cargo.toml index 5bd6742d9..a91cc9f27 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,6 +10,7 @@ members = [ "components/spider-task-executor", "components/spider-tdl", "components/spider-tdl-derive", + "components/spider-utils", "examples/huntsman/complex/tasks", "examples/huntsman/complex/types", "tests/huntsman/em-runtime", @@ -28,4 +29,5 @@ default-members = [ "components/spider-task-executor", "components/spider-tdl", "components/spider-tdl-derive", + "components/spider-utils", ] diff --git a/components/spider-execution-manager/src/client/grpc/storage.rs b/components/spider-execution-manager/src/client/grpc/storage.rs index 378de5d4f..4e0a4ef1e 100644 --- a/components/spider-execution-manager/src/client/grpc/storage.rs +++ b/components/spider-execution-manager/src/client/grpc/storage.rs @@ -5,7 +5,7 @@ use async_trait::async_trait; use spider_core::types::{ - id::{ExecutionManagerId, JobId, SessionId, TaskId}, + id::{ExecutionManagerId, JobId, SessionId, TaskId, TaskInstanceId}, io::ExecutionContext, }; use spider_proto_rust::storage::{ @@ -87,6 +87,7 @@ impl StorageClient for GrpcStorageClient { &self, job_id: JobId, task_id: TaskId, + task_instance_id: TaskInstanceId, em_id: ExecutionManagerId, session_id: SessionId, serialized_outputs: Option>, @@ -97,6 +98,7 @@ impl StorageClient for GrpcStorageClient { execution_manager_id: em_id.get(), session_id, serialized_outputs: serialized_outputs.unwrap_or_default(), + task_instance_id, }; let response = self .client @@ -113,6 +115,7 @@ impl StorageClient for GrpcStorageClient { &self, job_id: JobId, task_id: TaskId, + task_instance_id: TaskInstanceId, em_id: ExecutionManagerId, session_id: SessionId, error_message: String, @@ -123,6 +126,7 @@ impl StorageClient for GrpcStorageClient { execution_manager_id: em_id.get(), session_id, error_message, + task_instance_id, }; let response = self .client diff --git a/components/spider-execution-manager/src/client/storage.rs b/components/spider-execution-manager/src/client/storage.rs index 89732c8ce..b8e13a4d7 100644 --- a/components/spider-execution-manager/src/client/storage.rs +++ b/components/spider-execution-manager/src/client/storage.rs @@ -5,7 +5,7 @@ use async_trait::async_trait; use spider_core::types::{ - id::{ExecutionManagerId, JobId, SessionId, TaskId}, + id::{ExecutionManagerId, JobId, SessionId, TaskId, TaskInstanceId}, io::ExecutionContext, }; @@ -79,6 +79,7 @@ pub trait StorageClient: Send + Sync { /// /// * `job_id` - The owning job. /// * `task_id` - The task that ran. + /// * `task_instance_id` - The task instance that produced the outcome. /// * `em_id` - The identity of the calling execution manager. /// * `session_id` - The session id captured from the scheduler assignment. /// * `serialized_outputs` - The wire-format encoded task outputs buffer, forwarded verbatim to @@ -99,6 +100,7 @@ pub trait StorageClient: Send + Sync { &self, job_id: JobId, task_id: TaskId, + task_instance_id: TaskInstanceId, em_id: ExecutionManagerId, session_id: SessionId, serialized_outputs: Option>, @@ -110,6 +112,7 @@ pub trait StorageClient: Send + Sync { /// /// * `job_id` - The owning job. /// * `task_id` - The task that ran. + /// * `task_instance_id` - The task instance that produced the outcome. /// * `em_id` - The identity of the calling execution manager. /// * `session_id` - The session id captured from the scheduler assignment. /// * `error_message` - The formatted error message. @@ -127,6 +130,7 @@ pub trait StorageClient: Send + Sync { &self, job_id: JobId, task_id: TaskId, + task_instance_id: TaskInstanceId, em_id: ExecutionManagerId, session_id: SessionId, error_message: String, diff --git a/components/spider-execution-manager/src/runtime.rs b/components/spider-execution-manager/src/runtime.rs index 03d0dc76f..91d011395 100644 --- a/components/spider-execution-manager/src/runtime.rs +++ b/components/spider-execution-manager/src/runtime.rs @@ -5,7 +5,7 @@ use std::{collections::VecDeque, net::IpAddr, path::PathBuf, sync::Arc, time::Du use spider_core::{ session::SessionTracker, types::{ - id::{ExecutionManagerId, JobId, SessionId, TaskId}, + id::{ExecutionManagerId, JobId, SessionId, TaskId, TaskInstanceId}, io::ExecutionContext, scheduler::TaskAssignmentRecord, }, @@ -306,6 +306,7 @@ impl< let Some(execution_context) = self.register_task_instance(response).await? else { continue; }; + let task_instance_id = execution_context.task_instance_id; let hard_timeout = Duration::from_millis(execution_context.timeout_policy.hard_timeout_ms); @@ -348,6 +349,7 @@ impl< em: self.em_id, job: response.task_assignment.job_id, task: response.task_assignment.task_id, + task_instance_id, session: response.session_id, }, outcome, @@ -447,6 +449,7 @@ struct ReportTarget { em: ExecutionManagerId, job: JobId, task: TaskId, + task_instance_id: TaskInstanceId, session: SessionId, } @@ -532,17 +535,18 @@ impl Report { em, job, task, + task_instance_id, session, } = target; match self { Self::Success(outputs) => { storage_client - .report_task_success(job, task, em, session, outputs) + .report_task_success(job, task, task_instance_id, em, session, outputs) .await } Self::Failure(message) => { storage_client - .report_task_failure(job, task, em, session, message) + .report_task_failure(job, task, task_instance_id, em, session, message) .await } } diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index 78e88fc4f..67739ff57 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -7,8 +7,6 @@ pub struct SubmitJobRequest { pub serialized_task_graph: ::prost::alloc::vec::Vec, #[prost(bytes = "vec", tag = "3")] pub serialized_inputs: ::prost::alloc::vec::Vec, - #[prost(uint64, tag = "4")] - pub session_id: u64, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct SubmitJobResponse { @@ -81,40 +79,6 @@ pub mod job_error_response { } } #[derive(Clone, Copy, PartialEq, ::prost::Message)] -pub struct DeleteExpiredTerminatedJobsRequest { - #[prost(uint64, tag = "1")] - pub expire_after_sec: u64, - #[prost(uint64, tag = "2")] - pub session_id: u64, -} -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct DeleteExpiredTerminatedJobsResponse { - #[prost(oneof = "delete_expired_terminated_jobs_response::Result", tags = "1, 2")] - pub result: ::core::option::Option, -} -/// Nested message and enum types in `DeleteExpiredTerminatedJobsResponse`. -pub mod delete_expired_terminated_jobs_response { - #[derive(Clone, PartialEq, ::prost::Oneof)] - pub enum Result { - #[prost(message, tag = "1")] - DeletedJobs(super::DeletedJobs), - #[prost(message, tag = "2")] - Error(super::JobOrchestrationError), - } -} -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct DeletedJobs { - #[prost(uint64, repeated, tag = "1")] - pub job_ids: ::prost::alloc::vec::Vec, -} -#[derive(Clone, Copy, PartialEq, ::prost::Message)] -pub struct ResendReadyTasksRequest { - #[prost(uint64, optional, tag = "1")] - pub job_id: ::core::option::Option, - #[prost(uint64, tag = "2")] - pub session_id: u64, -} -#[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct PollReadyTasksRequest { #[prost(uint64, tag = "1")] pub max_items: u64, @@ -185,10 +149,12 @@ pub struct ReportTaskSuccessRequest { #[prost(message, optional, tag = "2")] pub task_id: ::core::option::Option, #[prost(uint64, tag = "3")] - pub execution_manager_id: u64, + pub task_instance_id: u64, #[prost(uint64, tag = "4")] + pub execution_manager_id: u64, + #[prost(uint64, tag = "5")] pub session_id: u64, - #[prost(bytes = "vec", tag = "5")] + #[prost(bytes = "vec", tag = "6")] pub serialized_outputs: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, ::prost::Message)] @@ -198,10 +164,12 @@ pub struct ReportTaskFailureRequest { #[prost(message, optional, tag = "2")] pub task_id: ::core::option::Option, #[prost(uint64, tag = "3")] - pub execution_manager_id: u64, + pub task_instance_id: u64, #[prost(uint64, tag = "4")] + pub execution_manager_id: u64, + #[prost(uint64, tag = "5")] pub session_id: u64, - #[prost(string, tag = "5")] + #[prost(string, tag = "6")] pub error_message: ::prost::alloc::string::String, } #[derive(Clone, PartialEq, ::prost::Message)] @@ -213,13 +181,6 @@ pub struct AddResourceGroupRequest { #[prost(uint64, tag = "3")] pub session_id: u64, } -#[derive(Clone, Copy, PartialEq, ::prost::Message)] -pub struct ResourceGroupIdRequest { - #[prost(uint64, tag = "1")] - pub resource_group_id: u64, - #[prost(uint64, tag = "2")] - pub session_id: u64, -} #[derive(Clone, PartialEq, ::prost::Message)] pub struct ResourceGroupIdResponse { #[prost(oneof = "resource_group_id_response::Result", tags = "1, 2")] @@ -1028,64 +989,6 @@ pub mod job_orchestration_service_client { ); self.inner.unary(req, path, codec).await } - pub async fn delete_expired_terminated_jobs( - &mut self, - request: impl tonic::IntoRequest, - ) -> std::result::Result< - tonic::Response, - tonic::Status, - > { - self.inner - .ready() - .await - .map_err(|e| { - tonic::Status::unknown( - format!("Service was not ready: {}", e.into()), - ) - })?; - let codec = tonic::codec::ProstCodec::default(); - let path = http::uri::PathAndQuery::from_static( - "/storage.JobOrchestrationService/DeleteExpiredTerminatedJobs", - ); - let mut req = request.into_request(); - req.extensions_mut() - .insert( - GrpcMethod::new( - "storage.JobOrchestrationService", - "DeleteExpiredTerminatedJobs", - ), - ); - self.inner.unary(req, path, codec).await - } - pub async fn resend_ready_tasks( - &mut self, - request: impl tonic::IntoRequest, - ) -> std::result::Result< - tonic::Response, - tonic::Status, - > { - self.inner - .ready() - .await - .map_err(|e| { - tonic::Status::unknown( - format!("Service was not ready: {}", e.into()), - ) - })?; - let codec = tonic::codec::ProstCodec::default(); - let path = http::uri::PathAndQuery::from_static( - "/storage.JobOrchestrationService/ResendReadyTasks", - ); - let mut req = request.into_request(); - req.extensions_mut() - .insert( - GrpcMethod::new( - "storage.JobOrchestrationService", - "ResendReadyTasks", - ), - ); - self.inner.unary(req, path, codec).await - } } } /// Generated client implementations. @@ -1598,35 +1501,6 @@ pub mod resource_group_management_service_client { ); self.inner.unary(req, path, codec).await } - pub async fn delete_resource_group( - &mut self, - request: impl tonic::IntoRequest, - ) -> std::result::Result< - tonic::Response, - tonic::Status, - > { - self.inner - .ready() - .await - .map_err(|e| { - tonic::Status::unknown( - format!("Service was not ready: {}", e.into()), - ) - })?; - let codec = tonic::codec::ProstCodec::default(); - let path = http::uri::PathAndQuery::from_static( - "/storage.ResourceGroupManagementService/DeleteResourceGroup", - ); - let mut req = request.into_request(); - req.extensions_mut() - .insert( - GrpcMethod::new( - "storage.ResourceGroupManagementService", - "DeleteResourceGroup", - ), - ); - self.inner.unary(req, path, codec).await - } } } /// Generated client implementations. @@ -2111,20 +1985,6 @@ pub mod job_orchestration_service_server { tonic::Response, tonic::Status, >; - async fn delete_expired_terminated_jobs( - &self, - request: tonic::Request, - ) -> std::result::Result< - tonic::Response, - tonic::Status, - >; - async fn resend_ready_tasks( - &self, - request: tonic::Request, - ) -> std::result::Result< - tonic::Response, - tonic::Status, - >; } #[derive(Debug)] pub struct JobOrchestrationServiceServer { @@ -2488,109 +2348,6 @@ pub mod job_orchestration_service_server { }; Box::pin(fut) } - "/storage.JobOrchestrationService/DeleteExpiredTerminatedJobs" => { - #[allow(non_camel_case_types)] - struct DeleteExpiredTerminatedJobsSvc( - pub Arc, - ); - impl< - T: JobOrchestrationService, - > tonic::server::UnaryService< - super::DeleteExpiredTerminatedJobsRequest, - > for DeleteExpiredTerminatedJobsSvc { - type Response = super::DeleteExpiredTerminatedJobsResponse; - type Future = BoxFuture< - tonic::Response, - tonic::Status, - >; - fn call( - &mut self, - request: tonic::Request< - super::DeleteExpiredTerminatedJobsRequest, - >, - ) -> Self::Future { - let inner = Arc::clone(&self.0); - let fut = async move { - ::delete_expired_terminated_jobs( - &inner, - request, - ) - .await - }; - Box::pin(fut) - } - } - let accept_compression_encodings = self.accept_compression_encodings; - let send_compression_encodings = self.send_compression_encodings; - let max_decoding_message_size = self.max_decoding_message_size; - let max_encoding_message_size = self.max_encoding_message_size; - let inner = self.inner.clone(); - let fut = async move { - let method = DeleteExpiredTerminatedJobsSvc(inner); - let codec = tonic::codec::ProstCodec::default(); - let mut grpc = tonic::server::Grpc::new(codec) - .apply_compression_config( - accept_compression_encodings, - send_compression_encodings, - ) - .apply_max_message_size_config( - max_decoding_message_size, - max_encoding_message_size, - ); - let res = grpc.unary(method, req).await; - Ok(res) - }; - Box::pin(fut) - } - "/storage.JobOrchestrationService/ResendReadyTasks" => { - #[allow(non_camel_case_types)] - struct ResendReadyTasksSvc(pub Arc); - impl< - T: JobOrchestrationService, - > tonic::server::UnaryService - for ResendReadyTasksSvc { - type Response = super::JobManagementOperationResponse; - type Future = BoxFuture< - tonic::Response, - tonic::Status, - >; - fn call( - &mut self, - request: tonic::Request, - ) -> Self::Future { - let inner = Arc::clone(&self.0); - let fut = async move { - ::resend_ready_tasks( - &inner, - request, - ) - .await - }; - Box::pin(fut) - } - } - let accept_compression_encodings = self.accept_compression_encodings; - let send_compression_encodings = self.send_compression_encodings; - let max_decoding_message_size = self.max_decoding_message_size; - let max_encoding_message_size = self.max_encoding_message_size; - let inner = self.inner.clone(); - let fut = async move { - let method = ResendReadyTasksSvc(inner); - let codec = tonic::codec::ProstCodec::default(); - let mut grpc = tonic::server::Grpc::new(codec) - .apply_compression_config( - accept_compression_encodings, - send_compression_encodings, - ) - .apply_max_message_size_config( - max_decoding_message_size, - max_encoding_message_size, - ); - let res = grpc.unary(method, req).await; - Ok(res) - }; - Box::pin(fut) - } _ => { Box::pin(async move { let mut response = http::Response::new(empty_body()); @@ -3255,13 +3012,6 @@ pub mod resource_group_management_service_server { tonic::Response, tonic::Status, >; - async fn delete_resource_group( - &self, - request: tonic::Request, - ) -> std::result::Result< - tonic::Response, - tonic::Status, - >; } #[derive(Debug)] pub struct ResourceGroupManagementServiceServer { @@ -3442,57 +3192,6 @@ pub mod resource_group_management_service_server { }; Box::pin(fut) } - "/storage.ResourceGroupManagementService/DeleteResourceGroup" => { - #[allow(non_camel_case_types)] - struct DeleteResourceGroupSvc( - pub Arc, - ); - impl< - T: ResourceGroupManagementService, - > tonic::server::UnaryService - for DeleteResourceGroupSvc { - type Response = super::ResourceGroupOperationResponse; - type Future = BoxFuture< - tonic::Response, - tonic::Status, - >; - fn call( - &mut self, - request: tonic::Request, - ) -> Self::Future { - let inner = Arc::clone(&self.0); - let fut = async move { - ::delete_resource_group( - &inner, - request, - ) - .await - }; - Box::pin(fut) - } - } - let accept_compression_encodings = self.accept_compression_encodings; - let send_compression_encodings = self.send_compression_encodings; - let max_decoding_message_size = self.max_decoding_message_size; - let max_encoding_message_size = self.max_encoding_message_size; - let inner = self.inner.clone(); - let fut = async move { - let method = DeleteResourceGroupSvc(inner); - let codec = tonic::codec::ProstCodec::default(); - let mut grpc = tonic::server::Grpc::new(codec) - .apply_compression_config( - accept_compression_encodings, - send_compression_encodings, - ) - .apply_max_message_size_config( - max_decoding_message_size, - max_encoding_message_size, - ); - let res = grpc.unary(method, req).await; - Ok(res) - }; - Box::pin(fut) - } _ => { Box::pin(async move { let mut response = http::Response::new(empty_body()); diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index 020ab5124..d9a5845bd 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -9,9 +9,6 @@ service JobOrchestrationService { rpc GetJobState(JobIdRequest) returns (JobStateResponse); rpc GetJobOutputs(JobIdRequest) returns (JobOutputsResponse); rpc GetJobError(JobIdRequest) returns (JobErrorResponse); - rpc DeleteExpiredTerminatedJobs(DeleteExpiredTerminatedJobsRequest) - returns (DeleteExpiredTerminatedJobsResponse); - rpc ResendReadyTasks(ResendReadyTasksRequest) returns (JobManagementOperationResponse); } service TaskInstanceManagementService { @@ -29,7 +26,6 @@ service InboundQueueService { service ResourceGroupManagementService { rpc AddResourceGroup(AddResourceGroupRequest) returns (ResourceGroupIdResponse); rpc VerifyResourceGroup(VerifyResourceGroupRequest) returns (ResourceGroupOperationResponse); - rpc DeleteResourceGroup(ResourceGroupIdRequest) returns (ResourceGroupOperationResponse); } service ExecutionManagerLivenessService { @@ -52,7 +48,6 @@ message SubmitJobRequest { uint64 resource_group_id = 1; bytes serialized_task_graph = 2; bytes serialized_inputs = 3; - uint64 session_id = 4; } message SubmitJobResponse { @@ -91,27 +86,6 @@ message JobErrorResponse { } } -message DeleteExpiredTerminatedJobsRequest { - uint64 expire_after_sec = 1; - uint64 session_id = 2; -} - -message DeleteExpiredTerminatedJobsResponse { - oneof result { - DeletedJobs deleted_jobs = 1; - JobOrchestrationError error = 2; - } -} - -message DeletedJobs { - repeated uint64 job_ids = 1; -} - -message ResendReadyTasksRequest { - optional uint64 job_id = 1; - uint64 session_id = 2; -} - message PollReadyTasksRequest { uint64 max_items = 1; uint64 wait_ms = 2; @@ -152,17 +126,19 @@ message RegisterTaskInstanceResponse { message ReportTaskSuccessRequest { uint64 job_id = 1; TaskId task_id = 2; - uint64 execution_manager_id = 3; - uint64 session_id = 4; - bytes serialized_outputs = 5; + uint64 task_instance_id = 3; + uint64 execution_manager_id = 4; + uint64 session_id = 5; + bytes serialized_outputs = 6; } message ReportTaskFailureRequest { uint64 job_id = 1; TaskId task_id = 2; - uint64 execution_manager_id = 3; - uint64 session_id = 4; - string error_message = 5; + uint64 task_instance_id = 3; + uint64 execution_manager_id = 4; + uint64 session_id = 5; + string error_message = 6; } message AddResourceGroupRequest { @@ -171,11 +147,6 @@ message AddResourceGroupRequest { uint64 session_id = 3; } -message ResourceGroupIdRequest { - uint64 resource_group_id = 1; - uint64 session_id = 2; -} - message ResourceGroupIdResponse { oneof result { uint64 resource_group_id = 1; diff --git a/components/spider-storage/Cargo.toml b/components/spider-storage/Cargo.toml index 1c81604d3..bb3947012 100644 --- a/components/spider-storage/Cargo.toml +++ b/components/spider-storage/Cargo.toml @@ -8,6 +8,10 @@ autotests = false name = "spider_storage" path = "src/lib.rs" +[[bin]] +name = "spider_storage_grpc_server" +path = "src/bin/grpc_server.rs" + [[test]] name = "test_spider_storage" path = "tests/test_spider_storage.rs" @@ -15,6 +19,8 @@ path = "tests/test_spider_storage.rs" [dependencies] async-channel = "2.3.1" async-trait = "0.1.89" +bincode = "1.3.3" +clap = { version = "4.6.1", features = ["derive"] } const_format = "0.2.35" dashmap = "6.1.0" rmp-serde = "1.3.1" @@ -22,16 +28,20 @@ secrecy = { version = "0.10.3", features = ["serde"] } serde = { version = "1.0.228", features = ["derive"] } spider-core = { path = "../spider-core" } spider-derive = { path = "../spider-derive" } +spider-proto-rust = { path = "../spider-proto-rust" } spider-tdl = { path = "../spider-tdl" } +spider-utils = { path = "../spider-utils" } sqlx = { version = "0.8.6", features = ["mysql", "runtime-tokio"] } subtle = "2.6.1" thiserror = "2.0.18" tokio = { version = "1.50.0", - features = ["macros", "rt-multi-thread", "sync", "time"] + features = ["macros", "rt-multi-thread", "signal", "sync", "time"] } tokio-util = { version = "0.7.18", features = ["rt"] } +tonic = "0.12.3" tracing = { version = "0.1.44", features = ["attributes"] } +yaml_serde = "0.10.4" [dev-dependencies] anyhow = "1.0.98" diff --git a/components/spider-storage/src/bin/grpc_server.rs b/components/spider-storage/src/bin/grpc_server.rs new file mode 100644 index 000000000..7f6acda29 --- /dev/null +++ b/components/spider-storage/src/bin/grpc_server.rs @@ -0,0 +1,74 @@ +//! Command-line entrypoint for the storage gRPC server. + +use std::{error::Error, net::SocketAddr, path::PathBuf}; + +use clap::Parser; +use spider_proto_rust::storage::{ + execution_manager_liveness_service_server::ExecutionManagerLivenessServiceServer, + inbound_queue_service_server::InboundQueueServiceServer, + job_orchestration_service_server::JobOrchestrationServiceServer, + resource_group_management_service_server::ResourceGroupManagementServiceServer, + scheduler_registration_service_server::SchedulerRegistrationServiceServer, + session_management_service_server::SessionManagementServiceServer, + task_instance_management_service_server::TaskInstanceManagementServiceServer, +}; +use spider_storage::{ServerConfig, grpc::GrpcServiceState, state::runtime::create_runtime}; +use spider_utils::logging::set_up_logging; +use tokio::select; +use tonic::transport::Server; + +/// Command-line arguments for the storage gRPC server. +#[derive(Debug, Parser)] +#[command(about = "Run the Spider storage gRPC server.")] +struct Cli { + /// Path to the YAML server configuration file. + #[arg(short, long, value_name = "PATH")] + config: PathBuf, +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let _log_guard = set_up_logging(); + let cli = Cli::parse(); + let server_config = ServerConfig::from_yaml_file(&cli.config)?; + let listen_addr = SocketAddr::new(server_config.host, server_config.port); + let (runtime, cancellation_token) = create_runtime(&server_config.runtime).await?; + let grpc_service = GrpcServiceState::new(runtime.get_service_state()); + tracing::info!(listen_addr = % listen_addr, "Starting storage gRPC server."); + + let serve_result = Server::builder() + .add_service(JobOrchestrationServiceServer::new(grpc_service.clone())) + .add_service(TaskInstanceManagementServiceServer::new( + grpc_service.clone(), + )) + .add_service(InboundQueueServiceServer::new(grpc_service.clone())) + .add_service(ResourceGroupManagementServiceServer::new( + grpc_service.clone(), + )) + .add_service(ExecutionManagerLivenessServiceServer::new( + grpc_service.clone(), + )) + .add_service(SchedulerRegistrationServiceServer::new( + grpc_service.clone(), + )) + .add_service(SessionManagementServiceServer::new(grpc_service)) + .serve_with_shutdown(listen_addr, async move { + select! { + () = cancellation_token.cancelled() => { + tracing::info!("Shutting down storage gRPC server."); + } + result = tokio::signal::ctrl_c() => { + if let Err(error) = result { + tracing::error!(error = % error, "Failed to listen for Ctrl-C."); + } + cancellation_token.cancel(); + } + } + }) + .await; + + let stop_result = runtime.stop().await; + serve_result?; + stop_result?; + Ok(()) +} diff --git a/components/spider-storage/src/config.rs b/components/spider-storage/src/config.rs index 30dbb8491..946247fac 100644 --- a/components/spider-storage/src/config.rs +++ b/components/spider-storage/src/config.rs @@ -1,5 +1,57 @@ +use std::{fs::File, io, net::IpAddr, path::Path}; + use secrecy::SecretString; use serde::{Deserialize, Serialize}; +use thiserror::Error; + +use crate::state::runtime::RuntimeConfig; + +/// Top-level configuration for the storage gRPC server. +/// +/// Pairs the server's listening endpoint with the [`RuntimeConfig`] used to build the storage +/// runtime. +#[derive(Clone, Debug, Deserialize)] +pub struct ServerConfig { + /// The IP address the gRPC server listens on. + pub host: IpAddr, + + /// The port the gRPC server listens on. + pub port: u16, + + /// Configuration for the storage runtime. + pub runtime: RuntimeConfig, +} + +impl ServerConfig { + /// Loads a [`ServerConfig`] from the YAML file at the given path. + /// + /// # Returns + /// + /// The parsed [`ServerConfig`] on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`yaml_serde::from_reader`]'s return values on failure. + pub fn from_yaml_file(path: &Path) -> Result { + let file = File::open(path)?; + let config = yaml_serde::from_reader(file)?; + Ok(config) + } +} + +/// Errors returned while loading a [`ServerConfig`]. +#[derive(Debug, Error)] +pub enum ConfigError { + /// Forwards an error from opening the configuration file. + #[error("failed to open config file: {0}")] + Io(#[from] io::Error), + + /// Forwards an error from deserializing the YAML configuration. + #[error("failed to parse config file: {0}")] + Parse(#[from] yaml_serde::Error), +} /// Configuration parameters for connecting to the database. #[derive(Clone, Debug, Deserialize, Serialize)] diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs new file mode 100644 index 000000000..2590a455c --- /dev/null +++ b/components/spider-storage/src/grpc.rs @@ -0,0 +1,252 @@ +//! gRPC service adapters for the storage runtime. + +use async_trait::async_trait; +use spider_proto_rust::storage::{ + self, + execution_manager_liveness_service_server::ExecutionManagerLivenessService, + inbound_queue_service_server::InboundQueueService, + job_orchestration_service_server::JobOrchestrationService, + resource_group_management_service_server::ResourceGroupManagementService, + scheduler_registration_service_server::SchedulerRegistrationService, + session_management_service_server::SessionManagementService, + task_instance_management_service_server::TaskInstanceManagementService, +}; +use tonic::{Request, Response, Status}; + +use crate::{ + db::DbStorage, + ready_queue::ReadyQueueSender, + state::ServiceState, + task_instance_pool::TaskInstancePoolConnector, +}; + +/// gRPC adapter over a storage [`ServiceState`]. +/// +/// # Type Parameters +/// +/// * `ReadyQueueSenderType` - The ready queue sender type. +/// * `DbConnectorType` - The database connector type. +/// * `TaskInstancePoolConnectorType` - The task instance pool connector type. +#[derive(Clone)] +pub struct GrpcServiceState< + ReadyQueueSenderType: ReadyQueueSender, + DbConnectorType: DbStorage, + TaskInstancePoolConnectorType: TaskInstancePoolConnector, +> { + _inner: ServiceState, +} + +impl< + ReadyQueueSenderType: ReadyQueueSender, + DbConnectorType: DbStorage, + TaskInstancePoolConnectorType: TaskInstancePoolConnector, +> GrpcServiceState +{ + /// Factory function. + /// + /// # Returns + /// + /// A new [`GrpcServiceState`] wrapping [`ServiceState`]. + #[must_use] + pub const fn new( + inner: ServiceState, + ) -> Self { + Self { _inner: inner } + } +} + +#[async_trait] +impl< + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: DbStorage + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, +> JobOrchestrationService + for GrpcServiceState +{ + async fn submit_job( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } + + async fn start_job( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } + + async fn cancel_job( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } + + async fn get_job_state( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } + + async fn get_job_outputs( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } + + async fn get_job_error( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } +} + +#[async_trait] +impl< + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: DbStorage + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, +> TaskInstanceManagementService + for GrpcServiceState +{ + async fn register_task_instance( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } + + async fn report_task_success( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } + + async fn report_task_failure( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } +} + +#[async_trait] +impl< + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: DbStorage + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, +> InboundQueueService + for GrpcServiceState +{ + async fn poll_ready_tasks( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } + + async fn poll_ready_commit_tasks( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } + + async fn poll_ready_cleanup_tasks( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } +} + +#[async_trait] +impl< + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: DbStorage + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, +> ResourceGroupManagementService + for GrpcServiceState +{ + async fn add_resource_group( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } + + async fn verify_resource_group( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } +} + +#[async_trait] +impl< + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: DbStorage + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, +> ExecutionManagerLivenessService + for GrpcServiceState +{ + async fn register_execution_manager( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } + + async fn update_execution_manager_heartbeat( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } +} + +#[async_trait] +impl< + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: DbStorage + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, +> SchedulerRegistrationService + for GrpcServiceState +{ + async fn register_scheduler( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } + + async fn get_schedulers( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } +} + +#[async_trait] +impl< + ReadyQueueSenderType: ReadyQueueSender + 'static, + DbConnectorType: DbStorage + 'static, + TaskInstancePoolConnectorType: TaskInstancePoolConnector + 'static, +> SessionManagementService + for GrpcServiceState +{ + async fn get_session( + &self, + _request: Request, + ) -> Result, Status> { + todo!("Not implemented") + } +} diff --git a/components/spider-storage/src/lib.rs b/components/spider-storage/src/lib.rs index 921e148e9..6ea141c9e 100644 --- a/components/spider-storage/src/lib.rs +++ b/components/spider-storage/src/lib.rs @@ -1,8 +1,9 @@ pub mod cache; mod config; pub mod db; +pub mod grpc; pub mod ready_queue; pub mod state; pub mod task_instance_pool; -pub use config::DatabaseConfig; +pub use config::{ConfigError, DatabaseConfig, ServerConfig}; diff --git a/components/spider-storage/src/ready_queue.rs b/components/spider-storage/src/ready_queue.rs index 1917b9e30..db453fb71 100644 --- a/components/spider-storage/src/ready_queue.rs +++ b/components/spider-storage/src/ready_queue.rs @@ -13,6 +13,7 @@ use std::time::Duration; use async_channel::{Receiver, Sender}; use async_trait::async_trait; +use serde::Deserialize; use spider_core::{ task::TaskIndex, types::id::{JobId, ResourceGroupId}, @@ -47,7 +48,8 @@ pub struct ReadyQueueEntry { } /// Configuration of a ready queue. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, Deserialize)] +#[serde(default)] pub struct ReadyQueueConfig { /// The capacity of the task lane. Must be greater than zero. pub task_capacity: usize, diff --git a/components/spider-storage/src/state/job_cache_gc.rs b/components/spider-storage/src/state/job_cache_gc.rs index 6d99ce3d8..916737597 100644 --- a/components/spider-storage/src/state/job_cache_gc.rs +++ b/components/spider-storage/src/state/job_cache_gc.rs @@ -5,6 +5,7 @@ use std::{ time::{Duration, Instant}, }; +use serde::Deserialize; use spider_core::types::id::JobId; use tokio::{ sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel}, @@ -21,7 +22,8 @@ use crate::{ }; /// Configuration for the job-cache GC actor. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, Deserialize)] +#[serde(default)] pub struct JobCacheGcConfig { /// Seconds to keep a terminated job in the in-memory cache before GC can remove it. pub terminated_job_retention_sec: u64, diff --git a/components/spider-storage/src/state/runtime.rs b/components/spider-storage/src/state/runtime.rs index 7c72318fe..8d718e3ef 100644 --- a/components/spider-storage/src/state/runtime.rs +++ b/components/spider-storage/src/state/runtime.rs @@ -1,5 +1,6 @@ use std::time::Duration; +use serde::Deserialize; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; @@ -21,10 +22,14 @@ use crate::{ }; /// Runtime configuration for the storage service. +#[derive(Clone, Debug, Deserialize)] pub struct RuntimeConfig { pub db_config: DatabaseConfig, + #[serde(default)] pub ready_queue_config: ReadyQueueConfig, + #[serde(default)] pub task_instance_pool_config: TaskInstancePoolConfig, + #[serde(default)] pub job_cache_gc_config: JobCacheGcConfig, } @@ -74,10 +79,10 @@ impl< tracing::info!("Task instance pool stopped."); } Ok(Err(e)) => { - tracing::error!(error = ? e, "Task instance pool exited on error."); + tracing::error!(error = % e, "Task instance pool exited on error."); } Err(e) => { - tracing::error!(error = ? e, "Task instance pool exited on panic."); + tracing::error!(error = % e, "Task instance pool exited on panic."); } } }; @@ -88,10 +93,10 @@ impl< tracing::info!("Job cache GC stopped."); } Ok(Err(e)) => { - tracing::error!(error = ? e, "Job cache GC exited on error."); + tracing::error!(error = % e, "Job cache GC exited on error."); } Err(e) => { - tracing::error!(error = ? e, "Job cache GC exited on panic."); + tracing::error!(error = % e, "Job cache GC exited on panic."); } } }; diff --git a/components/spider-storage/src/task_instance_pool.rs b/components/spider-storage/src/task_instance_pool.rs index 4b06dee53..955d0ef69 100644 --- a/components/spider-storage/src/task_instance_pool.rs +++ b/components/spider-storage/src/task_instance_pool.rs @@ -23,6 +23,7 @@ use std::{ }; use async_trait::async_trait; +use serde::Deserialize; use spider_core::types::id::{ExecutionManagerId, JobId, ResourceGroupId, TaskId, TaskInstanceId}; use tokio::{sync::mpsc, task::JoinHandle}; use tokio_util::sync::CancellationToken; @@ -39,7 +40,8 @@ use crate::{ /// Configuration for a task instance pool actor. /// /// Controls GC timing, channel buffering, and execution manager staleness detection. -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, Deserialize)] +#[serde(default)] pub struct TaskInstancePoolConfig { /// Seconds without a heartbeat after which an execution manager is considered stale. pub execution_manager_stale_cutoff_sec: u64, diff --git a/components/spider-utils/Cargo.toml b/components/spider-utils/Cargo.toml new file mode 100644 index 000000000..763b8ac34 --- /dev/null +++ b/components/spider-utils/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "spider-utils" +version = "0.1.0" +edition = "2024" + +[lib] +name = "spider_utils" +path = "src/lib.rs" + +[dependencies] +tracing-appender = "0.2.5" +tracing-subscriber = { + version = "0.3.23", + default-features = false, + features = ["env-filter", "fmt", "json"] +} diff --git a/components/spider-utils/src/lib.rs b/components/spider-utils/src/lib.rs new file mode 100644 index 000000000..300d05fdb --- /dev/null +++ b/components/spider-utils/src/lib.rs @@ -0,0 +1,3 @@ +//! Shared utilities for Spider crates. + +pub mod logging; diff --git a/components/spider-utils/src/logging.rs b/components/spider-utils/src/logging.rs new file mode 100644 index 000000000..4d8617c16 --- /dev/null +++ b/components/spider-utils/src/logging.rs @@ -0,0 +1,35 @@ +//! Logging setup for Spider binaries. + +use tracing_appender::non_blocking::{NonBlockingBuilder, WorkerGuard}; +use tracing_subscriber::EnvFilter; + +/// Initializes the global tracing subscriber with JSON-formatted, non-blocking output to `stderr`. +/// +/// Log events are forwarded to a dedicated background worker so that emitting a log does not block +/// the async runtime on the write syscall. The buffer is lossless: when it fills, the emitting +/// thread blocks until space frees up rather than dropping events. The returned [`WorkerGuard`] +/// must be held for the lifetime of the program; dropping it flushes any buffered logs and shuts +/// the worker down. +/// +/// # Returns +/// +/// The [`WorkerGuard`] for the non-blocking writer. +pub fn set_up_logging() -> WorkerGuard { + let (non_blocking_writer, guard) = NonBlockingBuilder::default() + .lossy(false) + .finish(std::io::stderr()); + tracing_subscriber::fmt() + .event_format( + tracing_subscriber::fmt::format() + .with_level(true) + .with_target(false) + .with_file(true) + .with_line_number(true) + .json(), + ) + .with_env_filter(EnvFilter::from_default_env()) + .with_ansi(false) + .with_writer(non_blocking_writer) + .init(); + guard +} diff --git a/tests/huntsman/em-runtime/tests/test_runtime.rs b/tests/huntsman/em-runtime/tests/test_runtime.rs index 79d1b5d01..93f268cdc 100644 --- a/tests/huntsman/em-runtime/tests/test_runtime.rs +++ b/tests/huntsman/em-runtime/tests/test_runtime.rs @@ -493,6 +493,7 @@ async fn success_outcome_reports_outputs() -> anyhow::Result<()> { let report = &reports[0]; assert_eq!(report.job_id, response.task_assignment.job_id); assert_eq!(report.task_id, response.task_assignment.task_id); + assert_eq!(report.task_instance_id, 1); assert_eq!(report.em_id, em_id); assert_eq!(report.session_id, SESSION_ID); let outputs = report diff --git a/tests/huntsman/test-utils/src/mock.rs b/tests/huntsman/test-utils/src/mock.rs index 47b61fcef..f120bf817 100644 --- a/tests/huntsman/test-utils/src/mock.rs +++ b/tests/huntsman/test-utils/src/mock.rs @@ -20,7 +20,7 @@ use std::{ use async_trait::async_trait; use dashmap::DashSet; use spider_core::types::{ - id::{ExecutionManagerId, JobId, SessionId, TaskId}, + id::{ExecutionManagerId, JobId, SessionId, TaskId, TaskInstanceId}, io::ExecutionContext, scheduler::TaskAssignmentRecord, }; @@ -161,6 +161,7 @@ pub struct RegisterCall { pub struct SuccessReport { pub job_id: JobId, pub task_id: TaskId, + pub task_instance_id: TaskInstanceId, pub em_id: ExecutionManagerId, pub session_id: SessionId, pub serialized_outputs: Option>, @@ -171,6 +172,7 @@ pub struct SuccessReport { pub struct FailureReport { pub job_id: JobId, pub task_id: TaskId, + pub task_instance_id: TaskInstanceId, pub em_id: ExecutionManagerId, pub session_id: SessionId, pub error_message: String, @@ -308,6 +310,7 @@ impl StorageClient for MockStorage { &self, job_id: JobId, task_id: TaskId, + task_instance_id: TaskInstanceId, em_id: ExecutionManagerId, session_id: SessionId, serialized_outputs: Option>, @@ -315,6 +318,7 @@ impl StorageClient for MockStorage { lock(&self.inner.success_reports).push(SuccessReport { job_id, task_id, + task_instance_id, em_id, session_id, serialized_outputs, @@ -329,6 +333,7 @@ impl StorageClient for MockStorage { &self, job_id: JobId, task_id: TaskId, + task_instance_id: TaskInstanceId, em_id: ExecutionManagerId, session_id: SessionId, error_message: String, @@ -336,6 +341,7 @@ impl StorageClient for MockStorage { lock(&self.inner.failure_reports).push(FailureReport { job_id, task_id, + task_instance_id, em_id, session_id, error_message,