BE-629: Add spherical k-means entity clustering endpoint via /entities/embeddings/clusters

.github/workflows/codspeed.yml

@@ -70,6 +70,9 @@ jobs:
         with:
           scope: ${{ matrix.name }}
 
+      - name: Warm up repository
+        uses: ./.github/actions/warm-up-repo
+
       - name: Build the benchmark target
         run: turbo run build:codspeed --filter=${{ matrix.name }}
 

Cargo.toml

@@ -224,6 +224,7 @@ quote                              = { version = "1.0.41", default-features = fa
 rand                               = { version = "0.10.0", default-features = false }
 rand_core                          = { version = "0.10.0", default-features = false }
 rand_distr                         = { version = "0.6.0", default-features = false }
+rand_xoshiro                       = { version = "0.8.1" }
 rapidfuzz                          = { version = "0.5.0", default-features = false }
 ratatui                            = { version = "0.30.0" }
 rayon                              = { version = "1.11.0", default-features = false }

apps/hash-graph/Cargo.toml

@@ -34,6 +34,7 @@ futures = { workspace = true }
 jsonwebtoken = { workspace = true }
 mimalloc = { workspace = true }
 multiaddr = { workspace = true }
+rayon = { workspace = true }
 regex = { workspace = true }
 reqwest = { workspace = true, features = ["rustls"] }
 simple-mermaid = { workspace = true }

apps/hash-graph/src/args.rs

@@ -7,7 +7,7 @@ use clap::{
 };
 use hash_telemetry::TracingConfig;
 
-use crate::subcommand::Subcommand;
+use crate::subcommand::{Subcommand, WorkerThreads};
 
 /// Arguments passed to the program.
 #[derive(Debug, Parser)]
@@ -16,6 +16,19 @@ pub struct Args {
     #[clap(flatten)]
     pub tracing_config: TracingConfig,
 
+    /// Number of threads in the global worker pool used for CPU-bound work such as entity
+    /// clustering.
+    ///
+    /// Accepts a fixed count (e.g. `4`) or a count relative to the available CPU cores: `n` for
+    /// all cores, `n/2` for half, `n/4` for a quarter, and so on.
+    #[clap(
+        long,
+        global = true,
+        default_value_t,
+        env = "HASH_GRAPH_WORKER_THREADS"
+    )]
+    pub worker_threads: WorkerThreads,
+
     /// Specify a subcommand to run.
     #[command(subcommand)]
     pub subcommand: Subcommand,

apps/hash-graph/src/main.rs

@@ -30,9 +30,10 @@ fn main() -> Result<(), Report<GraphError>> {
     let Args {
         subcommand,
         tracing_config,
+        worker_threads,
     } = Args::parse_args();
 
     let _sentry_guard = init(&tracing_config.sentry, release_name!());
 
-    subcommand.execute(tracing_config)
+    subcommand.execute(tracing_config, worker_threads)
 }

apps/hash-graph/src/subcommand/mod.rs

@@ -6,15 +6,28 @@ mod server;
 mod snapshot;
 mod type_fetcher;
 
-use core::time::Duration;
-use std::time::Instant;
+use core::{fmt, num::NonZero, str::FromStr, time::Duration};
+use std::{sync::Once, thread::available_parallelism, time::Instant};
 
 use clap::Parser;
 use error_stack::{Report, ensure};
 use hash_telemetry::{TracingConfig, init_tracing};
 use tokio::time::sleep;
 use tokio_util::{sync::CancellationToken, task::TaskTracker};
 
+pub use self::{
+    admin_server::{AdminServerArgs, admin_server},
+    completions::{CompletionsArgs, completions},
+    migrate::{MigrateArgs, migrate},
+    server::{ServerArgs, server},
+    snapshot::{SnapshotArgs, snapshot},
+    type_fetcher::{TypeFetcherArgs, type_fetcher},
+};
+use crate::{
+    error::{GraphError, HealthcheckError},
+    subcommand::reindex_cache::{ReindexCacheArgs, reindex_cache},
+};
+
 /// Drop guard that fires the `abort` token when a server task exits unexpectedly.
 ///
 /// "Unexpectedly" means the `shutdown` token has not been cancelled yet. This covers both
