new commit because lots of changes, generated protos and added metadata, support for lifecycle management

Author: MarcusSorealheis

nativelink-proto/BUILD.bazel

@@ -11,6 +11,7 @@ load(
 PROTO_NAMES = [
     "build.bazel.remote.asset.v1",
     "build.bazel.remote.execution.v2",
+    "build.bazel.remote.execution.worker_protocol",
     "build.bazel.semver",
     "com.github.trace_machina.nativelink.remote_execution",
     "com.github.trace_machina.nativelink.events",
@@ -80,6 +81,7 @@ genrule(
     srcs = [
         "build/bazel/remote/asset/v1/remote_asset.proto",
         "build/bazel/remote/execution/v2/remote_execution.proto",
+        "build/bazel/remote/execution/worker_protocol.proto",
         "build/bazel/semver/semver.proto",
         "com/github/trace_machina/nativelink/remote_execution/events.proto",
         "com/github/trace_machina/nativelink/remote_execution/worker_api.proto",

nativelink-proto/genproto/build.bazel.remote.execution.worker_protocol.pb.rs

@@ -0,0 +1,63 @@
+// Copyright 2022 The NativeLink Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// This file is @generated by prost-build.
+/// Work request sent to persistent worker
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct WorkRequest {
+    /// Command-line arguments for the worker
+    #[prost(string, repeated, tag = "1")]
+    pub arguments: ::prost::alloc::vec::Vec<::prost::alloc::string::String>,
+    /// Input files for this request
+    #[prost(message, repeated, tag = "2")]
+    pub inputs: ::prost::alloc::vec::Vec<Input>,
+    /// Unique request ID for this work item
+    #[prost(string, tag = "3")]
+    pub request_id: ::prost::alloc::string::String,
+    /// Whether to cancel this request
+    #[prost(bool, tag = "4")]
+    pub cancel: bool,
+    /// Verbosity level for logging
+    #[prost(int32, tag = "5")]
+    pub verbosity: i32,
+    /// Sandbox directory for this request
+    #[prost(string, tag = "6")]
+    pub sandbox_dir: ::prost::alloc::string::String,
+}
+/// Input file for work request
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct Input {
+    /// Path to the input file
+    #[prost(string, tag = "1")]
+    pub path: ::prost::alloc::string::String,
+    /// Digest of the file content
+    #[prost(bytes = "bytes", tag = "2")]
+    pub digest: ::prost::bytes::Bytes,
+}
+/// Work response from persistent worker
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct WorkResponse {
+    /// Exit code of the work (0 = success)
+    #[prost(int32, tag = "1")]
+    pub exit_code: i32,
+    /// Output from the worker
+    #[prost(string, tag = "2")]
+    pub output: ::prost::alloc::string::String,
+    /// Request ID this response corresponds to
+    #[prost(string, tag = "3")]
+    pub request_id: ::prost::alloc::string::String,
+    /// Whether the request was cancelled
+    #[prost(bool, tag = "4")]
+    pub was_cancelled: bool,
+}

nativelink-proto/genproto/lib.rs

@@ -45,6 +45,9 @@ pub mod build {
                 pub mod v2 {
                     include!("build.bazel.remote.execution.v2.pb.rs");
                 }
+                pub mod worker_protocol {
+                    include!("build.bazel.remote.execution.worker_protocol.pb.rs");
+                }
             }
         }
         pub mod semver {

nativelink-scheduler/src/api_worker_scheduler.rs

@@ -151,13 +151,18 @@ impl ApiWorkerSchedulerImpl {
         let worker_id = worker.id.clone();
 
         // Check if this worker supports persistent worker operations
-        if let Some(persistent_key) = worker.platform_properties.get("persistentWorkerKey") {
+        if let Some(persistent_key) = worker
+            .platform_properties
+            .properties
+            .get("persistentWorkerKey")
+        {
+            let key_str = format!("{persistent_key:?}");
             self.persistent_worker_pools
-                .entry(persistent_key.clone())
-                .or_insert_with(Vec::new)
+                .entry(key_str)
+                .or_default()
                 .push(worker_id.clone());
             tracing::info!(
-                "Registered worker {} for persistent key: {}",
+                "Registered worker {} for persistent key: {:?}",
                 worker_id,
                 persistent_key
             );
@@ -190,7 +195,7 @@ impl ApiWorkerSchedulerImpl {
         let result = self.workers.pop(worker_id);
 
         // Remove from persistent worker pools if applicable
-        for (_, workers) in self.persistent_worker_pools.iter_mut() {
+        for workers in self.persistent_worker_pools.values_mut() {
             workers.retain(|id| id != worker_id);
         }
 
@@ -219,18 +224,22 @@ impl ApiWorkerSchedulerImpl {
         platform_properties: &PlatformProperties,
     ) -> Option<WorkerId> {
         // First check if this is a persistent worker request
-        if let Some(persistent_key) = platform_properties.get("persistentWorkerKey") {
-            tracing::trace!("Looking for persistent worker with key: {}", persistent_key);
+        if let Some(persistent_key) = platform_properties.properties.get("persistentWorkerKey") {
+            tracing::trace!(
+                "Looking for persistent worker with key: {:?}",
+                persistent_key
+            );
 
+            let key_str = format!("{persistent_key:?}");
             // Try to find a dedicated persistent worker
-            if let Some(worker_ids) = self.persistent_worker_pools.get(persistent_key) {
+            if let Some(worker_ids) = self.persistent_worker_pools.get(&key_str) {
                 for worker_id in worker_ids {
                     if let Some(worker) = self.workers.peek(worker_id) {
                         if worker.can_accept_work()
                             && platform_properties.is_satisfied_by(&worker.platform_properties)
                         {
                             tracing::debug!(
-                                "Found persistent worker {} for key {}",
+                                "Found persistent worker {} for key {:?}",
                                 worker_id,
                                 persistent_key
                             );
@@ -241,7 +250,7 @@ impl ApiWorkerSchedulerImpl {
             }
 
             tracing::info!(
-                "No available persistent worker found for key: {}, will spawn new one",
+                "No available persistent worker found for key: {:?}, will spawn new one",
                 persistent_key
             );
         }

nativelink-worker/BUILD.bazel

@@ -12,6 +12,8 @@ rust_library(
     srcs = [
         "src/lib.rs",
         "src/local_worker.rs",
+        "src/persistent_worker.rs",
+        "src/persistent_worker_runner.rs",
         "src/running_actions_manager.rs",
         "src/worker_api_client_wrapper.rs",
         "src/worker_utils.rs",
@@ -52,6 +54,8 @@ rust_test_suite(
     timeout = "short",
     srcs = [
         "tests/local_worker_test.rs",
+        "tests/persistent_worker_runner_test.rs",
+        "tests/persistent_worker_test.rs",
         "tests/running_actions_manager_test.rs",
     ],
     compile_data = [