build/rbe/rustc_remote_wrapper.py - fuchsia - Git at Google

 #!/usr/bin/env python3.8
 # Copyright 2022 The Fuchsia Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 _MAIN_HELP = """Wraps a Rust compile command for remote execution (reclient, RBE).

 Given a Rust compile command, this script
 1) identifies all inputs needed to execute the command hermetically (remotely).
 2) identifies all outputs to be retrieved from remote execution.
 3) composes a `rewrapper` (reclient) command to request remote execution.
    This includes downloading remotely produced artifacts.
 4) forwards stdout/stderr back to the local environment

 This script was ported over from bin/rbe/rustc-remote-wrapper.sh.

 Usage:
   rustc_remote_wrapper.py [options] -- RUST-COMPILE-COMMAND

 Options:
   --help | -h : print help and exit
   --local : run the command locally
   --dry-run : print diagnostics and exit without running
   --verbose : print additional diagnostics
   --fsatrace : trace file access (works locally and remotely)

   All unknown options are forwarded to rewrapper.

 The RUST-COMPILE-COMMAND supports the following pseudo-flags, which are
 filtered out prior to execution, and forwarded to rewrapper:

   --remote-disable : same as --local
   --remote-inputs FILE,... : forwarded as --inputs to rewrapper
   --remote-outputs FILE,... : forwarded as --output_files to rewrapper
   --remote-flag OPT : forwarded as OPT (can be flag) to rewrapper

 This allows rustflags in GN to influence remote execution parameters.
 """

 import argparse
 import os
 import sys
 from typing import Iterable, Sequence, Tuple


 # This is provided as a function so that tests can fake it.
 def _dependent_globals(this_script, cwd=os.curdir):
     """Compute set of related globals that depend on this script's path."""
     script_dir = os.path.dirname(this_script)

     # This script lives under 'build/rbe', so the path to the root is '../..'.
     default_project_root = os.path.realpath(
         os.path.join(script_dir, '..', '..'))

     # This is the relative path to the project root dir from the build output dir.
     project_root_rel = os.path.relpath(default_project_root, start=cwd)

     # This is the relative path to the build output dir from the project root dir.
     build_subdir = os.path.relpath(cwd, start=default_project_root)

     return argparse.Namespace(
         script_dir=script_dir,
         default_project_root=default_project_root,
         project_root_rel=project_root_rel,
         build_subdir=build_subdir,

         # This is the script that eventually calls 'rewrapper' (reclient).
         generic_remote_action_wrapper=os.path.join(
             script_dir, 'fuchsia-rbe-action.sh'),

         # This command is used to check local determinism.
         check_determinism_command=[
             os.path.join(
                 default_project_root, 'build', 'tracer', 'output_cacher.py'),
             '--check-repeatability',
         ],

         # The path to the prebuilt fsatrace in Fuchsia's project tree.
         fsatrace_path=os.path.join(
             project_root_rel, 'prebuilt', 'fsatrace', 'fsatrace'),
         detail_diff=os.path.join(script_dir, 'detail-diff.sh'),
     )


 # Global variables
 _ARGV = sys.argv
 _SCRIPT_PATH = _ARGV[0]
 _GLOBALS = _dependent_globals(_SCRIPT_PATH)

 # Global constants

 # This is a known path where remote execution occurs.
 _REMOTE_PROJECT_ROOT = '/b/f/w'

 # Use this env both locally and remotely.
 _ENV = '/usr/bin/env'


 def apply_remote_flags_from_pseudo_flags(
         main_config: argparse.Namespace, command_params: argparse.Namespace):
     """Apply some flags from the command to the rewrapper configuration.

     Args:
       main_config: main configuration (modified-by-reference).
       command_params: command parameters inferred from parsing the command.
     """
     if command_params.remote_disable:
         main_config.local = True


 def parse_main_args(
         command: Sequence[str]) -> Tuple[argparse.Namespace, Sequence[str]]:
     """Scan a the main args for parameters.

     Interface matches that of argparse.ArgumentParser.parse_known_args().
     All unhandled tokens are intended to be forwarded to rewrapper.

     Args:
       command: command, like sys.argv

     Returns:
       A namespace struct with parameters, and a sequence of unhandled tokens.
     """
     params = argparse.Namespace(
         help=None,
         local=False,
         dry_run=False,
         verbose=False,
         fsatrace=False,
         command=[],  # compile command
     )
     forward_to_rewrapper = []

     opt_arg_func = None
     for i, token in enumerate(command):
         # Handle detached --option argument
         if opt_arg_func is not None:
             opt_arg_func(token)
             opt_arg_func = None
             continue

         opt, sep, arg = token.partition('=')

         if token in {'--help', '-h'}:
             params.help = _MAIN_HELP
             break
         if token == '--local':
             params.local = True
         elif token == '--dry-run':
             params.dry_run = True
         elif token == '--verbose':
             params.verbose = True
         elif token == '--fsatrace':
             params.fsatrace = True
         elif token == '--':  # stop option processing
             params.command = command[i + 1:]
             break
         # TODO: if needed, add --project-root override
         # TODO: if needed, add --source override
         # TODO: if needed, add --depfile override
         else:
             forward_to_rewrapper.append(token)

     return params, forward_to_rewrapper


 def filter_compile_command(
         command: Sequence[str]) -> Tuple[argparse.Namespace, Sequence[str]]:
     """Scan a command for remote execution parameters, filter out pseudo-flags.

     Interface matches that of argparse.ArgumentParser.parse_known_args().

     Args:
       command: command to scan

     Returns:
       a namespace struct with parameters, and a filtered version of the
       command to remotely execute (with pseudo-flags removed).
     """
     params = argparse.Namespace(
         remote_disable=False,
         remote_inputs=[],
         remote_outputs=[],
         remote_flags=[],
     )
     forward_as_compile_command = []

     opt_arg_func = None
     for token in command:
         # Handle detached --option argument
         if opt_arg_func is not None:
             opt_arg_func(token)
             opt_arg_func = None
             continue

         opt, sep, arg = token.partition('=')

         if token == '--remote-disable':
             params.remote_disable = True
         elif token == '--remote-inputs':
             opt_arg_func = lambda x: setattr(params, 'remote_inputs', x)
         elif opt == '--remote-inputs' and sep == '=':
             params.remote_inputs = arg
         elif token == '--remote-outputs':
             opt_arg_func = lambda x: setattr(params, 'remote_outputs', x)
         elif opt == '--remote-outputs' and sep == '=':
             params.remote_outputs = arg
         elif token == '--remote-flag':
             opt_arg_func = lambda x: params.remote_flags.append(x)
         elif opt == '--remote-flag' and sep == '=':
             params.remote_flags.append(arg)
         else:
             forward_as_compile_command.append(token)

     return params, forward_as_compile_command


 # string.removeprefix() only appeared in python 3.9
 # This is needed in some places to workaround b/203540556 (reclient).
 def remove_dot_slash_prefix(text: str) -> str:
     if text.startswith('./'):
         return text[2:]
     return text


 # string.removesuffix() only appeared in python 3.9
 def remove_suffix(text: str, suffix: str) -> str:
     if text.endswith(suffix):
         return text[:-len(suffix)]
     return text


 def parse_rust_compile_command(
     compile_command: Sequence[str],
     globals: argparse.Namespace,
 ) -> argparse.Namespace:
     """Scans a Rust compile command for remote execution parameters.

     Args:
       compile_command: the full (local) Rust compile command, which
         may contain environment variable prefixes.
       globals: variables that depend on the current working directory.

     Returns:
       A namespace of variables containing remote execution information.
     """
     params = argparse.Namespace(
         depfile=None,
         dep_only_command=[_ENV],  # a modified copy of compile_command
         emit_llvm_ir = False,
         emit_llvm_bc = False,
         output=None,
         extra_filename='',
         target_triple='',
     )
     opt_arg_func = None
     for token in compile_command:
         if opt_arg_func is not None:
             opt_arg_func(token)
             opt_arg_func = None
             params.dep_only_command.append(token)
             continue

         opt, sep, arg = token.partition('=')

         if token == '-o':
             opt_arg_func = lambda x: setattr(
                 params, 'output', remove_dot_slash_prefix(x))

         elif token == '--target':
             opt_arg_func = lambda x: setattr(params, 'target_triple', x)

         # Create a modified copy of the compile command that will be
         # used to only generate a depfile.
         # Rewrite the --emit token to do exactly this, ignoring
         # all other requested emit outputs.
         elif opt == '--emit' and sep == '=':
             emit_args = arg.split(',')
             for emit_arg in emit_args:
                 emit_key, emit_sep, emit_value = emit_arg.partition('=')
                 if emit_key == 'dep-info' and emit_sep == '=':
                     params.depfile = remove_dot_slash_prefix(emit_value)
                 elif emit_arg == 'llvm-ir':
                     params.emit_llvm_ir = True
                 elif emit_arg == 'llvm-bc':
                     params.emit_llvm_bc = True

             # Tell rustc to report all transitive *library* dependencies,
             # not just the sources, because these all need to be uploaded.
             # This includes (prebuilt) system libraries as well.
             # TODO(https://fxbug.dev/78292): this -Z flag is not known to be stable yet.
             params.dep_only_command += [
                 '-Zbinary-dep-depinfo',
                 f'--emit=dep-info={params.depfile}.nolink'
             ]
             continue

         elif token == '-Cextra-filename':
             opt_arg_func = lambda x: setattr(params, 'extra_filename', x)
         elif opt == '-Cextra-filename' and sep == '=':
             params.extra_filename = arg

         # By default, copy over most tokens for depfile generation.
         params.dep_only_command.append(token)

     return params


 def main(argv: Sequence[str]):
     # Parse flags, and forward all unhandled flags to rewrapper.
     main_config, forwarded_rewrapper_args = parse_main_args(
         argv[1:])  # drop argv[0], which is this script

     # Exit on --help.
     if main_config.help is not None:
         print(main_config.help)
         return 0

     # The command to run remotely is in main_config.command.
     # Remove pseudo flags from the remote command.
     remote_params, filtered_command = filter_compile_command(
         main_config.command)
     forwarded_rewrapper_args.extend(remote_params.remote_flags)

     # Import some remote parameters back to the main_config.
     apply_remote_flags_from_pseudo_flags(main_config, remote_params)

     compile_params = parse_rust_compile_command(filtered_command, _GLOBALS)

     # TODO: infer inputs and outputs for remote execution
     remote_inputs = remote_params.remote_inputs
     remote_outputs = remote_params.remote_outputs

     if compile_params.output is None:
         return 1

     output_base = remove_suffix(compile_params.output, '.rlib')

     if compile_params.emit_llvm_ir:
         remote_outputs.append(
             f'{output_base}{compile_params.extra_filename}.ll')

     if compile_params.emit_llvm_bc:
         remote_outputs.append(
             f'{output_base}{compile_params.extra_filename}.bc')

     # Use the dep-scanning command from compile_params.dep_only_command

     # TODO: construct remote execution command and run it
     # TODO: post-execution diagnostics

     return 0


 if __name__ == "__main__":
     sys.exit(main(_ARGV))
	#!/usr/bin/env python3.8
	# Copyright 2022 The Fuchsia Authors. All rights reserved.
	# Use of this source code is governed by a BSD-style license that can be
	# found in the LICENSE file.
	_MAIN_HELP = """Wraps a Rust compile command for remote execution (reclient, RBE).

	Given a Rust compile command, this script
	1) identifies all inputs needed to execute the command hermetically (remotely).
	2) identifies all outputs to be retrieved from remote execution.
	3) composes a `rewrapper` (reclient) command to request remote execution.
	This includes downloading remotely produced artifacts.
	4) forwards stdout/stderr back to the local environment

	This script was ported over from bin/rbe/rustc-remote-wrapper.sh.

	Usage:
	rustc_remote_wrapper.py [options] -- RUST-COMPILE-COMMAND

	Options:
	--help \| -h : print help and exit
	--local : run the command locally
	--dry-run : print diagnostics and exit without running
	--verbose : print additional diagnostics
	--fsatrace : trace file access (works locally and remotely)

	All unknown options are forwarded to rewrapper.

	The RUST-COMPILE-COMMAND supports the following pseudo-flags, which are
	filtered out prior to execution, and forwarded to rewrapper:

	--remote-disable : same as --local
	--remote-inputs FILE,... : forwarded as --inputs to rewrapper
	--remote-outputs FILE,... : forwarded as --output_files to rewrapper
	--remote-flag OPT : forwarded as OPT (can be flag) to rewrapper

	This allows rustflags in GN to influence remote execution parameters.
	"""

	import argparse
	import os
	import sys
	from typing import Iterable, Sequence, Tuple


	# This is provided as a function so that tests can fake it.
	def _dependent_globals(this_script, cwd=os.curdir):
	"""Compute set of related globals that depend on this script's path."""
	script_dir = os.path.dirname(this_script)

	# This script lives under 'build/rbe', so the path to the root is '../..'.
	default_project_root = os.path.realpath(
	os.path.join(script_dir, '..', '..'))

	# This is the relative path to the project root dir from the build output dir.
	project_root_rel = os.path.relpath(default_project_root, start=cwd)

	# This is the relative path to the build output dir from the project root dir.
	build_subdir = os.path.relpath(cwd, start=default_project_root)

	return argparse.Namespace(
	script_dir=script_dir,
	default_project_root=default_project_root,
	project_root_rel=project_root_rel,
	build_subdir=build_subdir,

	# This is the script that eventually calls 'rewrapper' (reclient).
	generic_remote_action_wrapper=os.path.join(
	script_dir, 'fuchsia-rbe-action.sh'),

	# This command is used to check local determinism.
	check_determinism_command=[
	os.path.join(
	default_project_root, 'build', 'tracer', 'output_cacher.py'),
	'--check-repeatability',
	],

	# The path to the prebuilt fsatrace in Fuchsia's project tree.
	fsatrace_path=os.path.join(
	project_root_rel, 'prebuilt', 'fsatrace', 'fsatrace'),
	detail_diff=os.path.join(script_dir, 'detail-diff.sh'),
	)


	# Global variables
	_ARGV = sys.argv
	_SCRIPT_PATH = _ARGV[0]
	_GLOBALS = _dependent_globals(_SCRIPT_PATH)

	# Global constants

	# This is a known path where remote execution occurs.
	_REMOTE_PROJECT_ROOT = '/b/f/w'

	# Use this env both locally and remotely.
	_ENV = '/usr/bin/env'


	def apply_remote_flags_from_pseudo_flags(
	main_config: argparse.Namespace, command_params: argparse.Namespace):
	"""Apply some flags from the command to the rewrapper configuration.

	Args:
	main_config: main configuration (modified-by-reference).
	command_params: command parameters inferred from parsing the command.
	"""
	if command_params.remote_disable:
	main_config.local = True


	def parse_main_args(
	command: Sequence[str]) -> Tuple[argparse.Namespace, Sequence[str]]:
	"""Scan a the main args for parameters.

	Interface matches that of argparse.ArgumentParser.parse_known_args().
	All unhandled tokens are intended to be forwarded to rewrapper.

	Args:
	command: command, like sys.argv

	Returns:
	A namespace struct with parameters, and a sequence of unhandled tokens.
	"""
	params = argparse.Namespace(
	help=None,
	local=False,
	dry_run=False,
	verbose=False,
	fsatrace=False,
	command=[], # compile command
	)
	forward_to_rewrapper = []

	opt_arg_func = None
	for i, token in enumerate(command):
	# Handle detached --option argument
	if opt_arg_func is not None:
	opt_arg_func(token)
	opt_arg_func = None
	continue

	opt, sep, arg = token.partition('=')

	if token in {'--help', '-h'}:
	params.help = _MAIN_HELP
	break
	if token == '--local':
	params.local = True
	elif token == '--dry-run':
	params.dry_run = True
	elif token == '--verbose':
	params.verbose = True
	elif token == '--fsatrace':
	params.fsatrace = True
	elif token == '--': # stop option processing
	params.command = command[i + 1:]
	break
	# TODO: if needed, add --project-root override
	# TODO: if needed, add --source override
	# TODO: if needed, add --depfile override
	else:
	forward_to_rewrapper.append(token)

	return params, forward_to_rewrapper


	def filter_compile_command(
	command: Sequence[str]) -> Tuple[argparse.Namespace, Sequence[str]]:
	"""Scan a command for remote execution parameters, filter out pseudo-flags.

	Interface matches that of argparse.ArgumentParser.parse_known_args().

	Args:
	command: command to scan

	Returns:
	a namespace struct with parameters, and a filtered version of the
	command to remotely execute (with pseudo-flags removed).
	"""
	params = argparse.Namespace(
	remote_disable=False,
	remote_inputs=[],
	remote_outputs=[],
	remote_flags=[],
	)
	forward_as_compile_command = []

	opt_arg_func = None
	for token in command:
	# Handle detached --option argument
	if opt_arg_func is not None:
	opt_arg_func(token)
	opt_arg_func = None
	continue

	opt, sep, arg = token.partition('=')

	if token == '--remote-disable':
	params.remote_disable = True
	elif token == '--remote-inputs':
	opt_arg_func = lambda x: setattr(params, 'remote_inputs', x)
	elif opt == '--remote-inputs' and sep == '=':
	params.remote_inputs = arg
	elif token == '--remote-outputs':
	opt_arg_func = lambda x: setattr(params, 'remote_outputs', x)
	elif opt == '--remote-outputs' and sep == '=':
	params.remote_outputs = arg
	elif token == '--remote-flag':
	opt_arg_func = lambda x: params.remote_flags.append(x)
	elif opt == '--remote-flag' and sep == '=':
	params.remote_flags.append(arg)
	else:
	forward_as_compile_command.append(token)

	return params, forward_as_compile_command


	# string.removeprefix() only appeared in python 3.9
	# This is needed in some places to workaround b/203540556 (reclient).
	def remove_dot_slash_prefix(text: str) -> str:
	if text.startswith('./'):
	return text[2:]
	return text


	# string.removesuffix() only appeared in python 3.9
	def remove_suffix(text: str, suffix: str) -> str:
	if text.endswith(suffix):
	return text[:-len(suffix)]
	return text


	def parse_rust_compile_command(
	compile_command: Sequence[str],
	globals: argparse.Namespace,
	) -> argparse.Namespace:
	"""Scans a Rust compile command for remote execution parameters.

	Args:
	compile_command: the full (local) Rust compile command, which
	may contain environment variable prefixes.
	globals: variables that depend on the current working directory.

	Returns:
	A namespace of variables containing remote execution information.
	"""
	params = argparse.Namespace(
	depfile=None,
	dep_only_command=[_ENV], # a modified copy of compile_command
	emit_llvm_ir = False,
	emit_llvm_bc = False,
	output=None,
	extra_filename='',
	target_triple='',
	)
	opt_arg_func = None
	for token in compile_command:
	if opt_arg_func is not None:
	opt_arg_func(token)
	opt_arg_func = None
	params.dep_only_command.append(token)
	continue

	opt, sep, arg = token.partition('=')

	if token == '-o':
	opt_arg_func = lambda x: setattr(
	params, 'output', remove_dot_slash_prefix(x))

	elif token == '--target':
	opt_arg_func = lambda x: setattr(params, 'target_triple', x)

	# Create a modified copy of the compile command that will be
	# used to only generate a depfile.
	# Rewrite the --emit token to do exactly this, ignoring
	# all other requested emit outputs.
	elif opt == '--emit' and sep == '=':
	emit_args = arg.split(',')
	for emit_arg in emit_args:
	emit_key, emit_sep, emit_value = emit_arg.partition('=')
	if emit_key == 'dep-info' and emit_sep == '=':
	params.depfile = remove_dot_slash_prefix(emit_value)
	elif emit_arg == 'llvm-ir':
	params.emit_llvm_ir = True
	elif emit_arg == 'llvm-bc':
	params.emit_llvm_bc = True

	# Tell rustc to report all transitive library dependencies,
	# not just the sources, because these all need to be uploaded.
	# This includes (prebuilt) system libraries as well.
	# TODO(https://fxbug.dev/78292): this -Z flag is not known to be stable yet.
	params.dep_only_command += [
	'-Zbinary-dep-depinfo',
	f'--emit=dep-info={params.depfile}.nolink'
	]
	continue

	elif token == '-Cextra-filename':
	opt_arg_func = lambda x: setattr(params, 'extra_filename', x)
	elif opt == '-Cextra-filename' and sep == '=':
	params.extra_filename = arg

	# By default, copy over most tokens for depfile generation.
	params.dep_only_command.append(token)

	return params


	def main(argv: Sequence[str]):
	# Parse flags, and forward all unhandled flags to rewrapper.
	main_config, forwarded_rewrapper_args = parse_main_args(
	argv[1:]) # drop argv[0], which is this script

	# Exit on --help.
	if main_config.help is not None:
	print(main_config.help)
	return 0

	# The command to run remotely is in main_config.command.
	# Remove pseudo flags from the remote command.
	remote_params, filtered_command = filter_compile_command(
	main_config.command)
	forwarded_rewrapper_args.extend(remote_params.remote_flags)

	# Import some remote parameters back to the main_config.
	apply_remote_flags_from_pseudo_flags(main_config, remote_params)

	compile_params = parse_rust_compile_command(filtered_command, _GLOBALS)

	# TODO: infer inputs and outputs for remote execution
	remote_inputs = remote_params.remote_inputs
	remote_outputs = remote_params.remote_outputs

	if compile_params.output is None:
	return 1

	output_base = remove_suffix(compile_params.output, '.rlib')

	if compile_params.emit_llvm_ir:
	remote_outputs.append(
	f'{output_base}{compile_params.extra_filename}.ll')

	if compile_params.emit_llvm_bc:
	remote_outputs.append(
	f'{output_base}{compile_params.extra_filename}.bc')

	# Use the dep-scanning command from compile_params.dep_only_command

	# TODO: construct remote execution command and run it
	# TODO: post-execution diagnostics

	return 0


	if __name__ == "__main__":
	sys.exit(main(_ARGV))