Route company tasks by GPU worker tags

Cargo.toml

@@ -17,7 +17,7 @@ spz-lib = { git = "https://github.com/404-Repo/spz-rs.git", rev = "227c6cb6cdc25
     "async",
 ] }
 multer = "3.1.0"
-tokio = { version = "1.52.2", features = [
+tokio = { version = "1.52.3", features = [
     "rt-multi-thread",
     "macros",
     "sync",
@@ -43,14 +43,14 @@ async-trait = "0.1.89"
 quinn = { version = "0.11.9", features = ["rustls", "runtime-tokio", "log"] }
 h3 = { version = "0.0.8", features = ["tracing"] }
 h3-quinn = { version = "0.0.10" }
-http = "1.4.1"
+http = "1.4.2"
 backon = "1.6.0"
 openraft = { version = "0.9.24", features = ["storage-v2", "serde"] }
 tokio-postgres-rustls = "0.14.0"
 rustls = { version = "0.23.40" }
 rustls-platform-verifier = "0.7.0"
 toml = "1.1.2"
-uuid = { version = "1.23.0", features = ["v4", "serde"] }
+uuid = { version = "1.23.3", features = ["v4", "serde"] }
 serde = { version = "1.0.228", features = ["derive"] }
 serde_json = "1.0.150"
 rmp-serde = "1.3.1"
@@ -59,7 +59,7 @@ bytes = "1.11.1"
 rcgen = { version = "0.14.8", features = ["pem"] }
 foldhash = "0.2.0"
 portable-atomic = "1.13.1"
-scc = "3.7.1"
+scc = "3.7.2"
 sdd = "4.8.6"
 schnorrkel = "0.11.5"
 base64-simd = "0.8"
@@ -69,17 +69,18 @@ prometheus = "0.14.0"
 blake3 = "1.8.5"
 moka = { version = "0.12.15", features = ["sync"] }
 mimalloc = { version = "0.1.52", default-features = false }
-regex = "1.12.3"
+regex = "1.12.4"
 image = "0.25.10"
 chrono = { version = "0.4", features = ["serde", "clock"] }
 itoa = "1.0.18"
 zstd = "0.13"
 tempfile = "3.27.0"
 notify = "8.2.0"
 rand = "0.10.1"
+crossbeam-skiplist = "0.1.3"
 
 [dev-dependencies]
-hyper = "1.9.0"
+hyper = "1.10.1"
 http-body-util = "0.1.3"
 futures-lite = { version = "2.6.1", default-features = false, features = [
     "std",
@@ -107,7 +108,7 @@ required-features = ["test-support"]
 
 [build-dependencies]
 anyhow = "1.0.100"
-vergen-gitcl = { version = "9.1.0", features = ["build", "cargo", "rustc"] }
+vergen-gitcl = { version = "10.0.0", features = ["build", "cargo", "rustc"] }
 
 [profile.release]
 codegen-units = 1

dev-env/init-scripts/init-schema.sql

@@ -249,6 +249,7 @@ CREATE TABLE IF NOT EXISTS companies (
   task_limit_concurrent INTEGER NOT NULL DEFAULT 1 CHECK (task_limit_concurrent >= 0),
   task_limit_daily INTEGER NOT NULL CHECK (task_limit_daily >= 0),
   ownership_state TEXT NOT NULL DEFAULT 'unassigned' CHECK (ownership_state IN ('unassigned', 'owned')),
+  worker_tags TEXT[] NOT NULL DEFAULT '{}'::TEXT[],
   created_by_user_id BIGINT REFERENCES users(id) ON DELETE SET NULL,
   created_at BIGINT NOT NULL,
   updated_at BIGINT NOT NULL

src/api/request.rs

@@ -1,4 +1,5 @@
 use bytes::Bytes;
+use foldhash::{HashSet, fast::RandomState};
 use serde::Deserialize;
 use serde_json::Value;
 use std::sync::Arc;
@@ -8,6 +9,36 @@ use uuid::Uuid;
 pub use super::gateway_info::{GatewayInfoExt, GatewayInfoExtRef};
 use crate::crypto::hotkey::Hotkey;
 
+pub const MAX_WORKER_TAG_LENGTH: usize = 32;
+
+pub fn normalize_worker_tags(tags: &[String]) -> Result<Vec<String>, String> {
+    let mut normalized = Vec::with_capacity(tags.len());
+    let mut seen = HashSet::with_capacity_and_hasher(tags.len(), RandomState::default());
+    for tag in tags {
+        let tag = tag.trim().to_ascii_lowercase();
+        if tag.is_empty() {
+            return Err("worker_tags cannot contain empty tags".to_string());
+        }
+        if tag.len() > MAX_WORKER_TAG_LENGTH {
+            return Err(format!(
+                "worker_tags entries cannot exceed {MAX_WORKER_TAG_LENGTH} characters"
+            ));
+        }
+        if !tag.bytes().all(|byte| {
+            byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'_' || byte == b'-'
+        }) {
+            return Err(
+                "worker_tags entries may only contain lowercase letters, numbers, underscores, or dashes"
+                    .to_string(),
+            );
+        }
+        if seen.insert(tag.clone()) {
+            normalized.push(tag);
+        }
+    }
+    Ok(normalized)
+}
+
 #[derive(Debug, Clone, Deserialize)]
 pub struct AddTaskRequest {
     pub seed: Option<SeedValue>,
@@ -57,6 +88,8 @@ pub struct GetTasksRequest {
     pub timestamp: String,
     pub requested_task_count: usize,
     pub model: ModelFilter,
+    #[serde(default)]
+    pub worker_tags: Vec<String>,
 }
 
 #[derive(Debug, Deserialize)]

src/build.rs

@@ -1,12 +1,17 @@
 use anyhow::Result;
 
-use vergen_gitcl::{BuildBuilder, CargoBuilder, Emitter, GitclBuilder, RustcBuilder};
+use vergen_gitcl::{Build, Cargo, Emitter, Gitcl, Rustc};
 
 fn main() -> Result<()> {
+    let build = Build::all_build();
+    let cargo = Cargo::all_cargo();
+    let gitcl = Gitcl::all_git();
+    let rustc = Rustc::all_rustc();
+
     Emitter::default()
-        .add_instructions(&BuildBuilder::all_build()?)?
-        .add_instructions(&CargoBuilder::all_cargo()?)?
-        .add_instructions(&GitclBuilder::all_git()?)?
-        .add_instructions(&RustcBuilder::all_rustc()?)?
+        .add_instructions(&build)?
+        .add_instructions(&cargo)?
+        .add_instructions(&gitcl)?
+        .add_instructions(&rustc)?
         .emit()
 }