Merge branch 'main' into remote-persistent-safe

Author: MarcusSorealheis

Cargo.lock

@@ -16,6 +16,7 @@ rust_library(
         "src/schedulers.rs",
         "src/serde_utils.rs",
         "src/stores.rs",
+        "src/warm_worker_pools.rs",
     ],
     compile_data = [
         "README.md",

nativelink-config/BUILD.bazel

@@ -10,7 +10,7 @@ version = "0.7.6"
 nativelink-error = { path = "../nativelink-error" }
 
 byte-unit = { version = "5.1.6", default-features = false, features = ["byte"] }
-humantime = { version = "2.2.0", default-features = false }
+humantime = { version = "2.3.0", default-features = false }
 rand = { version = "0.9.0", default-features = false, features = [
   "thread_rng",
 ] }

nativelink-config/Cargo.toml

@@ -17,7 +17,4 @@ pub mod cas_server;
 pub mod schedulers;
 pub mod serde_utils;
 pub mod stores;
-
-// Warm worker pools configuration (optional feature)
-#[cfg(feature = "warm-worker-pools")]
 pub mod warm_worker_pools;

nativelink-config/src/lib.rs

@@ -21,7 +21,6 @@ use crate::serde_utils::{
     convert_numeric_with_shellexpand,
 };
 use crate::stores::{GrpcEndpoint, Retry, StoreRefName};
-
 // Import warm worker pool configuration
 #[cfg(feature = "warm-worker-pools")]
 use crate::warm_worker_pools::WarmWorkerPoolsConfig;

nativelink-config/src/schedulers.rs

@@ -1,26 +1,98 @@
+load("@bazel_skylib//lib:selects.bzl", "selects")
 load("@rules_rust//rust:defs.bzl", "rust_library", "rust_test")
 
+# Platform configurations for protoc
+PLATFORM_OS_ARCH = [
+    ("linux", "aarch64"),
+    ("linux", "x86_64"),
+    ("macos", "aarch64"),
+    ("macos", "x86_64"),
+    ("windows", "aarch64"),
+    ("windows", "x86_64"),
+]
+
+[
+    selects.config_setting_group(
+        name = "{}_{}".format(
+            os.replace("macos", "osx"),
+            arch.replace("aarch64", "aarch_64"),
+        ),
+        match_all = [
+            "@platforms//cpu:{}".format(arch),
+            "@platforms//os:{}".format(os),
+        ],
+    )
+    for (os, arch) in PLATFORM_OS_ARCH
+]
+
+PLATFORM_NAMES = [
+    "{}_{}".format(
+        os.replace("macos", "osx"),
+        arch.replace("aarch64", "aarch_64"),
+    )
+    for (os, arch) in PLATFORM_OS_ARCH
+]
+
+# Generate Rust code from CRI proto files
+genrule(
+    name = "gen_cri_protos",
+    srcs = ["proto/cri/api.proto"],
+    outs = ["runtime.v1.pb.rs"],
+    cmd = select({
+        platform: '''
+        set -e
+        export PROTOC=$(execpath @@toolchains_protoc++protoc+toolchains_protoc_hub.{}//:bin/protoc)
+
+        $(execpath //nativelink-proto:gen_protos_tool) $(SRCS) -o $(RULEDIR)
+
+        for file in $(RULEDIR)/*.rs; do
+            mv -- "$$file" "$${{file%.rs}}.pb.rs"
+        done
+        '''.format(platform)
+        for platform in PLATFORM_NAMES
+    }),
+    tools = [
+        "//nativelink-proto:gen_protos_tool",
+    ] + select({
+        platform: ["@@toolchains_protoc++protoc+toolchains_protoc_hub.{}//:bin/protoc".format(platform)]
+        for platform in PLATFORM_NAMES
+    }),
+)
+
 rust_library(
     name = "nativelink-crio-worker-pool",
     srcs = [
         "src/cache.rs",
         "src/config.rs",
         "src/cri_client.rs",
+        "src/cri_client_grpc.rs",
+        "src/isolation.rs",
         "src/lib.rs",
         "src/lifecycle.rs",
         "src/pool_manager.rs",
         "src/warmup.rs",
         "src/worker.rs",
+        ":gen_cri_protos",
+    ],
+    compile_data = [
+        ":gen_cri_protos",
     ],
+    tags = ["no-clippy"],
     visibility = ["//visibility:public"],
     deps = [
+        "//nativelink-config",
         "//nativelink-error",
+        "//nativelink-metric",
         "//nativelink-util",
+        "@crates//:hyper-util",
+        "@crates//:prost",
         "@crates//:serde",
         "@crates//:serde_json",
         "@crates//:serde_with",
         "@crates//:tempfile",
         "@crates//:tokio",
+        "@crates//:tonic",
+        "@crates//:tower",
         "@crates//:tracing",
         "@crates//:uuid",
     ],
@@ -29,4 +101,5 @@ rust_library(
 rust_test(
     name = "unit_tests",
     crate = ":nativelink-crio-worker-pool",
+    tags = ["no-clippy"],
 )

nativelink-crio-worker-pool/BUILD.bazel

@@ -3,6 +3,10 @@ edition = "2024"
 name = "nativelink-crio-worker-pool"
 version = "0.1.0"
 
+[features]
+cargo-build = []
+default = ["cargo-build"]
+
 [package.metadata.cargo-machete]
 ignored = ["prost", "tonic-build"]
 

nativelink-crio-worker-pool/Cargo.toml

@@ -29,6 +29,13 @@ use tower::service_fn;
 
 // Generated from CRI protocol buffers
 pub mod runtime {
+    #![allow(clippy::all)]
+    // When building with Bazel, include the generated proto directly
+    #[cfg(not(feature = "cargo-build"))]
+    include!("../runtime.v1.pb.rs");
+
+    // When building with Cargo, use tonic's include_proto macro
+    #[cfg(feature = "cargo-build")]
     tonic::include_proto!("runtime.v1");
 }
 

nativelink-crio-worker-pool/src/cri_client_grpc.rs

@@ -4,7 +4,7 @@ use std::collections::{HashMap, HashSet, VecDeque};
 use std::path::PathBuf;
 use std::sync::Arc;
 
-use nativelink_config::warm_worker_pools::IsolationStrategy;
+use nativelink_config::warm_worker_pools::{IsolationStrategy, WarmWorkerPoolsConfig};
 use nativelink_error::{Code, Error, ResultExt, make_err};
 use nativelink_metric::MetricsComponent;
 use nativelink_util::action_messages::WorkerId;
@@ -14,7 +14,6 @@ use uuid::Uuid;
 
 use crate::cache::CachePrimingAgent;
 use crate::config::WorkerPoolConfig;
-use nativelink_config::warm_worker_pools::WarmWorkerPoolsConfig;
 use crate::cri_client::{
     ContainerConfig, ContainerMetadata, CriClient, ImageSpec, KeyValue, LinuxContainerConfig,
     LinuxContainerResources, LinuxPodSandboxConfig, LinuxSandboxSecurityContext, NamespaceOptions,
@@ -83,7 +82,11 @@ impl WarmWorkerPoolManager {
     ///
     /// If the pool has isolation enabled, this will create an ephemeral COW clone
     /// of a warm template. Otherwise, it falls back to regular acquisition.
-    pub async fn acquire_isolated(&self, pool: &str, job_id: &str) -> Result<WarmWorkerLease, Error> {
+    pub async fn acquire_isolated(
+        &self,
+        pool: &str,
+        job_id: &str,
+    ) -> Result<WarmWorkerLease, Error> {
         let worker_pool = self
             .pools
             .get(pool)
@@ -159,7 +162,8 @@ impl WorkerPool {
         let template = self.get_or_create_template().await?;
 
         // Create isolated clone
-        self.clone_from_template(&template, job_id, isolation_config).await
+        self.clone_from_template(&template, job_id, isolation_config)
+            .await
     }
 
     /// Gets the existing template or creates a new one.
@@ -261,11 +265,7 @@ impl WorkerPool {
         job_id: &str,
         isolation_config: &nativelink_config::warm_worker_pools::IsolationConfig,
     ) -> Result<WarmWorkerLease, Error> {
-        let worker_id = WorkerId(format!(
-            "crio:{}:isolated:{}",
-            self.config.name,
-            job_id
-        ));
+        let worker_id = WorkerId(format!("crio:{}:isolated:{}", self.config.name, job_id));
 
         // Create OverlayFS mount structure
         let mount = OverlayFsMount::new(
@@ -305,7 +305,12 @@ impl WorkerPool {
             .runtime
             .create_container(&sandbox_id, &container_config, &sandbox_config)
             .await
-            .err_tip(|| format!("while creating container for isolated worker {}", worker_id.0))?;
+            .err_tip(|| {
+                format!(
+                    "while creating container for isolated worker {}",
+                    worker_id.0
+                )
+            })?;
 
         self.runtime
             .start_container(&container_id)
@@ -844,7 +849,13 @@ impl WarmWorkerLease {
             // For isolated workers, we need different cleanup
             if let Some(mount) = self.isolation_mount.take() {
                 self.pool
-                    .release_isolated_worker(worker_id, &self.container_id, &self.sandbox_id, mount, outcome)
+                    .release_isolated_worker(
+                        worker_id,
+                        &self.container_id,
+                        &self.sandbox_id,
+                        mount,
+                        outcome,
+                    )
                     .await
                     .err_tip(|| "while releasing isolated worker")
             } else {

nativelink-crio-worker-pool/src/pool_manager.rs

@@ -32,6 +32,7 @@ PROTO_NAMES = [
 rust_binary(
     name = "gen_protos_tool",
     srcs = ["gen_protos_tool.rs"],
+    visibility = ["//visibility:public"],
     deps = [
         "@crates//:clap",
         "@crates//:prost-build",