@@ -87,6 +100,85 @@ impl ServerLifecycle {
     }
 }
 
+/// Number of threads for the global worker pool used for CPU-bound work.
+///
+/// Parses either a fixed thread count (e.g. `4`) or a count relative to the number of available
+/// CPU cores: `n` for all cores, `n/2` for half of them, `n/4` for a quarter, and so on.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum WorkerThreads {
+    /// The available CPU cores divided by the given divisor (`n`, `n/2`, `n/4`, ...).
+    Cores { divisor: NonZero<usize> },
+    /// A fixed number of threads.
+    Fixed(NonZero<usize>),
+}
+
+impl WorkerThreads {
+    /// Resolves to a concrete thread count, clamped to at least one thread.
+    #[expect(
+        clippy::integer_division,
+        reason = "Deriving a thread count from the core count is inherently lossy."
+    )]
+    fn resolve(self) -> NonZero<usize> {
+        match self {
+            Self::Fixed(threads) => threads,
+            Self::Cores { divisor } => available_parallelism()
+                .ok()
+                .and_then(|cores| NonZero::new(cores.get() / divisor))
+                .unwrap_or(NonZero::<usize>::MIN),
+        }
+    }
+}
+
+impl Default for WorkerThreads {
+    fn default() -> Self {
+        const HALF: NonZero<usize> = NonZero::new(2).expect("two should be non-zero");
+        Self::Cores { divisor: HALF }
+    }
+}
+
+impl fmt::Display for WorkerThreads {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match *self {
+            Self::Cores { divisor } if divisor == NonZero::<usize>::MIN => fmt.write_str("n"),
+            Self::Cores { divisor } => write!(fmt, "n/{divisor}"),
+            Self::Fixed(threads) => write!(fmt, "{threads}"),
+        }
+    }
+}
+
+/// Error returned when parsing a [`WorkerThreads`] value fails.
+#[derive(Debug)]
+pub struct ParseWorkerThreadsError;
+
+impl fmt::Display for ParseWorkerThreadsError {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt.write_str("expected a positive integer, `n`, or `n/<divisor>` (e.g. `4`, `n`, `n/2`)")
+    }
+}
+
+impl core::error::Error for ParseWorkerThreadsError {}
+
+impl FromStr for WorkerThreads {
+    type Err = ParseWorkerThreadsError;
+
+    fn from_str(value: &str) -> Result<Self, Self::Err> {
+        match value.strip_prefix(['n', 'N']) {
+            Some("") => Ok(Self::Cores {
+                divisor: NonZero::<usize>::MIN,
+            }),
+            Some(rest) => rest
+                .strip_prefix('/')
+                .and_then(|divisor| divisor.parse().ok())
+                .map(|divisor| Self::Cores { divisor })
+                .ok_or(ParseWorkerThreadsError),
+            None => value
+                .parse()
+                .map(Self::Fixed)
+                .map_err(|_error: core::num::ParseIntError| ParseWorkerThreadsError),
+        }
+    }
+}
+
 /// Shared healthcheck arguments for all server subcommands.
 #[derive(Debug, Clone, Parser)]
 pub(crate) struct HealthcheckArgs {
@@ -103,19 +195,6 @@ pub(crate) struct HealthcheckArgs {
     pub timeout: Option<u64>,
 }
 
-pub use self::{
-    admin_server::{AdminServerArgs, admin_server},
-    completions::{CompletionsArgs, completions},
-    migrate::{MigrateArgs, migrate},
-    server::{ServerArgs, server},
-    snapshot::{SnapshotArgs, snapshot},
-    type_fetcher::{TypeFetcherArgs, type_fetcher},
-};
-use crate::{
-    error::{GraphError, HealthcheckError},
-    subcommand::reindex_cache::{ReindexCacheArgs, reindex_cache},
-};
-
 /// Subcommand for the program.
 #[derive(Debug, clap::Subcommand)]
 pub enum Subcommand {
@@ -145,7 +224,17 @@ fn block_on(
     future: impl Future<Output = Result<(), Report<GraphError>>>,
     service_name: &'static str,
     tracing_config: TracingConfig,
+    worker_threads: WorkerThreads,
 ) -> Result<(), Report<GraphError>> {
+    static THREAD_POOL: Once = Once::new();
+    THREAD_POOL.call_once(|| {
+        rayon::ThreadPoolBuilder::new()
+            .num_threads(worker_threads.resolve().get())
+            .thread_name(|index| format!("rayon-{index}"))
+            .build_global()
+            .expect("rayon pool should be initialized exactly once");
+    });
+
     tokio::runtime::Builder::new_multi_thread()
         .enable_all()
         .build()
@@ -159,24 +248,49 @@ fn block_on(
 }
 
 impl Subcommand {
-    pub(crate) fn execute(self, tracing_config: TracingConfig) -> Result<(), Report<GraphError>> {
+    pub(crate) fn execute(
+        self,
+        tracing_config: TracingConfig,
+        worker_threads: WorkerThreads,
+    ) -> Result<(), Report<GraphError>> {
         match self {
-            Self::Server(args) => block_on(server(*args), "Graph API", tracing_config),
-            Self::AdminServer(args) => {
-                block_on(admin_server(*args), "Graph Admin API", tracing_config)
-            }
-            Self::Migrate(args) => block_on(migrate(*args), "Graph Migrations", tracing_config),
-            Self::TypeFetcher(args) => {
-                block_on(type_fetcher(*args), "Type Fetcher", tracing_config)
+            Self::Server(args) => {
+                block_on(server(*args), "Graph API", tracing_config, worker_threads)
             }
+            Self::AdminServer(args) => block_on(
+                admin_server(*args),
+                "Graph Admin API",
+                tracing_config,
+                worker_threads,
+            ),
+            Self::Migrate(args) => block_on(
+                migrate(*args),
+                "Graph Migrations",
+                tracing_config,
+                worker_threads,
+            ),
+            Self::TypeFetcher(args) => block_on(
+                type_fetcher(*args),
+                "Type Fetcher",
+                tracing_config,
+                worker_threads,
+            ),
             Self::Completions(ref args) => {
                 completions(args);
                 Ok(())
             }
-            Self::Snapshot(args) => block_on(snapshot(*args), "Graph Snapshot", tracing_config),
-            Self::ReindexCache(args) => {
-                block_on(reindex_cache(*args), "Graph Indexer", tracing_config)
-            }
+            Self::Snapshot(args) => block_on(
+                snapshot(*args),
+                "Graph Snapshot",
+                tracing_config,
+                worker_threads,
+            ),
+            Self::ReindexCache(args) => block_on(
+                reindex_cache(*args),
+                "Graph Indexer",
+                tracing_config,
+                worker_threads,
+            ),
         }
     }
 }

apps/hash-graph/src/subcommand/server.rs

@@ -16,8 +16,8 @@ use harpc_server::Server;
 use hash_codec::bytes::JsonLinesEncoder;
 use hash_graph_api::{
     rest::{
-        ApiConfig, QueryLogger, RestApiStore, RestRouterDependencies, hashql::CompilerContext,
-        rest_api_router,
+        ApiConfig, QueryLogger, RestApiStore, RestRouterDependencies, entity::ClusteringContext,
+        hashql::CompilerContext, rest_api_router,
     },
     rpc::Dependencies,
 };
@@ -239,6 +239,13 @@ pub struct ServerConfig {
     #[clap(flatten)]
     pub compiler: CompilerConfig,
 
+    /// Maximum number of entity-clustering requests processed at the same time.
+    ///
+    /// Excess requests wait until a slot frees up. If not set, the number of concurrent
+    /// clustering requests is unbounded.
+    #[clap(long, env = "HASH_GRAPH_CLUSTERING_CONCURRENCY_LIMIT")]
+    pub clustering_concurrency_limit: Option<NonZero<usize>>,
+
     /// Outputs the queries made to the graph to the specified file.
     #[clap(long)]
     pub log_queries: Option<PathBuf>,
@@ -457,6 +464,7 @@ where
         query_logger,
         api_config: config.api_config,
         compiler,
+        clustering: Arc::new(ClusteringContext::new(config.clustering_concurrency_limit)),
     });
     start_rest_server(router, config.http_address, lifecycle);