diff --git a/Cargo.lock b/Cargo.lock index 6941a5a57..33a28bb43 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -832,6 +832,40 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "huntsman-clp-search-client" +version = "0.1.0" +dependencies = [ + "anyhow", + "clap", + "rand 0.9.4", + "rmp-serde", + "spider-client", + "spider-core", + "tokio", + "tonic", +] + +[[package]] +name = "huntsman-clp-search-pool-ref" +version = "0.1.0" +dependencies = [ + "anyhow", + "clap", + "rand 0.9.4", + "tokio", +] + +[[package]] +name = "huntsman-clp-search-tasks" +version = "0.1.0" +dependencies = [ + "serde", + "spider-tdl", + "tracing", + "tracing-subscriber", +] + [[package]] name = "huntsman-complex" version = "0.1.0" @@ -841,6 +875,20 @@ dependencies = [ "spider-tdl", ] +[[package]] +name = "huntsman-complex-client" +version = "0.1.0" +dependencies = [ + "anyhow", + "clap", + "huntsman-complex-types", + "rmp-serde", + "spider-client", + "spider-core", + "tokio", + "tonic", +] + [[package]] name = "huntsman-complex-types" version = "0.1.0" @@ -849,6 +897,27 @@ dependencies = [ "spider-tdl", ] +[[package]] +name = "huntsman-nn-bench" +version = "0.1.0" +dependencies = [ + "anyhow", + "clap", + "rmp-serde", + "spider-client", + "spider-core", + "tokio", + "tonic", +] + +[[package]] +name = "huntsman-nn-bench-tasks" +version = "0.1.0" +dependencies = [ + "serde", + "spider-tdl", +] + [[package]] name = "hyper" version = "1.10.1" @@ -1912,6 +1981,18 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "spider-client" +version = "0.1.0" +dependencies = [ + "spider-core", + "spider-proto-rust", + "spider-utils", + "thiserror", + "tokio", + "tonic", +] + [[package]] name = "spider-core" version = "0.1.0" @@ -1986,6 +2067,7 @@ dependencies = [ "anyhow", "async-channel", "async-trait", + "clap", "dashmap", "serde", "spider-core", diff --git a/Cargo.toml b/Cargo.toml index a91cc9f27..636549051 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,7 @@ [workspace] resolver = "3" members = [ + "components/spider-client", "components/spider-core", "components/spider-derive", "components/spider-execution-manager", @@ -11,8 +12,14 @@ members = [ "components/spider-tdl", "components/spider-tdl-derive", "components/spider-utils", + "examples/huntsman/clp-search/client", + "examples/huntsman/clp-search/pool-ref", + "examples/huntsman/clp-search/tasks", + "examples/huntsman/complex/client", "examples/huntsman/complex/tasks", "examples/huntsman/complex/types", + "examples/huntsman/nn/client", + "examples/huntsman/nn/tasks", "tests/huntsman/em-runtime", "tests/huntsman/integration-test-tasks", "tests/huntsman/task-executor", @@ -20,6 +27,7 @@ members = [ "tests/huntsman/test-utils", ] default-members = [ + "components/spider-client", "components/spider-core", "components/spider-derive", "components/spider-execution-manager", diff --git a/components/spider-client/Cargo.toml b/components/spider-client/Cargo.toml new file mode 100644 index 000000000..4f7fc796d --- /dev/null +++ b/components/spider-client/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "spider-client" +version = "0.1.0" +edition = "2024" + +[lib] +name = "spider_client" +path = "src/lib.rs" + +[dependencies] +spider-core = { path = "../spider-core" } +spider-proto-rust = { path = "../spider-proto-rust" } +spider-utils = { path = "../spider-utils" } +thiserror = "2.0.18" +tokio = { version = "1.52.3", features = ["macros"] } +tonic = "0.14.6" diff --git a/components/spider-client/src/client.rs b/components/spider-client/src/client.rs new file mode 100644 index 000000000..cba01ec12 --- /dev/null +++ b/components/spider-client/src/client.rs @@ -0,0 +1,184 @@ +//! [`SpiderClient`] — the top-level handle holding the gRPC connection pools. + +use std::num::NonZeroUsize; + +use spider_core::{ + job::JobState, + task::TaskGraph, + types::{ + id::{JobId, ResourceGroupId}, + io::{TaskInput, TaskOutput}, + }, +}; +use tonic::transport::Endpoint; + +use crate::{ + error::ClientError, + grpc::{job::JobOrchestrationClient, resource_group::ResourceGroupManagementClient}, +}; + +/// User-facing client for the Spider storage gRPC services. +/// +/// Wraps a [`JobOrchestrationClient`] and a [`ResourceGroupManagementClient`] against the same +/// storage endpoint, so callers who need both job-lifecycle and resource-group operations get a +/// single handle and one [`SpiderClient::connect`] call. Callers who need only one service may +/// construct the inner client directly. +#[derive(Debug, Clone)] +pub struct SpiderClient { + job_orchestration: JobOrchestrationClient, + resource_group: ResourceGroupManagementClient, +} + +impl SpiderClient { + /// Connects pools of `pool_size` connections to the storage gRPC endpoint. + /// + /// Both the job-orchestration and resource-group-management services are reached through the + /// same `endpoint`. + /// + /// # Returns + /// + /// A new [`SpiderClient`] connected to `endpoint` on success. + /// + /// # Errors + /// + /// Returns [`ClientError::Transport`] if tonic cannot establish a connection to `endpoint`. + pub async fn connect(endpoint: Endpoint, pool_size: NonZeroUsize) -> Result { + let (job_orchestration, resource_group) = tokio::try_join!( + JobOrchestrationClient::connect(endpoint.clone(), pool_size), + ResourceGroupManagementClient::connect(endpoint, pool_size), + )?; + + Ok(Self { + job_orchestration, + resource_group, + }) + } + + /// Serializes and zstd-compresses the task graph and inputs, registers the job, and returns + /// its assigned id. Delegates to [`JobOrchestrationClient::submit_job`]. + /// + /// # Returns + /// + /// The [`JobId`] the storage server assigned to the registered job on success. + /// + /// # Errors + /// + /// See [`JobOrchestrationClient::submit_job`]. + pub async fn submit_job( + &self, + resource_group_id: ResourceGroupId, + task_graph: &TaskGraph, + inputs: Vec, + ) -> Result { + self.job_orchestration + .submit_job(resource_group_id, task_graph, inputs) + .await + } + + /// Starts a registered job. Delegates to [`JobOrchestrationClient::start_job`]. + /// + /// # Returns + /// + /// The job's [`JobState`] after the start request is accepted on success. + /// + /// # Errors + /// + /// See [`JobOrchestrationClient::start_job`]. + pub async fn start_job(&self, job_id: JobId) -> Result { + self.job_orchestration.start_job(job_id).await + } + + /// Cancels a job. Delegates to [`JobOrchestrationClient::cancel_job`]. + /// + /// # Returns + /// + /// The job's [`JobState`] after the cancellation request is accepted on success. + /// + /// # Errors + /// + /// See [`JobOrchestrationClient::cancel_job`]. + pub async fn cancel_job(&self, job_id: JobId) -> Result { + self.job_orchestration.cancel_job(job_id).await + } + + /// Gets the current state of a job. Delegates to [`JobOrchestrationClient::get_job_state`]. + /// + /// # Returns + /// + /// The job's current [`JobState`] on success. + /// + /// # Errors + /// + /// See [`JobOrchestrationClient::get_job_state`]. + pub async fn get_job_state(&self, job_id: JobId) -> Result { + self.job_orchestration.get_job_state(job_id).await + } + + /// Gets a job's task outputs. Delegates to [`JobOrchestrationClient::get_job_outputs`]. + /// + /// # Returns + /// + /// The job's outputs, deserialized from the storage wire format into opaque msgpack payloads, + /// on success. + /// + /// # Errors + /// + /// See [`JobOrchestrationClient::get_job_outputs`]. + pub async fn get_job_outputs(&self, job_id: JobId) -> Result, ClientError> { + self.job_orchestration.get_job_outputs(job_id).await + } + + /// Gets a job's error message. Delegates to [`JobOrchestrationClient::get_job_error`]. + /// + /// # Returns + /// + /// The job's error message on success. + /// + /// # Errors + /// + /// See [`JobOrchestrationClient::get_job_error`]. + pub async fn get_job_error(&self, job_id: JobId) -> Result { + self.job_orchestration.get_job_error(job_id).await + } + + /// Registers an external resource group and returns its server-assigned id. Delegates to + /// [`ResourceGroupManagementClient::add_resource_group`]. + /// + /// # Returns + /// + /// The [`ResourceGroupId`] the storage server assigned to the registered resource group on + /// success. + /// + /// # Errors + /// + /// See [`ResourceGroupManagementClient::add_resource_group`]. + pub async fn add_resource_group( + &self, + external_resource_group_id: String, + password: Vec, + ) -> Result { + self.resource_group + .add_resource_group(external_resource_group_id, password) + .await + } + + /// Verifies a resource group's password. Delegates to + /// [`ResourceGroupManagementClient::verify_resource_group`]. + /// + /// # Returns + /// + /// `Ok(())` on success. + /// + /// # Errors + /// + /// See [`ResourceGroupManagementClient::verify_resource_group`]. + pub async fn verify_resource_group( + &self, + resource_group_id: ResourceGroupId, + password: Vec, + ) -> Result<(), ClientError> { + self.resource_group + .verify_resource_group(resource_group_id, password) + .await + } +} diff --git a/components/spider-client/src/error.rs b/components/spider-client/src/error.rs new file mode 100644 index 000000000..8759dbff4 --- /dev/null +++ b/components/spider-client/src/error.rs @@ -0,0 +1,110 @@ +//! Client error type for the Spider storage gRPC services. +//! +//! [`ClientError`] is the single error type returned by [`crate::client::SpiderClient`] methods. +//! It folds transport failures, tonic error status, and payload serialization and +//! deserialization failures into one concrete enum. See [`ClientError`] for the variants and when +//! each arises. + +use spider_core::types::id::JobId; +use tonic::{Code, Status}; + +/// Errors returned by [`crate::client::SpiderClient`] operations. +#[derive(Debug, thiserror::Error)] +pub enum ClientError { + /// The gRPC transport failed or the connection was lost or unestablished. + #[error("transport error: {0}")] + Transport(String), + + /// The storage server returned an otherwise-uncategorized error. + #[error("storage server error: {0}")] + Server(String), + + /// No job with the requested identifier exists. + #[error("job not found: {0:?}")] + JobNotFound(JobId), + + /// The job is not in a state that allows the requested operation. + #[error("invalid job state: {0}")] + InvalidJobState(String), + + /// The storage server rejected the request as invalid. + #[error("invalid argument: {0}")] + InvalidArgument(String), + + /// The resource group or password was rejected. + #[error("unauthenticated: {0}")] + Unauthenticated(String), + + /// A failure to serialize, compress, or wire-frame a request payload. + #[error("serialization error: {0}")] + Serialization(String), + + /// A failure to deserialize, decompress, or wire-frame a response payload. + #[error("deserialization error: {0}")] + Deserialization(String), + + /// The server returned an unspecified job state that has no core representation. + #[error("job state is unspecified")] + UnspecifiedJobState, +} + +/// Maps a job-orchestration gRPC [`Status`] to a [`ClientError`]. +/// +/// `job_id` is the job the call targeted; it is attached to [`ClientError::JobNotFound`] when the +/// server reports `NOT_FOUND`. +/// +/// # Returns +/// +/// The [`ClientError`] for `status`'s code: +/// +/// * [`ClientError::JobNotFound`] for `NOT_FOUND`. +/// * [`ClientError::InvalidJobState`] for `FAILED_PRECONDITION`. +/// * [`ClientError::InvalidArgument`] for `INVALID_ARGUMENT`. +/// * [`ClientError::Unauthenticated`] for `UNAUTHENTICATED`. +/// * [`ClientError::Transport`] for `UNAVAILABLE` (a lost or unestablished connection). +/// * [`ClientError::Server`] for any other code. +pub(crate) fn job_status_to_error(status: &Status, job_id: JobId) -> ClientError { + match status.code() { + Code::NotFound => ClientError::JobNotFound(job_id), + Code::FailedPrecondition => ClientError::InvalidJobState(status.message().to_owned()), + Code::InvalidArgument => ClientError::InvalidArgument(status.message().to_owned()), + Code::Unauthenticated => ClientError::Unauthenticated(status.message().to_owned()), + Code::Unavailable => ClientError::Transport(status.message().to_owned()), + _ => ClientError::Server(status.message().to_owned()), + } +} + +/// Maps a resource-group-management gRPC [`Status`] to a [`ClientError`]. +/// +/// # Returns +/// +/// The [`ClientError`] for `status`'s code: +/// +/// * [`ClientError::InvalidArgument`] for `INVALID_ARGUMENT`. +/// * [`ClientError::Unauthenticated`] for `UNAUTHENTICATED` (an unknown or unauthorized resource +/// group, or an invalid password). +/// * [`ClientError::Transport`] for `UNAVAILABLE` (a lost or unestablished connection). +/// * [`ClientError::Server`] for any other code (including `NOT_FOUND`, `FAILED_PRECONDITION`, and +/// `INTERNAL`). +pub(crate) fn resource_group_status_to_error(status: &Status) -> ClientError { + match status.code() { + Code::InvalidArgument => ClientError::InvalidArgument(status.message().to_owned()), + Code::Unauthenticated => ClientError::Unauthenticated(status.message().to_owned()), + Code::Unavailable => ClientError::Transport(status.message().to_owned()), + _ => ClientError::Server(status.message().to_owned()), + } +} + +/// Converts a displayable transport-layer error into [`ClientError::Transport`]. +/// +/// Used by the `connect` methods of [`crate::client::SpiderClient`], +/// [`crate::job::JobOrchestrationClient`], +/// and [`crate::resource_group::ResourceGroupManagementClient`] to fold `spider_utils::grpc::Error` +/// into [`ClientError`]. +/// +/// # Returns +/// +/// A [`ClientError::Transport`] containing `error`'s display string. +pub(crate) fn to_transport_error(error: impl std::fmt::Display) -> ClientError { + ClientError::Transport(error.to_string()) +} diff --git a/components/spider-client/src/grpc/job.rs b/components/spider-client/src/grpc/job.rs new file mode 100644 index 000000000..64054d931 --- /dev/null +++ b/components/spider-client/src/grpc/job.rs @@ -0,0 +1,332 @@ +//! [`JobOrchestrationClient`] — gRPC client for the storage job-orchestration service. + +use std::num::NonZeroUsize; + +use spider_core::{ + compression::encode_zstd_bytes, + job::JobState, + task::TaskGraph, + types::{ + id::{JobId, ResourceGroupId}, + io::{SerializedTaskOutputs, TaskInput, TaskInputsSerializer, TaskOutput}, + }, +}; +use spider_proto_rust::{ + error::Error as ProtoError, + storage::{self, job_orchestration_service_client::JobOrchestrationServiceClient}, +}; +use spider_utils::grpc::client::ConnectionPool; +use tonic::{ + Code, + Status, + transport::{Channel, Endpoint}, +}; + +use crate::error::{ClientError, job_status_to_error, to_transport_error}; + +/// gRPC client for the storage server's job-orchestration service. +/// +/// Holds a round-robin pool of connections and exposes the job-lifecycle methods (submit, start, +/// cancel, get state, get outputs, get error). Build one with [`JobOrchestrationClient::connect`]. +/// [`crate::client::SpiderClient`] wraps one of these alongside a +/// [`crate::resource_group::ResourceGroupManagementClient`] for callers who need both services +/// behind a single handle. +#[derive(Debug, Clone)] +pub struct JobOrchestrationClient { + connection_pool: ConnectionPool>, +} + +impl JobOrchestrationClient { + /// Connects a pool of `pool_size` connections to the job-orchestration gRPC endpoint. + /// + /// # Returns + /// + /// A new [`JobOrchestrationClient`] connected to `endpoint` on success. + /// + /// # Errors + /// + /// Returns [`ClientError::Transport`] if tonic cannot establish a connection to `endpoint`. + pub(crate) async fn connect( + endpoint: Endpoint, + pool_size: NonZeroUsize, + ) -> Result { + let connection_pool = ConnectionPool::connect(endpoint, pool_size, |channel| { + JobOrchestrationServiceClient::new(channel) + }) + .await + .map_err(to_transport_error)?; + + Ok(Self { connection_pool }) + } + + /// Serializes and zstd-compresses the task graph and inputs, registers the job, and returns + /// its assigned id. + /// + /// # Returns + /// + /// The [`JobId`] the storage server assigned to the registered job on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`ClientError::Serialization`] if the task graph or inputs cannot be serialized or + /// compressed. + /// * [`ClientError::InvalidArgument`] if the storage server rejects the task graph or inputs. + /// * [`ClientError::Unauthenticated`] if the resource group is unknown or unauthorized. + /// * [`ClientError::Transport`] if the gRPC transport fails or the connection is lost. + /// * [`ClientError::Server`] for any other server-reported error. + /// + /// A freshly registered job has no id yet, so the server-reported `NOT_FOUND` and + /// `FAILED_PRECONDITION` codes (which a job id would otherwise attach) cannot arise for + /// registration and are folded into [`ClientError::Server`]. + pub(crate) async fn submit_job( + &self, + resource_group_id: ResourceGroupId, + task_graph: &TaskGraph, + inputs: Vec, + ) -> Result { + let compressed_serialized_task_graph = task_graph + .to_zstd_compressed_json() + .map_err(|error| ClientError::Serialization(error.to_string()))?; + let compressed_serialized_inputs = serialize_inputs(inputs)?; + let request = storage::RegisterJobRequest { + resource_group_id: resource_group_id.get(), + compressed_serialized_task_graph, + compressed_serialized_inputs, + }; + let response = self + .connection_pool + .get_client() + .register_job(request) + .await + .map_err(|status| submit_status_to_error(&status))? + .into_inner(); + + Ok(JobId::from(response.job_id)) + } + + /// Starts a registered job. + /// + /// # Returns + /// + /// The job's [`JobState`] after the start request is accepted on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`ClientError::JobNotFound`] if no job with `job_id` exists. + /// * [`ClientError::InvalidJobState`] if the job is not in a state that allows starting. + /// * [`ClientError::UnspecifiedJobState`] if the server reports an unspecified job state. + /// * [`ClientError::Transport`] if the gRPC transport fails, the connection is lost, or the + /// server reports an unrecognized job state. + /// * [`ClientError::Server`] for any other server-reported error. + pub(crate) async fn start_job(&self, job_id: JobId) -> Result { + let request = storage::JobIdRequest { + job_id: job_id.get(), + }; + let response = self + .connection_pool + .get_client() + .start_job(request) + .await + .map_err(|status| job_status_to_error(&status, job_id))? + .into_inner(); + + job_state_response_to_result(response) + } + + /// Cancels a job. + /// + /// # Returns + /// + /// The job's [`JobState`] after the cancellation request is accepted on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`ClientError::JobNotFound`] if no job with `job_id` exists. + /// * [`ClientError::InvalidJobState`] if the job is not in a state that allows cancellation. + /// * [`ClientError::UnspecifiedJobState`] if the server reports an unspecified job state. + /// * [`ClientError::Transport`] if the gRPC transport fails, the connection is lost, or the + /// server reports an unrecognized job state. + /// * [`ClientError::Server`] for any other server-reported error. + pub(crate) async fn cancel_job(&self, job_id: JobId) -> Result { + let request = storage::JobIdRequest { + job_id: job_id.get(), + }; + let response = self + .connection_pool + .get_client() + .cancel_job(request) + .await + .map_err(|status| job_status_to_error(&status, job_id))? + .into_inner(); + + job_state_response_to_result(response) + } + + /// Gets the current state of a job. + /// + /// # Returns + /// + /// The job's current [`JobState`] on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`ClientError::JobNotFound`] if no job with `job_id` exists. + /// * [`ClientError::UnspecifiedJobState`] if the server reports an unspecified job state. + /// * [`ClientError::Transport`] if the gRPC transport fails, the connection is lost, or the + /// server reports an unrecognized job state. + /// * [`ClientError::Server`] for any other server-reported error. + pub(crate) async fn get_job_state(&self, job_id: JobId) -> Result { + let request = storage::JobIdRequest { + job_id: job_id.get(), + }; + let response = self + .connection_pool + .get_client() + .get_job_state(request) + .await + .map_err(|status| job_status_to_error(&status, job_id))? + .into_inner(); + + job_state_response_to_result(response) + } + + /// Gets a job's task outputs. + /// + /// # Returns + /// + /// The job's outputs, deserialized from the storage wire format into opaque msgpack payloads, + /// on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`ClientError::JobNotFound`] if no job with `job_id` exists. + /// * [`ClientError::InvalidJobState`] if the job has not yet succeeded. + /// * [`ClientError::Deserialization`] if the returned outputs cannot be decompressed or + /// unframed. + /// * [`ClientError::Transport`] if the gRPC transport fails or the connection is lost. + /// * [`ClientError::Server`] for any other server-reported error. + pub(crate) async fn get_job_outputs( + &self, + job_id: JobId, + ) -> Result, ClientError> { + let request = storage::JobIdRequest { + job_id: job_id.get(), + }; + let response = self + .connection_pool + .get_client() + .get_job_outputs(request) + .await + .map_err(|status| job_status_to_error(&status, job_id))? + .into_inner(); + + SerializedTaskOutputs::deserialize_from_raw(&response.serialized_outputs) + .map_err(|error| ClientError::Deserialization(error.to_string())) + } + + /// Gets a job's error message. + /// + /// # Returns + /// + /// The job's error message on success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`ClientError::JobNotFound`] if no job with `job_id` exists. + /// * [`ClientError::InvalidJobState`] if the job has not yet failed. + /// * [`ClientError::Transport`] if the gRPC transport fails or the connection is lost. + /// * [`ClientError::Server`] for any other server-reported error. + pub(crate) async fn get_job_error(&self, job_id: JobId) -> Result { + let request = storage::JobIdRequest { + job_id: job_id.get(), + }; + let response = self + .connection_pool + .get_client() + .get_job_error(request) + .await + .map_err(|status| job_status_to_error(&status, job_id))? + .into_inner(); + + Ok(response.error_message) + } +} + +/// Serializes and zstd-compresses a job's task inputs for the `RegisterJob` request. +/// +/// # Returns +/// +/// The zstd-compressed wire-format input bytes on success. +/// +/// # Errors +/// +/// Returns [`ClientError::Serialization`] if an input cannot be framed or the wire buffer cannot +/// be compressed. +fn serialize_inputs(inputs: Vec) -> Result, ClientError> { + let mut serializer = TaskInputsSerializer::new(); + for input in inputs { + serializer + .append(input) + .map_err(|error| ClientError::Serialization(error.to_string()))?; + } + encode_zstd_bytes(&serializer.release()) + .map_err(|error| ClientError::Serialization(error.to_string())) +} + +/// Converts a `RegisterJob` gRPC [`Status`] to a [`ClientError`]. +/// +/// Registration has no job id yet, so the `NOT_FOUND` and `FAILED_PRECONDITION` codes that +/// [`job_status_to_error`] would attach a job id to cannot arise here and fall back to +/// [`ClientError::Server`]. The remaining arms match [`job_status_to_error`]. +/// +/// # Returns +/// +/// The [`ClientError`] for `status`'s code: +/// +/// * [`ClientError::InvalidArgument`] for `INVALID_ARGUMENT`. +/// * [`ClientError::Unauthenticated`] for `UNAUTHENTICATED`. +/// * [`ClientError::Transport`] for `UNAVAILABLE` (a lost or unestablished connection). +/// * [`ClientError::Server`] for any other code. +fn submit_status_to_error(status: &Status) -> ClientError { + match status.code() { + Code::InvalidArgument => ClientError::InvalidArgument(status.message().to_owned()), + Code::Unauthenticated => ClientError::Unauthenticated(status.message().to_owned()), + Code::Unavailable => ClientError::Transport(status.message().to_owned()), + _ => ClientError::Server(status.message().to_owned()), + } +} + +/// Converts a `JobStateResponse` into a [`JobState`]. +/// +/// # Returns +/// +/// The [`JobState`] carried by `response` on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * [`ClientError::UnspecifiedJobState`] if the server reports an unspecified job state. +/// * [`ClientError::Transport`] if `response` carries an unrecognized job state. +fn job_state_response_to_result( + response: storage::JobStateResponse, +) -> Result { + let proto_state = storage::JobState::try_from(response.state) + .map_err(|error| ClientError::Transport(error.to_string()))?; + JobState::try_from(proto_state).map_err(|error| match error { + ProtoError::JobStateUnspecified => ClientError::UnspecifiedJobState, + other => ClientError::Transport(other.to_string()), + }) +} diff --git a/components/spider-client/src/grpc/mod.rs b/components/spider-client/src/grpc/mod.rs new file mode 100644 index 000000000..e3b9561d5 --- /dev/null +++ b/components/spider-client/src/grpc/mod.rs @@ -0,0 +1,2 @@ +pub mod job; +pub mod resource_group; diff --git a/components/spider-client/src/grpc/resource_group.rs b/components/spider-client/src/grpc/resource_group.rs new file mode 100644 index 000000000..8fc4b033a --- /dev/null +++ b/components/spider-client/src/grpc/resource_group.rs @@ -0,0 +1,117 @@ +//! [`ResourceGroupManagementClient`] — gRPC client for the storage resource-group-management +//! service. + +use std::num::NonZeroUsize; + +use spider_core::types::id::ResourceGroupId; +use spider_proto_rust::storage::{ + self, + resource_group_management_service_client::ResourceGroupManagementServiceClient, +}; +use spider_utils::grpc::client::ConnectionPool; +use tonic::transport::{Channel, Endpoint}; + +use crate::error::{ClientError, resource_group_status_to_error, to_transport_error}; + +/// gRPC client for the storage server's resource-group-management service. +/// +/// Holds a round-robin pool of connections and exposes the resource-group operations (add, verify). +/// Build one with [`ResourceGroupManagementClient::connect`]. [`crate::client::SpiderClient`] wraps +/// one of these alongside a [`crate::job::JobOrchestrationClient`] for callers who need both +/// services behind a single handle. +#[derive(Debug, Clone)] +pub struct ResourceGroupManagementClient { + connection_pool: ConnectionPool>, +} + +impl ResourceGroupManagementClient { + /// Connects a pool of `pool_size` connections to the resource-group-management gRPC endpoint. + /// + /// # Returns + /// + /// A new [`ResourceGroupManagementClient`] connected to `endpoint` on success. + /// + /// # Errors + /// + /// Returns [`ClientError::Transport`] if tonic cannot establish a connection to `endpoint`. + pub(crate) async fn connect( + endpoint: Endpoint, + pool_size: NonZeroUsize, + ) -> Result { + let connection_pool = ConnectionPool::connect(endpoint, pool_size, |channel| { + ResourceGroupManagementServiceClient::new(channel) + }) + .await + .map_err(to_transport_error)?; + + Ok(Self { connection_pool }) + } + + /// Registers an external resource group and returns its server-assigned id. + /// + /// # Returns + /// + /// The [`ResourceGroupId`] the storage server assigned to the registered resource group on + /// success. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`ClientError::InvalidArgument`] if the storage server rejects the request as invalid. + /// * [`ClientError::Unauthenticated`] if the resource group is unknown or the password is + /// invalid. + /// * [`ClientError::Transport`] if the gRPC transport fails or the connection is lost. + /// * [`ClientError::Server`] for any other server-reported error. + pub(crate) async fn add_resource_group( + &self, + external_resource_group_id: String, + password: Vec, + ) -> Result { + let request = storage::AddResourceGroupRequest { + external_resource_group_id, + password, + }; + let response = self + .connection_pool + .get_client() + .add_resource_group(request) + .await + .map_err(|status| resource_group_status_to_error(&status))? + .into_inner(); + + Ok(ResourceGroupId::from(response.resource_group_id)) + } + + /// Verifies a resource group's password. + /// + /// # Returns + /// + /// `Ok(())` on success — the storage server's response is empty, so success is implicit. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`ClientError::InvalidArgument`] if the storage server rejects the request as invalid. + /// * [`ClientError::Unauthenticated`] if the resource group is unknown or the password is + /// invalid. + /// * [`ClientError::Transport`] if the gRPC transport fails or the connection is lost. + /// * [`ClientError::Server`] for any other server-reported error. + pub(crate) async fn verify_resource_group( + &self, + resource_group_id: ResourceGroupId, + password: Vec, + ) -> Result<(), ClientError> { + let request = storage::VerifyResourceGroupRequest { + resource_group_id: resource_group_id.get(), + password, + }; + self.connection_pool + .get_client() + .verify_resource_group(request) + .await + .map_err(|status| resource_group_status_to_error(&status))?; + Ok(()) + } +} diff --git a/components/spider-client/src/lib.rs b/components/spider-client/src/lib.rs new file mode 100644 index 000000000..5e32240ff --- /dev/null +++ b/components/spider-client/src/lib.rs @@ -0,0 +1,21 @@ +//! User-facing client library for the Spider storage gRPC services. +//! +//! [`SpiderClient`] wraps the storage server's job-orchestration and resource-group-management +//! gRPC services and exposes high-level typed methods that operate on `spider-core` types. It +//! hides proto-level concerns — task-graph and input serialization, zstd compression, the task +//! output wire format, and `tonic::Status` mapping — behind an ergonomic async API: +//! +//! * Job lifecycle: [`SpiderClient::submit_job`], [`SpiderClient::start_job`], +//! [`SpiderClient::cancel_job`], [`SpiderClient::get_job_state`], +//! [`SpiderClient::get_job_outputs`], [`SpiderClient::get_job_error`]. +//! * Resource group operations: [`SpiderClient::add_resource_group`], +//! [`SpiderClient::verify_resource_group`]. +//! +//! Each service is also available as a standalone client — [`JobOrchestrationClient`] and +//! [`ResourceGroupManagementClient`] — for callers who need only one of the two. + +pub mod client; +pub mod error; +pub(crate) mod grpc; + +pub use client::SpiderClient; diff --git a/components/spider-execution-manager/src/process_pool.rs b/components/spider-execution-manager/src/process_pool.rs index da55168ba..5408725dd 100644 --- a/components/spider-execution-manager/src/process_pool.rs +++ b/components/spider-execution-manager/src/process_pool.rs @@ -228,6 +228,13 @@ impl ProcessPool { .stdout(Stdio::piped()) .stderr(Stdio::from(log_file)) .kill_on_drop(true); + // Propagate the execution manager's `RUST_LOG` to the executor so both the executor's own + // tracing subscriber and the dlopened TDL package's (package-local) subscriber honor the + // same log level. `Command` inherits the parent environment, but forwarding explicitly + // keeps the filter reproducible regardless of how the executor is spawned. + if let Ok(rust_log) = std::env::var("RUST_LOG") { + command.env("RUST_LOG", rust_log); + } let mut child = command.spawn()?; let stdin = child .stdin diff --git a/components/spider-execution-manager/src/runtime.rs b/components/spider-execution-manager/src/runtime.rs index b3e7c0d5a..cafac86f0 100644 --- a/components/spider-execution-manager/src/runtime.rs +++ b/components/spider-execution-manager/src/runtime.rs @@ -1,6 +1,11 @@ //! Runtime — the execution manager's main loop. -use std::{collections::VecDeque, net::IpAddr, path::PathBuf, time::Duration}; +use std::{ + collections::VecDeque, + net::IpAddr, + path::PathBuf, + time::{Duration, Instant}, +}; use spider_core::{ session::SessionTracker, @@ -271,8 +276,13 @@ impl< /// /// * Forwards [`Self::register_task_instance`]'s return values on failure. /// * Forwards [`ProcessPool::execute`]'s return values on failure. + // The benchmark timing instrumentation pushes this event loop one line past the pedantic + // limit; the loop reads as a single linear pipeline and is not worth splitting. + #[allow(clippy::too_many_lines)] async fn main_loop(&mut self) -> Result<(), RuntimeError> { loop { + // Benchmark: time the scheduler poll (logged below only when a task is returned). + let next_task_start = Instant::now(); let response = tokio::select! { biased; () = self.cancellation_token.cancelled() => return Ok(()), @@ -292,6 +302,7 @@ impl< }; tracing::info!( + scheduler_next_task_us = elapsed_us(next_task_start), bundle_session = response.session_id, job_id = ? response.task_assignment.job_id, task_id = ? response.task_assignment.task_id, @@ -331,6 +342,8 @@ impl< resource_group_id: response.task_assignment.resource_group_id, ctx: execution_context, }; + // Benchmark: time the EM-side task-executor round-trip (includes the clp-s subprocess). + let execute_start = Instant::now(); let outcome = self .process_pool .execute(request, hard_timeout) @@ -343,6 +356,12 @@ impl< "Process pool failed to dispatch task. Bailing out." ); })?; + tracing::info!( + task_executor_execute_us = elapsed_us(execute_start), + job_id = ? response.task_assignment.job_id, + task_id = ? response.task_assignment.task_id, + "Task executor returned an outcome." + ); let current_session = self.session_tracker.current(); if response.session_id < current_session { @@ -396,6 +415,8 @@ impl< &mut self, response: SchedulerResponse, ) -> Result, RuntimeError> { + // Benchmark instrumentation: time the register-task-instance gRPC call to storage. + let register_start = Instant::now(); let register_result = tokio::select! { biased; () = self.cancellation_token.cancelled() => return Ok(None), @@ -409,6 +430,12 @@ impl< match register_result { Ok(execution_context) => { + tracing::info!( + register_task_instance_us = elapsed_us(register_start), + job_id = ? response.task_assignment.job_id, + task_id = ? response.task_assignment.task_id, + "Registered task instance with storage." + ); self.mark_consume(&response); Ok(Some(execution_context)) } @@ -458,6 +485,15 @@ impl< } } +/// Benchmark instrumentation helper: microseconds elapsed since `start`, saturating on overflow. +/// +/// # Returns +/// +/// The elapsed time since `start` in microseconds, clamped to [`u64::MAX`]. +fn elapsed_us(start: Instant) -> u64 { + u64::try_from(start.elapsed().as_micros()).unwrap_or(u64::MAX) +} + /// Identifies a single task-instance attempt that an outcome report belongs to. #[derive(Debug, Clone, Copy)] struct ReportTarget { diff --git a/components/spider-proto-rust/src/assignment.rs b/components/spider-proto-rust/src/assignment.rs index 0c0700e43..e3df93b6e 100644 --- a/components/spider-proto-rust/src/assignment.rs +++ b/components/spider-proto-rust/src/assignment.rs @@ -24,6 +24,15 @@ impl From for ProtoTaskAssignmentRecord { } } +impl From for TaskAssignmentRecord { + fn from(record: ProtoTaskAssignmentRecord) -> Self { + Self { + id: TaskAssignmentId::from(record.id), + from: SchedulerId::from(record.from), + } + } +} + impl TryFrom for Option { type Error = Error; @@ -67,6 +76,14 @@ mod tests { assert_eq!(record.from, 9); } + #[test] + fn protocol_assignment_record_converts_to_core() { + let record = TaskAssignmentRecord::from(ProtoTaskAssignmentRecord { id: 7, from: 9 }); + + assert_eq!(record.id, TaskAssignmentId::from(7)); + assert_eq!(record.from, SchedulerId::from(9)); + } + #[test] fn next_task_response_converts_assignment() { let response = NextTaskResponse { diff --git a/components/spider-proto-rust/src/generated/storage.rs b/components/spider-proto-rust/src/generated/storage.rs index b4d4fcc5e..1276939e6 100644 --- a/components/spider-proto-rust/src/generated/storage.rs +++ b/components/spider-proto-rust/src/generated/storage.rs @@ -152,6 +152,13 @@ pub struct VerifyResourceGroupRequest { pub password: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct DeleteResourceGroupRequest { + #[prost(uint64, tag = "1")] + pub resource_group_id: u64, + #[prost(bytes = "vec", tag = "2")] + pub password: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct RegisterExecutionManagerRequest { #[prost(string, tag = "1")] pub ip_address: ::prost::alloc::string::String, @@ -1635,6 +1642,32 @@ pub mod inbound_queue_service_client { ); 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_prost::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/storage.InboundQueueService/ResendReadyTasks", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert( + GrpcMethod::new("storage.InboundQueueService", "ResendReadyTasks"), + ); + self.inner.unary(req, path, codec).await + } } } /// Generated server implementations. @@ -1671,6 +1704,13 @@ pub mod inbound_queue_service_server { tonic::Response, tonic::Status, >; + async fn resend_ready_tasks( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; } #[derive(Debug)] pub struct InboundQueueServiceServer { @@ -1895,6 +1935,55 @@ pub mod inbound_queue_service_server { }; Box::pin(fut) } + "/storage.InboundQueueService/ResendReadyTasks" => { + #[allow(non_camel_case_types)] + struct ResendReadyTasksSvc(pub Arc); + impl< + T: InboundQueueService, + > tonic::server::UnaryService + for ResendReadyTasksSvc { + type Response = super::super::common::Void; + 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_prost::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( @@ -2086,6 +2175,35 @@ 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_prost::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 server implementations. @@ -2115,6 +2233,13 @@ 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 { @@ -2295,6 +2420,57 @@ 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::super::common::Void; + 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_prost::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( diff --git a/components/spider-proto-rust/src/lib.rs b/components/spider-proto-rust/src/lib.rs index 6016d9045..e19242bb7 100644 --- a/components/spider-proto-rust/src/lib.rs +++ b/components/spider-proto-rust/src/lib.rs @@ -5,6 +5,7 @@ pub mod error; pub mod id; pub mod io; pub mod job; +pub mod scheduler_registration; pub mod unpack; #[allow(clippy::all, clippy::nursery, clippy::pedantic)] diff --git a/components/spider-proto-rust/src/scheduler_registration.rs b/components/spider-proto-rust/src/scheduler_registration.rs new file mode 100644 index 000000000..e01403e0b --- /dev/null +++ b/components/spider-proto-rust/src/scheduler_registration.rs @@ -0,0 +1,40 @@ +//! Conversions between protobuf scheduler messages and their Spider core representations. + +use spider_core::types::scheduler::RegisteredScheduler; + +use crate::storage; + +impl From for storage::Scheduler { + fn from(scheduler: RegisteredScheduler) -> Self { + Self { + scheduler_id: scheduler.id.get(), + ip_address: scheduler.ip_address.to_string(), + port: u32::from(scheduler.port), + } + } +} + +#[cfg(test)] +mod tests { + use std::net::IpAddr; + + use spider_core::types::{id::SchedulerId, scheduler::RegisteredScheduler}; + + use crate::storage; + + #[test] + fn registered_scheduler_to_protocol_carries_id_ip_and_port() { + const SCHEDULER_ID: SchedulerId = SchedulerId::from(42); + const PORT: u16 = 5678; + + let scheduler = storage::Scheduler::from(RegisteredScheduler { + id: SCHEDULER_ID, + ip_address: IpAddr::V4("127.0.0.1".parse().expect("valid IP")), + port: PORT, + }); + + assert_eq!(scheduler.scheduler_id, SCHEDULER_ID.get()); + assert_eq!(scheduler.ip_address, "127.0.0.1"); + assert_eq!(scheduler.port, u32::from(PORT)); + } +} diff --git a/components/spider-proto-rust/src/unpack/mod.rs b/components/spider-proto-rust/src/unpack/mod.rs index 1da2c7af6..8f3425c00 100644 --- a/components/spider-proto-rust/src/unpack/mod.rs +++ b/components/spider-proto-rust/src/unpack/mod.rs @@ -5,8 +5,10 @@ //! //! * [`common`] — shared helpers for `common.proto` types (e.g. [`common::TaskId`]). //! * [`storage`] — request unpacking for `storage.proto`. +//! * [`scheduler`] — request unpacking for `scheduler.proto`. mod common; +mod scheduler; mod storage; use tonic::{Code, Status}; diff --git a/components/spider-proto-rust/src/unpack/scheduler.rs b/components/spider-proto-rust/src/unpack/scheduler.rs new file mode 100644 index 000000000..f167ccae4 --- /dev/null +++ b/components/spider-proto-rust/src/unpack/scheduler.rs @@ -0,0 +1,59 @@ +//! [`RequestUnpack`] implementations for `scheduler.proto` requests. + +use std::time::Duration; + +use spider_core::types::{id::ExecutionManagerId, scheduler::TaskAssignmentRecord}; + +use crate::{ + scheduler::{ + HeartbeatRequest, + NextTaskRequest, + ShutdownRequest, + TaskAssignmentRecord as ProtoTaskAssignmentRecord, + }, + unpack::{RequestUnpack, UnpackError}, +}; + +/// Unpacks [`NextTaskRequest`] into a tuple containing: +/// +/// * The execution manager ID. +/// * The previously consumed assignment record, if any. +/// * The maximum duration to wait for an assignment. +impl RequestUnpack for NextTaskRequest { + type Unpacked = (ExecutionManagerId, Option, Duration); + + fn unpack(self) -> Result { + Ok(( + ExecutionManagerId::from(self.execution_manager_id), + self.prev_assignment.map(ProtoTaskAssignmentRecord::into), + Duration::from_millis(self.wait_time_ms), + )) + } +} + +/// Unpacks [`HeartbeatRequest`] into an [`ExecutionManagerId`]. +impl RequestUnpack for HeartbeatRequest { + type Unpacked = ExecutionManagerId; + + fn unpack(self) -> Result { + Ok(ExecutionManagerId::from(self.execution_manager_id)) + } +} + +/// Unpacks [`ShutdownRequest`] into a tuple containing: +/// +/// * The execution manager ID. +/// * The previously consumed assignment records. +impl RequestUnpack for ShutdownRequest { + type Unpacked = (ExecutionManagerId, Vec); + + fn unpack(self) -> Result { + Ok(( + ExecutionManagerId::from(self.execution_manager_id), + self.prev_assignments + .into_iter() + .map(ProtoTaskAssignmentRecord::into) + .collect(), + )) + } +} diff --git a/components/spider-proto-rust/src/unpack/storage.rs b/components/spider-proto-rust/src/unpack/storage.rs index ded53f6d2..f5218637c 100644 --- a/components/spider-proto-rust/src/unpack/storage.rs +++ b/components/spider-proto-rust/src/unpack/storage.rs @@ -1,5 +1,7 @@ //! [`RequestUnpack`] implementations for `storage.proto` requests. +use std::{net::IpAddr, time::Duration}; + use spider_core::types::id::{ ExecutionManagerId, JobId, @@ -8,14 +10,22 @@ use spider_core::types::id::{ TaskId, TaskInstanceId, }; +use tonic::Code; use crate::{ storage::{ + AddResourceGroupRequest, + DeleteResourceGroupRequest, + ExecutionManagerIdRequest, JobIdRequest, + PollReadyTasksRequest, + RegisterExecutionManagerRequest, RegisterJobRequest, + RegisterSchedulerRequest, RegisterTaskInstanceRequest, ReportTaskFailureRequest, ReportTaskSuccessRequest, + VerifyResourceGroupRequest, }, unpack::{RequestUnpack, UnpackError, common::unpack_task_id}, }; @@ -134,3 +144,107 @@ impl RequestUnpack for ReportTaskFailureRequest { )) } } + +/// Unpacks [`AddResourceGroupRequest`] into a tuple containing: +/// +/// * The external resource group ID. +/// * The password. +impl RequestUnpack for AddResourceGroupRequest { + type Unpacked = (String, Vec); + + fn unpack(self) -> Result { + Ok((self.external_resource_group_id, self.password)) + } +} + +/// Unpacks [`VerifyResourceGroupRequest`] into a tuple containing: +/// +/// * The resource group ID. +/// * The password. +impl RequestUnpack for VerifyResourceGroupRequest { + type Unpacked = (ResourceGroupId, Vec); + + fn unpack(self) -> Result { + Ok((ResourceGroupId::from(self.resource_group_id), self.password)) + } +} + +/// Unpacks [`DeleteResourceGroupRequest`] into a tuple containing: +/// +/// * The resource group ID. +/// * The password proving ownership of the resource group. +impl RequestUnpack for DeleteResourceGroupRequest { + type Unpacked = (ResourceGroupId, Vec); + + fn unpack(self) -> Result { + Ok((ResourceGroupId::from(self.resource_group_id), self.password)) + } +} + +/// Unpacks [`RegisterExecutionManagerRequest`] into the execution manager's IP address. +impl RequestUnpack for RegisterExecutionManagerRequest { + type Unpacked = IpAddr; + + fn unpack(self) -> Result { + self.ip_address + .parse::() + .map_err(|error| invalid_argument(format!("invalid IP address: {error}"))) + } +} + +/// Unpacks [`ExecutionManagerIdRequest`] into an [`ExecutionManagerId`]. +impl RequestUnpack for ExecutionManagerIdRequest { + type Unpacked = ExecutionManagerId; + + fn unpack(self) -> Result { + Ok(ExecutionManagerId::from(self.execution_manager_id)) + } +} + +/// Unpacks [`RegisterSchedulerRequest`] into a tuple containing: +/// +/// * The scheduler IP address. +/// * The scheduler port. +impl RequestUnpack for RegisterSchedulerRequest { + type Unpacked = (IpAddr, u16); + + fn unpack(self) -> Result { + let ip_address = self + .ip_address + .parse::() + .map_err(|error| invalid_argument(format!("invalid IP address: {error}")))?; + let port = u16::try_from(self.port) + .map_err(|_| invalid_argument(format!("port does not fit in `u16`: {}", self.port)))?; + Ok((ip_address, port)) + } +} + +/// Unpacks [`PollReadyTasksRequest`] into a tuple containing: +/// +/// * The maximum number of entries to return. +/// * The maximum duration to block waiting for entries. +impl RequestUnpack for PollReadyTasksRequest { + type Unpacked = (usize, Duration); + + fn unpack(self) -> Result { + let max_items = usize::try_from(self.max_items).map_err(|_| { + invalid_argument(format!( + "max_items does not fit in `usize`: {}", + self.max_items + )) + })?; + Ok((max_items, Duration::from_millis(self.wait_ms))) + } +} + +/// Builds an [`UnpackError`] carrying [`Code::InvalidArgument`] and the given message. +/// +/// # Returns +/// +/// An [`UnpackError`] whose [`Code`] is [`Code::InvalidArgument`] and whose message is `message`. +const fn invalid_argument(message: String) -> UnpackError { + UnpackError { + code: Code::InvalidArgument, + message, + } +} diff --git a/components/spider-proto/storage/storage.proto b/components/spider-proto/storage/storage.proto index 11caf6e3f..0654dd3f8 100644 --- a/components/spider-proto/storage/storage.proto +++ b/components/spider-proto/storage/storage.proto @@ -23,11 +23,13 @@ service InboundQueueService { rpc PollReadyTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); rpc PollReadyCommitTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); rpc PollReadyCleanupTasks(PollReadyTasksRequest) returns (PollReadyTasksResponse); + rpc ResendReadyTasks(common.Void) returns (common.Void); } service ResourceGroupManagementService { rpc AddResourceGroup(AddResourceGroupRequest) returns (ResourceGroupIdResponse); rpc VerifyResourceGroup(VerifyResourceGroupRequest) returns (common.Void); + rpc DeleteResourceGroup(DeleteResourceGroupRequest) returns (common.Void); } service ExecutionManagerLivenessService { @@ -152,6 +154,11 @@ message VerifyResourceGroupRequest { bytes password = 2; } +message DeleteResourceGroupRequest { + uint64 resource_group_id = 1; + bytes password = 2; +} + message RegisterExecutionManagerRequest { string ip_address = 1; } diff --git a/components/spider-scheduler/Cargo.toml b/components/spider-scheduler/Cargo.toml index 8944a11ed..ed2cc8598 100644 --- a/components/spider-scheduler/Cargo.toml +++ b/components/spider-scheduler/Cargo.toml @@ -7,15 +7,23 @@ edition = "2024" name = "spider_scheduler" path = "src/lib.rs" +[[bin]] +name = "spider_scheduler_grpc_server" +path = "src/bin/grpc_server.rs" + [dependencies] async-channel = "2.3.1" async-trait = "0.1.89" +clap = { version = "4.6.1", features = ["derive"] } serde = { version = "1.0.228", features = ["derive"] } spider-core = { path = "../spider-core" } spider-proto-rust = { path = "../spider-proto-rust" } spider-utils = { path = "../spider-utils" } thiserror = "2.0.18" -tokio = { version = "1.52.3", features = ["macros", "rt", "sync", "time"] } +tokio = { + version = "1.52.3", + features = ["macros", "rt-multi-thread", "signal", "sync", "time"] +} tokio-util = "0.7.18" tonic = "0.14.6" tracing = { version = "0.1.41", default-features = false, features = ["std"] } diff --git a/components/spider-scheduler/src/bin/grpc_server.rs b/components/spider-scheduler/src/bin/grpc_server.rs new file mode 100644 index 000000000..4968508c7 --- /dev/null +++ b/components/spider-scheduler/src/bin/grpc_server.rs @@ -0,0 +1,65 @@ +//! Command-line entrypoint for the scheduler gRPC server. + +use std::{error::Error, net::SocketAddr, path::PathBuf}; + +use clap::Parser; +use spider_proto_rust::scheduler::scheduler_service_server::SchedulerServiceServer; +use spider_scheduler::{ + GrpcSchedulerStorageClient, + ServerConfig, + create_runtime, + grpc::GrpcSchedulerService, +}; +use spider_utils::{config::YamlConfig, logging::set_up_logging}; +use tokio::select; +use tonic::transport::Server; + +/// Command-line arguments for the scheduler gRPC server. +#[derive(Debug, Parser)] +#[command(about = "Run the Spider scheduler 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.runtime.host, server_config.runtime.port); + + let storage_client = GrpcSchedulerStorageClient::connect( + server_config.storage_endpoint.endpoint()?, + server_config.storage_connection_pool_size, + ) + .await?; + + let (runtime, service, cancellation_token) = + create_runtime(server_config.runtime, storage_client).await?; + let grpc_service = GrpcSchedulerService::new(service, cancellation_token.clone()); + tracing::info!(listen_addr = % listen_addr, "Starting scheduler gRPC server."); + + let serve_result = Server::builder() + .add_service(SchedulerServiceServer::new(grpc_service)) + .serve_with_shutdown(listen_addr, async move { + select! { + () = cancellation_token.cancelled() => { + tracing::info!("Shutting down scheduler 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-scheduler/src/config.rs b/components/spider-scheduler/src/config.rs index b056627b8..d37c19643 100644 --- a/components/spider-scheduler/src/config.rs +++ b/components/spider-scheduler/src/config.rs @@ -1,14 +1,36 @@ -//! The scheduler core configuration that selects and configures the scheduling algorithm. +//! Scheduler configuration: the top-level server configuration and the scheduler core configuration +//! that selects the scheduling algorithm. + +use std::num::NonZeroUsize; use serde::Deserialize; +use spider_utils::config::EndpointConfig; use crate::{ core::SchedulerCore, core_impl::RoundRobinConfig, dispatch_queue::DispatchQueueSink, + runtime::RuntimeConfig, storage_client::SchedulerStorageClient, }; +/// Top-level configuration for the scheduler gRPC server. +/// +/// Wraps the scheduler [`RuntimeConfig`] (which also supplies the gRPC listen `host`/`port`) and +/// adds the storage endpoint and connection-pool size that [`crate::create_runtime`] deliberately +/// leaves out, since it receives the storage client as a parameter. +#[derive(Clone, Debug, Deserialize)] +pub struct ServerConfig { + /// The storage gRPC endpoint the scheduler registers with and polls for ready tasks. + pub storage_endpoint: EndpointConfig, + + /// The number of connections per pool used to reach the storage service. + pub storage_connection_pool_size: NonZeroUsize, + + /// The scheduler runtime configuration (also supplies the gRPC listen host/port). + pub runtime: RuntimeConfig, +} + /// The configuration that selects and configures the scheduler core's scheduling algorithm. #[derive(Clone, Debug, Deserialize)] #[serde(rename_all = "snake_case")] diff --git a/components/spider-scheduler/src/core_impl/round_robin/tests.rs b/components/spider-scheduler/src/core_impl/round_robin/tests.rs index a98228013..b8a259442 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/tests.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/tests.rs @@ -178,6 +178,10 @@ impl SchedulerStorageClient for MockStorageClient { async fn job_state(&self, _job_id: JobId) -> Result { Ok(JobState::Running) } + + async fn resend_ready_tasks(&self) -> Result<(), StorageClientError> { + Ok(()) + } } /// # Returns diff --git a/components/spider-scheduler/src/error.rs b/components/spider-scheduler/src/error.rs index a10689f41..c0d61762a 100644 --- a/components/spider-scheduler/src/error.rs +++ b/components/spider-scheduler/src/error.rs @@ -11,13 +11,6 @@ pub enum StorageClientError { #[error("job not found: {0:?}")] JobNotFound(JobId), - /// The scheduler's storage session is stale. - #[error("stale storage session: {storage_session:?}")] - StaleSession { - /// Storage's current session ID. - storage_session: SessionId, - }, - /// The storage server returned an invalid input error. #[error("invalid storage request: {0}")] InvalidInput(String), diff --git a/components/spider-scheduler/src/execution_manager_registry.rs b/components/spider-scheduler/src/execution_manager_registry.rs index e9954b261..4937380f1 100644 --- a/components/spider-scheduler/src/execution_manager_registry.rs +++ b/components/spider-scheduler/src/execution_manager_registry.rs @@ -26,7 +26,7 @@ pub enum ExecutionManagerRegistryError { EmNotFound(ExecutionManagerId), } -#[derive(Clone, Deserialize)] +#[derive(Clone, Debug, Deserialize)] pub struct ExecutionManagerRegistryConfig { /// The time, in seconds, that an execution manager is considered dead without receiving any /// heartbeat. diff --git a/components/spider-scheduler/src/grpc.rs b/components/spider-scheduler/src/grpc.rs new file mode 100644 index 000000000..ff3cf0172 --- /dev/null +++ b/components/spider-scheduler/src/grpc.rs @@ -0,0 +1,244 @@ +//! gRPC service adapter for the scheduler service. +//! +//! [`GrpcSchedulerService`] wraps a [`SchedulerServiceState`] and implements the generated +//! [`SchedulerService`] trait, translating inbound protobuf requests into domain calls and mapping +//! [`SchedulerServiceError`]s back to [`tonic::Status`]. It owns the runtime [`CancellationToken`] +//! so a fatal internal error can cancel the scheduler runtime, mirroring the split in +//! `spider-storage` between the domain [`SchedulerServiceState`] and its gRPC adapter. + +use async_trait::async_trait; +use spider_core::types::{ + id::{SchedulerId, SessionId}, + scheduler::TaskAssignment, +}; +use spider_proto_rust::{ + common, + scheduler::{ + self, + NextTaskResponse, + SchedulerAssignment, + next_task_response, + scheduler_service_server::SchedulerService, + }, + unpack::RequestUnpack, +}; +use tokio_util::sync::CancellationToken; +use tonic::{Request, Response, Status}; + +use crate::{ + dispatch_queue::DispatchQueueSource, + error::{SchedulerError, SchedulerServiceError}, + execution_manager_registry::ExecutionManagerRegistryError, + service::SchedulerServiceState, +}; + +/// gRPC adapter over a [`SchedulerServiceState`]. +/// +/// # Type Parameters +/// +/// * `DispatchQueueSourceType` - The reader side of the dispatching queue the underlying service +/// drains. +#[derive(Clone)] +pub struct GrpcSchedulerService { + inner: SchedulerServiceState, + cancellation_token: CancellationToken, +} + +impl + GrpcSchedulerService +{ + /// Factory function. + /// + /// # Returns + /// + /// A new [`GrpcSchedulerService`] wrapping [`SchedulerServiceState`]. + #[must_use] + pub const fn new( + inner: SchedulerServiceState, + cancellation_token: CancellationToken, + ) -> Self { + Self { + inner, + cancellation_token, + } + } + + /// Error handler for scheduler service errors. + /// + /// This function maps the given [`SchedulerServiceError`] to a [`Status`] that can be sent to + /// the client. The errors are logged for observability. + /// + /// # Returns + /// + /// The [`Status`] to send to the client: + /// + /// * `NOT_FOUND` for an unknown execution manager or task assignment. + /// * `FAILED_PRECONDITION` for an invalid storage session. + /// * `INTERNAL` when the dispatching queue is closed (the scheduler is shutting down), for a + /// fatal internal error (the service will be cancelled), and any other otherwise unexpected + /// error. + pub fn service_error_handler(&self, error: SchedulerServiceError, tag: &'static str) -> Status { + const SERVICE_NAME: &str = "Scheduler"; + match error { + SchedulerServiceError::Scheduler(SchedulerError::DispatchQueueClosed) => { + tracing::warn!( + error = %error, + service = SERVICE_NAME, + tag, + "Dispatch queue is closed." + ); + Status::internal("scheduler is shutting down") + } + + SchedulerServiceError::Scheduler(SchedulerError::InvalidSessionId(session_id)) => { + tracing::warn!( + error = %error, + service = SERVICE_NAME, + tag, + session_id, + "Invalid session ID." + ); + Status::failed_precondition(error.to_string()) + } + + SchedulerServiceError::EMRegistry(ExecutionManagerRegistryError::EmNotFound(em_id)) => { + tracing::warn!( + error = %error, + service = SERVICE_NAME, + tag, + em_id = %em_id, + "Execution manager not found." + ); + Status::not_found("execution manager not found") + } + + SchedulerServiceError::EMRegistry( + ExecutionManagerRegistryError::TaskAssignmentNotFound(em_id, assignment_id), + ) => { + tracing::warn!( + error = %error, + service = SERVICE_NAME, + tag, + em_id = %em_id, + assignment_id = %assignment_id, + "Task assignment not found." + ); + Status::not_found("task assignment not found") + } + + SchedulerServiceError::Scheduler(SchedulerError::Internal(e)) => { + tracing::error!( + error = %e, + service = SERVICE_NAME, + tag, + "Internal error. Cancelling service." + ); + self.cancellation_token.cancel(); + Status::internal("scheduler service unavailable") + } + + error => { + tracing::error!( + error = %error, + service = SERVICE_NAME, + tag, + "Unexpected internal error." + ); + Status::internal("internal error") + } + } + } +} + +/// Implementation of [`SchedulerService`]. +/// +/// All possible errors that can occur during scheduling can be found in +/// [`GrpcSchedulerService::service_error_handler`]. +#[async_trait] +impl SchedulerService + for GrpcSchedulerService +{ + async fn next_task( + &self, + request: Request, + ) -> Result, Status> { + const TAG: &str = "next_task"; + + let (em_id, prev_assignment, wait_time) = request.into_inner().unpack()?; + tracing::info!(em_id = em_id.get(), "NextTask request received."); + + match self + .inner + .next_task(em_id, prev_assignment, wait_time) + .await + { + Ok(Some((session_id, assignment))) => { + // Benchmark instrumentation: the log timestamp marks when the scheduler dispatches + // this task assignment to an execution manager. + tracing::info!( + job_id = assignment.job_id.get(), + task_id = ? assignment.task_id, + "Dispatched a task assignment to an execution manager." + ); + Ok(Response::new(make_next_task_response( + assignment, + self.inner.scheduler_id(), + session_id, + ))) + } + Ok(None) => Ok(Response::new(NextTaskResponse { + result: Some(next_task_response::Result::NoTask(common::Void {})), + })), + Err(error) => Err(self.service_error_handler(error, TAG)), + } + } + + async fn heartbeat( + &self, + request: Request, + ) -> Result, Status> { + let em_id = request.into_inner().unpack()?; + tracing::info!(em_id = em_id.get(), "Heartbeat request received."); + + match self.inner.heartbeat(em_id).await { + Ok(()) => Ok(Response::new(common::Void {})), + Err(error) => Err(self.service_error_handler(error, "heartbeat")), + } + } + + async fn shutdown( + &self, + request: Request, + ) -> Result, Status> { + let (em_id, prev_assignments) = request.into_inner().unpack()?; + tracing::info!(em_id = em_id.get(), "Shutdown request received."); + + match self.inner.shutdown(em_id, prev_assignments).await { + Ok(()) => Ok(Response::new(common::Void {})), + Err(error) => Err(self.service_error_handler(error, "shutdown")), + } + } +} + +/// # Returns +/// +/// A [`NextTaskResponse`] carrying the given assignment, stamped with `scheduler_id` and paired +/// with `session_id`. +fn make_next_task_response( + assignment: TaskAssignment, + scheduler_id: SchedulerId, + session_id: SessionId, +) -> NextTaskResponse { + NextTaskResponse { + result: Some(next_task_response::Result::Assignment( + SchedulerAssignment { + id: assignment.id.get(), + resource_group_id: assignment.resource_group_id.get(), + job_id: assignment.job_id.get(), + task_id: Some(common::TaskId::from(assignment.task_id)), + scheduler_id: scheduler_id.get(), + session_id, + }, + )), + } +} diff --git a/components/spider-scheduler/src/lib.rs b/components/spider-scheduler/src/lib.rs index 65a3932af..b5edc09be 100644 --- a/components/spider-scheduler/src/lib.rs +++ b/components/spider-scheduler/src/lib.rs @@ -37,13 +37,14 @@ pub mod core_impl; pub mod dispatch_queue; pub mod error; pub mod execution_manager_registry; +pub mod grpc; pub mod runtime; pub mod service; pub mod storage_client; pub mod types; pub use crate::{ - config::SchedulerConfig, + config::{SchedulerConfig, ServerConfig}, core::SchedulerCore, dispatch_queue::{ DispatchQueueReader, diff --git a/components/spider-scheduler/src/runtime.rs b/components/spider-scheduler/src/runtime.rs index d514498a7..2267d4f43 100644 --- a/components/spider-scheduler/src/runtime.rs +++ b/components/spider-scheduler/src/runtime.rs @@ -23,7 +23,7 @@ use crate::{ }; /// Runtime configuration for the scheduler service. -#[derive(Deserialize)] +#[derive(Clone, Debug, Deserialize)] pub struct RuntimeConfig { /// The scheduler core configuration that selects and configures the scheduling algorithm. pub scheduler: SchedulerConfig, @@ -235,6 +235,10 @@ mod tests { async fn job_state(&self, _job_id: JobId) -> Result { Ok(JobState::Running) } + + async fn resend_ready_tasks(&self) -> Result<(), StorageClientError> { + Ok(()) + } } /// # Returns diff --git a/components/spider-scheduler/src/storage_client/grpc.rs b/components/spider-scheduler/src/storage_client/grpc.rs index 7c5239698..c9ca8c2a0 100644 --- a/components/spider-scheduler/src/storage_client/grpc.rs +++ b/components/spider-scheduler/src/storage_client/grpc.rs @@ -7,11 +7,14 @@ use spider_core::{ job::JobState, types::id::{JobId, ResourceGroupId, SchedulerId, SessionId, TaskId}, }; -use spider_proto_rust::storage::{ - self, - inbound_queue_service_client::InboundQueueServiceClient, - job_orchestration_service_client::JobOrchestrationServiceClient, - scheduler_registration_service_client::SchedulerRegistrationServiceClient, +use spider_proto_rust::{ + common, + storage::{ + self, + inbound_queue_service_client::InboundQueueServiceClient, + job_orchestration_service_client::JobOrchestrationServiceClient, + scheduler_registration_service_client::SchedulerRegistrationServiceClient, + }, }; use spider_utils::grpc::client::ConnectionPool; use tonic::{ @@ -167,6 +170,15 @@ impl SchedulerStorageClient for GrpcSchedulerStorageClient { .into_inner(); job_state_response_to_result(response) } + + async fn resend_ready_tasks(&self) -> Result<(), StorageClientError> { + self.inbound_queue + .get_client() + .resend_ready_tasks(common::Void {}) + .await + .map_err(|status| inbound_status_to_error(&status))?; + Ok(()) + } } /// Maps an inbound-queue gRPC [`Status`] to a [`StorageClientError`]. diff --git a/components/spider-scheduler/src/storage_client/mod.rs b/components/spider-scheduler/src/storage_client/mod.rs index 3f3f75246..5627f4584 100644 --- a/components/spider-scheduler/src/storage_client/mod.rs +++ b/components/spider-scheduler/src/storage_client/mod.rs @@ -145,4 +145,19 @@ pub trait SchedulerStorageClient: Send + Sync + Clone { /// * [`StorageClientError::Server`] if the storage server returns an error. /// * [`StorageClientError::Transport`] if the storage server returns malformed data. async fn job_state(&self, job_id: JobId) -> Result; + + /// Asks storage to re-enqueue the ready tasks of every cached job back onto the inbound queue. + /// + /// Used after a storage session change (e.g. a scheduler reconnect) to recover tasks that were + /// drained but not yet placed. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`StorageClientError::Server`] if the inbound queue is closed and can no longer yield + /// entries, or the storage server returns another error. + /// * [`StorageClientError::Transport`] if the storage transport fails or returns malformed + /// data. + async fn resend_ready_tasks(&self) -> Result<(), StorageClientError>; } diff --git a/components/spider-storage/src/cache/job.rs b/components/spider-storage/src/cache/job.rs index c5373297d..4703eae7f 100644 --- a/components/spider-storage/src/cache/job.rs +++ b/components/spider-storage/src/cache/job.rs @@ -200,6 +200,12 @@ impl< self.inner.id } + /// Returns the resource group that owns this job. + #[must_use] + pub fn resource_group_id(&self) -> ResourceGroupId { + self.inner.owner_id + } + /// # Returns /// /// The current job state. diff --git a/components/spider-storage/src/db/mariadb.rs b/components/spider-storage/src/db/mariadb.rs index cac6d9ea2..babc17cd1 100644 --- a/components/spider-storage/src/db/mariadb.rs +++ b/components/spider-storage/src/db/mariadb.rs @@ -464,8 +464,41 @@ impl ResourceGroupManagement for MariaDbStorageConnector { } } - async fn delete(&self, _resource_group_id: ResourceGroupId) -> Result<(), DbError> { - todo!("not implemented") + async fn delete(&self, resource_group_id: ResourceGroupId) -> Result<(), DbError> { + const SELECT_FOR_UPDATE_QUERY: &str = formatcp!( + "SELECT `id` FROM `{table}` WHERE `id` = ? FOR UPDATE;", + table = RESOURCE_GROUPS_TABLE_NAME, + ); + const DELETE_JOBS_QUERY: &str = formatcp!( + "DELETE FROM `{table}` WHERE `resource_group_id` = ?;", + table = JOBS_TABLE_NAME, + ); + const DELETE_RESOURCE_GROUP_QUERY: &str = formatcp!( + "DELETE FROM `{table}` WHERE `id` = ?;", + table = RESOURCE_GROUPS_TABLE_NAME, + ); + + let mut tx = self.pool.begin().await?; + + let Some(_): Option = sqlx::query_scalar(SELECT_FOR_UPDATE_QUERY) + .bind(resource_group_id) + .fetch_optional(&mut *tx) + .await? + else { + return Err(DbError::ResourceGroupNotFound(resource_group_id)); + }; + + sqlx::query(DELETE_JOBS_QUERY) + .bind(resource_group_id) + .execute(&mut *tx) + .await?; + sqlx::query(DELETE_RESOURCE_GROUP_QUERY) + .bind(resource_group_id) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + Ok(()) } } diff --git a/components/spider-storage/src/grpc.rs b/components/spider-storage/src/grpc.rs index ea0a145ed..cd599d01f 100644 --- a/components/spider-storage/src/grpc.rs +++ b/components/spider-storage/src/grpc.rs @@ -1,7 +1,10 @@ //! gRPC service adapters for the storage runtime. use async_trait::async_trait; -use spider_core::types::{id::TaskId, io::SerializedTaskOutputs}; +use spider_core::types::{ + id::{SessionId, TaskId}, + io::SerializedTaskOutputs, +}; use spider_proto_rust::{ common, storage::{ @@ -20,9 +23,9 @@ use tokio_util::sync::CancellationToken; use tonic::{Request, Response, Status}; use crate::{ - cache::error::CacheError, + cache::error::{CacheError, InternalError}, db::{DbError, DbStorage}, - ready_queue::ReadyQueueSender, + ready_queue::{ReadyQueueEntry, ReadyQueueSender}, state::{ServiceState, StorageServerError}, task_instance_pool::TaskInstancePoolConnector, }; @@ -78,7 +81,7 @@ impl< /// * `INTERNAL` for: /// * A fatal cache-internal error (the service will restart). /// * Any other (database or otherwise unexpected) error. - /// * `UNAUTHENTICATED` for an unknown or unauthorized resource group. + /// * `UNAUTHENTICATED` for an unknown or unauthorized resource group, or a wrong password. /// * `NOT_FOUND` for a missing job. /// * `FAILED_PRECONDITION` for operations on an invalid job state. /// * `INVALID_ARGUMENT` for a malformed task graph, inputs, or request. @@ -90,14 +93,7 @@ impl< const SERVICE_NAME: &str = "JobOrchestration"; match error { StorageServerError::Cache(CacheError::Internal(e)) => { - tracing::error!( - error = % e, - service = SERVICE_NAME, - tag, - "Internal error in the cache layer. Cancelling service." - ); - self.cancellation_token.cancel(); - Status::internal("storage service unavailable") + self.fatal_internal_status(SERVICE_NAME, tag, &e) } StorageServerError::Db(db_error) => match &db_error { @@ -159,15 +155,7 @@ impl< Status::invalid_argument(error.to_string()) } - _ => { - tracing::error!( - error = % error, - service = SERVICE_NAME, - tag, - "Unexpected internal error." - ); - Status::internal("internal error") - } + _ => self.unexpected_internal_status(SERVICE_NAME, tag, &error), } } @@ -194,14 +182,7 @@ impl< const SERVICE_NAME: &str = "TaskInstanceManagement"; match error { StorageServerError::Cache(CacheError::Internal(e)) => { - tracing::error!( - error = % e, - service = SERVICE_NAME, - tag, - "Internal error in the cache layer. Cancelling service." - ); - self.cancellation_token.cancel(); - Status::internal("storage service unavailable") + self.fatal_internal_status(SERVICE_NAME, tag, &e) } StorageServerError::StaleSession(storage_session) => { @@ -247,18 +228,227 @@ impl< Status::invalid_argument(error.to_string()) } - _ => { - tracing::error!( + _ => self.unexpected_internal_status(SERVICE_NAME, tag, &error), + } + } + + /// Error handler for inbound queue service errors. + /// + /// This function maps the given [`StorageServerError`] to a [`Status`] that can be sent to the + /// client. The errors are logged for observability. + /// + /// # Returns + /// + /// The [`Status`] to send to the client: + /// + /// * `INTERNAL` when the ready-queue channel is closed (the inbound queue can no longer yield + /// entries), for a fatal cache-internal error (the service will restart), or for any other + /// unexpected failure. + pub fn inbound_queue_service_error_handler( + &self, + error: StorageServerError, + tag: &'static str, + ) -> Status { + const SERVICE_NAME: &str = "InboundQueue"; + match error { + StorageServerError::Cache(CacheError::Internal( + InternalError::ReadyQueueChannelClosed, + )) => { + tracing::warn!( + service = SERVICE_NAME, + tag, + "Inbound queue channel is closed." + ); + Status::internal("inbound queue is closed") + } + + StorageServerError::Cache(CacheError::Internal(e)) => { + self.fatal_internal_status(SERVICE_NAME, tag, &e) + } + + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error), + } + } + + /// Error handler for resource group management service errors. + /// + /// This function maps the given [`StorageServerError`] to a [`Status`] that can be sent to the + /// client. The errors are logged for observability. + /// + /// # Returns + /// + /// The [`Status`] to send to the client: + /// + /// * `UNAUTHENTICATED` for an unknown resource group or a wrong password. + /// * `ALREADY_EXISTS` for a duplicate external resource group ID. + /// * `INTERNAL` for: + /// * A fatal cache-internal error (the service will restart). + /// * Any other unexpected failure. + pub fn resource_group_management_service_error_handler( + &self, + error: StorageServerError, + tag: &'static str, + ) -> Status { + const SERVICE_NAME: &str = "ResourceGroupManagement"; + match error { + error @ StorageServerError::Db( + DbError::ResourceGroupNotFound(_) | DbError::InvalidPassword(_), + ) => { + tracing::warn!( + error = % error, + service = SERVICE_NAME, + tag, + "Invalid resource group." + ); + Status::unauthenticated("invalid resource group") + } + + error @ StorageServerError::Db(DbError::ResourceGroupAlreadyExists(_)) => { + tracing::warn!( + error = % error, + service = SERVICE_NAME, + tag, + "Resource group already exists." + ); + Status::already_exists(error.to_string()) + } + + StorageServerError::Cache(CacheError::Internal(e)) => { + self.fatal_internal_status(SERVICE_NAME, tag, &e) + } + + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error), + } + } + + /// Error handler for execution manager liveness service errors. + /// + /// This function maps the given [`StorageServerError`] to a [`Status`] that can be sent to the + /// client. The errors are logged for observability. + /// + /// # Returns + /// + /// The [`Status`] to send to the client: + /// + /// * `FAILED_PRECONDITION` when the execution manager has already been reaped. + /// * `INVALID_ARGUMENT` for an illegal execution manager ID. + /// * `INTERNAL` for: + /// * A fatal cache-internal error (the service will restart). + /// * Any other unexpected failure. + pub fn execution_manager_liveness_service_error_handler( + &self, + error: StorageServerError, + tag: &'static str, + ) -> Status { + const SERVICE_NAME: &str = "ExecutionManagerLiveness"; + match error { + error @ StorageServerError::Db(DbError::ExecutionManagerAlreadyDead(_)) => { + tracing::warn!( error = % error, service = SERVICE_NAME, tag, - "Unexpected internal error. Cancelling service to avoid cache corruption." + "Execution manager already marked dead." ); - self.cancellation_token.cancel(); - Status::internal("internal error") + Status::failed_precondition(error.to_string()) + } + + error @ StorageServerError::Db(DbError::IllegalExecutionManagerId(_)) => { + tracing::warn!( + error = % error, + service = SERVICE_NAME, + tag, + "Illegal execution manager id." + ); + Status::invalid_argument(error.to_string()) + } + + StorageServerError::Cache(CacheError::Internal(e)) => { + self.fatal_internal_status(SERVICE_NAME, tag, &e) + } + + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error), + } + } + + /// Error handler for scheduler registration service errors. + /// + /// This function maps the given [`StorageServerError`] to a [`Status`] that can be sent to the + /// client. The errors are logged for observability. + /// + /// # Returns + /// + /// The [`Status`] to send to the client: + /// + /// * `INTERNAL` for: + /// * A fatal cache-internal error (the service will restart). + /// * Any other failure; scheduler registration currently has no caller-visible error + /// classification beyond a generic server error. + #[must_use] + pub fn scheduler_registration_service_error_handler( + &self, + error: StorageServerError, + tag: &'static str, + ) -> Status { + const SERVICE_NAME: &str = "SchedulerRegistration"; + match error { + StorageServerError::Cache(CacheError::Internal(e)) => { + self.fatal_internal_status(SERVICE_NAME, tag, &e) } + + error => self.unexpected_internal_status(SERVICE_NAME, tag, &error), } } + + /// Logs a fatal cache-internal error, cancels the service, and returns an `INTERNAL` status. + /// + /// Shared by every service error handler's `Cache(CacheError::Internal)` arm. A fatal + /// cache-internal error is unrecoverable, so the whole storage service is cancelled to avoid + /// cache corruption. It is reported as `INTERNAL` rather than `UNAVAILABLE`, which is reserved + /// for transport-level unavailability such as a dropped connection. + /// + /// # Returns + /// + /// `Status::internal("storage service unavailable")`. + fn fatal_internal_status( + &self, + service_name: &'static str, + tag: &'static str, + error: &InternalError, + ) -> Status { + tracing::error!( + error = % error, + service = service_name, + tag, + "Internal error in the cache layer. Cancelling service." + ); + self.cancellation_token.cancel(); + Status::internal("storage service unavailable") + } + + /// Logs an unexpected error, cancels the service, and returns an `INTERNAL` status. + /// + /// Shared by every service error handler for the catch-all fallback arm. An unmapped error is + /// treated as unrecoverable, so the whole storage service is cancelled to avoid cache + /// corruption. + /// + /// # Returns + /// + /// `Status::internal("internal error")`. + fn unexpected_internal_status( + &self, + service_name: &'static str, + tag: &'static str, + error: &StorageServerError, + ) -> Status { + tracing::error!( + error = % error, + service = service_name, + tag, + "Unexpected internal error. Cancelling service to avoid cache corruption." + ); + self.cancellation_token.cancel(); + Status::internal("internal error") + } } /// Implementation of [`JobOrchestrationService`]. @@ -495,23 +685,81 @@ impl< { async fn poll_ready_tasks( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (max_items, wait) = request.into_inner().unpack()?; + tracing::debug!(max_items, ?wait, "Poll ready tasks request received."); + let entries = self + .inner + .poll_ready_tasks(max_items, wait) + .await + .map_err(|error| self.inbound_queue_service_error_handler(error, "poll_ready_tasks"))?; + Ok(Response::new(storage::PollReadyTasksResponse { + tasks: Some(build_ready_tasks( + self.inner.session_id(), + entries, + |task_index| common::TaskId::from(TaskId::Index(task_index)), + )), + })) } async fn poll_ready_commit_tasks( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (max_items, wait) = request.into_inner().unpack()?; + tracing::debug!( + max_items, + ?wait, + "Poll ready commit tasks request received." + ); + let entries = self + .inner + .poll_commit_ready_tasks(max_items, wait) + .await + .map_err(|error| { + self.inbound_queue_service_error_handler(error, "poll_ready_commit_tasks") + })?; + Ok(Response::new(storage::PollReadyTasksResponse { + tasks: Some(build_ready_tasks(self.inner.session_id(), entries, |_| { + common::TaskId::from(TaskId::Commit) + })), + })) } async fn poll_ready_cleanup_tasks( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (max_items, wait) = request.into_inner().unpack()?; + tracing::debug!( + max_items, + ?wait, + "Poll ready cleanup tasks request received." + ); + let entries = self + .inner + .poll_cleanup_ready_tasks(max_items, wait) + .await + .map_err(|error| { + self.inbound_queue_service_error_handler(error, "poll_ready_cleanup_tasks") + })?; + Ok(Response::new(storage::PollReadyTasksResponse { + tasks: Some(build_ready_tasks(self.inner.session_id(), entries, |_| { + common::TaskId::from(TaskId::Cleanup) + })), + })) + } + + async fn resend_ready_tasks( + &self, + _request: Request, + ) -> Result, Status> { + tracing::info!("Resend ready tasks request received."); + self.inner.resend_ready_tasks().await.map_err(|error| { + self.inbound_queue_service_error_handler(error, "resend_ready_tasks") + })?; + Ok(Response::new(common::Void {})) } } @@ -525,16 +773,56 @@ impl< { async fn add_resource_group( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (external_id, password) = request.into_inner().unpack()?; + tracing::info!(external_id = % external_id, "Add resource group request received."); + let rg_id = self + .inner + .add_resource_group(external_id, password) + .await + .map_err(|error| { + self.resource_group_management_service_error_handler(error, "add_resource_group") + })?; + Ok(Response::new(storage::ResourceGroupIdResponse { + resource_group_id: rg_id.get(), + })) } async fn verify_resource_group( &self, - _request: Request, + request: Request, + ) -> Result, Status> { + let (rg_id, password) = request.into_inner().unpack()?; + tracing::info!( + rg_id = rg_id.get(), + "Verify resource group request received." + ); + self.inner + .verify_resource_group(rg_id, &password) + .await + .map_err(|error| { + self.resource_group_management_service_error_handler(error, "verify_resource_group") + })?; + Ok(Response::new(common::Void {})) + } + + async fn delete_resource_group( + &self, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (rg_id, password) = request.into_inner().unpack()?; + tracing::info!( + rg_id = rg_id.get(), + "Delete resource group request received." + ); + self.inner + .delete_resource_group(rg_id, &password) + .await + .map_err(|error| { + self.resource_group_management_service_error_handler(error, "delete_resource_group") + })?; + Ok(Response::new(common::Void {})) } } @@ -548,16 +836,51 @@ impl< { async fn register_execution_manager( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let ip_address = request.into_inner().unpack()?; + tracing::info!(% ip_address, "Execution manager registration request received."); + let em_id = self + .inner + .register_execution_manager(ip_address) + .await + .map_err(|error| { + self.execution_manager_liveness_service_error_handler( + error, + "register_execution_manager", + ) + })?; + Ok(Response::new(storage::RegisterExecutionManagerResponse { + registration: Some(storage::ExecutionManagerRegistration { + execution_manager_id: em_id.get(), + session_id: self.inner.session_id(), + }), + })) } async fn update_execution_manager_heartbeat( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let em_id = request.into_inner().unpack()?; + tracing::info!( + em_id = em_id.get(), + "Execution manager heartbeat request received." + ); + self.inner + .update_execution_manager_heartbeat(em_id) + .await + .map_err(|error| { + self.execution_manager_liveness_service_error_handler( + error, + "update_execution_manager_heartbeat", + ) + })?; + Ok(Response::new( + storage::UpdateExecutionManagerHeartbeatResponse { + session_id: self.inner.session_id(), + }, + )) } } @@ -571,16 +894,41 @@ impl< { async fn register_scheduler( &self, - _request: Request, + request: Request, ) -> Result, Status> { - todo!("Not implemented") + let (ip_address, port) = request.into_inner().unpack()?; + tracing::info!(% ip_address, port, "Scheduler registration request received."); + let scheduler_id = self + .inner + .register_scheduler(ip_address, port) + .await + .map_err(|error| { + self.scheduler_registration_service_error_handler(error, "register_scheduler") + })?; + Ok(Response::new(storage::RegisterSchedulerResponse { + registration: Some(storage::SchedulerRegistration { + scheduler_id: scheduler_id.get(), + session_id: self.inner.session_id(), + }), + })) } async fn get_schedulers( &self, _request: Request, ) -> Result, Status> { - todo!("Not implemented") + tracing::info!("Get schedulers request received."); + let schedulers = self.inner.get_schedulers().await.map_err(|error| { + self.scheduler_registration_service_error_handler(error, "get_schedulers") + })?; + Ok(Response::new(storage::GetSchedulersResponse { + schedulers: Some(storage::SchedulerRegistrations { + schedulers: schedulers + .into_iter() + .map(storage::Scheduler::from) + .collect(), + }), + })) } } @@ -596,10 +944,49 @@ impl< &self, _request: Request, ) -> Result, Status> { - todo!("Not implemented") + let session_id = self.inner.session_id(); + tracing::info!(session_id, "Get session request received."); + Ok(Response::new(storage::GetSessionResponse { session_id })) } } +/// Builds a [`storage::ReadyTasks`] message from a batch of ready-queue entries. +/// +/// # Type Parameters +/// +/// * `TaskKindType` - The kind of ready-queue task carried by each entry +/// ([`spider_core::task::TaskIndex`] for the regular lane, +/// [`crate::ready_queue::CommitTaskMarker`] for the commit lane, or +/// [`crate::ready_queue::CleanupTaskMarker`] for the cleanup lane). +/// +/// # Arguments +/// +/// * `to_task_id` - Converts each entry's lane-specific task kind into its protobuf task ID. +/// +/// # Returns +/// +/// A [`storage::ReadyTasks`] carrying the storage session and the flattened ready tasks. +fn build_ready_tasks( + session_id: SessionId, + entries: Vec>, + to_task_id: impl Fn(TaskKindType) -> common::TaskId, +) -> storage::ReadyTasks { + let tasks = entries + .into_iter() + .map(|entry| { + let resource_group_id = entry.resource_group_id.get(); + let job_id = entry.job_id.get(); + let task_id = to_task_id(entry.task_kind); + storage::ReadyTask { + resource_group_id, + job_id, + task_id: Some(task_id), + } + }) + .collect(); + storage::ReadyTasks { session_id, tasks } +} + /// # Returns /// /// A [`storage::JobStateResponse`] carrying the given job state. @@ -610,3 +997,332 @@ fn make_job_state_response( state: storage::JobState::from(state).into(), }) } + +#[cfg(test)] +mod tests { + use spider_core::types::id::{ExecutionManagerId, JobId, ResourceGroupId, SessionId}; + use tokio_util::sync::CancellationToken; + use tonic::{Code, Request}; + + use super::*; + use crate::{ + ready_queue::{ReadyQueueConfig, ReadyQueueSenderHandle, create_ready_queue}, + state::{ + JobCache, + JobCacheGcHandle, + test_utils::{MockDbConnector, MockReadyQueueSender, MockTaskInstancePoolConnector}, + }, + }; + + type TestGrpcState = + GrpcServiceState; + + type TestGrpcStateWithReadyQueue = + GrpcServiceState; + + const TEST_SESSION_ID: SessionId = 1; + + /// # Returns + /// + /// A [`TestGrpcState`] backed by a default mock DB connector. + fn create_grpc_service() -> TestGrpcState { + create_grpc_service_with_db(MockDbConnector::default()) + } + + /// # Returns + /// + /// A [`TestGrpcState`] backed by `db` and a [`MockReadyQueueSender`]. + fn create_grpc_service_with_db(db: MockDbConnector) -> TestGrpcState { + let (_sender, receiver) = + create_ready_queue(&ReadyQueueConfig::default()).expect("ready queue creation"); + let service = ServiceState::new( + db, + TEST_SESSION_ID, + JobCache::new(), + MockReadyQueueSender, + receiver, + MockTaskInstancePoolConnector, + JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), + ); + GrpcServiceState::new(service, CancellationToken::new()) + } + + /// # Returns + /// + /// A [`TestGrpcStateWithReadyQueue`] wired to a real ready queue, plus the queue's sender + /// handle so tests can enqueue entries. + fn create_grpc_service_with_ready_queue( + db: MockDbConnector, + ) -> (TestGrpcStateWithReadyQueue, ReadyQueueSenderHandle) { + let (sender, receiver) = + create_ready_queue(&ReadyQueueConfig::default()).expect("ready queue creation"); + let service = ServiceState::new( + db, + TEST_SESSION_ID, + JobCache::new(), + sender.clone(), + receiver, + MockTaskInstancePoolConnector, + JobCacheGcHandle::new(tokio::sync::mpsc::unbounded_channel().0), + ); + ( + GrpcServiceState::new(service, CancellationToken::new()), + sender, + ) + } + + #[tokio::test] + async fn get_session_returns_service_session_id() -> anyhow::Result<()> { + let service = create_grpc_service(); + let response = service + .get_session(Request::new(common::Void {})) + .await? + .into_inner(); + assert_eq!(response.session_id, TEST_SESSION_ID); + Ok(()) + } + + #[tokio::test] + async fn add_verify_delete_resource_group_round_trip() -> anyhow::Result<()> { + let service = create_grpc_service(); + let password = b"secret".to_vec(); + + let add_response = service + .add_resource_group(Request::new(storage::AddResourceGroupRequest { + external_resource_group_id: "external-rg".to_owned(), + password: password.clone(), + })) + .await? + .into_inner(); + let rg_id = add_response.resource_group_id; + + service + .verify_resource_group(Request::new(storage::VerifyResourceGroupRequest { + resource_group_id: rg_id, + password: password.clone(), + })) + .await?; + + service + .delete_resource_group(Request::new(storage::DeleteResourceGroupRequest { + resource_group_id: rg_id, + password: password.clone(), + })) + .await?; + + let verify_after_delete = service + .verify_resource_group(Request::new(storage::VerifyResourceGroupRequest { + resource_group_id: rg_id, + password, + })) + .await; + let status = verify_after_delete.expect_err("verify should fail after delete"); + assert_eq!(status.code(), Code::Unauthenticated); + Ok(()) + } + + #[tokio::test] + async fn delete_resource_group_rejects_wrong_password_as_unauthenticated() -> anyhow::Result<()> + { + let service = create_grpc_service(); + let password = b"secret".to_vec(); + let rg_id = service + .add_resource_group(Request::new(storage::AddResourceGroupRequest { + external_resource_group_id: "external-rg".to_owned(), + password: password.clone(), + })) + .await? + .into_inner() + .resource_group_id; + + let result = service + .delete_resource_group(Request::new(storage::DeleteResourceGroupRequest { + resource_group_id: rg_id, + password: b"wrong".to_vec(), + })) + .await; + let status = result.expect_err("a wrong password should be rejected"); + assert_eq!(status.code(), Code::Unauthenticated); + Ok(()) + } + + #[tokio::test] + async fn verify_resource_group_rejects_wrong_password_as_unauthenticated() -> anyhow::Result<()> + { + let service = create_grpc_service(); + let password = b"secret".to_vec(); + let rg_id = service + .add_resource_group(Request::new(storage::AddResourceGroupRequest { + external_resource_group_id: "external-rg".to_owned(), + password: password.clone(), + })) + .await? + .into_inner() + .resource_group_id; + + let result = service + .verify_resource_group(Request::new(storage::VerifyResourceGroupRequest { + resource_group_id: rg_id, + password: b"wrong".to_vec(), + })) + .await; + let status = result.expect_err("a wrong password should be rejected"); + assert_eq!(status.code(), Code::Unauthenticated); + Ok(()) + } + + #[tokio::test] + async fn register_task_instance_reports_missing_job_as_failed_precondition() + -> anyhow::Result<()> { + let service = create_grpc_service(); + let result = service + .register_task_instance(Request::new(storage::RegisterTaskInstanceRequest { + job_id: JobId::random().get(), + task_id: Some(common::TaskId::from(TaskId::Index(0))), + execution_manager_id: ExecutionManagerId::from(1).get(), + session_id: TEST_SESSION_ID, + })) + .await; + let status = result.expect_err("an unknown job should be rejected"); + assert_eq!(status.code(), Code::FailedPrecondition); + Ok(()) + } + + #[test] + fn job_orchestration_maps_unknown_resource_group_to_unauthenticated() { + let service = create_grpc_service(); + let status = service.job_orchestration_service_error_handler( + StorageServerError::Db(DbError::ResourceGroupNotFound(ResourceGroupId::from(7))), + "test", + ); + assert_eq!(status.code(), Code::Unauthenticated); + } + + #[test] + fn job_orchestration_maps_wrong_password_to_unauthenticated() { + let service = create_grpc_service(); + let status = service.job_orchestration_service_error_handler( + StorageServerError::Db(DbError::InvalidPassword(ResourceGroupId::from(7))), + "test", + ); + assert_eq!(status.code(), Code::Unauthenticated); + } + + #[test] + fn job_orchestration_maps_fatal_cache_internal_to_internal() { + let service = create_grpc_service(); + let status = service.job_orchestration_service_error_handler( + StorageServerError::Cache(CacheError::Internal(InternalError::TaskNotRunning)), + "test", + ); + assert_eq!(status.code(), Code::Internal); + } + + #[test] + fn resource_group_management_maps_unknown_resource_group_to_unauthenticated() { + let service = create_grpc_service(); + let status = service.resource_group_management_service_error_handler( + StorageServerError::Db(DbError::ResourceGroupNotFound(ResourceGroupId::from(7))), + "test", + ); + assert_eq!(status.code(), Code::Unauthenticated); + } + + #[test] + fn resource_group_management_maps_fatal_cache_internal_to_internal() { + let service = create_grpc_service(); + let status = service.resource_group_management_service_error_handler( + StorageServerError::Cache(CacheError::Internal(InternalError::TaskNotRunning)), + "test", + ); + assert_eq!(status.code(), Code::Internal); + } + + #[tokio::test] + async fn register_execution_manager_returns_id_and_session() -> anyhow::Result<()> { + let service = create_grpc_service(); + let response = service + .register_execution_manager(Request::new(storage::RegisterExecutionManagerRequest { + ip_address: "127.0.0.1".to_owned(), + })) + .await? + .into_inner(); + let registration = response + .registration + .expect("registration should be present"); + assert_eq!(registration.session_id, TEST_SESSION_ID); + assert_ne!(registration.execution_manager_id, 0); + Ok(()) + } + + #[tokio::test] + async fn heartbeat_returns_session_for_registered_em() -> anyhow::Result<()> { + let service = create_grpc_service(); + let em_id = service + .register_execution_manager(Request::new(storage::RegisterExecutionManagerRequest { + ip_address: "127.0.0.1".to_owned(), + })) + .await? + .into_inner() + .registration + .expect("registration should be present") + .execution_manager_id; + + let response = service + .update_execution_manager_heartbeat(Request::new(storage::ExecutionManagerIdRequest { + execution_manager_id: em_id, + })) + .await? + .into_inner(); + assert_eq!(response.session_id, TEST_SESSION_ID); + Ok(()) + } + + #[tokio::test] + async fn heartbeat_rejects_unknown_em() -> anyhow::Result<()> { + let service = create_grpc_service(); + let result = service + .update_execution_manager_heartbeat(Request::new(storage::ExecutionManagerIdRequest { + execution_manager_id: ExecutionManagerId::from(999).get(), + })) + .await; + let status = result.expect_err("an unknown em id should be rejected"); + assert_eq!(status.code(), Code::InvalidArgument); + Ok(()) + } + + #[tokio::test] + async fn resend_ready_tasks_succeeds() -> anyhow::Result<()> { + let service = create_grpc_service(); + service + .resend_ready_tasks(Request::new(common::Void {})) + .await?; + Ok(()) + } + + #[tokio::test] + async fn poll_ready_tasks_returns_entries() -> anyhow::Result<()> { + const TASK_INDEX: usize = 3; + let (service, sender) = create_grpc_service_with_ready_queue(MockDbConnector::default()); + let rg_id = ResourceGroupId::from(7); + let job_id = JobId::from(11); + sender + .send_task_ready(rg_id, job_id, vec![TASK_INDEX]) + .await + .expect("send_task_ready should succeed"); + + let response = service + .poll_ready_tasks(Request::new(storage::PollReadyTasksRequest { + max_items: 10, + wait_ms: 100, + })) + .await? + .into_inner(); + let tasks = response.tasks.expect("ready tasks should be present"); + assert_eq!(tasks.session_id, TEST_SESSION_ID); + assert_eq!(tasks.tasks.len(), 1); + assert_eq!(tasks.tasks[0].resource_group_id, rg_id.get()); + assert_eq!(tasks.tasks[0].job_id, job_id.get()); + Ok(()) + } +} diff --git a/components/spider-storage/src/state.rs b/components/spider-storage/src/state.rs index 4b76ec055..3ea42a3b5 100644 --- a/components/spider-storage/src/state.rs +++ b/components/spider-storage/src/state.rs @@ -11,4 +11,4 @@ pub use runtime::{Runtime, create_runtime}; pub use service::ServiceState; #[cfg(test)] -mod test_utils; +pub(crate) mod test_utils; diff --git a/components/spider-storage/src/state/job_cache.rs b/components/spider-storage/src/state/job_cache.rs index c01a2d74c..5b9a63f62 100644 --- a/components/spider-storage/src/state/job_cache.rs +++ b/components/spider-storage/src/state/job_cache.rs @@ -3,7 +3,7 @@ use std::{ sync::Arc, }; -use spider_core::types::id::JobId; +use spider_core::types::id::{JobId, ResourceGroupId}; use tokio::sync::RwLock; use crate::{ @@ -119,6 +119,27 @@ impl< .count() } + /// Removes every job control block belonging to the given resource group from the cache. + /// + /// # Returns + /// + /// The number of job control blocks that existed and were removed. + pub async fn remove_by_resource_group(&self, resource_group_id: ResourceGroupId) -> usize { + let victim_ids: Vec = self + .jobs + .read() + .await + .iter() + .filter(|(_, jcb)| jcb.resource_group_id() == resource_group_id) + .map(|(job_id, _)| *job_id) + .collect(); + let mut jobs = self.jobs.write().await; + victim_ids + .iter() + .filter(|job_id| jobs.remove(job_id).is_some()) + .count() + } + /// Resends all ready tasks for every job in the cache to the ready queue. /// /// # Errors @@ -177,9 +198,23 @@ mod tests { state::test_utils::{MockDbConnector, MockReadyQueueSender, MockTaskInstancePoolConnector}, }; + /// # Returns + /// + /// A test job control block owned by a random resource group. async fn create_test_jcb( job_id: JobId, ) -> SharedJobControlBlock + { + create_test_jcb_with_resource_group(job_id, ResourceGroupId::random()).await + } + + /// # Returns + /// + /// A test job control block owned by `resource_group_id`. + async fn create_test_jcb_with_resource_group( + job_id: JobId, + resource_group_id: ResourceGroupId, + ) -> SharedJobControlBlock { let bytes_type = DataTypeDescriptor::Value(ValueTypeDescriptor::bytes()); let mut submitted = @@ -201,7 +236,7 @@ mod tests { create_validated_submission(submitted, vec![TaskInput::ValuePayload(vec![0u8; 4])]); SharedJobControlBlock::create( job_id, - spider_core::types::id::ResourceGroupId::random(), + resource_group_id, job_submission, MockReadyQueueSender, MockDbConnector::default(), @@ -272,6 +307,49 @@ mod tests { Ok(()) } + #[tokio::test] + async fn job_cache_remove_by_resource_group_evicts_only_owned_jobs() -> anyhow::Result<()> { + let cache: JobCache = + JobCache::new(); + let target_group = ResourceGroupId::from(42); + let other_group = ResourceGroupId::from(7); + let first_job_id = JobId::random(); + let second_job_id = JobId::random(); + let unrelated_job_id = JobId::random(); + + cache + .insert(create_test_jcb_with_resource_group(first_job_id, target_group).await) + .await?; + cache + .insert(create_test_jcb_with_resource_group(second_job_id, target_group).await) + .await?; + cache + .insert(create_test_jcb_with_resource_group(unrelated_job_id, other_group).await) + .await?; + + let num_removed_jobs = cache.remove_by_resource_group(target_group).await; + + assert_eq!( + num_removed_jobs, 2, + "only the resource group's jobs should be removed" + ); + assert_eq!( + cache.get(first_job_id).await.map(|_| ()), + None, + "first owned job should be removed" + ); + assert_eq!( + cache.get(second_job_id).await.map(|_| ()), + None, + "second owned job should be removed" + ); + assert!( + cache.get(unrelated_job_id).await.is_some(), + "unrelated job should remain" + ); + Ok(()) + } + #[tokio::test] async fn job_cache_get_returns_none_for_nonexistent_job() -> anyhow::Result<()> { let cache: JobCache = diff --git a/components/spider-storage/src/state/service.rs b/components/spider-storage/src/state/service.rs index 6d4d5d306..f916ebe4d 100644 --- a/components/spider-storage/src/state/service.rs +++ b/components/spider-storage/src/state/service.rs @@ -536,6 +536,38 @@ impl< .map_err(StorageServerError::from) } + /// Deletes a resource group and all of its jobs from database and cache. + /// + /// The caller must supply the resource group's password, which is verified before any deletion + /// occurs so that only an authenticated owner can remove a group and its jobs. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * Forwards [`ResourceGroupManagement::verify`]'s return values on failure (wrong password or + /// unknown resource group). + /// * Forwards [`ResourceGroupManagement::delete`]'s return values on failure. + pub async fn delete_resource_group( + &self, + resource_group_id: ResourceGroupId, + password: &[u8], + ) -> Result<(), StorageServerError> { + self.inner.db.verify(resource_group_id, password).await?; + self.inner.db.delete(resource_group_id).await?; + let evicted_jobs = self + .inner + .job_cache + .remove_by_resource_group(resource_group_id) + .await; + tracing::info!( + rg_id = ? resource_group_id, + evicted_jobs, + "Resource group deleted.", + ); + Ok(()) + } + /// Polls the ready queue for task entries. /// /// # Returns @@ -1594,6 +1626,62 @@ mod tests { Ok(()) } + #[tokio::test] + async fn delete_resource_group_succeeds_for_existing() -> anyhow::Result<()> { + let service = create_test_service(); + let password = vec![1, 2, 3]; + let rg_id = service + .add_resource_group("external_123".to_owned(), password.clone()) + .await?; + service.delete_resource_group(rg_id, &password).await?; + let result = service.verify_resource_group(rg_id, &[1, 2, 3]).await; + assert!( + result.is_err(), + "verify should fail after the resource group is deleted" + ); + Ok(()) + } + + #[tokio::test] + async fn delete_resource_group_rejects_wrong_password() -> anyhow::Result<()> { + let service = create_test_service(); + let rg_id = service + .add_resource_group("external_123".to_owned(), vec![1, 2, 3]) + .await?; + let result = service.delete_resource_group(rg_id, &[4, 5, 6]).await; + assert!( + matches!( + result, + Err(StorageServerError::Db(DbError::InvalidPassword(_))) + ), + "delete_resource_group should return InvalidPassword for a wrong password" + ); + assert!( + service + .verify_resource_group(rg_id, &[1, 2, 3]) + .await + .is_ok(), + "resource group should still exist when the password is wrong" + ); + Ok(()) + } + + #[tokio::test] + async fn delete_resource_group_returns_error_for_unknown() -> anyhow::Result<()> { + let service = create_test_service(); + let result = service + .delete_resource_group(ResourceGroupId::random(), &[1, 2, 3]) + .await; + assert!( + matches!( + result, + Err(StorageServerError::Db(DbError::ResourceGroupNotFound(_))) + ), + "delete_resource_group should return ResourceGroupNotFound for an unknown id" + ); + Ok(()) + } + #[tokio::test] async fn poll_ready_tasks_returns_entries_from_ready_queue() -> anyhow::Result<()> { const TASK_INDEX: TaskIndex = 0; diff --git a/components/spider-storage/tests/mariadb_test.rs b/components/spider-storage/tests/mariadb_test.rs index 91e7bfe1e..412b0abc5 100644 --- a/components/spider-storage/tests/mariadb_test.rs +++ b/components/spider-storage/tests/mariadb_test.rs @@ -615,6 +615,46 @@ async fn test_verify_nonexistent_resource_group() { ); } +#[tokio::test] +#[ignore = "requires MariaDB"] +async fn test_delete_resource_group_removes_group_and_its_jobs() { + let storage = create_mariadb_connector().await; + let rg_id = create_test_resource_group(&storage).await; + let (graph, inputs) = single_task_graph(); + let job_submission = create_validated_submission(graph, inputs); + let job_id = storage + .register(rg_id, &job_submission) + .await + .expect("register should succeed"); + + storage.delete(rg_id).await.expect("delete should succeed"); + + let verify_result = storage.verify(rg_id, b"test-password").await; + assert!( + matches!(verify_result, Err(DbError::ResourceGroupNotFound(_))), + "verify should fail after the resource group is deleted, got {verify_result:?}" + ); + + let job_state_result = storage.get_state(job_id).await; + assert!( + matches!(job_state_result, Err(DbError::JobNotFound(_))), + "the resource group's jobs should be removed, got {job_state_result:?}" + ); +} + +#[tokio::test] +#[ignore = "requires MariaDB"] +async fn test_delete_nonexistent_resource_group() { + let storage = create_mariadb_connector().await; + let fake_rg_id = ResourceGroupId::random(); + + let result = storage.delete(fake_rg_id).await; + assert!( + matches!(result, Err(DbError::ResourceGroupNotFound(_))), + "expected ResourceGroupNotFound, got {result:?}" + ); +} + #[tokio::test] #[ignore = "requires MariaDB"] async fn test_start_job_not_found() { diff --git a/examples/huntsman/clp-search/README.md b/examples/huntsman/clp-search/README.md new file mode 100644 index 000000000..21b11cb40 --- /dev/null +++ b/examples/huntsman/clp-search/README.md @@ -0,0 +1,147 @@ +# CLP-search-over-Spider benchmark harness + +A small benchmark harness that runs a **multi-worker CLP search** over a directory of CLP archives +by fanning the work out across a live Spider stack. Each archive is searched by an independent +Spider task that shells out to the `clp-s` binary; the client collects and prints the results. + +It exists to benchmark Spider on a real workload (embarrassingly-parallel search over hundreds of +archives) and to compare it against non-Spider baselines. + +## Crates + +| Path | Crate / artifact | What it is | +|-------------|--------------------------------------|-------------------------------------------------------------------| +| `client/` | `huntsman-clp-search-client` (bin) | The harness: builds a Spider job from a query + archive dir, runs it, prints results. | +| `tasks/` | `huntsman-clp-search-tasks` (cdylib) | The TDL package `clp_search` exposing the single task `clp_search::search`. Staged as `libclp_search.so`. | +| `pool-ref/` | `huntsman-clp-search-pool-ref` (bin) | Baseline: the same fan-out + `clp-s` work through a local process pool, **no Spider** — used to isolate Spider's scheduling overhead. | + +## Prerequisites + +* The `clp-s` binary. Defaults to `/home/lzh/dev/clp/build/core/clp-s`; override with the + `CLP_S_BIN` environment variable (read by the task and the pool-ref binary). +* A directory of CLP archives, where **every immediate subdirectory is one archive** (e.g. + `~/dev/clp/build/clp-package/var/data/archives/default`). +* The workspace built and the `clp_search` package staged so task executors can `dlopen` it: + + ```shell + task build:rust + task build:packages + ``` + + `build:packages` builds the example binaries and stages `libclp_search.so` into + `build/tdl_packages/clp_search/`. +* A running Spider stack. See `../../../stack-doc.md`. The stack must have the `clp_search` package + staged (the step above) before you submit a job that references it. + +## Running + +In one shell, bring up the stack (defaults are 16 workers and a large scheduler ready-task queue, +which the concurrent-job case needs): + +```shell +uv run --script tools/scripts/stack/run.py +``` + +In another shell, once the stack reports it is up, run the client: + +```shell +build/rust-targets/release/huntsman-clp-search-client \ + --input ~/dev/clp/build/clp-package/var/data/archives/default \ + --query '*NonDFS*' +``` + +### Client arguments + +| Flag | Default | Meaning | +|----------------|---------------------------------------|----------------------------------------------------------------| +| `--input` | (required) | Directory whose immediate subdirectories are CLP archives. | +| `--query` | (required) | The KQL search query, applied to every archive. | +| `--endpoint` | `http://127.0.0.1:50051` | Spider storage gRPC endpoint. | +| `--pool-size` | `4` | `SpiderClient` gRPC connection-pool size. | +| `--output-dir` | `build/spider-run/clp-search-results` | Base dir; each run creates a unique `run-` subdir. | +| `--no-shuffle` | off (archives are shuffled) | Submit archives in sorted order instead of shuffling them. | + +By default the client **shuffles** the archive order before submission, so heavy archives are spread +across workers rather than concentrated on one straggler. For a cheap query whose per-task cost is +uniform this mostly adds run-to-run variance (the assignment is random each run), so pass +`--no-shuffle` for reproducible measurements; keep the shuffle when tasks are expensive and a sorted +order would systematically pile heavy archives onto a few workers. + +**Output convention:** search results (JSONL) go to **stdout**; the submit line and the per-phase +timing breakdown go to **stderr**. So `client ... > results.jsonl` captures just the results. + +## How the client works + +1. **Discover archives** — every immediate subdirectory of `--input`, canonicalized to an absolute + path and sorted by name, then shuffled (unless `--no-shuffle`) to randomize task-to-worker + assignment. +2. **Prepare outputs + graph** — create a unique run directory and assign each archive a unique + output file `/-.jsonl` (tasks never share an output path). Build a + **flat** `TaskGraph`: one `clp_search::search` task per archive, each with three graph inputs + (archive path, query, output path) and **no outputs**, and no inter-task dependencies — so the + scheduler can run all of them in parallel across the execution managers. +3. **Build inputs** — flatten the per-task inputs in insertion/position order (archive path, query, + output path), matching `TaskGraph::get_task_graph_input_indices`. Each string is msgpack-encoded + into a `TaskInput::ValuePayload`; the graph declares each input as an opaque `bytes` type. +4. **Submit → start → poll** — register a per-run resource group, `submit_job`, `start_job`, then + poll `get_job_state` (every 10 ms) until the job reaches a terminal state. +5. **Print results** — on success, read each per-task output file in archive order and concatenate + them to stdout. On failure, fetch and report the job error. +6. **Report timing** — print a phase breakdown to stderr: `query_processing` (discovery + graph + + inputs + connect + submit/start), `spider_execution` (the poll-to-terminal wait — this is where + nearly all the wall-clock lives), and `post_processing` (reading + printing results). + +## The search task (`clp_search::search`) + +A single TDL task with signature `search(archive_path, query, output_path) -> Result<(), TdlError>` +(no output). It runs `clp-s s ` with the child's stdout redirected into +`output_path`, and returns an error if `clp-s` cannot be spawned or exits non-zero. The binary is +resolved from `CLP_S_BIN` (default `/home/lzh/dev/clp/build/core/clp-s`). + +## Baseline (`pool-ref`) + +`clp-search-pool-ref` runs the identical workload — same discovery, same per-archive output files, +same `clp-s s ` invocation — but through a local `tokio` process pool bounded by +`--pool-size` (default 16) instead of Spider. It emits the same phase-timing labels, so the gap +between its execution time and Spider's `spider_execution` is Spider's scheduling/coordination +overhead. Run it exactly like the client but without `--endpoint`: + +```shell +build/rust-targets/release/clp-search-pool-ref \ + --input ~/dev/clp/build/clp-package/var/data/archives/default --query '*NonDFS*' +``` + +The pool-ref also prints a `[clp_s]` line to stderr with the per-archive `clp-s` execution +distribution (count/sum/mean/median/min/max/p95), directly comparable to Spider's `clp_s_elapsed_us` +metric below — useful for separating `clp-s` execution cost (and CPU contention) from Spider's +coordination overhead. + +## Benchmark metrics + +Beyond the client's own stderr phase breakdown (`query_processing` / `spider_execution` / +`post_processing`, with `submit_job` and `start_job` timed separately, and a +`client_job_start_epoch_us` line for the scheduling-overhead calc), the stack emits per-task metrics +as structured JSON logs when it runs at `info` (set `log_level: "info"` in `spider.yaml`, or +`RUST_LOG=info`). Collect them from `build/spider-run/`: + +| Metric (JSON field) | Where | Meaning | +|---------------------|-------|---------| +| `scheduler_next_task_us` | `em-*.log` | EM time to fetch a task from the scheduler (excl. the first, idle long-poll). | +| `register_task_instance_us` | `em-*.log` | EM time to register the task instance in storage. | +| `task_executor_execute_us` | `em-*.log` | EM-side task-executor round-trip (includes `clp-s`). | +| `clp_s_elapsed_us` | `em-logs/-.log` | `clp-s` subprocess wall time, measured inside the task. | +| `Dispatched a task assignment` (log timestamp) | `scheduler.log` | When the scheduler handed a task to an EM. | + +Aggregate by filtering log lines on `job_id`. Scheduling overhead per job = the scheduler's first +dispatch timestamp minus the client's `client_job_start_epoch_us`. + +## Notes + +* Because everything runs on one host, the task processes and the client share a filesystem, so the + client reads back exactly the files the tasks wrote. All paths are made absolute for this reason. +* Prefer a **freshly launched** stack per benchmark: a stack left running can drift into a bad state + and fail on the next submission. +* The scheduler's `ready_task_capacity` (in `tools/scripts/stack/spider.yaml`) must exceed the total + number of simultaneously-ready tasks when submitting **concurrent** jobs (N jobs × M archives). + Otherwise concurrent jobs backpressure the scheduler and become slow and uneven. The default is + sized generously for this. diff --git a/examples/huntsman/clp-search/client/Cargo.toml b/examples/huntsman/clp-search/client/Cargo.toml new file mode 100644 index 000000000..708e0a7c0 --- /dev/null +++ b/examples/huntsman/clp-search/client/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "huntsman-clp-search-client" +version = "0.1.0" +edition = "2024" +publish = false + +[[bin]] +name = "huntsman-clp-search-client" +path = "src/main.rs" + +[dependencies] +anyhow = "1.0.98" +clap = { version = "4.6.1", features = ["derive"] } +rand = "0.9.1" +rmp-serde = "1.3.1" +spider-client = { path = "../../../../components/spider-client" } +spider-core = { path = "../../../../components/spider-core" } +tokio = { version = "1.50.0", features = ["macros", "rt-multi-thread"] } +tonic = "0.14.6" diff --git a/examples/huntsman/clp-search/client/src/main.rs b/examples/huntsman/clp-search/client/src/main.rs new file mode 100644 index 000000000..2dbe2628a --- /dev/null +++ b/examples/huntsman/clp-search/client/src/main.rs @@ -0,0 +1,426 @@ +//! Spider client that runs a multi-worker CLP search across a directory of CLP archives. +//! +//! The client discovers every immediate subdirectory of `--input` as one CLP archive, then builds +//! a FLAT Spider task graph with one `clp_search::search` task per archive. All tasks are +//! independent (no data-flow dependencies), so the scheduler is free to run them in parallel across +//! every available execution manager for maximum throughput. +//! +//! Each task writes its per-archive matches to a unique output file under a per-run output +//! directory. After the job succeeds, the client concatenates those files to STDOUT (the search +//! results); all progress and the final end-to-end latency summary are written to STDERR. + +use std::{ + fs, + num::NonZeroUsize, + path::{Path, PathBuf}, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use anyhow::{Context, anyhow}; +use clap::Parser; +use rand::seq::SliceRandom; +use spider_client::SpiderClient; +use spider_core::{ + job::JobState, + task::{DataTypeDescriptor, TaskDescriptor, TaskGraph, TdlContext, ValueTypeDescriptor}, + types::{id::JobId, io::TaskInput}, +}; +use tonic::transport::Endpoint; + +/// TDL package and task function each search task drives. +const PACKAGE: &str = "clp_search"; +const TASK_FUNC: &str = "clp_search::search"; + +/// Password used when registering the per-run resource group. +const RESOURCE_GROUP_PASSWORD: &[u8] = b"huntsman-clp-search-client"; + +/// Command-line arguments for the CLP search client. +#[derive(Debug, Parser)] +#[command(about = "Run a multi-worker CLP search over a directory of CLP archives via Spider.")] +struct Cli { + /// Directory whose immediate subdirectories are each a CLP archive to search. + #[arg(long, value_name = "PATH")] + input: PathBuf, + + /// KQL search query to run against every archive (e.g. `*NonDFS*`). + #[arg(long, value_name = "STRING")] + query: String, + + /// Spider storage gRPC endpoint to connect to. + #[arg(long, value_name = "URL", default_value = "http://127.0.0.1:50051")] + endpoint: String, + + /// `SpiderClient` gRPC connection pool size. + #[arg(long, default_value_t = 4)] + pool_size: usize, + + /// Base directory under which this run creates a unique subdirectory to hold its outputs. + #[arg( + long, + value_name = "PATH", + default_value = "build/spider-run/clp-search-results" + )] + output_dir: PathBuf, + + /// Process archives in sorted order instead of shuffling them before submission. + #[arg(long)] + no_shuffle: bool, +} + +/// Discovers the CLP archives directly under `input`. +/// +/// Every immediate subdirectory of `input` is treated as one CLP archive. The returned paths are +/// absolute (canonicalized) and sorted by directory name so the task ordering is deterministic. +/// +/// # Returns +/// +/// The absolute archive directory paths, sorted by name, on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * No archive subdirectory is found under `input`. +/// * Forwards [`fs::read_dir`]'s return values on failure. +/// * Forwards [`fs::canonicalize`]'s return values on failure. +fn discover_archives(input: &Path) -> anyhow::Result> { + let entries = fs::read_dir(input) + .with_context(|| format!("failed to read --input {}", input.display()))?; + let mut archives = Vec::new(); + for entry in entries { + let entry = entry.context("failed to read a directory entry under --input")?; + let path = entry.path(); + if path.is_dir() { + let absolute = fs::canonicalize(&path).with_context(|| { + format!("failed to canonicalize archive path {}", path.display()) + })?; + archives.push(absolute); + } + } + if archives.is_empty() { + return Err(anyhow!( + "no archive subdirectory found under --input {}", + input.display() + )); + } + archives.sort(); + Ok(archives) +} + +/// Builds a flat task graph with one independent `clp_search::search` task per archive. +/// +/// Every task has three graph inputs (archive path, query, output path) and no outputs, so all +/// tasks are input tasks with no dependencies between them. +/// +/// # Returns +/// +/// The assembled task graph on success. +/// +/// # Errors +/// +/// Forwards [`TaskGraph::new`]'s return values on failure. +/// Forwards [`TaskGraph::insert_task`]'s return values on failure. +fn build_graph(num_tasks: usize) -> anyhow::Result { + let bytes_type = DataTypeDescriptor::Value(ValueTypeDescriptor::bytes()); + let tdl_context = TdlContext { + package: PACKAGE.to_owned(), + task_func: TASK_FUNC.to_owned(), + }; + + let mut graph = TaskGraph::new(None, None)?; + for _ in 0..num_tasks { + graph.insert_task(TaskDescriptor { + tdl_context: tdl_context.clone(), + execution_policy: None, + inputs: vec![bytes_type.clone(); 3], + outputs: vec![], + input_sources: None, + })?; + } + Ok(graph) +} + +/// Serializes a string value into a msgpack task-input payload. +/// +/// # Returns +/// +/// The [`TaskInput::ValuePayload`] carrying the msgpack-encoded string on success. +/// +/// # Errors +/// +/// Forwards [`rmp_serde::to_vec`]'s return values on failure. +fn value_input(value: &str) -> anyhow::Result { + let payload = rmp_serde::to_vec(value).context("failed to serialize a task input")?; + Ok(TaskInput::ValuePayload(payload)) +} + +/// Prepares the per-run outputs, the flat task graph, and the flattened task inputs. +/// +/// Creates the unique run output directory `/run-`, assigns each archive a +/// unique output file path under it, builds the flat `clp_search::search` task graph, and flattens +/// the graph inputs so that, for archive `i`, positions 0,1,2 are the archive path, the query, and +/// the output path -- matching the order of [`TaskGraph::get_task_graph_input_indices`]. +/// +/// # Returns +/// +/// A tuple of the per-task absolute output paths, the assembled task graph, and the flattened task +/// inputs on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * An archive path has no final path component. +/// * An archive path or a computed output path is not valid UTF-8. +/// * Forwards [`fs::create_dir_all`]'s return values on failure. +/// * Forwards [`fs::canonicalize`]'s return values on failure. +/// * Forwards [`build_graph`]'s return values on failure. +/// * Forwards [`value_input`]'s return values on failure. +fn prepare_job( + archives: &[PathBuf], + query: &str, + output_dir: &Path, + nanos: u128, +) -> anyhow::Result<(Vec, TaskGraph, Vec)> { + let run_dir = output_dir.join(format!("run-{nanos}")); + fs::create_dir_all(&run_dir).with_context(|| { + format!( + "failed to create run output directory {}", + run_dir.display() + ) + })?; + let run_dir = fs::canonicalize(&run_dir).with_context(|| { + format!( + "failed to canonicalize run output directory {}", + run_dir.display() + ) + })?; + + let index_width = archives.len().to_string().len(); + let mut output_paths = Vec::with_capacity(archives.len()); + for (i, archive) in archives.iter().enumerate() { + let archive_name = archive + .file_name() + .context("archive path has no final component")? + .to_string_lossy(); + output_paths.push(run_dir.join(format!("{i:0index_width$}-{archive_name}.jsonl"))); + } + + let graph = build_graph(archives.len())?; + + let mut task_inputs: Vec = Vec::with_capacity(archives.len() * 3); + for (archive, output_path) in archives.iter().zip(output_paths.iter()) { + let archive_str = archive + .to_str() + .context("archive path is not valid UTF-8")?; + let output_str = output_path + .to_str() + .context("output path is not valid UTF-8")?; + task_inputs.push(value_input(archive_str)?); + task_inputs.push(value_input(query)?); + task_inputs.push(value_input(output_str)?); + } + + Ok((output_paths, graph, task_inputs)) +} + +/// Polls the job state until it reaches a terminal state. +/// +/// # Returns +/// +/// The terminal [`JobState`] on success. +/// +/// # Errors +/// +/// Forwards [`SpiderClient::get_job_state`]'s return values on failure. +async fn poll_until_terminal(client: &SpiderClient, job_id: JobId) -> anyhow::Result { + loop { + let state = client + .get_job_state(job_id) + .await + .context("get_job_state")?; + if state.is_terminal() { + return Ok(state); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } +} + +/// Prints a single labeled phase-timing line to STDERR, in milliseconds. +/// +/// The label is left-padded so the colons of consecutive lines align. +fn print_timing(label: &str, duration: Duration) { + eprintln!( + "[timing] {label:<26}: {:.1} ms", + duration.as_secs_f64() * 1000.0 + ); +} + +/// End-to-end per-phase timings for a single CLP search run. +struct PhaseTimings { + discovery: Duration, + graph_and_inputs: Duration, + connect_and_resource_group: Duration, + submit_job: Duration, + start_job: Duration, + spider_execution: Duration, + post_processing: Duration, + total: Duration, +} + +impl PhaseTimings { + /// Prints the per-phase breakdown followed by the three headline rollups and the total to + /// STDERR. + /// + /// `query_processing` aggregates the discovery, graph/input construction, connection, and + /// job-submission phases (everything before the distributed execution begins). + fn print(&self) { + let query_processing = self.discovery + + self.graph_and_inputs + + self.connect_and_resource_group + + self.submit_job + + self.start_job; + print_timing("discovery", self.discovery); + print_timing("graph_and_inputs", self.graph_and_inputs); + print_timing( + "connect_and_resource_group", + self.connect_and_resource_group, + ); + print_timing("submit_job (register)", self.submit_job); + print_timing("start_job", self.start_job); + print_timing("spider_execution", self.spider_execution); + print_timing("post_processing", self.post_processing); + print_timing("== query_processing", query_processing); + print_timing("== spider_execution", self.spider_execution); + print_timing("== post_processing", self.post_processing); + print_timing("== total", self.total); + } +} + +// A linear benchmark pipeline (discover -> submit -> poll -> read) with inline phase timing; the +// `--no-shuffle` branch pushes it one line past the pedantic limit, but splitting it would hurt +// readability more than help. +#[allow(clippy::too_many_lines)] +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let cli = Cli::parse(); + let start = Instant::now(); + + let pool_size = NonZeroUsize::new(cli.pool_size).context("--pool-size must be >= 1")?; + + let phase_start = Instant::now(); + let mut archives = discover_archives(&cli.input)?; + if !cli.no_shuffle { + // Shuffle so heavy archives spread across workers instead of concentrating on a straggler. + archives.shuffle(&mut rand::rng()); + } + let discovery_duration = phase_start.elapsed(); + + // Use a unique run id per run so repeated runs do not collide, both for the output directory + // and for the resource-group external id (checked against a persistent MariaDB). + let phase_start = Instant::now(); + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .context("system clock before UNIX epoch")? + .as_nanos(); + let (output_paths, graph, task_inputs) = + prepare_job(&archives, &cli.query, &cli.output_dir, nanos)?; + let graph_and_inputs_duration = phase_start.elapsed(); + + let phase_start = Instant::now(); + let endpoint: Endpoint = cli + .endpoint + .parse() + .with_context(|| format!("invalid --endpoint {:?}", cli.endpoint))?; + let client = SpiderClient::connect(endpoint, pool_size) + .await + .context("failed to connect to the Spider storage service")?; + + let resource_group_id = client + .add_resource_group( + format!("clp-search-{nanos}"), + RESOURCE_GROUP_PASSWORD.to_vec(), + ) + .await + .context("add_resource_group")?; + let connect_and_resource_group_duration = phase_start.elapsed(); + + let phase_start = Instant::now(); + let job_id = client + .submit_job(resource_group_id, &graph, task_inputs) + .await + .context("submit_job")?; + let submit_job_duration = phase_start.elapsed(); + + let phase_start = Instant::now(); + client.start_job(job_id).await.context("start_job")?; + let start_job_duration = phase_start.elapsed(); + + // Client-side wall-clock (epoch microseconds) at the moment the job is started. The scheduling + // overhead is the scheduler's first-task-dispatch timestamp minus this value. + let client_job_start_epoch_us = SystemTime::now() + .duration_since(UNIX_EPOCH) + .context("system clock before UNIX epoch")? + .as_micros(); + eprintln!( + "[metric] client_job_start_epoch_us={client_job_start_epoch_us} job_id={}", + job_id.get() + ); + + eprintln!( + "Submitted CLP search job: archives={}, tasks={}, query={:?}, job_id={}", + archives.len(), + archives.len(), + cli.query, + job_id.get() + ); + + let phase_start = Instant::now(); + let state = poll_until_terminal(&client, job_id).await?; + let spider_execution_duration = phase_start.elapsed(); + + let phase_start = Instant::now(); + match state { + JobState::Succeeded => { + for output_path in &output_paths { + let contents = fs::read_to_string(output_path).with_context(|| { + format!("failed to read output file {}", output_path.display()) + })?; + print!("{contents}"); + } + } + JobState::Failed => { + let message = client + .get_job_error(job_id) + .await + .context("get_job_error")?; + return Err(anyhow!("job failed: {message}")); + } + other => { + return Err(anyhow!("job ended in unexpected state {other:?}")); + } + } + let post_processing_duration = phase_start.elapsed(); + + PhaseTimings { + discovery: discovery_duration, + graph_and_inputs: graph_and_inputs_duration, + connect_and_resource_group: connect_and_resource_group_duration, + submit_job: submit_job_duration, + start_job: start_job_duration, + spider_execution: spider_execution_duration, + post_processing: post_processing_duration, + total: start.elapsed(), + } + .print(); + + eprintln!( + "Job succeeded: archives={}, tasks={}, job_id={}, elapsed={:.3?}", + archives.len(), + archives.len(), + job_id.get(), + start.elapsed() + ); + + Ok(()) +} diff --git a/examples/huntsman/clp-search/pool-ref/Cargo.toml b/examples/huntsman/clp-search/pool-ref/Cargo.toml new file mode 100644 index 000000000..b2f8fba86 --- /dev/null +++ b/examples/huntsman/clp-search/pool-ref/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "huntsman-clp-search-pool-ref" +version = "0.1.0" +edition = "2024" +publish = false + +[[bin]] +name = "clp-search-pool-ref" +path = "src/main.rs" + +[dependencies] +anyhow = "1.0.98" +clap = { version = "4.6.1", features = ["derive"] } +rand = "0.9.1" +tokio = { + version = "1.50.0", + features = ["macros", "process", "rt-multi-thread", "sync"] +} diff --git a/examples/huntsman/clp-search/pool-ref/src/main.rs b/examples/huntsman/clp-search/pool-ref/src/main.rs new file mode 100644 index 000000000..7ef297d2a --- /dev/null +++ b/examples/huntsman/clp-search/pool-ref/src/main.rs @@ -0,0 +1,424 @@ +//! Local process-pool reference binary that runs the SAME CLP search workload as the Spider +//! client, but WITHOUT Spider. +//! +//! Instead of submitting a Spider job, this binary drives one `clp-s` process per discovered CLP +//! archive through a bounded local pool that keeps at most `--pool-size` processes running +//! concurrently. It exists to establish a baseline end-to-end latency so Spider's +//! scheduling/coordination overhead can be isolated by comparing the two. +//! +//! Archive discovery, the per-run output directory, the per-archive output-file naming, the +//! `clp-s` invocation, and the phase-timing labels are all mirrored from the Spider client so the +//! two runs are directly comparable. All search results are written to STDOUT; all progress and the +//! per-phase timing breakdown are written to STDERR. + +use std::{ + env, + fs, + fs::File, + path::{Path, PathBuf}, + process::Stdio, + sync::Arc, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use anyhow::{Context, anyhow}; +use clap::Parser; +use rand::seq::SliceRandom; +use tokio::{process::Command, sync::Semaphore, task::JoinSet}; + +/// Environment variable that overrides the `clp-s` binary path. +const CLP_S_BIN_ENV: &str = "CLP_S_BIN"; + +/// Default `clp-s` binary path used when [`CLP_S_BIN_ENV`] is unset. +const DEFAULT_CLP_S_BIN: &str = "/home/lzh/dev/clp/build/core/clp-s"; + +/// Command-line arguments for the CLP search process-pool reference binary. +#[derive(Debug, Parser)] +#[command( + about = "Run a CLP search over a directory of CLP archives via a local clp-s process pool (no \ + Spider)." +)] +struct Cli { + /// Directory whose immediate subdirectories are each a CLP archive to search. + #[arg(long, value_name = "PATH")] + input: PathBuf, + + /// KQL search query to run against every archive (e.g. `*NonDFS*`). + #[arg(long, value_name = "STRING")] + query: String, + + /// Base directory under which this run creates a unique subdirectory to hold its outputs. + #[arg( + long, + value_name = "PATH", + default_value = "build/spider-run/clp-search-pool-results" + )] + output_dir: PathBuf, + + /// Maximum number of `clp-s` processes running concurrently. + #[arg(long, default_value_t = 16)] + pool_size: usize, +} + +/// Discovers the CLP archives directly under `input`. +/// +/// Every immediate subdirectory of `input` is treated as one CLP archive. The returned paths are +/// absolute (canonicalized) and sorted by directory name so the run ordering is deterministic. +/// +/// # Returns +/// +/// The absolute archive directory paths, sorted by name, on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * No archive subdirectory is found under `input`. +/// * Forwards [`fs::read_dir`]'s return values on failure. +/// * Forwards [`fs::canonicalize`]'s return values on failure. +fn discover_archives(input: &Path) -> anyhow::Result> { + let entries = fs::read_dir(input) + .with_context(|| format!("failed to read --input {}", input.display()))?; + let mut archives = Vec::new(); + for entry in entries { + let entry = entry.context("failed to read a directory entry under --input")?; + let path = entry.path(); + if path.is_dir() { + let absolute = fs::canonicalize(&path).with_context(|| { + format!("failed to canonicalize archive path {}", path.display()) + })?; + archives.push(absolute); + } + } + if archives.is_empty() { + return Err(anyhow!( + "no archive subdirectory found under --input {}", + input.display() + )); + } + archives.sort(); + Ok(archives) +} + +/// Prepares the per-run output directory and assigns each archive a unique output-file path. +/// +/// Creates the unique run output directory `/run-` and, for archive `i`, assigns +/// the absolute output path `/-.jsonl`. This is the local +/// analog of the Spider client's graph/input construction, so it is timed under the same +/// `graph_and_inputs` label. +/// +/// # Returns +/// +/// The per-archive absolute output paths, in archive order, on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * An archive path has no final path component. +/// * Forwards [`fs::create_dir_all`]'s return values on failure. +/// * Forwards [`fs::canonicalize`]'s return values on failure. +fn prepare_outputs( + archives: &[PathBuf], + output_dir: &Path, + nanos: u128, +) -> anyhow::Result> { + let run_dir = output_dir.join(format!("run-{nanos}")); + fs::create_dir_all(&run_dir).with_context(|| { + format!( + "failed to create run output directory {}", + run_dir.display() + ) + })?; + let run_dir = fs::canonicalize(&run_dir).with_context(|| { + format!( + "failed to canonicalize run output directory {}", + run_dir.display() + ) + })?; + + let index_width = archives.len().to_string().len(); + let mut output_paths = Vec::with_capacity(archives.len()); + for (i, archive) in archives.iter().enumerate() { + let archive_name = archive + .file_name() + .context("archive path has no final component")? + .to_string_lossy(); + output_paths.push(run_dir.join(format!("{i:0index_width$}-{archive_name}.jsonl"))); + } + Ok(output_paths) +} + +/// Runs the KQL `query` over the CLP archive at `archive_path`, writing the matching records as +/// JSONL to `output_path`. +/// +/// Invokes `clp-s` (resolved from `clp_s_bin`) as `clp-s s `, redirects its +/// stdout into the freshly truncated file at `output_path`, and captures its stderr so a failing +/// archive can be reported. This mirrors the Spider TDL search task exactly. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// # Returns +/// +/// The wall-clock duration of the `clp-s` subprocess on success (for execution-time analysis). +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * The output file cannot be created. +/// * The `clp-s` process cannot be spawned or waited on. +/// * The `clp-s` process exits with a non-success status. +async fn search_archive( + clp_s_bin: &str, + archive_path: &Path, + query: &str, + output_path: &Path, +) -> anyhow::Result { + let output_file = File::create(output_path).with_context(|| { + format!( + "failed to create output file `{}` for archive `{}`", + output_path.display(), + archive_path.display() + ) + })?; + + // Time only the `clp-s` subprocess, matching the Spider task's `clp_s_elapsed_us` metric. + let clp_s_start = Instant::now(); + let child = Command::new(clp_s_bin) + .arg("s") + .arg(archive_path) + .arg(query) + .stdout(Stdio::from(output_file)) + .stderr(Stdio::piped()) + .spawn() + .with_context(|| { + format!( + "failed to spawn `{clp_s_bin}` for archive `{}`", + archive_path.display() + ) + })?; + + let output = child.wait_with_output().await.with_context(|| { + format!( + "failed to wait on `{clp_s_bin}` for archive `{}`", + archive_path.display() + ) + })?; + let clp_s_elapsed = clp_s_start.elapsed(); + + if output.status.success() { + return Ok(clp_s_elapsed); + } + + Err(anyhow!( + "`clp-s` failed for archive `{}` with status {}: {}", + archive_path.display(), + output.status, + String::from_utf8_lossy(&output.stderr), + )) +} + +/// Runs one `clp-s` search per archive through a bounded local process pool. +/// +/// A [`Semaphore`] with `pool_size` permits caps the number of `clp-s` processes running at once: +/// every spawned task must acquire a permit before it spawns its child process and releases it when +/// the child completes. All archives are spawned up front and awaited to completion; the failures +/// of all archives are collected so a single failed archive does not abort the others (mirroring +/// how a failed Spider job is only reported once all tasks settle). +/// +/// # Returns +/// +/// The per-archive `clp-s` subprocess durations (in completion order) on success. +/// +/// # Errors +/// +/// Returns an error if: +/// +/// * Any archive's `clp-s` process fails to spawn or exits with a non-success status. +/// * A search task panics. +async fn run_pool( + clp_s_bin: String, + query: String, + archives: &[PathBuf], + output_paths: &[PathBuf], + pool_size: usize, +) -> anyhow::Result> { + let semaphore = Arc::new(Semaphore::new(pool_size)); + let clp_s_bin = Arc::new(clp_s_bin); + let query = Arc::new(query); + + let mut join_set = JoinSet::new(); + for (archive, output_path) in archives.iter().zip(output_paths.iter()) { + let permit_source = Arc::clone(&semaphore); + let clp_s_bin = Arc::clone(&clp_s_bin); + let query = Arc::clone(&query); + let archive = archive.clone(); + let output_path = output_path.clone(); + join_set.spawn(async move { + let _permit = permit_source + .acquire_owned() + .await + .expect("pool semaphore closed unexpectedly"); + search_archive(&clp_s_bin, &archive, &query, &output_path).await + }); + } + + let mut failures = Vec::new(); + let mut clp_s_times = Vec::with_capacity(archives.len()); + while let Some(joined) = join_set.join_next().await { + match joined { + Ok(Ok(elapsed)) => clp_s_times.push(elapsed), + Ok(Err(error)) => failures.push(format!("{error:#}")), + Err(join_error) => failures.push(format!("a search task panicked: {join_error}")), + } + } + + if failures.is_empty() { + return Ok(clp_s_times); + } + + for failure in &failures { + eprintln!("clp-s failure: {failure}"); + } + Err(anyhow!( + "the pool run failed: {} archive(s) failed", + failures.len() + )) +} + +/// Prints summary statistics of the per-archive `clp-s` execution times to STDERR. +/// +/// Reports count, total, mean, median, min, max, and p95 (all in milliseconds) so the pool's +/// `clp-s` execution distribution can be compared against Spider's `clp_s_elapsed_us` metric. +fn print_clp_s_stats(times: &[Duration]) { + if times.is_empty() { + return; + } + let mut ms: Vec = times.iter().map(|d| d.as_secs_f64() * 1000.0).collect(); + ms.sort_by(f64::total_cmp); + let n = ms.len(); + let sum: f64 = ms.iter().sum(); + let mean = sum / f64::from(u32::try_from(n).unwrap_or(u32::MAX)); + eprintln!( + "[clp_s] n={n} sum={sum:.1}ms mean={mean:.3}ms median={:.3}ms min={:.3}ms max={:.3}ms \ + p95={:.3}ms", + ms[n / 2], + ms[0], + ms[n - 1], + ms[(n * 95) / 100], + ); +} + +/// Prints a single labeled phase-timing line to STDERR, in milliseconds. +/// +/// The label is left-padded so the colons of consecutive lines align. +fn print_timing(label: &str, duration: Duration) { + eprintln!( + "[timing] {label:<26}: {:.1} ms", + duration.as_secs_f64() * 1000.0 + ); +} + +/// End-to-end per-phase timings for a single CLP search pool run. +struct PhaseTimings { + discovery: Duration, + graph_and_inputs: Duration, + spider_execution: Duration, + post_processing: Duration, + total: Duration, +} + +impl PhaseTimings { + /// Prints the per-phase breakdown followed by the three headline rollups and the total to + /// STDERR. + /// + /// `query_processing` aggregates the discovery and output-preparation phases (everything before + /// the pool execution begins), matching the Spider client's rollup layout. + fn print(&self) { + let query_processing = self.discovery + self.graph_and_inputs; + print_timing("discovery", self.discovery); + print_timing("graph_and_inputs", self.graph_and_inputs); + print_timing("spider_execution", self.spider_execution); + print_timing("post_processing", self.post_processing); + print_timing("== query_processing", query_processing); + print_timing("== spider_execution", self.spider_execution); + print_timing("== post_processing", self.post_processing); + print_timing("== total", self.total); + } +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let cli = Cli::parse(); + let start = Instant::now(); + + if cli.pool_size == 0 { + return Err(anyhow!("--pool-size must be >= 1")); + } + + let phase_start = Instant::now(); + let mut archives = discover_archives(&cli.input)?; + // Randomize archive order so heavy archives are spread across the pool instead of clustering + // (matches the Spider client's shuffle, keeping the two comparable). + archives.shuffle(&mut rand::rng()); + let discovery_duration = phase_start.elapsed(); + + let phase_start = Instant::now(); + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .context("system clock before UNIX epoch")? + .as_nanos(); + let output_paths = prepare_outputs(&archives, &cli.output_dir, nanos)?; + let graph_and_inputs_duration = phase_start.elapsed(); + + let clp_s_bin = env::var(CLP_S_BIN_ENV).unwrap_or_else(|_| DEFAULT_CLP_S_BIN.to_owned()); + + eprintln!( + "Starting CLP search pool run: archives={}, pool_size={}, query={:?}, clp_s_bin={:?}", + archives.len(), + cli.pool_size, + cli.query, + clp_s_bin + ); + + let phase_start = Instant::now(); + let clp_s_times = run_pool( + clp_s_bin, + cli.query.clone(), + &archives, + &output_paths, + cli.pool_size, + ) + .await?; + let spider_execution_duration = phase_start.elapsed(); + print_clp_s_stats(&clp_s_times); + + let phase_start = Instant::now(); + for output_path in &output_paths { + let contents = fs::read_to_string(output_path) + .with_context(|| format!("failed to read output file {}", output_path.display()))?; + print!("{contents}"); + } + let post_processing_duration = phase_start.elapsed(); + + PhaseTimings { + discovery: discovery_duration, + graph_and_inputs: graph_and_inputs_duration, + spider_execution: spider_execution_duration, + post_processing: post_processing_duration, + total: start.elapsed(), + } + .print(); + + eprintln!( + "Pool run succeeded: archives={}, pool_size={}, elapsed={:.3?}", + archives.len(), + cli.pool_size, + start.elapsed() + ); + + Ok(()) +} diff --git a/examples/huntsman/clp-search/tasks/Cargo.toml b/examples/huntsman/clp-search/tasks/Cargo.toml new file mode 100644 index 000000000..7d6ffbb1b --- /dev/null +++ b/examples/huntsman/clp-search/tasks/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "huntsman-clp-search-tasks" +version = "0.1.0" +edition = "2024" +publish = false + +[lib] +crate-type = ["cdylib"] +name = "clp_search" +path = "src/lib.rs" + +[dependencies] +serde = { version = "1.0.228", features = ["derive"] } +spider-tdl = { + path = "../../../../components/spider-tdl", + features = ["derive"] +} +tracing = { version = "0.1.41", default-features = false, features = ["std"] } +tracing-subscriber = { + version = "0.3.23", + default-features = false, + features = ["env-filter", "fmt", "json"] +} diff --git a/examples/huntsman/clp-search/tasks/src/lib.rs b/examples/huntsman/clp-search/tasks/src/lib.rs new file mode 100644 index 000000000..1304f9b62 --- /dev/null +++ b/examples/huntsman/clp-search/tasks/src/lib.rs @@ -0,0 +1,113 @@ +//! TDL package that runs KQL searches over CLP archives via the `clp-s` binary. + +mod task_decl { + use std::{ + env, + fs::File, + process::{Command, Stdio}, + sync::Once, + time::Instant, + }; + + use spider_tdl::{TaskContext, TdlError, task}; + + /// Environment variable that overrides the `clp-s` binary path. + const CLP_S_BIN_ENV: &str = "CLP_S_BIN"; + + /// Default `clp-s` binary path used when [`CLP_S_BIN_ENV`] is unset. + const DEFAULT_CLP_S_BIN: &str = "/home/lzh/dev/clp/build/core/clp-s"; + + /// Guards one-time installation of this package's tracing subscriber. + static LOG_INIT: Once = Once::new(); + + /// Installs a package-local tracing subscriber exactly once. + /// + /// This TDL package is a `cdylib` with its own copy of `tracing`'s global dispatcher, distinct + /// from the task executor that `dlopen`s it, so the executor's subscriber never observes events + /// emitted here. This installs a subscriber owned by the package that writes JSON to stderr -- + /// the same stream the executor redirects to `em-logs/-.log` -- and honors + /// `RUST_LOG` (propagated from the execution manager). `try_init` makes a redundant call on a + /// later task invocation a no-op. + fn init_task_logging() { + LOG_INIT.call_once(|| { + let _ = tracing_subscriber::fmt() + .event_format(tracing_subscriber::fmt::format().with_target(false).json()) + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_ansi(false) + .with_writer(std::io::stderr) + .try_init(); + }); + } + + /// Runs the KQL `query` over the CLP archive at `archive_path`, writing the matching records as + /// JSONL to `output_path`. + /// + /// Invokes the `clp-s` binary (resolved from the `CLP_S_BIN` environment variable, or + /// [`DEFAULT_CLP_S_BIN`] when unset) as `clp-s s ` and redirects its + /// stdout into the freshly truncated file at `output_path`. The parent directory of + /// `output_path` is assumed to already exist. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`TdlError::ExecutionError`] if the output file cannot be created, the `clp-s` process + /// cannot be spawned, the process cannot be waited on, or the process exits with a + /// non-success status. + #[task(name = "clp_search::search")] + pub fn search( + ctx: TaskContext, + archive_path: String, + query: String, + output_path: String, + ) -> Result<(), TdlError> { + init_task_logging(); + let clp_s_bin = env::var(CLP_S_BIN_ENV).unwrap_or_else(|_| DEFAULT_CLP_S_BIN.to_owned()); + + let output_file = File::create(&output_path).map_err(|error| { + TdlError::ExecutionError(format!( + "failed to create output file `{output_path}` for archive `{archive_path}`: \ + {error}" + )) + })?; + + // Benchmark instrumentation: time only the `clp-s` subprocess. + let clp_s_start = Instant::now(); + let output = Command::new(&clp_s_bin) + .arg("s") + .arg(&archive_path) + .arg(&query) + .stdout(Stdio::from(output_file)) + .stderr(Stdio::piped()) + .output() + .map_err(|error| { + TdlError::ExecutionError(format!( + "failed to run `{clp_s_bin}` for archive `{archive_path}`: {error}" + )) + })?; + let clp_s_elapsed_us = u64::try_from(clp_s_start.elapsed().as_micros()).unwrap_or(u64::MAX); + tracing::info!( + clp_s_elapsed_us, + job_id = ? ctx.job_id, + task_id = ? ctx.task_id, + "clp-s subprocess finished." + ); + + if output.status.success() { + return Ok(()); + } + + Err(TdlError::ExecutionError(format!( + "`clp-s` failed for archive `{archive_path}` with status {}: {}", + output.status, + String::from_utf8_lossy(&output.stderr), + ))) + } +} + +spider_tdl::register_tdl_package! { + package_name: "clp_search", + tasks: [ + task_decl::search + ], +} diff --git a/examples/huntsman/complex/client/Cargo.toml b/examples/huntsman/complex/client/Cargo.toml new file mode 100644 index 000000000..4e8306503 --- /dev/null +++ b/examples/huntsman/complex/client/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "huntsman-complex-client" +version = "0.1.0" +edition = "2024" +publish = false + +[[bin]] +name = "huntsman-complex-client" +path = "src/main.rs" + +[dependencies] +anyhow = "1.0.98" +clap = { version = "4.6.1", features = ["derive"] } +huntsman-complex-types = { path = "../types" } +rmp-serde = "1.3.1" +spider-client = { path = "../../../../components/spider-client" } +spider-core = { path = "../../../../components/spider-core" } +tokio = { version = "1.50.0", features = ["macros", "rt-multi-thread"] } +tonic = "0.14.6" \ No newline at end of file diff --git a/examples/huntsman/complex/client/src/main.rs b/examples/huntsman/complex/client/src/main.rs new file mode 100644 index 000000000..a83d2359e --- /dev/null +++ b/examples/huntsman/complex/client/src/main.rs @@ -0,0 +1,340 @@ +//! Simple Spider client that builds a layered `complex::add` task graph and runs it against a live +//! Spider stack. +//! +//! The graph is "neural-network-shaped": `--level` layers of `--width` tasks each. Layer 0 takes +//! its two inputs from the graph inputs; every inner task adds two outputs from the previous layer +//! (a fixed fan-in pattern, `prev[i]` and `prev[(i + 1) % width]`). Tasks within a layer are +//! independent, so `--width` controls how much parallelism the scheduler can exploit -- e.g. +//! `--width 16` keeps a 16-worker stack busy. +//! +//! After the job finishes, the client decodes the final layer's outputs and compares them against +//! an in-process simulation of the same DAG, proving the stack executed the graph correctly. + +use std::{ + num::NonZeroUsize, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use anyhow::{Context, anyhow}; +use clap::Parser; +use huntsman_complex_types::{Complex, ComplexVec}; +use spider_client::SpiderClient; +use spider_core::{ + job::JobState, + task::{ + DataTypeDescriptor, + TaskDescriptor, + TaskGraph, + TaskIndex, + TaskInputOutputIndex, + TdlContext, + ValueTypeDescriptor, + }, + types::{id::JobId, io::TaskInput}, +}; +use tonic::transport::Endpoint; + +/// TDL package and task function the graph drives. +const PACKAGE: &str = "complex"; +const TASK_FUNC: &str = "complex::add"; + +/// Password used when registering the per-run resource group. +const RESOURCE_GROUP_PASSWORD: &[u8] = b"huntsman-complex-client"; + +/// Command-line arguments for the example client. +#[derive(Debug, Parser)] +#[command(about = "Build a layered complex::add task graph and run it against the Spider stack.")] +struct Cli { + /// Spider storage gRPC endpoint to connect to. + #[arg(long, value_name = "URL", default_value = "http://127.0.0.1:50051")] + endpoint: String, + + /// Number of layers (graph depth). + #[arg(long, default_value_t = 10)] + level: usize, + + /// Tasks per layer. Tasks within a layer are independent, so this controls parallelism + /// (e.g. `--width 16` saturates a 16-worker stack). + #[arg(long, default_value_t = 4)] + width: usize, + + /// `SpiderClient` gRPC connection pool size. + #[arg(long, default_value_t = 4)] + pool_size: usize, + + /// Print each final-layer output's complex values alongside the simulation, for inspection. + #[arg(long)] + print_outputs: bool, +} + +/// Builds the layered `complex::add` task graph. +/// +/// Layer 0 has `width` tasks whose two inputs come from the graph inputs. Each subsequent layer +/// has `width` tasks whose two inputs are outputs `prev[i]` and `prev[(i + 1) % width]` from the +/// previous layer. +/// +/// # Returns +/// +/// The assembled task graph on success. +/// +/// # Errors +/// +/// Forwards [`TaskGraph::new`]'s return values on failure. +/// Forwards [`TaskGraph::insert_task`]'s return values on failure. +fn build_graph(level: usize, width: usize) -> anyhow::Result { + let bytes_type = DataTypeDescriptor::Value(ValueTypeDescriptor::bytes()); + let tdl_context = TdlContext { + package: PACKAGE.to_owned(), + task_func: TASK_FUNC.to_owned(), + }; + + let mut graph = TaskGraph::new(None, None)?; + + // Layer 0: both inputs come from the graph inputs. + let mut prev_layer: Vec = Vec::with_capacity(width); + for _ in 0..width { + let task_idx = graph.insert_task(TaskDescriptor { + tdl_context: tdl_context.clone(), + execution_policy: None, + inputs: vec![bytes_type.clone(); 2], + outputs: vec![bytes_type.clone()], + input_sources: None, + })?; + prev_layer.push(task_idx); + } + + // Inner layers: both inputs come from the previous layer's outputs. + for _ in 1..level { + let mut cur_layer = Vec::with_capacity(width); + for i in 0..width { + let lhs = prev_layer[i]; + let rhs = prev_layer[(i + 1) % width]; + let task_idx = graph.insert_task(TaskDescriptor { + tdl_context: tdl_context.clone(), + execution_policy: None, + inputs: vec![bytes_type.clone(); 2], + outputs: vec![bytes_type.clone()], + input_sources: Some(vec![ + TaskInputOutputIndex { + task_idx: lhs, + position: 0, + }, + TaskInputOutputIndex { + task_idx: rhs, + position: 0, + }, + ]), + })?; + cur_layer.push(task_idx); + } + prev_layer = cur_layer; + } + + Ok(graph) +} + +/// Builds the graph inputs: `2 * width` short `ComplexVec` values, one per positional input of the +/// layer-0 tasks (insertion order: task `i` consumes inputs `2 * i` and `2 * i + 1`). +/// +/// # Returns +/// +/// The graph input vectors. +fn build_graph_inputs(width: usize) -> Vec { + let count = 2usize.checked_mul(width).expect("width too large"); + let mut values = Vec::with_capacity(count); + // `k` is an f64 counter so no integer->float cast (which would trip clippy::cast_precision_loss + // under -D pedantic); the values are deterministic test data either way. + let mut k = 0.0_f64; + for _ in 0..count { + values.push(ComplexVec { + items: vec![ + Complex { + re: k + 1.0, + im: 0.0, + }, + Complex { + re: 0.0, + im: k + 1.0, + }, + ], + }); + k += 1.0; + } + values +} + +/// Element-wise complex addition of two equal-length vectors. +/// +/// # Panics +/// +/// Panics if `a` and `b` differ in length (the `complex::add` task would reject the same). +fn complex_add(a: &ComplexVec, b: &ComplexVec) -> ComplexVec { + assert_eq!( + a.items.len(), + b.items.len(), + "complex_add: vector length mismatch" + ); + ComplexVec { + items: a + .items + .iter() + .zip(b.items.iter()) + .map(|(x, y)| Complex { + re: x.re + y.re, + im: x.im + y.im, + }) + .collect(), + } +} + +/// Simulates the layered DAG in-process using the same connection pattern as [`build_graph`], so +/// the retrieved outputs can be checked against an independent reference. +/// +/// # Returns +/// +/// The final layer's output vectors. +fn simulate(level: usize, width: usize, inputs: &[ComplexVec]) -> Vec { + let mut prev: Vec = (0..width) + .map(|i| complex_add(&inputs[2 * i], &inputs[2 * i + 1])) + .collect(); + for _ in 1..level { + let cur: Vec = (0..width) + .map(|i| complex_add(&prev[i], &prev[(i + 1) % width])) + .collect(); + prev = cur; + } + prev +} + +/// Polls the job state until it reaches a terminal state. +/// +/// # Returns +/// +/// The terminal [`JobState`] on success. +/// +/// # Errors +/// +/// Forwards [`SpiderClient::get_job_state`]'s return values on failure. +async fn poll_until_terminal(client: &SpiderClient, job_id: JobId) -> anyhow::Result { + loop { + let state = client + .get_job_state(job_id) + .await + .context("get_job_state")?; + if state.is_terminal() { + return Ok(state); + } + tokio::time::sleep(Duration::from_millis(500)).await; + } +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let cli = Cli::parse(); + if cli.level == 0 || cli.width == 0 { + return Err(anyhow!("--level and --width must be >= 1")); + } + let pool_size = NonZeroUsize::new(cli.pool_size).context("--pool-size must be >= 1")?; + + let endpoint: Endpoint = cli + .endpoint + .parse() + .with_context(|| format!("invalid --endpoint {:?}", cli.endpoint))?; + let client = SpiderClient::connect(endpoint, pool_size) + .await + .context("failed to connect to the Spider storage service")?; + + let graph_inputs = build_graph_inputs(cli.width); + let expected = simulate(cli.level, cli.width, &graph_inputs); + let graph = build_graph(cli.level, cli.width)?; + + let task_inputs: Vec = graph_inputs + .iter() + .map(|value| { + Ok::(TaskInput::ValuePayload( + rmp_serde::to_vec(value).context("failed to serialize a graph input")?, + )) + }) + .collect::>()?; + + // Use a unique external id per run so repeated runs (against a persistent MariaDB) do not + // collide with the resource-group-already-exists error. + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .context("system clock before UNIX epoch")? + .as_nanos(); + let resource_group_id = client + .add_resource_group( + format!("huntsman-nn-{nanos}"), + RESOURCE_GROUP_PASSWORD.to_vec(), + ) + .await + .context("add_resource_group")?; + + let job_id = client + .submit_job(resource_group_id, &graph, task_inputs) + .await + .context("submit_job")?; + client.start_job(job_id).await.context("start_job")?; + + println!( + "Submitted layered complex::add job: level={}, width={}, tasks={}, job_id={}", + cli.level, + cli.width, + cli.level * cli.width, + job_id.get() + ); + + let state = poll_until_terminal(&client, job_id).await?; + match state { + JobState::Succeeded => { + let outputs = client + .get_job_outputs(job_id) + .await + .context("get_job_outputs")?; + anyhow::ensure!( + outputs.len() == cli.width, + "expected {} graph outputs, got {}", + cli.width, + outputs.len() + ); + let mut mismatches = 0; + for (i, output) in outputs.iter().enumerate() { + let got: ComplexVec = rmp_serde::from_slice(output) + .with_context(|| format!("failed to decode output {i}"))?; + if got != expected[i] { + mismatches += 1; + } + if cli.print_outputs { + println!( + "output[{i}] got={:?} expected={:?}", + got.items, expected[i].items + ); + } + } + if mismatches == 0 { + println!( + "Job succeeded; all {} final-layer outputs match the local simulation.", + outputs.len() + ); + } else { + return Err(anyhow!( + "job succeeded but {mismatches}/{} outputs mismatched the simulation", + outputs.len() + )); + } + } + JobState::Failed => { + let message = client + .get_job_error(job_id) + .await + .context("get_job_error")?; + return Err(anyhow!("job failed: {message}")); + } + other => { + return Err(anyhow!("job ended in unexpected state {other:?}")); + } + } + + Ok(()) +} diff --git a/examples/huntsman/nn/README.md b/examples/huntsman/nn/README.md new file mode 100644 index 000000000..1e108cc6c --- /dev/null +++ b/examples/huntsman/nn/README.md @@ -0,0 +1,137 @@ +# Neural-network-shaped Spider benchmark + +A small benchmark harness that runs a **layered, neural-network-shaped task graph** over a live +Spider stack. Each task is a simulated neuron -- it consumes 25 128-byte inputs, sleeps a fixed +10 ms to model compute cost, and emits a fixed 128-byte output. The graph is `--level` layers of +`--width` tasks each (default `--level 10 --width 1000` => 10,000 tasks), where every inner task +draws its 25 inputs from distinct random outputs of the previous layer. The 25 inputs per task is +fixed by the `nn_bench::sleep` task's signature, not a tunable. + +It exists to benchmark Spider on a structured, dependency-heavy workload and to compare Spider's +actual execution time against an **analytic ideal-runtime lower bound** computed from the generated +graph shape. + +## Crates + +| Path | Crate / artifact | What it is | +|-------------|-----------------------------------|-------------------------------------------------------------------| +| `client/` | `huntsman-nn-bench` (bin) | The harness: builds the layered `nn_bench::sleep` graph, runs it, verifies outputs, reports timing + ideal runtime. | +| `tasks/` | `huntsman-nn-bench-tasks` (cdylib) | The TDL package `nn_bench` exposing the single task `nn_bench::sleep`. Staged as `libnn_bench.so`. | + +## Prerequisites + +* The workspace built and the `nn_bench` package staged so task executors can `dlopen` it: + + ```shell + task build:rust + task build:packages + ``` + + `build:packages` builds the example binaries and stages `libnn_bench.so` into + `build/tdl_packages/nn_bench/`. +* A running Spider stack. See `../../../stack-doc.md`. The stack must have the `nn_bench` package + staged (the step above) before you submit a job that references it. + +## Running + +In one shell, bring up the stack (defaults are 16 workers and a large scheduler ready-task queue, +which the concurrent-job case needs): + +```shell +uv run --script tools/scripts/stack/run.py +``` + +In another shell, once the stack reports it is up, run the client: + +```shell +build/rust-targets/release/huntsman-nn-bench --level 10 --width 32 +``` + +A width-32 graph keeps all 16 execution managers busy within each layer (and the task's 25-input +signature requires `--width >= 25` so each task can draw distinct previous-layer outputs). For a +larger workload, raise `--width` (and the stack's `--workers` to match, so a layer can run in +parallel). The default shape (`--level 10 --width 1000`) is 10,000 tasks and is meant for a +many-worker stack. + +### Client arguments + +| Flag | Default | Meaning | +|-------------------|--------------------------|----------------------------------------------------------------| +| `--endpoint` | `http://127.0.0.1:50051` | Spider storage gRPC endpoint. | +| `--level` | `10` | Number of layers (graph depth). | +| `--width` | `1000` | Tasks per layer; must be >= 25 (the task's fixed input count) so each task can sample distinct previous-layer outputs. Controls the parallelism the scheduler can exploit. | +| `--input-bytes` | `128` | Size in bytes of each input and of each task's fixed output. | +| `--seed` | `0x517_d3ad` | Seed for the random input-selection topology (deterministic across runs). | +| `--pool-size` | `4` | `SpiderClient` gRPC connection-pool size. | +| `--print-outputs` | off | Print each final-layer output + whether it matches the expected payload. | + +**Output convention:** per-output inspection (with `--print-outputs`) goes to **stdout**; the +submit line, the per-phase timing breakdown, the ideal-runtime figures, and the success summary go +to **stderr**. So `client ... --print-outputs > outputs.txt` captures just the inspection output. + +## How the client works + +1. **Build graph** — construct the layered `nn_bench::sleep` `TaskGraph`: layer 0's `width` tasks + take their 25 inputs from the graph inputs; every inner task's 25 inputs are distinct random + outputs from the previous layer (seeded by `--seed`). The builder also records the graph's task + count and depth (the longest dependency chain) for the ideal-runtime step. +2. **Build inputs** — `width * 25` byte vectors of `input_bytes` bytes, one per positional input + of the layer-0 tasks, each msgpack-encoded into a `TaskInput::ValuePayload`. +3. **Submit → start → poll** — register a per-run resource group, `submit_job`, `start_job`, then + poll `get_job_state` (every 500 ms) until the job reaches a terminal state. +4. **Verify outputs** — on success, decode each final-layer output and check it equals the fixed + 128-byte payload the task emits, proving every task ran. On failure, fetch and report the job + error. +5. **Report timing** — print a phase breakdown to stderr: `graph_and_inputs`, `connect_and_resource_group`, + `submit_and_start`, `spider_execution` (the poll-to-terminal wait — where nearly all the + wall-clock lives), `post_processing`, plus the `query_processing` / `spider_execution` / + `post_processing` / `total` rollups. +6. **Report ideal runtime** — print the analytic lower bound for 16, 32, 64, and 128 workers (see + below). + +## The task (`nn_bench::sleep`) + +A single TDL task with signature `sleep(ctx, i0..i24) -> Result, TdlError>` — 25 `bytes` +positional inputs modeling the neuron's 25 incoming data-flow edges. It sums the input lengths +(only to prove the inputs were delivered), sleeps 10 ms, and returns the fixed 128-byte payload. +Each invocation logs a START and END line to stderr carrying the job/task ids and a nanosecond +timestamp, so per-task start/end times can be recovered from the executor logs. + +## Ideal runtime + +Because every task costs a fixed 10 ms (a simulated sleep), the workload's ideal makespan is +analytic — no separate baseline run is needed. For a DAG of equal-duration tasks scheduled on `W` +workers, no schedule can beat + +``` +ideal(W) = max(critical_path, total_work / W) +``` + +where `total_work = total_tasks * 10 ms` (the perfect-parallelism bound) and `critical_path = +depth * 10 ms` (the longest dependency chain, which must run serially). The client computes this +from the generated graph's task count and depth and prints it for 16, 32, 64, and 128 workers: + +``` +[timing] == ideal (lower bound): +[timing] ideal 16 workers: 6250.0 ms +[timing] ideal 32 workers: 3125.0 ms +[timing] ideal 64 workers: 1562.5 ms +[timing] ideal 128 workers: 781.2 ms +``` + +(For the default `--level 10 --width 1000`: 10,000 tasks, depth 10, so `critical_path = 100 ms` +and `total_work = 100 s`.) Compare the measured `[timing] == spider_execution` against the line +matching the stack's worker count — the gap is Spider's scheduling/coordination overhead. As `W` +grows the `total_work / W` term shrinks until the `critical_path` floor dominates (here at ~1000 +workers); below that the ideal halves each time `W` doubles. + +## Notes + +* Prefer a **freshly launched** stack per benchmark: a stack left running can drift into a bad state + and fail on the next submission. +* The scheduler's `ready_task_capacity` (in `tools/scripts/stack/spider.yaml`) must exceed the total + number of simultaneously-ready tasks when submitting **concurrent** jobs. Otherwise concurrent + jobs backpressure the scheduler and become slow and uneven. The default is sized generously for a + single wide job. +* Per-task START/END timestamps are written by the task itself to stderr, captured by the execution + manager into `build/spider-run/em-logs/-.log`. \ No newline at end of file diff --git a/examples/huntsman/nn/client/Cargo.toml b/examples/huntsman/nn/client/Cargo.toml new file mode 100644 index 000000000..5393f36fa --- /dev/null +++ b/examples/huntsman/nn/client/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "huntsman-nn-bench" +version = "0.1.0" +edition = "2024" +publish = false + +[[bin]] +name = "huntsman-nn-bench" +path = "src/main.rs" + +[dependencies] +anyhow = "1.0.98" +clap = { version = "4.6.1", features = ["derive"] } +rmp-serde = "1.3.1" +spider-client = { path = "../../../../components/spider-client" } +spider-core = { path = "../../../../components/spider-core" } +tokio = { version = "1.50.0", features = ["macros", "rt-multi-thread"] } +tonic = "0.14.6" \ No newline at end of file diff --git a/examples/huntsman/nn/client/src/main.rs b/examples/huntsman/nn/client/src/main.rs new file mode 100644 index 000000000..960ce140f --- /dev/null +++ b/examples/huntsman/nn/client/src/main.rs @@ -0,0 +1,548 @@ +//! Neural-network-shaped benchmark client for the Spider stack. +//! +//! Builds a layered task graph that mimics a neural network -- `--level` layers of `--width` tasks +//! each -- where every task consumes 25 128-byte inputs randomly selected from the previous +//! layer's outputs (layer 0 takes 25 graph inputs per task), sleeps 10 ms inside the +//! `nn_bench::sleep` task, and emits a fixed 128-byte output. The default shape is +//! `--level 10 --width 1000` => 10,000 tasks. The 25 inputs per task is fixed by the task's +//! signature (`nn_bench::sleep` takes 25 positional inputs), not a tunable. +//! +//! After the job finishes the client decodes the final layer's outputs and checks each against the +//! fixed 128-byte payload the task emits, proving every task ran. It reports a per-phase timing +//! breakdown to STDERR (the `[timing]` convention: sub-phases plus `query_processing` / +//! `spider_execution` / `post_processing` / `total` rollups) and, since the task cost is a fixed +//! simulated sleep, an **ideal runtime** lower bound computed from the generated graph shape for +//! 16, 32, 64, and 128 workers -- the floor to compare Spider's actual `spider_execution` against. +//! +//! Per-task START/END timestamps are written by the task itself to stderr, which the execution +//! manager captures into `build/spider-run/em-logs/-.log`. + +use std::{ + collections::HashSet, + num::NonZeroUsize, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use anyhow::{Context, anyhow}; +use clap::Parser; +use spider_client::SpiderClient; +use spider_core::{ + job::JobState, + task::{ + DataTypeDescriptor, + TaskDescriptor, + TaskGraph, + TaskIndex, + TaskInputOutputIndex, + TdlContext, + ValueTypeDescriptor, + }, + types::{id::JobId, io::TaskInput}, +}; +use tonic::transport::Endpoint; + +/// TDL package and task function the graph drives. The task sleeps 10 ms and emits a fixed +/// 128-byte payload. +const PACKAGE: &str = "nn_bench"; +const TASK_FUNC: &str = "nn_bench::sleep"; + +/// Fixed 128-byte payload every `nn_bench::sleep` invocation emits; the client checks every +/// final-layer output against this. +const OUTPUT_PAYLOAD: [u8; 128] = [0; 128]; + +/// Number of positional inputs each `nn_bench::sleep` task consumes. Mirrors the task's `i0..i24` +/// parameter list (in `examples/huntsman/nn/tasks/src/lib.rs`); keep the two in sync. The graph +/// declares this many inputs per task, and the runtime maps them positionally onto the task's +/// parameters, so this is a fixed property of the task -- not a tunable. +const NUM_TASK_INPUT: usize = 25; + +/// Per-task compute cost in milliseconds. Mirrors `nn_bench::sleep`'s `SLEEP_DURATION` (in +/// `examples/huntsman/nn/tasks/src/lib.rs`); keep the two in sync. Used to compute the ideal +/// runtime lower bound. +const TASK_DURATION_MS: f64 = 10.0; + +/// Worker counts for which the ideal runtime lower bound is reported, alongside the measured +/// `spider_execution`. The client does not know the live stack's worker count, so it reports a +/// fixed reference set. +const IDEAL_WORKER_COUNTS: [usize; 4] = [16, 32, 64, 128]; + +/// Password used when registering the per-run resource group. +const RESOURCE_GROUP_PASSWORD: &[u8] = b"huntsman-nn-bench"; + +/// Command-line arguments for the benchmark client. +#[derive(Debug, Parser)] +#[command( + about = "Build a neural-network-shaped nn_bench::sleep task graph and benchmark the Spider \ + stack." +)] +struct Cli { + /// Spider storage gRPC endpoint to connect to. + #[arg(long, value_name = "URL", default_value = "http://127.0.0.1:50051")] + endpoint: String, + + /// Number of layers (graph depth). + #[arg(long, default_value_t = 10)] + level: usize, + + /// Tasks per layer. Tasks within a layer are independent, so this controls the parallelism the + /// scheduler can exploit (e.g. `--width 1000` against a 32-worker stack). Must be >= the + /// task's input count (`NUM_TASK_INPUT`) so each task can draw distinct previous-layer + /// outputs. + #[arg(long, default_value_t = 1000)] + width: usize, + + /// Size in bytes of each input and of each task's output (the task emits a fixed payload of + /// this size). + #[arg(long, default_value_t = 128)] + input_bytes: usize, + + /// Seed for the random input-selection topology (deterministic across runs). + #[arg(long, default_value_t = 0x517_d3ad)] + seed: u64, + + /// `SpiderClient` gRPC connection pool size. + #[arg(long, default_value_t = 4)] + pool_size: usize, + + /// Print each final-layer output alongside the expected payload, for inspection (to STDOUT). + #[arg(long)] + print_outputs: bool, +} + +/// A tiny xorshift64 RNG seeded by `--seed` so the random input-selection topology is +/// deterministic. +struct Rng(u64); + +impl Rng { + const fn new(seed: u64) -> Self { + // Avoid a zero seed, which would xorshift-stick at 0. + Self(if seed == 0 { + 0x9e37_79b9_7f4a_7c15 + } else { + seed + }) + } + + const fn next_u64(&mut self) -> u64 { + let mut x = self.0; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.0 = x; + x + } +} + +/// Samples `count` distinct task indices in `[0, width)` using rejection sampling. +/// +/// # Panics +/// +/// Panics if `count > width` (cannot pick that many distinct indices). +fn sample_distinct(rng: &mut Rng, count: usize, width: usize) -> Vec { + assert!( + count <= width, + "count ({count}) must be <= width ({width}) for distinct sampling" + ); + let width_u64 = u64::try_from(width).expect("width fits in u64"); + let mut picked: HashSet = HashSet::with_capacity(count); + while picked.len() < count { + let idx = usize::try_from(rng.next_u64() % width_u64).expect("index fits in usize"); + picked.insert(idx); + } + picked.into_iter().collect() +} + +/// Structural shape of the generated graph, used to compute the ideal runtime lower bound. +struct GraphShape { + /// Total number of tasks in the graph. + total_tasks: usize, + /// Graph depth in tasks -- the longest dependency chain. Every inner-layer task depends on a + /// previous-layer task, so a chain spans all `level` layers; the chain length equals the layer + /// count. + depth: usize, +} + +/// Builds the layered `nn_bench::sleep` task graph. +/// +/// Layer 0 has `width` tasks whose `NUM_TASK_INPUT` inputs come from the graph inputs. Each +/// subsequent layer has `width` tasks whose `NUM_TASK_INPUT` inputs are distinct random outputs +/// from the previous layer, drawn from a seeded RNG so the topology is reproducible. +/// +/// # Returns +/// +/// A tuple of the assembled task graph and its [`GraphShape`] (task count + depth, derived from +/// the build so the ideal-runtime calculation tracks the graph that was actually generated). +/// +/// # Errors +/// +/// Forwards [`TaskGraph::new`]'s return values on failure. +/// Forwards [`TaskGraph::insert_task`]'s return values on failure. +/// +/// # Panics +/// +/// Panics (via [`sample_distinct`]) if `width < NUM_TASK_INPUT`, since distinct sampling needs at +/// least that many previous-layer outputs to choose from. +fn build_graph(level: usize, width: usize, seed: u64) -> anyhow::Result<(TaskGraph, GraphShape)> { + let bytes_type = DataTypeDescriptor::Value(ValueTypeDescriptor::bytes()); + let tdl_context = TdlContext { + package: PACKAGE.to_owned(), + task_func: TASK_FUNC.to_owned(), + }; + + let mut graph = TaskGraph::new(None, None)?; + let mut total_tasks = 0usize; + + // Layer 0: all inputs come from the graph inputs. + let mut prev_layer: Vec = Vec::with_capacity(width); + for _ in 0..width { + let task_idx = graph.insert_task(TaskDescriptor { + tdl_context: tdl_context.clone(), + execution_policy: None, + inputs: vec![bytes_type.clone(); NUM_TASK_INPUT], + outputs: vec![bytes_type.clone()], + input_sources: None, + })?; + prev_layer.push(task_idx); + total_tasks += 1; + } + + // Inner layers: `NUM_TASK_INPUT` distinct random outputs from the previous layer per task. + let mut rng = Rng::new(seed); + for _ in 1..level { + let mut cur_layer = Vec::with_capacity(width); + for _ in 0..width { + let sources: Vec = + sample_distinct(&mut rng, NUM_TASK_INPUT, width) + .into_iter() + .map(|src| TaskInputOutputIndex { + task_idx: prev_layer[src], + position: 0, + }) + .collect(); + let task_idx = graph.insert_task(TaskDescriptor { + tdl_context: tdl_context.clone(), + execution_policy: None, + inputs: vec![bytes_type.clone(); NUM_TASK_INPUT], + outputs: vec![bytes_type.clone()], + input_sources: Some(sources), + })?; + cur_layer.push(task_idx); + total_tasks += 1; + } + prev_layer = cur_layer; + } + + let shape = GraphShape { + total_tasks, + depth: level, + }; + Ok((graph, shape)) +} + +/// Builds the graph inputs: `width * NUM_TASK_INPUT` byte vectors of `input_bytes` bytes each, one +/// per positional input of the layer-0 tasks (insertion order: task `i` consumes inputs +/// `NUM_TASK_INPUT * i .. NUM_TASK_INPUT * i + NUM_TASK_INPUT`). +/// +/// # Returns +/// +/// The graph input byte vectors. +/// +/// # Panics +/// +/// Panics on arithmetic overflow if `width * NUM_TASK_INPUT` overflows `usize`. +fn build_graph_inputs(width: usize, input_bytes: usize) -> Vec> { + let count = width + .checked_mul(NUM_TASK_INPUT) + .expect("width * NUM_TASK_INPUT overflow"); + vec![vec![0u8; input_bytes]; count] +} + +/// Polls the job state until it reaches a terminal state. +/// +/// # Returns +/// +/// The terminal [`JobState`] on success. +/// +/// # Errors +/// +/// Forwards [`SpiderClient::get_job_state`]'s return values on failure. +async fn poll_until_terminal(client: &SpiderClient, job_id: JobId) -> anyhow::Result { + loop { + let state = client + .get_job_state(job_id) + .await + .context("get_job_state")?; + if state.is_terminal() { + return Ok(state); + } + tokio::time::sleep(Duration::from_millis(500)).await; + } +} + +/// Builds the graph inputs and msgpack-encodes each into a [`TaskInput`]. +/// +/// # Returns +/// +/// The encoded task inputs on success. +/// +/// # Errors +/// +/// Forwards `rmp_serde::to_vec`'s return values on failure. +/// +/// # Panics +/// +/// Panics on arithmetic overflow if `width * NUM_TASK_INPUT` overflows `usize` (via +/// [`build_graph_inputs`]). +fn serialize_graph_inputs(width: usize, input_bytes: usize) -> anyhow::Result> { + let graph_inputs = build_graph_inputs(width, input_bytes); + graph_inputs + .iter() + .map(|value| { + Ok::(TaskInput::ValuePayload( + rmp_serde::to_vec(value).context("failed to serialize a graph input")?, + )) + }) + .collect() +} + +/// Verifies the final-layer outputs against the fixed payload, printing each to STDOUT when +/// `print_outputs` is set. +/// +/// # Errors +/// +/// Forwards [`SpiderClient::get_job_outputs`]'s return values on failure. +/// Returns an error if any output mismatches the expected payload. +async fn handle_succeeded( + client: &SpiderClient, + job_id: JobId, + width: usize, + print_outputs: bool, +) -> anyhow::Result<()> { + let outputs = client + .get_job_outputs(job_id) + .await + .context("get_job_outputs")?; + anyhow::ensure!( + outputs.len() == width, + "expected {width} graph outputs, got {}", + outputs.len(), + ); + + let mut mismatches = 0; + for (i, output) in outputs.iter().enumerate() { + let got: Vec = rmp_serde::from_slice(output) + .with_context(|| format!("failed to decode output {i}"))?; + if got != OUTPUT_PAYLOAD { + mismatches += 1; + } + if print_outputs { + // Inspection output goes to STDOUT (the results stream); everything else is on STDERR. + println!( + "output[{i}] len={} matches_expected={}", + got.len(), + got == OUTPUT_PAYLOAD, + ); + } + } + + eprintln!("outputs: {}/{} matched", width - mismatches, width); + if mismatches == 0 { + Ok(()) + } else { + Err(anyhow!( + "job succeeded but {mismatches}/{width} outputs mismatched the expected payload" + )) + } +} + +/// Prints a single labeled phase-timing line to STDERR, in milliseconds. +/// +/// The label is left-padded so the colons of consecutive lines align. +fn print_timing(label: &str, duration: Duration) { + eprintln!( + "[timing] {label:<26}: {:.1} ms", + duration.as_secs_f64() * 1000.0 + ); +} + +/// Prints a single labeled timing line to STDERR from a millisecond value (used for the analytic +/// ideal-runtime figures, which are not a measured `Duration`). +fn print_timing_ms(label: &str, ms: f64) { + eprintln!("[timing] {label:<26}: {ms:.1} ms"); +} + +/// End-to-end per-phase timings for a single nn benchmark run. +struct PhaseTimings { + graph_and_inputs: Duration, + connect_and_resource_group: Duration, + submit_and_start: Duration, + spider_execution: Duration, + post_processing: Duration, + total: Duration, +} + +impl PhaseTimings { + /// Prints the per-phase breakdown followed by the three headline rollups and the total to + /// STDERR. + /// + /// `query_processing` aggregates the graph/input construction, connection, and job-submission + /// phases (everything before the distributed execution begins). + fn print(&self) { + let query_processing = + self.graph_and_inputs + self.connect_and_resource_group + self.submit_and_start; + print_timing("graph_and_inputs", self.graph_and_inputs); + print_timing( + "connect_and_resource_group", + self.connect_and_resource_group, + ); + print_timing("submit_and_start", self.submit_and_start); + print_timing("spider_execution", self.spider_execution); + print_timing("post_processing", self.post_processing); + print_timing("== query_processing", query_processing); + print_timing("== spider_execution", self.spider_execution); + print_timing("== post_processing", self.post_processing); + print_timing("== total", self.total); + } +} + +/// Ideal runtime lower bound (in milliseconds) for the generated graph on `workers` workers. +/// +/// For a DAG of equal-duration tasks, no schedule can beat `max(critical_path, total_work / W)`: +/// the `total_work / W` term is the perfect-parallelism bound, and the `critical_path` term is the +/// longest dependency chain that must run serially. +fn ideal_runtime_ms(shape: &GraphShape, workers: usize) -> f64 { + // Convert through `u32` so the `usize` -> `f64` widening is lossless (avoids + // `clippy::cast_precision_loss`); the counts are bounded well below `u32::MAX`. + let total_tasks = + f64::from(u32::try_from(shape.total_tasks).expect("total task count fits in u32")); + let depth = f64::from(u32::try_from(shape.depth).expect("graph depth fits in u32")); + let workers = f64::from(u32::try_from(workers).expect("worker count fits in u32")); + let total_work_ms = total_tasks * TASK_DURATION_MS; + let critical_path_ms = depth * TASK_DURATION_MS; + let work_bound_ms = total_work_ms / workers; + critical_path_ms.max(work_bound_ms) +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let cli = Cli::parse(); + if cli.level == 0 || cli.width == 0 { + return Err(anyhow!("--level and --width must be >= 1")); + } + if cli.width < NUM_TASK_INPUT { + return Err(anyhow!( + "--width ({}) must be >= {} (the task's positional input count) for distinct sampling", + cli.width, + NUM_TASK_INPUT, + )); + } + let pool_size = NonZeroUsize::new(cli.pool_size).context("--pool-size must be >= 1")?; + + eprintln!( + "NN benchmark: level={} width={} input_bytes={}", + cli.level, cli.width, cli.input_bytes, + ); + + let total_start = Instant::now(); + + // (a) Graph build + input serialization. + let phase_start = Instant::now(); + let (graph, shape) = build_graph(cli.level, cli.width, cli.seed)?; + let task_inputs = serialize_graph_inputs(cli.width, cli.input_bytes)?; + let graph_and_inputs_duration = phase_start.elapsed(); + + // (b) Connect + register a per-run resource group. Use a unique external id per run so repeated + // runs (against a persistent MariaDB) do not collide with the resource-group-already-exists + // error. + let phase_start = Instant::now(); + let endpoint: Endpoint = cli + .endpoint + .parse() + .with_context(|| format!("invalid --endpoint {:?}", cli.endpoint))?; + let client = SpiderClient::connect(endpoint, pool_size) + .await + .context("failed to connect to the Spider storage service")?; + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .context("system clock before UNIX epoch")? + .as_nanos(); + let resource_group_id = client + .add_resource_group( + format!("huntsman-nn-{nanos}"), + RESOURCE_GROUP_PASSWORD.to_vec(), + ) + .await + .context("add_resource_group")?; + let connect_and_resource_group_duration = phase_start.elapsed(); + + // (c) submit_job + start_job. + let phase_start = Instant::now(); + let job_id = client + .submit_job(resource_group_id, &graph, task_inputs) + .await + .context("submit_job")?; + client.start_job(job_id).await.context("start_job")?; + let submit_and_start_duration = phase_start.elapsed(); + + eprintln!( + "Submitted nn_bench job: tasks={}, job_id={}", + shape.total_tasks, + job_id.get(), + ); + + // (d) start_job -> terminal: the execution wall time the benchmark cares about. + let phase_start = Instant::now(); + let state = poll_until_terminal(&client, job_id).await?; + let spider_execution_duration = phase_start.elapsed(); + + // (e) Decode + verify the final-layer outputs. + let phase_start = Instant::now(); + match state { + JobState::Succeeded => { + handle_succeeded(&client, job_id, cli.width, cli.print_outputs).await?; + } + JobState::Failed => { + let message = client + .get_job_error(job_id) + .await + .context("get_job_error")?; + return Err(anyhow!("job failed: {message}")); + } + other => { + return Err(anyhow!("job ended in unexpected state {other:?}")); + } + } + let post_processing_duration = phase_start.elapsed(); + + PhaseTimings { + graph_and_inputs: graph_and_inputs_duration, + connect_and_resource_group: connect_and_resource_group_duration, + submit_and_start: submit_and_start_duration, + spider_execution: spider_execution_duration, + post_processing: post_processing_duration, + total: total_start.elapsed(), + } + .print(); + + // Ideal runtime lower bound for the generated graph at each reference worker count. These are + // the floors to compare the measured `spider_execution` against -- the gap is Spider's + // scheduling/coordination overhead. + eprintln!("[timing] == ideal (lower bound):"); + for workers in IDEAL_WORKER_COUNTS { + print_timing_ms( + &format!("ideal {workers} workers"), + ideal_runtime_ms(&shape, workers), + ); + } + + let throughput = + f64::from(u32::try_from(shape.total_tasks).expect("total task count fits in u32")) + / spider_execution_duration.as_secs_f64(); + eprintln!("throughput: {throughput:.1} tasks/s"); + eprintln!( + "Job succeeded; all {} final-layer outputs match the expected payload.", + cli.width, + ); + + Ok(()) +} diff --git a/examples/huntsman/nn/tasks/Cargo.toml b/examples/huntsman/nn/tasks/Cargo.toml new file mode 100644 index 000000000..cf7b7533e --- /dev/null +++ b/examples/huntsman/nn/tasks/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "huntsman-nn-bench-tasks" +version = "0.1.0" +edition = "2024" +publish = false + +[lib] +crate-type = ["cdylib"] +name = "nn_bench" +path = "src/lib.rs" + +[dependencies] +serde = { version = "1.0.228", features = ["derive"] } +spider-tdl = { + path = "../../../../components/spider-tdl", + features = ["derive"] +} \ No newline at end of file diff --git a/examples/huntsman/nn/tasks/src/lib.rs b/examples/huntsman/nn/tasks/src/lib.rs new file mode 100644 index 000000000..21f0d6631 --- /dev/null +++ b/examples/huntsman/nn/tasks/src/lib.rs @@ -0,0 +1,140 @@ +//! Benchmark TDL package for the neural-network-shaped task graph. +//! +//! Exposes a single task, [`task_decl::sleep`], that mirrors one neuron in a layered NN benchmark: +//! it consumes 25 `bytes` (128-byte) inputs, sleeps for a fixed 10 ms to simulate compute cost, and +//! emits a fixed 128-byte output. The inputs are intentionally ignored -- the task "does no work" +//! beyond the sleep -- so the benchmark measures scheduling/execution overhead rather than real +//! computation. +//! +//! Each invocation logs a START and END line to stderr (captured by the execution manager into +//! `build/spider-run/em-logs/-.log`) carrying the job/task ids and a nanosecond +//! timestamp, so per-task start/end times can be recovered from the executor logs. + +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +/// Fixed sleep duration simulating a neuron's compute cost. +const SLEEP_DURATION: Duration = Duration::from_millis(10); + +/// Fixed 128-byte payload every invocation emits. +const OUTPUT_PAYLOAD: [u8; 128] = [0; 128]; + +/// Nanoseconds since the UNIX epoch, for the per-task log lines. +/// +/// # Panics +/// +/// Panics only if the system clock is before the UNIX epoch -- never in practice. +fn now_unix_nanos() -> u128 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock before UNIX epoch") + .as_nanos() +} + +mod task_decl { + // The `#[task]` macro generates a wrapper (`__sleep`) whose positional arity mirrors the + // user task's signature, so clippy's `too_many_arguments` (threshold 7) fires on the + // expansion. The 25-arity is a fixed property of the neuron model, not a code smell. The + // macro does not forward a function-level `#[allow]` to the generated wrapper, so the allow + // must be scoped at the module (which contains only this one task). Suppress here rather than + // restructuring the task or touching the derive macro. + #![allow(clippy::too_many_arguments)] + + // Alias `std::thread::sleep` to avoid colliding with the `sleep` marker struct the + // `#[task]` macro generates below (the macro names the marker after the function). + use std::thread::sleep as thread_sleep; + + use spider_tdl::{TaskContext, TdlError, task}; + + use super::{OUTPUT_PAYLOAD, SLEEP_DURATION, now_unix_nanos}; + + /// NN benchmark neuron: consumes 25 `bytes` inputs, sleeps 10 ms, emits a fixed 128-byte + /// output, and logs START/END timestamps to stderr. + /// + /// The 25 positional inputs model the neuron's 25 incoming data-flow edges. Their content is + /// ignored (the task does no real work); only their total byte count is logged to prove the + /// inputs were delivered. + #[task(name = "nn_bench::sleep")] + pub fn sleep( + ctx: TaskContext, + i0: Vec, + i1: Vec, + i2: Vec, + i3: Vec, + i4: Vec, + i5: Vec, + i6: Vec, + i7: Vec, + i8: Vec, + i9: Vec, + i10: Vec, + i11: Vec, + i12: Vec, + i13: Vec, + i14: Vec, + i15: Vec, + i16: Vec, + i17: Vec, + i18: Vec, + i19: Vec, + i20: Vec, + i21: Vec, + i22: Vec, + i23: Vec, + i24: Vec, + ) -> Result, TdlError> { + // Touch every input so the wrapper's parameters are used (and so the log proves the inputs + // were delivered); the content itself is irrelevant to the benchmark. + let input_bytes = i0.len() + + i1.len() + + i2.len() + + i3.len() + + i4.len() + + i5.len() + + i6.len() + + i7.len() + + i8.len() + + i9.len() + + i10.len() + + i11.len() + + i12.len() + + i13.len() + + i14.len() + + i15.len() + + i16.len() + + i17.len() + + i18.len() + + i19.len() + + i20.len() + + i21.len() + + i22.len() + + i23.len() + + i24.len(); + + let start_ns = now_unix_nanos(); + eprintln!( + "[nn_bench::sleep] START job={:?} task={:?} instance={} input_bytes={} t_ns={}", + ctx.job_id, ctx.task_id, ctx.task_instance_id, input_bytes, start_ns, + ); + + thread_sleep(SLEEP_DURATION); + + let end_ns = now_unix_nanos(); + eprintln!( + "[nn_bench::sleep] END job={:?} task={:?} instance={} t_ns={} dur_ns={}", + ctx.job_id, + ctx.task_id, + ctx.task_instance_id, + end_ns, + end_ns - start_ns, + ); + + Ok(OUTPUT_PAYLOAD.to_vec()) + } +} + +spider_tdl::register_tdl_package! { + package_name: "nn_bench", + tasks: [ + task_decl::sleep, + ], +} diff --git a/stack-doc.md b/stack-doc.md new file mode 100644 index 000000000..69a496266 --- /dev/null +++ b/stack-doc.md @@ -0,0 +1,206 @@ +# Running the Spider stack + +`tools/scripts/stack/run.py` brings up the whole Spider service stack locally -- a MariaDB +container plus the storage, scheduler, and N execution-manager Rust binaries -- in dependency +order and supervises them in the foreground. It is a standalone `uv` script, so it has no +install step beyond [uv] itself. + +## Prerequisites + +* The Rust release binaries must already be built. From the repo root: + + ```shell + task build:rust + ``` + + The binaries land in `build/rust-targets/release/` (the workspace's `CARGO_TARGET_DIR`). The + run script fails fast if any required binary is missing. +* The compiled task libraries must be staged so task executors can `dlopen` them. From the repo + root: + + ```shell + task build:packages + ``` + + This builds the workspace and copies each task library into + `build/tdl_packages//lib.so` (the `package_dir` configured in `spider.yaml`). + Run it after `task build:rust`. It is only needed when the jobs you submit reference a task + package (e.g. the `complex` package used by the example client below). +* [Docker] must be runnable by your user (used to start the MariaDB container). + +## Configuration + +There is one hand-written config, `tools/scripts/stack/spider.yaml` -- the single source of truth +for the whole stack. It factors shared values (MariaDB, gRPC endpoints, binary/run/package paths) +to the top and gives each service its own section listing every knob for that binary; knobs with +`#[serde(default)]` in the Rust schema are commented out so serde applies its own defaults. + +Each Rust binary still consumes its own `--config ` with its own serde schema, so at launch +`run.py` calls `tools/scripts/stack/generate.py` to derive three per-service configs from +`spider.yaml` and write them into `run_dir`: + +| Generated file | Consumed by | +|----------------------|-------------------------------| +| `gen-storage.yaml` | `spider_storage_grpc_server` | +| `gen-scheduler.yaml` | `spider_scheduler_grpc_server`| +| `gen-em.yaml` | `spider_execution_manager` | + +`run.py` then passes each generated file to its binary via `--config`. The generated files are +kept on disk (under the gitignored `build/` tree) for debugging -- inspect them to see exactly +what each binary was handed. Do not edit them; edit `spider.yaml` and regenerate. + +`generate.py` is also runnable standalone, which is the easiest way to check what the derivation +produces without launching the stack: + +```shell +uv run --script tools/scripts/stack/generate.py +uv run --script tools/scripts/stack/generate.py --output-dir /tmp +``` + +Paths in `spider.yaml` are resolved relative to the current working directory `run.py` is +launched from. The defaults assume that is the repository root. + +To run more tasks in parallel, increase the worker count -- each execution manager runs one +task-executor at a time. Either edit `workers:` in `spider.yaml` or pass `--workers`. + +The log level for every Rust service in the stack (storage, scheduler, execution managers, and the +task-executors they spawn) is controlled by `log_level:` in `spider.yaml`. It is forwarded to each +binary as `RUST_LOG`, so any `tracing` filter syntax is accepted -- e.g. `info`, `debug`, or +`spider_scheduler=debug,info` to raise one component while keeping the rest at `info`. Either edit +`log_level:` in `spider.yaml` or pass `--log-level`. + +> **Note:** `spider.yaml` writes the round-robin scheduler as `!round_robin` (a YAML tag), not as +> a `{ round_robin: { ... } }` map. The `yaml_serde` crate the server uses deserializes serde +> externally-tagged enums this way; a plain map will fail to parse. `generate.py` registers a +> matching representer so the generated `gen-scheduler.yaml` emits the same tag. + +## Running the stack + +```shell +uv run --script tools/scripts/stack/run.py +``` + +This generates the per-service configs, starts MariaDB, waits for it to accept connections, then +starts storage, waits for its port (`50051`), starts the scheduler, waits for its port (`50052`), +then launches the configured number of execution-manager workers. Once everything is up, it +supervises the services in the foreground. + +Useful flags: + +| Flag | Default | Description | +|---------------------|-----------------------------|------------------------------------------------------------------| +| `--config` | `tools/scripts/stack/spider.yaml` | Path to the global stack config. | +| `--workers N` | from config | Override the execution-manager worker count. | +| `--log-level L` | from config (`info`) | Override the RUST_LOG level (e.g. `info`, `debug`). | +| `--skip-mariadb` | off | Assume MariaDB is already running; do not start it. | +| `--teardown` | off | Also stop the MariaDB container when the run ends. | +| `--start-timeout S` | `30` | Seconds to wait for each service to become ready. | + +## Running workers on separate nodes + +`tools/scripts/stack/run_em.py` launches only execution-manager workers from the global config -- +no MariaDB, storage, or scheduler. Use it to run task-executor workers on a node separate from the +storage/scheduler services (e.g. when distributing workers across multiple nodes). The storage and +scheduler services must already be running on the scheduler node -- start them there with +`run.py --workers 0` (or any worker count, if that node should also run workers): + +```shell +# On the scheduler/storage node: +uv run --script tools/scripts/stack/run.py --workers 0 + +# On each worker node: +uv run --script tools/scripts/stack/run_em.py +``` + +`run_em.py` reads the same `spider.yaml`, calls `generate.py` to derive `gen-em.yaml` (the +storage/scheduler configs generate.py also writes are unused on a worker node), and launches the +configured number of execution managers. Each EM spawns a task-executor and registers with the +scheduler at `scheduler_endpoint`. Workers self-differentiate by a generated execution-manager ID, +so launching several EMs from one config on one node does not collide. + +For multi-node deployments, edit `spider.yaml` on each worker node so that: + +* `storage_endpoint` and `scheduler_endpoint` point at the scheduler/storage node (not `127.0.0.1`). +* `execution_manager.host` is this worker node's reachable IP (it is advertised to the + scheduler/storage as the EM's own address). + +The scheduler must listen on an address reachable from the worker nodes -- set +`scheduler_endpoint.host` on the scheduler node to its external IP (or `0.0.0.0`), not `127.0.0.1`. + +Useful flags: + +| Flag | Default | Description | +|---------------------|-----------------------------|------------------------------------------------------------------| +| `--config` | `tools/scripts/stack/spider.yaml` | Path to the global stack config. | +| `--workers N` | from config | Override the execution-manager worker count. | +| `--log-level L` | from config (`info`) | Override the RUST_LOG level (e.g. `info`, `debug`). | + +Press `Ctrl-C` (or send `SIGTERM`) to stop the workers. They are torn down in reverse launch +order; any that do not exit in time are `SIGKILL`ed. Per-worker logs land in the same +`build/spider-run/` layout as `run.py` -- `em-.log` per worker and `em-logs/-.log` +per task-executor. + +## Stopping + +Press `Ctrl-C` (or send `SIGTERM`) to stop the run. Services are torn down in reverse launch +order (execution managers -> scheduler -> storage); any that do not exit in time are `SIGKILL`ed. + +By default the MariaDB container is **left running** between runs so database state persists. +Pass `--teardown` to also stop and remove the MariaDB container when the run ends: + +```shell +uv run --script tools/scripts/stack/run.py --teardown +``` + +## Logs and generated configs + +`build/spider-run/` holds the runtime artifacts: + +* `gen-storage.yaml`, `gen-scheduler.yaml`, `gen-em.yaml` -- the generated per-service configs. +* `storage.log`, `scheduler.log`, `em-0.log`, `em-1.log`, ... -- per-service stdout/stderr + (truncated on each launch). +* `em-logs/-.log` -- per task-executor subprocess logs. + +This directory lives under the gitignored `build/` tree, so none of it is committed. + +## Example: a layered task graph with `huntsman-complex-client` + +`examples/huntsman/complex/client` is a small client binary that builds a "neural-network-shaped" +task graph out of the `complex` package's `complex::add` task and runs it against a live stack. The +graph has `--level` layers of `--width` tasks each: layer 0 takes its two inputs from the graph +inputs, and every inner task adds two outputs from the previous layer. Tasks within a layer are +independent, so `--width` controls how much parallelism the scheduler can exploit. After the job +finishes, the client decodes the final layer's outputs and checks them against an in-process +simulation of the same DAG, so a successful run proves the stack executed the graph correctly. + +The binary is built by `task build:packages` (which builds the whole workspace, including the +example crates) and lands in `build/rust-targets/release/huntsman-complex-client`. Run it against a +stack that is already up. In one shell, start the stack with enough workers to run a layer in +parallel: + +```shell +uv run --script tools/scripts/stack/run.py --workers 16 +``` + +In another shell, once the stack reports it is up, run the client: + +```shell +build/rust-targets/release/huntsman-complex-client --level 10 --width 16 +``` + +A width-16 graph keeps all 16 execution managers busy within each layer. On success the client +prints that every final-layer output matched the local simulation. Defaults are `--level 10` and +`--width 4`; pass `--help` for the full flag list. This example was validated with `--workers 16` +and `--width 16`. + +## Notes + +* The storage service creates its own database tables on connect (`CREATE TABLE IF NOT + EXISTS`), so no schema initialization step is needed or run. Do not pre-create tables from + the older `tools/scripts/mariadb/wolf/` schema -- a stale schema will cause storage to fail + with column-mismatch errors. +* The MariaDB container is started with `--rm`, so stopping it also discards its data -- each + `--teardown` run starts from a fresh database. + +[uv]: https://docs.astral.sh/uv/ +[Docker]: https://docs.docker.com/get-started/ \ No newline at end of file diff --git a/taskfiles/build.yaml b/taskfiles/build.yaml index cf6058384..73776d112 100644 --- a/taskfiles/build.yaml +++ b/taskfiles/build.yaml @@ -44,6 +44,33 @@ tasks: . "{{.G_RUST_TOOLCHAIN_ENV_FILE}}" cargo build --release --all-features + # Builds the workspace, then stages the compiled task libraries (cdylibs) into the on-disk layout + # the task executor dlopens: `${G_TDL_PACKAGES_DIR}//lib.so`. Run this + # before launching the stack with `run.py` so task executors can resolve packages referenced by + # submitted jobs. + packages: + dir: "{{.ROOT_DIR}}" + vars: + G_TDL_PACKAGES_DIR: "{{.G_BUILD_DIR}}/tdl_packages" + G_RUST_RELEASE_DIR: "{{.G_RUST_BUILD_DIR}}/release" + deps: + - "toolchains:rust" + cmd: |- + . "{{.G_RUST_TOOLCHAIN_ENV_FILE}}" + cargo build --release --workspace --all-features + mkdir -p "{{.G_TDL_PACKAGES_DIR}}/complex" \ + "{{.G_TDL_PACKAGES_DIR}}/integration_test_tasks" \ + "{{.G_TDL_PACKAGES_DIR}}/clp_search" \ + "{{.G_TDL_PACKAGES_DIR}}/nn_bench" + cp "{{.G_RUST_RELEASE_DIR}}/libhuntsman_complex.so" \ + "{{.G_TDL_PACKAGES_DIR}}/complex/libcomplex.so" + cp "{{.G_RUST_RELEASE_DIR}}/libintegration_test_tasks.so" \ + "{{.G_TDL_PACKAGES_DIR}}/integration_test_tasks/libintegration_test_tasks.so" + cp "{{.G_RUST_RELEASE_DIR}}/libclp_search.so" \ + "{{.G_TDL_PACKAGES_DIR}}/clp_search/libclp_search.so" + cp "{{.G_RUST_RELEASE_DIR}}/libnn_bench.so" \ + "{{.G_TDL_PACKAGES_DIR}}/nn_bench/libnn_bench.so" + tdl-generate-parsers: vars: CHECKSUM_FILE: "{{.G_BUILD_DIR}}/{{.TASK}}.md5" diff --git a/taskfiles/lint.yaml b/taskfiles/lint.yaml index 195473453..fc42b3e74 100644 --- a/taskfiles/lint.yaml +++ b/taskfiles/lint.yaml @@ -169,7 +169,8 @@ tasks: .github/ \ docs/ \ taskfile.yaml \ - taskfiles/ + taskfiles/ \ + tools/scripts/stack/ toml-check: cmds: diff --git a/tools/scripts/stack/generate.py b/tools/scripts/stack/generate.py new file mode 100755 index 000000000..225c80a0f --- /dev/null +++ b/tools/scripts/stack/generate.py @@ -0,0 +1,219 @@ +#!/usr/bin/env -S uv run --script +# /// script +# dependencies = ["pyyaml>=6.0"] +# /// +""" +Generate per-service Spider configs from a single global config. + +Reads the global ``spider.yaml`` and writes three serde-shaped per-service YAML files +(``gen-storage.yaml``, ``gen-scheduler.yaml``, ``gen-em.yaml``) that are each passed to their +Rust binary via ``--config``. Shared values (MariaDB, gRPC endpoints, binary/run/package paths) +are factored into the global config and defined once; this script derives each binary's exact +serde schema from them. + +Runnable standalone for debugging:: + + uv run --script tools/scripts/stack/generate.py + uv run --script tools/scripts/stack/generate.py --config path/to/spider.yaml --output-dir /tmp +""" + +import argparse +import logging +import sys +from pathlib import Path + +import yaml + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +logger = logging.getLogger(__name__) + +# Top-level fields the global config must provide. Missing any of these is a fail-fast error. +_REQUIRED_TOP_LEVEL = ( + "binary_dir", + "package_dir", + "run_dir", + "mariadb", + "storage_endpoint", + "scheduler_endpoint", + "storage", + "scheduler", + "execution_manager", +) + + +# ``yaml_serde`` (used by the Rust binaries) deserializes serde externally-tagged enums with a +# YAML ``!tag``. Register the scheduler's ``!round_robin`` variant on the safe loader so this +# script can read the global config, plus a marker dict + representer so it can emit it back into +# the generated scheduler config. +class _RoundRobin(dict): + """Mapping rendered as a ``!round_robin`` YAML tag for ``yaml_serde``.""" + + +yaml.SafeLoader.add_constructor( + "!round_robin", + lambda loader, node: loader.construct_mapping(node, deep=True), +) +yaml.SafeDumper.add_representer( + _RoundRobin, + lambda dumper, data: dumper.represent_mapping("!round_robin", data), +) + + +def _resolve(path: str) -> Path: + """Resolve ``path`` relative to the current working directory (absolute paths pass through).""" + return Path(path).resolve() if Path(path).is_absolute() else (Path.cwd() / path).resolve() + + +def _load_yaml(path: Path) -> dict: + """Load a YAML file with the ``!round_robin`` tag registered on the safe loader.""" + with path.open() as file: + return yaml.load(file, Loader=yaml.SafeLoader) + + +def _require(mapping: dict, key: str, where: str) -> object: + """Return ``mapping[key]`` or fail fast with a clear message about the missing field.""" + if key not in mapping: + logger.error("Missing required field '%s' in %s.", key, where) + sys.exit(1) + return mapping[key] + + +def _build_storage_config(config: dict) -> dict: + """Derive ``spider_storage_grpc_server``'s ``ServerConfig`` from the global config.""" + mariadb = config["mariadb"] + storage = config["storage"] + endpoint = config["storage_endpoint"] + runtime: dict = { + "db_config": { + "host": _require(mariadb, "host", "mariadb"), + "port": _require(mariadb, "port", "mariadb"), + "name": _require(mariadb, "database", "mariadb"), + "username": _require(mariadb, "username", "mariadb"), + "password": _require(mariadb, "password", "mariadb"), + "max_connections": _require(storage, "max_connections", "storage"), + }, + } + # These three sub-configs have #[serde(default)] in the Rust schema; only emit the ones the + # global config actually sets (the rest are commented out) so serde applies its own defaults. + for optional in ("ready_queue_config", "task_instance_pool_config", "job_cache_gc_config"): + if optional in storage: + runtime[optional] = storage[optional] + return { + "host": _require(endpoint, "host", "storage_endpoint"), + "port": _require(endpoint, "port", "storage_endpoint"), + "runtime": runtime, + } + + +def _build_scheduler_config(config: dict) -> dict: + """Derive ``spider_scheduler_grpc_server``'s ``ServerConfig`` from the global config.""" + scheduler = config["scheduler"] + runtime = _require(scheduler, "runtime", "scheduler") + storage_endpoint = config["storage_endpoint"] + scheduler_endpoint = config["scheduler_endpoint"] + out_runtime: dict = { + "scheduler": _RoundRobin(_require(runtime, "scheduler", "scheduler.runtime")), + "host": _require(scheduler_endpoint, "host", "scheduler_endpoint"), + "port": _require(scheduler_endpoint, "port", "scheduler_endpoint"), + } + # em_registry and stop_timeout_sec have #[serde(default)]; only emit if set. + if "em_registry" in runtime: + out_runtime["em_registry"] = runtime["em_registry"] + if "stop_timeout_sec" in runtime: + out_runtime["stop_timeout_sec"] = runtime["stop_timeout_sec"] + return { + "storage_endpoint": { + "host": _require(storage_endpoint, "host", "storage_endpoint"), + "port": _require(storage_endpoint, "port", "storage_endpoint"), + }, + "storage_connection_pool_size": _require( + scheduler, + "storage_connection_pool_size", + "scheduler", + ), + "runtime": out_runtime, + } + + +def _build_em_config(config: dict) -> dict: + """Derive ``spider_execution_manager``'s ``Config`` from the global config.""" + em = config["execution_manager"] + storage_endpoint = config["storage_endpoint"] + scheduler_endpoint = config["scheduler_endpoint"] + return { + "host": _require(em, "host", "execution_manager"), + "storage": { + "host": _require(storage_endpoint, "host", "storage_endpoint"), + "port": _require(storage_endpoint, "port", "storage_endpoint"), + }, + "scheduler": { + "host": _require(scheduler_endpoint, "host", "scheduler_endpoint"), + "port": _require(scheduler_endpoint, "port", "scheduler_endpoint"), + }, + "liveness": _require(em, "liveness", "execution_manager"), + # bin_path / log_dir are derived from the shared binary_dir / run_dir so they stay + # consistent with where run.py actually puts the binaries and logs. They are relative + # strings, resolved by the execution manager against its own working directory. + "task_executor": { + "bin_path": f"{config['binary_dir']}/spider-task-executor", + "package_dir": config["package_dir"], + "log_dir": f"{config['run_dir']}/em-logs", + }, + "connection_pool_size": _require(em, "connection_pool_size", "execution_manager"), + "scheduler_poll_wait_ms": _require(em, "scheduler_poll_wait_ms", "execution_manager"), + } + + +def main() -> int: + """Generate the three per-service configs and write them to the output directory.""" + parser = argparse.ArgumentParser( + description="Generate per-service Spider configs from a single global config.", + ) + parser.add_argument( + "--config", + type=str, + default="tools/scripts/stack/spider.yaml", + help="Path to the global stack config (default: %(default)s)", + ) + parser.add_argument( + "--output-dir", + type=str, + default=None, + help="Directory to write generated configs (default: the config's run_dir)", + ) + args = parser.parse_args() + + config = _load_yaml(_resolve(args.config)) + for section in _REQUIRED_TOP_LEVEL: + if section not in config: + logger.error("Missing required top-level field '%s' in global config.", section) + sys.exit(1) + + output_dir = _resolve(args.output_dir) if args.output_dir else _resolve(config["run_dir"]) + output_dir.mkdir(parents=True, exist_ok=True) + + services = { + "gen-storage.yaml": _build_storage_config(config), + "gen-scheduler.yaml": _build_scheduler_config(config), + "gen-em.yaml": _build_em_config(config), + } + for filename, data in services.items(): + path = output_dir / filename + with path.open("w") as file: + yaml.dump( + data, + file, + Dumper=yaml.SafeDumper, + default_flow_style=False, + sort_keys=False, + ) + logger.info("Wrote %s.", path) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/scripts/stack/run.py b/tools/scripts/stack/run.py new file mode 100755 index 000000000..221a2df6e --- /dev/null +++ b/tools/scripts/stack/run.py @@ -0,0 +1,375 @@ +#!/usr/bin/env -S uv run --script +# /// script +# dependencies = ["pyyaml>=6.0"] +# /// +""" +Run the Spider stack: MariaDB, storage, scheduler, and N execution managers. + +Services are launched in dependency order (storage -> scheduler -> execution managers), and each +is waited on until it accepts connections before the next is started. The script then supervises +the services in the foreground. + +Per-service configs are generated from a single global config (``spider.yaml``) by ``generate.py`` +at launch and written into ``run_dir``; each binary is then passed its generated file via +``--config``. + +Ctrl-C / SIGTERM tears down the services in reverse order. MariaDB is left running by default so +database state persists across run cycles; pass ``--teardown`` to also stop the MariaDB container +when the run ends. + +The Rust binaries must already be built -- run ``task build:rust`` first. The script fails fast +if any required binary is missing. +""" + +import argparse +import contextlib +import logging +import os +import signal +import socket +import subprocess +import sys +import threading +import time +from pathlib import Path + +import yaml + +# To silence Ruff S607: the absolute path of this executable may vary depending on the +# installation method. +_uv_executable = "uv" + +# The MariaDB helper scripts live next to this script (under tools/scripts/mariadb), so locate +# them relative to this file rather than the current working directory. +_MARIADB_SCRIPTS_DIR = Path(__file__).resolve().parent.parent / "mariadb" + +# generate.py lives next to this script and derives the per-service configs from the global +# config; locate it relative to this file rather than the current working directory. +_STACK_SCRIPTS_DIR = Path(__file__).resolve().parent + +_REQUIRED_BINARIES = ( + "spider_storage_grpc_server", + "spider_scheduler_grpc_server", + "spider_execution_manager", + "spider-task-executor", +) + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +logger = logging.getLogger(__name__) + +# (role, Popen) pairs, in launch order. Torn down in reverse. +_procs: list[tuple[str, subprocess.Popen]] = [] +# Set by the signal handler so the supervise loop breaks and the finally block tears down. +_exit_event = threading.Event() + +# ``yaml_serde`` (used by the Rust binaries) deserializes serde externally-tagged enums with a +# YAML ``!tag``. The global config contains the scheduler's ``!round_robin`` tag, so register it +# on the safe loader to let this script read the global config with PyYAML. +yaml.SafeLoader.add_constructor( + "!round_robin", + lambda loader, node: loader.construct_mapping(node, deep=True), +) + + +def _resolve(path: str) -> Path: + """Resolves a path from the config relative to the current working directory.""" + return Path(path).resolve() if Path(path).is_absolute() else (Path.cwd() / path).resolve() + + +def _load_yaml(path: Path) -> dict: + """Loads a YAML config file with the ``!round_robin`` tag registered on the safe loader.""" + with path.open() as file: + return yaml.load(file, Loader=yaml.SafeLoader) + + +def _port_open(host: str, port: int) -> bool: + """Returns whether ``host:port`` currently accepts a TCP connection.""" + try: + with socket.create_connection((host, port), timeout=1.0): + return True + except OSError: + return False + + +def _wait_for_port(host: str, port: int, timeout: float) -> bool: + """Blocks until ``host:port`` accepts a TCP connection or ``timeout`` elapses.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if _port_open(host, port): + return True + time.sleep(0.5) + return False + + +def _check_binaries(binary_dir: Path) -> None: + """Fails fast if any required release binary is missing.""" + missing = [name for name in _REQUIRED_BINARIES if not (binary_dir / name).exists()] + if missing: + logger.error( + "Missing binaries in %s: %s. Run `task build:rust` first.", + binary_dir, + ", ".join(missing), + ) + sys.exit(1) + + +def _ensure_mariadb(mariadb: dict, skip: bool) -> None: + """ + Starts the MariaDB container if it is not already running. + + The container creates the configured database on first start; the storage service creates + its own tables on connect, so no schema initialization is needed here. + """ + if skip: + logger.info("Skipping MariaDB startup (--skip-mariadb).") + return + + common_args = [ + "--port", + str(mariadb["port"]), + "--username", + mariadb["username"], + "--password", + mariadb["password"], + "--database", + mariadb["database"], + ] + + start_cmd = [ + _uv_executable, + "run", + "--script", + str(_MARIADB_SCRIPTS_DIR / "start.py"), + "--name", + mariadb["name"], + *common_args, + ] + result = subprocess.run(start_cmd, check=False) + # mariadb/start.py returns 1 when the container already exists; treat that as success. + if result.returncode == 1: + logger.info("MariaDB container %s already running.", mariadb["name"]) + elif result.returncode != 0: + logger.error("Failed to start MariaDB container (exit %d).", result.returncode) + sys.exit(1) + logger.info("MariaDB is ready.") + + +def _stop_mariadb(name: str) -> None: + """Stops the MariaDB container via the existing mariadb script.""" + result = subprocess.run( + [_uv_executable, "run", "--script", str(_MARIADB_SCRIPTS_DIR / "stop.py"), "--name", name], + check=False, + ) + if result.returncode != 0: + logger.error("MariaDB stop script exited with code %d.", result.returncode) + else: + logger.info("MariaDB container stopped.") + + +def _generate_configs(global_config: Path, run_dir: Path) -> None: + """Run ``generate.py`` to materialize the per-service configs into ``run_dir``.""" + result = subprocess.run( + [ + _uv_executable, + "run", + "--script", + str(_STACK_SCRIPTS_DIR / "generate.py"), + "--config", + str(global_config), + "--output-dir", + str(run_dir), + ], + check=False, + ) + if result.returncode != 0: + logger.error("generate.py failed (exit %d).", result.returncode) + sys.exit(1) + logger.info("Generated per-service configs in %s.", run_dir) + + +def _launch(role: str, args: list[str], log_file: Path, log_level: str) -> subprocess.Popen: + """Launches a service process in a new session and tees its stderr to a log file.""" + log_file.parent.mkdir(parents=True, exist_ok=True) + # Truncate per launch so each run's log reflects only the current attempt, not stale output + # from earlier failed runs. + log = log_file.open("wb") + # The Rust services read their log level from RUST_LOG; inject the resolved level so the stack + # is observable without forcing the caller to set the env var. Any value already present in + # os.environ is overwritten -- the config/CLI value is the single source of truth. The + # task-executor child processes inherit this env from their execution manager, so the level + # propagates to every binary in the stack. + env = {**os.environ, "RUST_LOG": log_level} + proc = subprocess.Popen( + args, + stdout=log, + stderr=subprocess.STDOUT, + start_new_session=True, + env=env, + ) + _procs.append((role, proc)) + logger.info("Started %s (pid %d): %s", role, proc.pid, " ".join(args)) + return proc + + +def _teardown(mariadb: dict, stop_mariadb: bool) -> None: + """SIGTERMs every running service in reverse launch order, then SIGKILLs stragglers.""" + for role, proc in reversed(_procs): + if proc.poll() is not None: + continue + logger.info("Stopping %s (pid %d).", role, proc.pid) + with contextlib.suppress(ProcessLookupError): + os.killpg(os.getpgid(proc.pid), signal.SIGTERM) + + deadline = time.monotonic() + 10.0 + for role, proc in reversed(_procs): + if proc.poll() is not None: + continue + remaining = max(0.0, deadline - time.monotonic()) + try: + proc.wait(timeout=remaining) + except subprocess.TimeoutExpired: + logger.warning("%s did not exit in time; sending SIGKILL.", role) + with contextlib.suppress(ProcessLookupError): + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + logger.info("All services stopped.") + + if stop_mariadb: + _stop_mariadb(mariadb["name"]) + else: + logger.info("Leaving MariaDB running. Pass --teardown to stop it on exit.") + + +def _on_signal(_signum: int, _frame: object) -> None: + """Sets the exit event so the supervise loop breaks and the finally block tears down.""" + _exit_event.set() + + +def _parse_args() -> argparse.Namespace: + """Builds and parses the command-line arguments for the stack runner.""" + parser = argparse.ArgumentParser(description="Run the Spider stack.") + parser.add_argument( + "--config", + type=str, + default="tools/scripts/stack/spider.yaml", + help="Path to the top-level stack config (default: %(default)s)", + ) + parser.add_argument( + "--workers", + type=int, + default=None, + help="Override the worker count from the config", + ) + parser.add_argument( + "--log-level", + type=str, + default=None, + help="Override the RUST_LOG level from the config (e.g. info, debug)", + ) + parser.add_argument( + "--skip-mariadb", + action="store_true", + help="Assume MariaDB is already running and initialized; do not start it", + ) + parser.add_argument( + "--teardown", + action="store_true", + help="Also stop the MariaDB container when the run ends", + ) + parser.add_argument( + "--start-timeout", + type=float, + default=30.0, + help="Seconds to wait for each service to become ready (default: %(default)s)", + ) + return parser.parse_args() + + +def main() -> int: + """Main.""" + args = _parse_args() + + global_config = _resolve(args.config) + config = _load_yaml(global_config) + mariadb = config["mariadb"] + workers = args.workers if args.workers is not None else config["workers"] + log_level = args.log_level if args.log_level is not None else config.get("log_level", "info") + logger.info("Log level: %s", log_level) + binary_dir = _resolve(config["binary_dir"]) + run_dir = _resolve(config["run_dir"]) + run_dir.mkdir(parents=True, exist_ok=True) + storage_endpoint = config["storage_endpoint"] + scheduler_endpoint = config["scheduler_endpoint"] + + _check_binaries(binary_dir) + _generate_configs(global_config, run_dir) + _ensure_mariadb(mariadb, args.skip_mariadb) + + storage_cfg_path = run_dir / "gen-storage.yaml" + scheduler_cfg_path = run_dir / "gen-scheduler.yaml" + em_cfg_path = run_dir / "gen-em.yaml" + + storage_args = [ + str(binary_dir / "spider_storage_grpc_server"), + "--config", + str(storage_cfg_path), + ] + _launch("storage", storage_args, run_dir / "storage.log", log_level) + if not _wait_for_port( + str(storage_endpoint["host"]), + storage_endpoint["port"], + args.start_timeout, + ): + logger.error("Storage did not become ready in %ss.", args.start_timeout) + _teardown(mariadb, args.teardown) + return 1 + logger.info("Storage is ready.") + + scheduler_args = [ + str(binary_dir / "spider_scheduler_grpc_server"), + "--config", + str(scheduler_cfg_path), + ] + _launch("scheduler", scheduler_args, run_dir / "scheduler.log", log_level) + if not _wait_for_port( + str(scheduler_endpoint["host"]), + scheduler_endpoint["port"], + args.start_timeout, + ): + logger.error("Scheduler did not become ready in %ss.", args.start_timeout) + _teardown(mariadb, args.teardown) + return 1 + logger.info("Scheduler is ready.") + + em_args = [ + str(binary_dir / "spider_execution_manager"), + "--config", + str(em_cfg_path), + ] + for i in range(workers): + _launch(f"em-{i}", em_args, run_dir / f"em-{i}.log", log_level) + # Give each EM a moment to register before launching the next, so the scheduler + # sees them arrive in order. + time.sleep(1.0) + logger.info("Launched %d execution-manager worker(s).", workers) + logger.info("Stack is up. Press Ctrl-C to stop.") + + signal.signal(signal.SIGINT, _on_signal) + signal.signal(signal.SIGTERM, _on_signal) + try: + while not _exit_event.is_set(): + for role, proc in _procs: + if proc.poll() is not None: + logger.error("%s exited unexpectedly (code %d).", role, proc.returncode) + return 1 + _exit_event.wait(1.0) + finally: + _teardown(mariadb, args.teardown) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/scripts/stack/run_em.py b/tools/scripts/stack/run_em.py new file mode 100755 index 000000000..0781d1fd3 --- /dev/null +++ b/tools/scripts/stack/run_em.py @@ -0,0 +1,245 @@ +#!/usr/bin/env -S uv run --script +# /// script +# dependencies = ["pyyaml>=6.0"] +# /// +""" +Run Spider execution-manager workers only. + +Used to run task-executor workers on a node separate from the storage/scheduler services (e.g. when +distributing workers across multiple nodes). Launches the configured number of execution managers +from the global stack config; each spawns a task-executor and registers with the scheduler at +``scheduler_endpoint``. The storage and scheduler services must already be running -- start them +on the scheduler node via ``run.py`` first. + +The per-service EM config is generated from the global config by ``generate.py`` at launch (only +``gen-em.yaml`` is consumed here; the storage/scheduler configs it also writes are unused on a +worker node). Ctrl-C / SIGTERM tears down the workers in reverse launch order. + +For multi-node deployments, point the global config's ``storage_endpoint``/``scheduler_endpoint`` +at the scheduler/storage node and set ``execution_manager.host`` to this node's reachable IP before +running this script. Workers self-differentiate by a generated execution-manager ID, so launching +several EMs from one config on one node does not collide. + +The Rust binaries must already be built -- run ``task build:rust`` first. The script fails fast if +any required binary is missing. +""" + +import argparse +import contextlib +import logging +import os +import signal +import subprocess +import sys +import threading +import time +from pathlib import Path + +import yaml + +# To silence Ruff S607: the absolute path of this executable may vary depending on the +# installation method. +_uv_executable = "uv" + +# generate.py lives next to this script and derives the per-service configs from the global +# config; locate it relative to this file rather than the current working directory. +_STACK_SCRIPTS_DIR = Path(__file__).resolve().parent + +# Only the execution manager and the task-executor it spawns are needed on a worker node. +_REQUIRED_BINARIES = ( + "spider_execution_manager", + "spider-task-executor", +) + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +logger = logging.getLogger(__name__) + +# (role, Popen) pairs, in launch order. Torn down in reverse. +_procs: list[tuple[str, subprocess.Popen]] = [] +# Set by the signal handler so the supervise loop breaks and the finally block tears down. +_exit_event = threading.Event() + +# ``yaml_serde`` (used by the Rust binaries) deserializes serde externally-tagged enums with a +# YAML ``!tag``. The global config contains the scheduler's ``!round_robin`` tag, so register it +# on the safe loader to let this script read the global config with PyYAML. +yaml.SafeLoader.add_constructor( + "!round_robin", + lambda loader, node: loader.construct_mapping(node, deep=True), +) + + +def _resolve(path: str) -> Path: + """Resolves a path from the config relative to the current working directory.""" + return Path(path).resolve() if Path(path).is_absolute() else (Path.cwd() / path).resolve() + + +def _load_yaml(path: Path) -> dict: + """Loads a YAML config file with the ``!round_robin`` tag registered on the safe loader.""" + with path.open() as file: + return yaml.load(file, Loader=yaml.SafeLoader) + + +def _check_binaries(binary_dir: Path) -> None: + """Fails fast if any required release binary is missing.""" + missing = [name for name in _REQUIRED_BINARIES if not (binary_dir / name).exists()] + if missing: + logger.error( + "Missing binaries in %s: %s. Run `task build:rust` first.", + binary_dir, + ", ".join(missing), + ) + sys.exit(1) + + +def _generate_configs(global_config: Path, run_dir: Path) -> None: + """ + Run ``generate.py`` to materialize the per-service configs into ``run_dir``. + + Only ``gen-em.yaml`` is consumed here; generate.py also writes the storage/scheduler configs, + which are harmless on a worker node. + """ + result = subprocess.run( + [ + _uv_executable, + "run", + "--script", + str(_STACK_SCRIPTS_DIR / "generate.py"), + "--config", + str(global_config), + "--output-dir", + str(run_dir), + ], + check=False, + ) + if result.returncode != 0: + logger.error("generate.py failed (exit %d).", result.returncode) + sys.exit(1) + logger.info("Generated per-service configs in %s.", run_dir) + + +def _launch(role: str, args: list[str], log_file: Path, log_level: str) -> subprocess.Popen: + """Launches a service process in a new session and tees its stderr to a log file.""" + log_file.parent.mkdir(parents=True, exist_ok=True) + # Truncate per launch so each run's log reflects only the current attempt, not stale output + # from earlier failed runs. + log = log_file.open("wb") + # The Rust services read their log level from RUST_LOG; inject the resolved level so the stack + # is observable without forcing the caller to set the env var. Any value already present in + # os.environ is overwritten -- the config/CLI value is the single source of truth. The + # task-executor child processes inherit this env from their execution manager, so the level + # propagates to every binary in the stack. + env = {**os.environ, "RUST_LOG": log_level} + proc = subprocess.Popen( + args, + stdout=log, + stderr=subprocess.STDOUT, + start_new_session=True, + env=env, + ) + _procs.append((role, proc)) + logger.info("Started %s (pid %d): %s", role, proc.pid, " ".join(args)) + return proc + + +def _teardown() -> None: + """SIGTERMs every running worker in reverse launch order, then SIGKILLs stragglers.""" + for role, proc in reversed(_procs): + if proc.poll() is not None: + continue + logger.info("Stopping %s (pid %d).", role, proc.pid) + with contextlib.suppress(ProcessLookupError): + os.killpg(os.getpgid(proc.pid), signal.SIGTERM) + + deadline = time.monotonic() + 10.0 + for role, proc in reversed(_procs): + if proc.poll() is not None: + continue + remaining = max(0.0, deadline - time.monotonic()) + try: + proc.wait(timeout=remaining) + except subprocess.TimeoutExpired: + logger.warning("%s did not exit in time; sending SIGKILL.", role) + with contextlib.suppress(ProcessLookupError): + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + logger.info("All workers stopped.") + + +def _on_signal(_signum: int, _frame: object) -> None: + """Sets the exit event so the supervise loop breaks and the finally block tears down.""" + _exit_event.set() + + +def _parse_args() -> argparse.Namespace: + """Builds and parses the command-line arguments for the worker launcher.""" + parser = argparse.ArgumentParser(description="Run Spider execution-manager workers only.") + parser.add_argument( + "--config", + type=str, + default="tools/scripts/stack/spider.yaml", + help="Path to the top-level stack config (default: %(default)s)", + ) + parser.add_argument( + "--workers", + type=int, + default=None, + help="Override the worker count from the config", + ) + parser.add_argument( + "--log-level", + type=str, + default=None, + help="Override the RUST_LOG level from the config (e.g. info, debug)", + ) + return parser.parse_args() + + +def main() -> int: + """Main.""" + args = _parse_args() + + global_config = _resolve(args.config) + config = _load_yaml(global_config) + workers = args.workers if args.workers is not None else config["workers"] + log_level = args.log_level if args.log_level is not None else config.get("log_level", "info") + logger.info("Log level: %s", log_level) + binary_dir = _resolve(config["binary_dir"]) + run_dir = _resolve(config["run_dir"]) + run_dir.mkdir(parents=True, exist_ok=True) + + _check_binaries(binary_dir) + _generate_configs(global_config, run_dir) + + em_cfg_path = run_dir / "gen-em.yaml" + em_args = [ + str(binary_dir / "spider_execution_manager"), + "--config", + str(em_cfg_path), + ] + for i in range(workers): + _launch(f"em-{i}", em_args, run_dir / f"em-{i}.log", log_level) + # Give each EM a moment to register before launching the next, so the scheduler + # sees them arrive in order. + time.sleep(1.0) + logger.info("Launched %d execution-manager worker(s).", workers) + logger.info("Workers are up. Press Ctrl-C to stop.") + + signal.signal(signal.SIGINT, _on_signal) + signal.signal(signal.SIGTERM, _on_signal) + try: + while not _exit_event.is_set(): + for role, proc in _procs: + if proc.poll() is not None: + logger.error("%s exited unexpectedly (code %d).", role, proc.returncode) + return 1 + _exit_event.wait(1.0) + finally: + _teardown() + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/scripts/stack/spider.yaml b/tools/scripts/stack/spider.yaml new file mode 100644 index 000000000..9f067f8cd --- /dev/null +++ b/tools/scripts/stack/spider.yaml @@ -0,0 +1,114 @@ +# Global Spider stack configuration -- single source of truth. +# +# `tools/scripts/stack/generate.py` derives the three per-service configs consumed by each Rust +# binary (`gen-storage.yaml`, `gen-scheduler.yaml`, `gen-em.yaml`) from this file; `run.py` calls +# generate.py at launch and passes each generated file to its binary via `--config`. Edit this +# file to configure the stack; do not edit the generated files. +# +# Paths are resolved relative to the current working directory the stack is launched from. The +# defaults below assume that is the repository root. +# +# Shared values (MariaDB, gRPC endpoints, binary/run/package paths) are defined once here. Each +# per-service section lists every knob for that binary; the ones with `#[serde(default)]` in the +# Rust schema are commented out so serde applies its own defaults. Uncomment and edit to override. + +# Number of execution-manager workers to launch. Each execution manager runs one task-executor at +# a time, so increase this to run more tasks in parallel. +workers: 16 + +# Log level for every Rust service in the stack (storage, scheduler, execution managers, and the +# task-executors they spawn). Forwarded to each binary as RUST_LOG, so any tracing filter syntax is +# accepted (e.g. `info`, `debug`, `spider_scheduler=debug,info`). Override per run with --log-level. +log_level: "info" + +# Directory containing the built release binaries. Must already exist (run `task build:rust` +# first; output goes to the workspace CARGO_TARGET_DIR, `build/rust-targets`). run.py fails fast +# if any required binary is missing. +binary_dir: "build/rust-targets/release" + +# Directory for runtime artifacts (per-service logs, per-executor logs, generated configs). +# Created if missing; lives under the gitignored `build/` tree so artifacts are not committed. +run_dir: "build/spider-run" + +# Directory where `task build:packages` stages compiled task libraries as +# `/lib.so`. Task executors dlopen a package from here when a task referencing it +# runs. Run `task build:packages` before launching the stack so referenced packages are staged. +package_dir: "build/tdl_packages/" + +# MariaDB container settings. Feeds both the container start (run.py -> mariadb/start.py) and the +# storage service's `db_config`. The storage service creates its own tables on connect, so no +# schema initialization is performed. +mariadb: + name: "mariadb-spider-dev" + host: "127.0.0.1" + port: 3306 + username: "spider-user" + password: "spider-password" + database: "spider-db" + +# Shared gRPC endpoints. The storage service listens on `storage_endpoint` (the scheduler and the +# execution managers connect to it); the scheduler listens on `scheduler_endpoint` (the execution +# managers connect to it). Each per-service config derives its listen address and its client +# targets from these. +storage_endpoint: + host: "127.0.0.1" + port: 50051 +scheduler_endpoint: + host: "127.0.0.1" + port: 50052 + +# --- storage (spider_storage_grpc_server) --------------------------------------------- +# listen host/port <- storage_endpoint; db_config <- mariadb (+ storage.max_connections). The +# three sub-configs below have `#[serde(default)]`; they are commented out so serde applies its +# own defaults. Uncomment to override. +storage: + max_connections: 64 + ready_queue_config: + task_capacity: 1048576 + commit_capacity: 256 + cleanup_capacity: 256 + # task_instance_pool_config: + # execution_manager_stale_cutoff_sec: 30 + # gc_interval_sec: 5 + # message_channel_capacity: 1024 + # job_cache_gc_config: + # terminated_job_retention_sec: 300 + # gc_interval_sec: 10 + +# --- scheduler (spider_scheduler_grpc_server) ----------------------------------------- +# storage_endpoint <- shared; runtime.host/port <- scheduler_endpoint. `em_registry` and +# `stop_timeout_sec` have `#[serde(default)]`; commented out. +# +# NOTE: `runtime.scheduler` is a serde externally-tagged enum. The `yaml_serde` crate used by the +# server deserializes such enums using a YAML `!tag`, so the round-robin variant is written as +# `!round_robin`, not as a plain `{ round_robin: { ... } }` map. +scheduler: + storage_connection_pool_size: 4 + runtime: + # em_registry: + # dead_em_cutoff_sec: 60 + # liveness_tracking_interval_ms: 1000 + # stop_timeout_sec: 30 + scheduler: !round_robin + active_job_queue_capacity: 64 + dispatch_queue_capacity: 64 + ready_task_capacity: 65536 + commit_ready_task_capacity: 256 + cleanup_ready_task_capacity: 256 + storage_poll_timeout_ms: 10 + tick_interval_ms: 5 + finalizing_job_expiration_timeout_sec: 300 + +# --- execution_manager (spider_execution_manager) ------------------------------------- +# storage/scheduler client endpoints <- shared; task_executor.bin_path <- binary_dir + +# "/spider-task-executor"; task_executor.log_dir <- run_dir + "/em-logs"; task_executor.package_dir +# <- shared package_dir. This config has no `#[serde(default)]` fields -- every field is required. +# The same generated config is reused by every execution-manager worker; workers self-differentiate +# by a generated execution-manager ID, so identical configs do not collide. +execution_manager: + host: "127.0.0.1" + connection_pool_size: 4 + scheduler_poll_wait_ms: 1000 + liveness: + storage_heartbeat_interval_sec: 5 + scheduler_heartbeat_interval_sec: 5