diff --git a/src/python/pants/core/util_rules/environments.py b/src/python/pants/core/util_rules/environments.py index baf653244a5..6c5da245117 100644 --- a/src/python/pants/core/util_rules/environments.py +++ b/src/python/pants/core/util_rules/environments.py @@ -16,7 +16,7 @@ from pants.engine.internals.scheduler import SchedulerSession from pants.engine.internals.selectors import Params from pants.engine.platform import Platform -from pants.engine.rules import Get, MultiGet, QueryRule, collect_rules, rule +from pants.engine.rules import Get, MultiGet, QueryRule, collect_rules, rule, rule_helper from pants.engine.target import ( COMMON_TARGET_FIELDS, Field, @@ -85,6 +85,11 @@ class EnvironmentField(StringField): ) +class FallbackEnvironmentField(StringField): + alias = "fallback_environment" + default = None + + class CompatiblePlatformsField(StringSequenceField): alias = "compatible_platforms" default = tuple(plat.value for plat in Platform) @@ -101,9 +106,28 @@ class CompatiblePlatformsField(StringSequenceField): ) +class LocalFallbackEnvironmentField(FallbackEnvironmentField): + help = softwrap( + f""" + The environment to fallback to when this local environment cannot be used because the + field `{CompatiblePlatformsField.alias}` is not compatible with the local host. + + Must be an environment name from the option `[environments-preview].names`, the + special string `{LOCAL_ENVIRONMENT_MATCHER}` to use the relevant local environment, or the + Python value `None` to error when this specific local environment cannot be used. + + Tip: when targeting Linux, it can be particularly helpful to fallback to a + `docker_environment` or `remote_environment` target. That allows you to prefer using the + local host when possible, which often has less overhead (particularly compared to Docker). + If the local host is not compatible, then Pants will use Docker or remote execution to + still run in a similar environment. + """ + ) + + class LocalEnvironmentTarget(Target): alias = "_local_environment" - core_fields = (*COMMON_TARGET_FIELDS, CompatiblePlatformsField) + core_fields = (*COMMON_TARGET_FIELDS, CompatiblePlatformsField, LocalFallbackEnvironmentField) help = softwrap( """ Configuration of environment variables and search paths for running Pants locally. @@ -168,9 +192,28 @@ def docker_platform_field_default_factory( return FieldDefaultFactoryResult(lambda f: f.normalized_value) +class DockerFallbackEnvironmentField(FallbackEnvironmentField): + help = softwrap( + f""" + The environment to fallback to when this Docker environment cannot be used because the + field `{DockerPlatformField.alias}` is not compatible with the local host's CPU + architecture. (This is only an issue when the local host is Linux; macOS is fine.) + + Must be an environment name from the option `[environments-preview].names`, the + special string `{LOCAL_ENVIRONMENT_MATCHER}` to use the relevant local environment, or the + Python value `None` to error when this specific Docker environment cannot be used. + """ + ) + + class DockerEnvironmentTarget(Target): alias = "_docker_environment" - core_fields = (*COMMON_TARGET_FIELDS, DockerImageField, DockerPlatformField) + core_fields = ( + *COMMON_TARGET_FIELDS, + DockerImageField, + DockerPlatformField, + DockerFallbackEnvironmentField, + ) help = softwrap( """ Configuration of a Docker image used for building your code, including the environment @@ -204,9 +247,7 @@ class RemoteExtraPlatformPropertiesField(StringSequenceField): ) -class RemoteFallbackEnvironmentField(StringField): - alias = "fallback_environment" - default = None +class RemoteFallbackEnvironmentField(FallbackEnvironmentField): help = softwrap( f""" The environment to fallback to when remote execution is disabled via the global option @@ -273,7 +314,7 @@ class UnrecognizedEnvironmentError(Exception): pass -class RemoteExecutionDisabledError(Exception): +class NoFallbackEnvironmentError(Exception): pass @@ -405,6 +446,22 @@ async def determine_local_environment( return ChosenLocalEnvironmentName(result_name) +@rule_helper +async def _apply_fallback_environment(env_tgt: Target, error_msg: str) -> EnvironmentName: + fallback_field = env_tgt[FallbackEnvironmentField] + if fallback_field.value is None: + raise NoFallbackEnvironmentError(error_msg) + return await Get( + EnvironmentName, + EnvironmentNameRequest( + fallback_field.value, + description_of_origin=( + f"the `{fallback_field.alias}` field of the target {env_tgt.address}" + ), + ), + ) + + @rule async def resolve_environment_name( request: EnvironmentNameRequest, @@ -426,34 +483,74 @@ async def resolve_environment_name( """ ) ) + + # Get the target so that we can apply the environment_fallback field, if relevant. env_tgt = await Get(EnvironmentTarget, EnvironmentName(request.raw_value)) if env_tgt.val is None: raise AssertionError(f"EnvironmentTarget.val is None for the name `{request.raw_value}`") - # If remote execution is disabled and it's a remote environment, try falling back. - if not global_options.remote_execution and env_tgt.val.has_field( - RemoteFallbackEnvironmentField + if ( + env_tgt.val.has_field(RemoteFallbackEnvironmentField) + and not global_options.remote_execution ): - fallback_field = env_tgt.val[RemoteFallbackEnvironmentField] - if fallback_field.value is None: - raise RemoteExecutionDisabledError( - softwrap( - f""" - The global option `--remote-execution` is set to false, but the remote - environment `{request.raw_value}` is used in {request.description_of_origin}. - - Either enable the option `--remote-execution`, or set the field - `{fallback_field.alias}` for the target {env_tgt.val.address}. - """ - ) - ) - return await Get( - EnvironmentName, - EnvironmentNameRequest( - fallback_field.value, - description_of_origin=( - f"the `{fallback_field.alias}` field of the target {env_tgt.val.address}" - ), + return await _apply_fallback_environment( + env_tgt.val, + error_msg=softwrap( + f""" + The global option `--remote-execution` is set to false, but the remote + environment `{request.raw_value}` is used in {request.description_of_origin}. + + Either enable the option `--remote-execution`, or set the field + `{FallbackEnvironmentField.alias}` for the target {env_tgt.val.address}. + """ + ), + ) + + localhost_platform = Platform.create_for_localhost().value + + if ( + env_tgt.val.has_field(DockerFallbackEnvironmentField) + and localhost_platform in (Platform.linux_x86_64.value, Platform.linux_arm64.value) + and localhost_platform != env_tgt.val[DockerPlatformField].normalized_value.value + ): + return await _apply_fallback_environment( + env_tgt.val, + error_msg=softwrap( + f""" + The docker environment `{request.raw_value}` is specified in + {request.description_of_origin}, but it cannot be used because the local host has + the platform `{localhost_platform}` and the Docker environment has the platform + {env_tgt.val[DockerPlatformField].normalized_value}. + + Consider setting the field `{FallbackEnvironmentField.alias}` for the target + {env_tgt.val.address}, such as to a `docker_environment` target that sets + `{DockerPlatformField.alias}` to `{localhost_platform}`. Alternatively, consider + not explicitly setting the field `{DockerPlatformField.alias}` for the target + {env_tgt.val.address} because the default behavior is to use the CPU architecture + of the current host for the platform (although this requires the docker image + supports that CPU architecture). + """ + ), + ) + + if ( + env_tgt.val.has_field(LocalFallbackEnvironmentField) + and localhost_platform not in env_tgt.val[CompatiblePlatformsField].value + ): + return await _apply_fallback_environment( + env_tgt.val, + error_msg=softwrap( + f""" + The local environment `{request.raw_value}` was specified in + {request.description_of_origin}, but it is not compatible with the current + machine's platform: {localhost_platform}. The environment only works with the + platforms: {env_tgt.val[CompatiblePlatformsField].value} + + Consider setting the the field `{FallbackEnvironmentField.alias}` for the target + {env_tgt.val.address}, such as to a `docker_environment` or `remote_environment` + target. You can also set that field to another `local_environment` target, such as + one that is compatible with the current platform {localhost_platform}. + """ ), ) diff --git a/src/python/pants/core/util_rules/environments_test.py b/src/python/pants/core/util_rules/environments_test.py index cfaafd9b829..a1390e84f3e 100644 --- a/src/python/pants/core/util_rules/environments_test.py +++ b/src/python/pants/core/util_rules/environments_test.py @@ -15,25 +15,36 @@ AllEnvironmentTargets, AmbiguousEnvironmentError, ChosenLocalEnvironmentName, + CompatiblePlatformsField, DockerEnvironmentTarget, DockerImageField, + DockerPlatformField, EnvironmentField, EnvironmentName, EnvironmentNameRequest, + EnvironmentsSubsystem, EnvironmentTarget, + FallbackEnvironmentField, LocalEnvironmentTarget, NoCompatibleEnvironmentError, + NoFallbackEnvironmentError, RemoteEnvironmentTarget, - RemoteExecutionDisabledError, RemoteExtraPlatformPropertiesField, UnrecognizedEnvironmentError, extract_process_config_from_environment, + resolve_environment_name, ) from pants.engine.platform import Platform from pants.engine.target import FieldSet, OptionalSingleSourceField, Target from pants.option.global_options import GlobalOptions from pants.testutil.option_util import create_subsystem -from pants.testutil.rule_runner import QueryRule, RuleRunner, engine_error, run_rule_with_mocks +from pants.testutil.rule_runner import ( + MockGet, + QueryRule, + RuleRunner, + engine_error, + run_rule_with_mocks, +) @pytest.fixture @@ -195,10 +206,11 @@ def test_resolve_environment_name(rule_runner: RuleRunner) -> None: "BUILD": dedent( """\ _local_environment(name='local') - # Intentionally set this to no platforms so that it cannot be autodiscovered. - _local_environment(name='hardcoded', compatible_platforms=[]) + _local_environment( + name='local-fallback', compatible_platforms=[], fallback_environment='local' + ) _docker_environment(name='docker', image="centos6:latest") - _remote_environment(name='remote-no-fallback', fallback_environment=None) + _remote_environment(name='remote-no-fallback') _remote_environment(name='remote-fallback', fallback_environment="docker") _remote_environment(name='remote-bad-fallback', fallback_environment="fake") """ @@ -215,12 +227,12 @@ def get_name(v: str) -> EnvironmentName: assert get_name(LOCAL_ENVIRONMENT_MATCHER).val is None # Else, error for unrecognized names. with engine_error(UnrecognizedEnvironmentError): - get_name("hardcoded") + get_name("local") env_names_arg = ( "--environments-preview-names={" + "'local': '//:local', " - + "'hardcoded': '//:hardcoded', " + + "'local-fallback': '//:local-fallback', " + "'docker': '//:docker', " + "'remote-no-fallback': '//:remote-no-fallback', " + "'remote-fallback': '//:remote-fallback'," @@ -228,19 +240,120 @@ def get_name(v: str) -> EnvironmentName: ) rule_runner.set_options([env_names_arg, "--remote-execution"]) assert get_name(LOCAL_ENVIRONMENT_MATCHER).val == "local" - for name in ("hardcoded", "docker", "remote-no-fallback", "remote-fallback"): + for name in ("local", "docker", "remote-no-fallback", "remote-fallback"): assert get_name(name).val == name with engine_error(UnrecognizedEnvironmentError): get_name("fake") + assert get_name("local-fallback").val == "local" + rule_runner.set_options([env_names_arg, "--no-remote-execution"]) assert get_name("remote-fallback").val == "docker" - with engine_error(RemoteExecutionDisabledError): + with engine_error(NoFallbackEnvironmentError): get_name("remote-no-fallback") with engine_error(UnrecognizedEnvironmentError): get_name("remote-bad-fallback") +def test_resolve_environment_name_local_and_docker_fallbacks(monkeypatch) -> None: + # We can't monkeypatch the Platform with RuleRunner, so instead use run_run_with_mocks. + def get_env_name(env_tgt: Target, platform: Platform) -> str | None: + monkeypatch.setattr(Platform, "create_for_localhost", lambda: platform) + result = run_rule_with_mocks( + resolve_environment_name, + rule_args=[ + EnvironmentNameRequest("env", description_of_origin="foo"), + create_subsystem(EnvironmentsSubsystem, names={"env": "", "fallback": ""}), + create_subsystem(GlobalOptions, remote_execution=False), + ], + mock_gets=[ + MockGet( + output_type=ChosenLocalEnvironmentName, + input_types=(), + mock=lambda: ChosenLocalEnvironmentName(None), + ), + MockGet( + output_type=EnvironmentTarget, + input_types=(EnvironmentName,), + mock=lambda _: EnvironmentTarget(env_tgt), + ), + MockGet( + output_type=EnvironmentName, + input_types=(EnvironmentNameRequest,), + mock=lambda req: EnvironmentName(req.raw_value), + ), + ], + ).val + return result # type: ignore[no-any-return] + + def create_local_tgt( + *, compatible_platforms: list[Platform] | None = None, fallback: bool = False + ) -> LocalEnvironmentTarget: + return LocalEnvironmentTarget( + { + CompatiblePlatformsField.alias: [plat.value for plat in compatible_platforms] + if compatible_platforms + else None, + FallbackEnvironmentField.alias: "fallback" if fallback else None, + }, + Address("envs"), + ) + + assert get_env_name(create_local_tgt(), Platform.linux_arm64) == "env" + assert ( + get_env_name( + create_local_tgt(compatible_platforms=[Platform.linux_x86_64], fallback=True), + Platform.linux_arm64, + ) + == "fallback" + ) + with pytest.raises(NoFallbackEnvironmentError): + get_env_name( + create_local_tgt(compatible_platforms=[Platform.linux_x86_64]), + Platform.linux_arm64, + ) + + def create_docker_tgt( + *, platform: Platform | None = None, fallback: bool = False + ) -> DockerEnvironmentTarget: + return DockerEnvironmentTarget( + { + DockerImageField.alias: "image", + DockerPlatformField.alias: platform.value if platform else None, + FallbackEnvironmentField.alias: "fallback" if fallback else None, + }, + Address("envs"), + ) + + # If the Docker platform is not set, we default to the CPU arch. So regardless of localhost, + # the Docker environment should be used. + assert get_env_name(create_docker_tgt(), Platform.linux_arm64) == "env" + + # The Docker env can be used if we're on macOS, or on Linux and the CPU arch matches. + for plat in (Platform.macos_arm64, Platform.macos_x86_64, Platform.linux_x86_64): + assert get_env_name(create_docker_tgt(platform=Platform.linux_x86_64), plat) == "env" + for plat in (Platform.macos_arm64, Platform.macos_x86_64, Platform.linux_arm64): + assert get_env_name(create_docker_tgt(platform=Platform.linux_arm64), plat) == "env" + + # But if on Linux and a different CPU arch is used, fallback. + assert ( + get_env_name( + create_docker_tgt(platform=Platform.linux_x86_64, fallback=True), Platform.linux_arm64 + ) + == "fallback" + ) + assert ( + get_env_name( + create_docker_tgt(platform=Platform.linux_arm64, fallback=True), Platform.linux_x86_64 + ) + == "fallback" + ) + with pytest.raises(NoFallbackEnvironmentError): + get_env_name(create_docker_tgt(platform=Platform.linux_x86_64), Platform.linux_arm64) + with pytest.raises(NoFallbackEnvironmentError): + get_env_name(create_docker_tgt(platform=Platform.linux_arm64), Platform.linux_x86_64) + + def test_resolve_environment_tgt(rule_runner: RuleRunner) -> None: rule_runner.write_files({"BUILD": "_local_environment(name='env')"}) rule_runner.set_options(