Skip to content

Commit

Permalink
Add a new data field meta_internal_extra_params to CommandExecutorConfig
Browse files Browse the repository at this point in the history 
Summary:
To better support meta internal infra in buck without extending CommandExecutorConfig every time when a new data needs to add, add a generic data field meta_internal_extra_params to hold any future meta internal parameters.

Meta:
For RL physical devices (remote execution platform=riot), we need to pass in setup_preference_key to specify the device state, otherwise, the leased devices will be with random OS build/APK versions and random device state (wifi, nux, etc). This will leads to test flakiness.

The setup_preference_key is in [the thrift struct TExecutionPolicy](https://www.internalfb.com/code/fbsource/[da3da7c7fda1]/fbcode/remote_execution/lib/if/common.thrift?lines=528). Add the support for the remote execution policy in buck. This data is similar to the remote_execution_properties and remote_execution_use_case, which will be passed down from bzl for RL on device tests.

Reviewed By: JakobDegen

Differential Revision: D68185161

fbshipit-source-id: f5ce79b044c98383797b1fafa572856215ea645e
  • Loading branch information
@dongyansu @facebook-github-bot
dongyansu authored and facebook-github-bot committed Jan 24, 2025
1 parent becb5be commit 51314c5
Show file tree
Hide file tree
Showing 11 changed files with 149 additions and 2 deletions.
5 changes: 4 additions & 1 deletion app/buck2_action_impl/src/actions/impls/run.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ use buck2_build_api::interpreter::rule_defs::cmd_args::StarlarkCmdArgs;
use buck2_build_api::interpreter::rule_defs::provider::builtin::worker_info::FrozenWorkerInfo;
use buck2_build_api::interpreter::rule_defs::provider::builtin::worker_info::WorkerInfo;
use buck2_core::category::CategoryRef;
use buck2_core::execution_types::executor_config::MetaInternalExtraParams;
use buck2_core::execution_types::executor_config::RemoteExecutorCustomImage;
use buck2_core::execution_types::executor_config::RemoteExecutorDependency;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
Expand Down Expand Up @@ -171,6 +172,7 @@ pub(crate) struct UnregisteredRunAction {
pub(crate) unique_input_inodes: bool,
pub(crate) remote_execution_dependencies: Vec<RemoteExecutorDependency>,
pub(crate) remote_execution_custom_image: Option<RemoteExecutorCustomImage>,
pub(crate) meta_internal_extra_params: MetaInternalExtraParams,
}

impl UnregisteredAction for UnregisteredRunAction {
Expand Down Expand Up @@ -632,7 +634,8 @@ impl RunAction {
.with_force_full_hybrid_if_capable(self.inner.force_full_hybrid_if_capable)
.with_unique_input_inodes(self.inner.unique_input_inodes)
.with_remote_execution_dependencies(self.inner.remote_execution_dependencies.clone())
.with_remote_execution_custom_image(self.inner.remote_execution_custom_image.clone());
.with_remote_execution_custom_image(self.inner.remote_execution_custom_image.clone())
.with_meta_internal_extra_params(self.inner.meta_internal_extra_params.clone());

let (dep_file_bundle, req) = if let Some(visitor) = dep_file_visitor {
let bundle = make_dep_file_bundle(ctx, visitor, cmdline_digest, req.paths())?;
Expand Down
11 changes: 11 additions & 0 deletions app/buck2_action_impl/src/context/run.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ use buck2_build_api::interpreter::rule_defs::cmd_args::SimpleCommandLineArtifact
use buck2_build_api::interpreter::rule_defs::cmd_args::StarlarkCmdArgs;
use buck2_build_api::interpreter::rule_defs::cmd_args::StarlarkCommandLineValueUnpack;
use buck2_build_api::interpreter::rule_defs::command_executor_config::parse_custom_re_image;
use buck2_build_api::interpreter::rule_defs::command_executor_config::parse_meta_internal_extra_params;
use buck2_build_api::interpreter::rule_defs::context::AnalysisActions;
use buck2_build_api::interpreter::rule_defs::provider::builtin::run_info::RunInfo;
use buck2_build_api::interpreter::rule_defs::provider::builtin::worker_run_info::WorkerRunInfo;
Expand All @@ -35,6 +36,7 @@ use host_sharing::WeightPercentage;
use starlark::environment::MethodsBuilder;
use starlark::eval::Evaluator;
use starlark::starlark_module;
use starlark::values::dict::DictRef;
use starlark::values::dict::UnpackDictEntries;
use starlark::values::list::UnpackList;
use starlark::values::none::NoneOr;
Expand Down Expand Up @@ -131,6 +133,8 @@ pub(crate) fn analysis_actions_methods_run(methods: &mut MethodsBuilder) {
/// * `drop_host_mount_globs`: list of strings containing file
/// globs. Any mounts globs specified will not be bind mounted
/// from the host.
/// * `meta_internal_extra_params`: a dictionary to pass extra parameters to RE, can add more keys in the future:
/// * `remote_execution_policy`: refer to TExecutionPolicy.
///
/// When actions execute, they'll do so from the root of the repository. As they execute,
/// actions have exclusive access to their output directory.
Expand Down Expand Up @@ -193,6 +197,9 @@ pub(crate) fn analysis_actions_methods_run(methods: &mut MethodsBuilder) {
#[starlark(require = named, default=UnpackList::default())]
remote_execution_dependencies: UnpackList<SmallMap<&'v str, &'v str>>,
#[starlark(default = NoneType, require = named)] remote_execution_dynamic_image: Value<'v>,
#[starlark(require = named, default = NoneOr::None)] meta_internal_extra_params: NoneOr<
DictRef<'v>,
>,
) -> starlark::Result<NoneType> {
struct RunCommandArtifactVisitor {
inner: SimpleCommandLineArtifactVisitor,
Expand Down Expand Up @@ -378,6 +385,9 @@ pub(crate) fn analysis_actions_methods_run(methods: &mut MethodsBuilder) {
remote_execution_dynamic_image,
)?;

let extra_params =
parse_meta_internal_extra_params(meta_internal_extra_params.into_option())?;

let action = UnregisteredRunAction {
executor_preference,
always_print_stderr,
Expand All @@ -392,6 +402,7 @@ pub(crate) fn analysis_actions_methods_run(methods: &mut MethodsBuilder) {
unique_input_inodes,
remote_execution_dependencies: re_dependencies,
remote_execution_custom_image: re_custom_image,
meta_internal_extra_params: extra_params,
};
this.state()?.register_action(
artifacts.inputs,
Expand Down
59 changes: 58 additions & 1 deletion app/buck2_build_api/src/interpreter/rule_defs/command_executor_config.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,12 @@ use buck2_core::execution_types::executor_config::Executor;
use buck2_core::execution_types::executor_config::HybridExecutionLevel;
use buck2_core::execution_types::executor_config::ImagePackageIdentifier;
use buck2_core::execution_types::executor_config::LocalExecutorOptions;
use buck2_core::execution_types::executor_config::MetaInternalExtraParams;
use buck2_core::execution_types::executor_config::PathSeparatorKind;
use buck2_core::execution_types::executor_config::RePlatformFields;
use buck2_core::execution_types::executor_config::RemoteEnabledExecutor;
use buck2_core::execution_types::executor_config::RemoteEnabledExecutorOptions;
use buck2_core::execution_types::executor_config::RemoteExecutionPolicy;
use buck2_core::execution_types::executor_config::RemoteExecutorCustomImage;
use buck2_core::execution_types::executor_config::RemoteExecutorDependency;
use buck2_core::execution_types::executor_config::RemoteExecutorOptions;
Expand Down Expand Up @@ -54,6 +56,8 @@ enum CommandExecutorConfigErrors {
"executor config must specify at least `local_enabled = True` or `remote_enabled = True`"
)]
NoExecutor,
#[error("expected a dict, got `{0}` (type `{1}`)")]
RePolicyNotADict(String, String),
}

#[derive(Debug, Display, NoSerialize, ProvidesStaticType, Allocative)]
Expand Down Expand Up @@ -94,6 +98,7 @@ pub fn register_command_executor_config(builder: &mut GlobalsBuilder) {
/// * `remote_execution_resource_units`: The resources (eg. GPUs) to use for remote execution
/// * `remote_execution_dependencies`: Dependencies for remote execution for this platform
/// * `remote_execution_custom_image`: Custom Tupperware image for remote execution for this platform
/// * `meta_internal_extra_params`: Json dict of extra params to pass to RE related to Meta internal infra.
#[starlark(as_type = StarlarkCommandExecutorConfig)]
fn CommandExecutorConfig<'v>(
#[starlark(require = named)] local_enabled: bool,
Expand Down Expand Up @@ -124,9 +129,12 @@ pub fn register_command_executor_config(builder: &mut GlobalsBuilder) {
#[starlark(default=UnpackList::default(), require = named)]
remote_execution_dependencies: UnpackList<SmallMap<&'v str, &'v str>>,
#[starlark(default = NoneType, require = named)] remote_execution_dynamic_image: Value<'v>,
#[starlark(default = NoneOr::None, require = named)] meta_internal_extra_params: NoneOr<
DictRef<'v>,
>,
) -> starlark::Result<StarlarkCommandExecutorConfig> {
let command_executor_config = {
let remote_execution_max_input_files_mebibytes =
let remote_execution_max_input_files_mebibytes: Option<i32> =
remote_execution_max_input_files_mebibytes.into_option();

let remote_execution_queue_time_threshold_s =
Expand Down Expand Up @@ -167,6 +175,9 @@ pub fn register_command_executor_config(builder: &mut GlobalsBuilder) {
remote_execution_dynamic_image,
)?;

let extra_params =
parse_meta_internal_extra_params(meta_internal_extra_params.into_option())?;

let re_use_case = if remote_execution_use_case.is_none() {
None
} else {
Expand Down Expand Up @@ -280,6 +291,7 @@ pub fn register_command_executor_config(builder: &mut GlobalsBuilder) {
remote_dep_file_cache_enabled,
dependencies: re_dependencies,
custom_image: re_dynamic_image,
meta_internal_extra_params: extra_params,
})
}
(Some(local), None, true) => {
Expand All @@ -296,6 +308,7 @@ pub fn register_command_executor_config(builder: &mut GlobalsBuilder) {
remote_dep_file_cache_enabled,
dependencies: re_dependencies,
custom_image: re_dynamic_image,
meta_internal_extra_params: extra_params,
})
}
// If remote cache is disabled, also disable the remote dep file cache as well
Expand Down Expand Up @@ -415,3 +428,47 @@ pub fn parse_custom_re_image(
drop_host_mount_globs,
}))
}

fn parse_remote_execution_policy(
policy: Option<Value>,
) -> buck2_error::Result<RemoteExecutionPolicy> {
if policy.is_none() {
Ok(RemoteExecutionPolicy::default())
} else {
let re_policy_dict = DictRef::from_value(policy.unwrap().to_value()).ok_or_else(|| {
buck2_error::Error::from(CommandExecutorConfigErrors::RePolicyNotADict(
policy.unwrap().to_value().to_repr(),
policy.unwrap().to_value().get_type().to_owned(),
))
})?;

Ok(RemoteExecutionPolicy {
setup_preference_key: re_policy_dict
.get_str("setup_preference_key")
.and_then(|v| v.unpack_str())
.map(|s| s.to_owned()),
region_preference: re_policy_dict
.get_str("region_preference")
.and_then(|v| v.unpack_str())
.map(|s| s.to_owned()),
priority: re_policy_dict
.get_str("priority")
.and_then(|v| v.unpack_i32())
.map(|i| i.to_owned()),
})
}
}

pub fn parse_meta_internal_extra_params<'v>(
params: Option<DictRef<'v>>,
) -> buck2_error::Result<MetaInternalExtraParams> {
if let Some(params) = params {
Ok(MetaInternalExtraParams {
remote_execution_policy: parse_remote_execution_policy(
params.get_str("remote_execution_policy"),
)?,
})
} else {
Ok(MetaInternalExtraParams::default())
}
}
18 changes: 18 additions & 0 deletions app/buck2_core/src/execution_types/executor_config.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ pub struct RemoteEnabledExecutorOptions {
pub remote_dep_file_cache_enabled: bool,
pub dependencies: Vec<RemoteExecutorDependency>,
pub custom_image: Option<RemoteExecutorCustomImage>,
pub meta_internal_extra_params: MetaInternalExtraParams,
}

#[derive(Debug, buck2_error::Error)]
Expand Down Expand Up @@ -327,3 +328,20 @@ impl CommandExecutorConfig {
}
}
}

/// This struct is used to pass policy info about the action to RE, its data should
/// match the TExecutionPolicy in the RE thrift API.
/// affinity_keys is not defined here because it's already defined in ReActionIdentity
/// duration_ms is not supported because we can't unpack i64 from starlark easily
#[derive(Default, Debug, Clone, Eq, Hash, PartialEq, Allocative)]
pub struct RemoteExecutionPolicy {
pub priority: Option<i32>,
pub region_preference: Option<String>,
pub setup_preference_key: Option<String>,
}

/// This struct is used to pass meta internal params to RE
#[derive(Default, Debug, Clone, Eq, Hash, PartialEq, Allocative)]
pub struct MetaInternalExtraParams {
pub remote_execution_policy: RemoteExecutionPolicy,
}
16 changes: 16 additions & 0 deletions app/buck2_execute/src/execute/request.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ use allocative::Allocative;
use buck2_common::file_ops::FileMetadata;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_common::local_resource_state::LocalResourceState;
use buck2_core::execution_types::executor_config::MetaInternalExtraParams;
use buck2_core::execution_types::executor_config::RemoteExecutorCustomImage;
use buck2_core::execution_types::executor_config::RemoteExecutorDependency;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
Expand Down Expand Up @@ -328,6 +329,8 @@ pub struct CommandExecutionRequest {
remote_execution_dependencies: Vec<RemoteExecutorDependency>,
/// RE custom tupperware image.
remote_execution_custom_image: Option<RemoteExecutorCustomImage>,
/// RE execution policy.
meta_internal_extra_params: MetaInternalExtraParams,
}

impl CommandExecutionRequest {
Expand Down Expand Up @@ -358,6 +361,7 @@ impl CommandExecutionRequest {
remote_dep_file_key: None,
remote_execution_dependencies: Vec::new(),
remote_execution_custom_image: None,
meta_internal_extra_params: MetaInternalExtraParams::default(),
}
}

Expand Down Expand Up @@ -562,6 +566,18 @@ impl CommandExecutionRequest {
pub fn remote_execution_custom_image(&self) -> &Option<RemoteExecutorCustomImage> {
&self.remote_execution_custom_image
}

pub fn with_meta_internal_extra_params(
mut self,
meta_internal_extra_params: MetaInternalExtraParams,
) -> Self {
self.meta_internal_extra_params = meta_internal_extra_params;
self
}

pub fn meta_internal_extra_params(&self) -> &MetaInternalExtraParams {
&self.meta_internal_extra_params
}
}

/// Is an output a file or a directory
Expand Down
18 changes: 18 additions & 0 deletions app/buck2_execute/src/re/client.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ use std::time::Duration;
use allocative::Allocative;
use anyhow::Context;
use buck2_core::buck2_env;
use buck2_core::execution_types::executor_config::MetaInternalExtraParams;
use buck2_core::execution_types::executor_config::RemoteExecutorDependency;
use buck2_core::execution_types::executor_config::RemoteExecutorUseCase;
use buck2_core::fs::project::ProjectRoot;
Expand Down Expand Up @@ -231,6 +232,7 @@ impl RemoteExecutionClient {
re_max_queue_time: Option<Duration>,
re_resource_units: Option<i64>,
knobs: &ExecutorGlobalKnobs,
meta_internal_extra_params: &MetaInternalExtraParams,
) -> buck2_error::Result<ExecuteResponseOrCancelled> {
self.data
.executes
Expand All @@ -246,6 +248,7 @@ impl RemoteExecutionClient {
re_max_queue_time,
re_resource_units,
knobs,
meta_internal_extra_params,
))
.await
}
Expand Down Expand Up @@ -1068,6 +1071,7 @@ impl RemoteExecutionClientImpl {
re_max_queue_time: Option<Duration>,
re_resource_units: Option<i64>,
knobs: &ExecutorGlobalKnobs,
meta_internal_extra_params: &MetaInternalExtraParams,
) -> buck2_error::Result<ExecuteResponseOrCancelled> {
let metadata = RemoteExecutionMetadata {
platform: Some(re_platform(platform)),
Expand All @@ -1087,6 +1091,20 @@ impl RemoteExecutionClientImpl {
skip_cache_lookup: self.skip_remote_cache || skip_cache_read,
execution_policy: Some(TExecutionPolicy {
affinity_keys: vec![identity.affinity_key.clone()],
priority: meta_internal_extra_params
.remote_execution_policy
.priority
.unwrap_or_default(),
region_preference: meta_internal_extra_params
.remote_execution_policy
.region_preference
.clone()
.unwrap_or_default(),
setup_preference_key: meta_internal_extra_params
.remote_execution_policy
.setup_preference_key
.clone()
.unwrap_or_default(),
..Default::default()
}),
action_digest: action_digest.to_re(),
Expand Down
3 changes: 3 additions & 0 deletions app/buck2_execute/src/re/manager.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ use allocative::Allocative;
use async_trait::async_trait;
use buck2_core::async_once_cell::AsyncOnceCell;
use buck2_core::buck2_env;
use buck2_core::execution_types::executor_config::MetaInternalExtraParams;
use buck2_core::execution_types::executor_config::RemoteExecutorDependency;
use buck2_core::execution_types::executor_config::RemoteExecutorUseCase;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
Expand Down Expand Up @@ -405,6 +406,7 @@ impl ManagedRemoteExecutionClient {
re_max_queue_time: Option<Duration>,
re_resource_units: Option<i64>,
knobs: &ExecutorGlobalKnobs,
meta_internal_extra_params: &MetaInternalExtraParams,
) -> buck2_error::Result<ExecuteResponseOrCancelled> {
let use_case = self.re_use_case_override.unwrap_or(use_case);
self.lock()?
Expand All @@ -422,6 +424,7 @@ impl ManagedRemoteExecutionClient {
re_max_queue_time,
re_resource_units,
knobs,
meta_internal_extra_params,
)
.await
}
Expand Down
4 changes: 4 additions & 0 deletions app/buck2_execute_impl/src/executors/re.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ use std::sync::Arc;
use std::time::Duration;

use async_trait::async_trait;
use buck2_core::execution_types::executor_config::MetaInternalExtraParams;
use buck2_core::execution_types::executor_config::RemoteExecutorDependency;
use buck2_core::execution_types::executor_config::RemoteExecutorUseCase;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
Expand Down Expand Up @@ -153,6 +154,7 @@ impl ReExecutor {
digest_config: DigestConfig,
platform: &RE::Platform,
dependencies: impl IntoIterator<Item = &'a RemoteExecutorDependency>,
meta_internal_extra_params: &MetaInternalExtraParams,
) -> ControlFlow<CommandExecutionResult, (CommandExecutionManager, ExecuteResponse)> {
info!(
"RE command line:\n```\n$ {}\n```\n for action `{}`",
Expand All @@ -174,6 +176,7 @@ impl ReExecutor {
self.re_max_queue_time_ms.map(Duration::from_millis),
self.re_resource_units,
&self.knobs,
meta_internal_extra_params,
)
.await;

Expand Down Expand Up @@ -341,6 +344,7 @@ impl PreparedCommandExecutor for ReExecutor {
self.dependencies
.iter()
.chain(remote_execution_dependencies.iter()),
&command.request.meta_internal_extra_params(),
)
.await?;

Expand Down
2 changes: 2 additions & 0 deletions app/buck2_server/src/daemon/common.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ use buck2_core::execution_types::executor_config::CommandGenerationOptions;
use buck2_core::execution_types::executor_config::Executor;
use buck2_core::execution_types::executor_config::HybridExecutionLevel;
use buck2_core::execution_types::executor_config::LocalExecutorOptions;
use buck2_core::execution_types::executor_config::MetaInternalExtraParams;
use buck2_core::execution_types::executor_config::PathSeparatorKind;
use buck2_core::execution_types::executor_config::RePlatformFields;
use buck2_core::execution_types::executor_config::RemoteEnabledExecutor;
Expand Down Expand Up @@ -492,6 +493,7 @@ pub fn get_default_executor_config(host_platform: HostPlatformOverride) -> Comma
remote_dep_file_cache_enabled: false,
dependencies: vec![],
custom_image: None,
meta_internal_extra_params: MetaInternalExtraParams::default(),
})
};

Expand Down
Loading

0 comments on commit 51314c5

Please sign in to comment.