Property replace (TraceMachina#1976)

Currently the property modifier scheduler can add or remove properties.  It would be useful to replace a property too, for example Siso sets the label:action_large=1 property and we would like to switch that to cpu_count=2000 instead.

Implement an additional option to the property modifier to replace properties as well as add and remove.  This has the option to match values and the option to replace the value or pass it through.

Author: chrisstaite-menlo

nativelink-config/src/schedulers.rs

@@ -195,13 +195,30 @@ pub struct PlatformPropertyAddition {
     pub value: String,
 }
 
+#[derive(Deserialize, Serialize, Debug, Clone)]
+#[serde(deny_unknown_fields)]
+pub struct PlatformPropertyReplacement {
+    /// The name of the property to replace.
+    pub name: String,
+    /// The the value to match against, if unset then any instance matches.
+    #[serde(default)]
+    pub value: Option<String>,
+    /// The new name of the property.
+    pub new_name: String,
+    /// The value to assign to the property, if unset will remain the same.
+    #[serde(default)]
+    pub new_value: Option<String>,
+}
+
 #[derive(Deserialize, Serialize, Debug, Clone)]
 #[serde(rename_all = "snake_case")]
 pub enum PropertyModification {
     /// Add a property to the action properties.
     Add(PlatformPropertyAddition),
     /// Remove a named property from the action.
     Remove(String),
+    /// If a property is found, then replace it with another one.
+    Replace(PlatformPropertyReplacement),
 }
 
 #[derive(Deserialize, Serialize, Debug)]

nativelink-scheduler/src/property_modifier_scheduler.rs

@@ -16,7 +16,9 @@ use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use async_trait::async_trait;
-use nativelink_config::schedulers::{PropertyModification, PropertyModifierSpec};
+use nativelink_config::schedulers::{
+    PlatformPropertyReplacement, PropertyModification, PropertyModifierSpec,
+};
 use nativelink_error::{Error, ResultExt};
 use nativelink_metric::{MetricsComponent, RootMetricsComponent};
 use nativelink_util::action_messages::{ActionInfo, OperationId};
@@ -71,7 +73,8 @@ impl PropertyModifierScheduler {
         );
         for modification in &self.modifications {
             match modification {
-                PropertyModification::Remove(name) => {
+                PropertyModification::Remove(name)
+                | PropertyModification::Replace(PlatformPropertyReplacement { name, .. }) => {
                     known_properties.insert(name.clone());
                 }
                 PropertyModification::Add(_) => (),
@@ -93,13 +96,36 @@ impl PropertyModifierScheduler {
         let action_info_mut = Arc::make_mut(&mut action_info);
         for modification in &self.modifications {
             match modification {
-                PropertyModification::Add(addition) => action_info_mut
-                    .platform_properties
-                    .insert(addition.name.clone(), addition.value.clone()),
+                PropertyModification::Add(addition) => {
+                    action_info_mut
+                        .platform_properties
+                        .insert(addition.name.clone(), addition.value.clone());
+                }
                 PropertyModification::Remove(name) => {
-                    action_info_mut.platform_properties.remove(name)
+                    action_info_mut.platform_properties.remove(name);
                 }
-            };
+                PropertyModification::Replace(replacement) => {
+                    if let Some((existing_name, existing_value)) = action_info_mut
+                        .platform_properties
+                        .remove_entry(&replacement.name)
+                    {
+                        if replacement
+                            .value
+                            .as_ref()
+                            .is_none_or(|value| *value == existing_value)
+                        {
+                            action_info_mut.platform_properties.insert(
+                                replacement.new_name.clone(),
+                                replacement.new_value.clone().unwrap_or(existing_value),
+                            );
+                        } else {
+                            action_info_mut
+                                .platform_properties
+                                .insert(existing_name, existing_value);
+                        }
+                    }
+                }
+            }
         }
         self.scheduler
             .add_action(client_operation_id, action_info)

nativelink-scheduler/tests/property_modifier_scheduler_test.rs

@@ -22,7 +22,8 @@ mod utils {
 
 use futures::{StreamExt, join};
 use nativelink_config::schedulers::{
-    PlatformPropertyAddition, PropertyModification, PropertyModifierSpec, SchedulerSpec, SimpleSpec,
+    PlatformPropertyAddition, PlatformPropertyReplacement, PropertyModification,
+    PropertyModifierSpec, SchedulerSpec, SimpleSpec,
 };
 use nativelink_error::Error;
 use nativelink_macro::nativelink_test;
@@ -181,7 +182,7 @@ async fn add_action_property_remove_after_add() -> Result<(), Error> {
     let context = make_modifier_scheduler(vec![
         PropertyModification::Add(PlatformPropertyAddition {
             name: name.clone(),
-            value: value.clone(),
+            value,
         }),
         PropertyModification::Remove(name.clone()),
     ]);
@@ -210,6 +211,139 @@ async fn add_action_property_remove_after_add() -> Result<(), Error> {
     Ok(())
 }
 
+#[nativelink_test]
+async fn add_action_property_replace() -> Result<(), Error> {
+    let name = "name".to_string();
+    let new_name = "new_name".to_string();
+    let value = "value".to_string();
+    let context = make_modifier_scheduler(vec![PropertyModification::Replace(
+        PlatformPropertyReplacement {
+            name: name.clone(),
+            value: None,
+            new_name: new_name.clone(),
+            new_value: None,
+        },
+    )]);
+    let mut action_info = make_base_action_info(UNIX_EPOCH, DigestInfo::zero_digest());
+    Arc::make_mut(&mut action_info)
+        .platform_properties
+        .insert(name, value.clone());
+    let (_forward_watch_channel_tx, forward_watch_channel_rx) =
+        watch::channel(Arc::new(ActionState {
+            client_operation_id: OperationId::default(),
+            stage: ActionStage::Queued,
+            action_digest: action_info.unique_qualifier.digest(),
+        }));
+    let client_operation_id = OperationId::default();
+    let (_, (passed_client_operation_id, action_info)) = join!(
+        context
+            .modifier_scheduler
+            .add_action(client_operation_id.clone(), action_info.clone()),
+        context
+            .mock_scheduler
+            .expect_add_action(Ok(Box::new(TokioWatchActionStateResult::new(
+                client_operation_id.clone(),
+                action_info,
+                forward_watch_channel_rx
+            )))),
+    );
+    assert_eq!(client_operation_id, passed_client_operation_id);
+    assert_eq!(
+        HashMap::from([(new_name, value)]),
+        action_info.platform_properties
+    );
+    Ok(())
+}
+
+#[nativelink_test]
+async fn add_action_property_replace_match_value() -> Result<(), Error> {
+    let name = "name".to_string();
+    let new_name = "new_name".to_string();
+    let value = "value".to_string();
+    let context = make_modifier_scheduler(vec![PropertyModification::Replace(
+        PlatformPropertyReplacement {
+            name: name.clone(),
+            value: Some(value.clone()),
+            new_name: new_name.clone(),
+            new_value: None,
+        },
+    )]);
+    let mut action_info = make_base_action_info(UNIX_EPOCH, DigestInfo::zero_digest());
+    Arc::make_mut(&mut action_info)
+        .platform_properties
+        .insert(name.clone(), value.clone());
+    let (_forward_watch_channel_tx, forward_watch_channel_rx) =
+        watch::channel(Arc::new(ActionState {
+            client_operation_id: OperationId::default(),
+            stage: ActionStage::Queued,
+            action_digest: action_info.unique_qualifier.digest(),
+        }));
+    let client_operation_id = OperationId::default();
+    let (_, (passed_client_operation_id, action_info)) = join!(
+        context
+            .modifier_scheduler
+            .add_action(client_operation_id.clone(), action_info.clone()),
+        context
+            .mock_scheduler
+            .expect_add_action(Ok(Box::new(TokioWatchActionStateResult::new(
+                client_operation_id.clone(),
+                action_info,
+                forward_watch_channel_rx
+            )))),
+    );
+    assert_eq!(client_operation_id, passed_client_operation_id);
+    assert_eq!(
+        HashMap::from([(new_name, value)]),
+        action_info.platform_properties
+    );
+    Ok(())
+}
+
+#[nativelink_test]
+async fn add_action_property_replace_value() -> Result<(), Error> {
+    let name = "name".to_string();
+    let new_name = "new_name".to_string();
+    let value = "value".to_string();
+    let value_two = "value_two".to_string();
+    let context = make_modifier_scheduler(vec![PropertyModification::Replace(
+        PlatformPropertyReplacement {
+            name: name.clone(),
+            value: None,
+            new_name: new_name.clone(),
+            new_value: Some(value_two.clone()),
+        },
+    )]);
+    let mut action_info = make_base_action_info(UNIX_EPOCH, DigestInfo::zero_digest());
+    Arc::make_mut(&mut action_info)
+        .platform_properties
+        .insert(name.clone(), value.clone());
+    let (_forward_watch_channel_tx, forward_watch_channel_rx) =
+        watch::channel(Arc::new(ActionState {
+            client_operation_id: OperationId::default(),
+            stage: ActionStage::Queued,
+            action_digest: action_info.unique_qualifier.digest(),
+        }));
+    let client_operation_id = OperationId::default();
+    let (_, (passed_client_operation_id, action_info)) = join!(
+        context
+            .modifier_scheduler
+            .add_action(client_operation_id.clone(), action_info.clone()),
+        context
+            .mock_scheduler
+            .expect_add_action(Ok(Box::new(TokioWatchActionStateResult::new(
+                client_operation_id.clone(),
+                action_info,
+                forward_watch_channel_rx
+            )))),
+    );
+    assert_eq!(client_operation_id, passed_client_operation_id);
+    assert_eq!(
+        HashMap::from([(new_name, value_two)]),
+        action_info.platform_properties
+    );
+    Ok(())
+}
+
 #[nativelink_test]
 async fn add_action_property_remove() -> Result<(), Error> {
     let name = "name".to_string();
@@ -303,3 +437,26 @@ async fn remove_retains_type_in_underlying_manager() -> Result<(), Error> {
     assert_eq!(Ok(vec![name]), known_props);
     Ok(())
 }
+
+#[nativelink_test]
+async fn replace_retains_type_in_underlying_manager() -> Result<(), Error> {
+    let name = "name".to_string();
+    let context = make_modifier_scheduler(vec![PropertyModification::Replace(
+        PlatformPropertyReplacement {
+            name: name.clone(),
+            value: None,
+            new_name: "new_name".to_string(),
+            new_value: None,
+        },
+    )]);
+    let known_properties = vec![name.clone()];
+    let instance_name_fut = context
+        .mock_scheduler
+        .expect_get_known_properties(Ok(known_properties));
+    let known_props_fut = context
+        .modifier_scheduler
+        .get_known_properties(INSTANCE_NAME);
+    let (_, known_props) = join!(instance_name_fut, known_props_fut);
+    assert_eq!(Ok(vec![name]), known_props);
+    Ok(())
+}