def _GetToolEnv(env=None):
"""Generate the environment that should be used for the subprocess.
Args:
env: {str, str}, An existing environment to augment. If None, the current
environment will be cloned and used as the base for the subprocess.
Returns:
The modified env.
"""
if env is None:
env = dict(os.environ)
env = encoding.EncodeEnv(env)
encoding.SetEncodedValue(env, 'CLOUDSDK_WRAPPER', '1')
# Flags can set properties which override the properties file and the existing
# env vars. We need to propagate them to children processes through the
# environment so that those commands will use the same settings.
for s in properties.VALUES:
for p in s:
encoding.SetEncodedValue(env, p.EnvironmentName(),
p.Get(required=False, validate=False))
# Configuration needs to be handled separately because it's not a real
# property (although it behaves like one).
encoding.SetEncodedValue(
env, config.CLOUDSDK_ACTIVE_CONFIG_NAME,
named_configs.ConfigurationStore.ActiveConfig().name)
return env
def ExecuteScript(cls,
script_name,
args,
timeout=None,
stdin=None,
env=None,
add_python_paths=True):
"""Execute the given wrapper script with the given args.
This wrapper must be a sh script on non-windows, or a .cmd file, without the
'.cmd' extension (so it has the same script_name as non-windows), on
windows.
Args:
script_name: str, The script to run.
args: [str], The arguments for the script.
timeout: int, The number of seconds to wait before killing the process.
stdin: str, Optional input for the script.
env: {str: str}, Optional environment variables for the script.
add_python_paths: bool, Add PYTHONPATH=sys.path to env if True.
Returns:
An ExecutionResult object.
"""
args = exec_utils.GetArgsForScript(script_name, args)
env = encoding.EncodeEnv(dict(env if env else os.environ.copy()))
if add_python_paths:
env = cls._AddPythonPathsToEnv(env)
# pylint: disable=protected-access
runner = exec_utils._ProcessRunner(args,
timeout=timeout,
stdin=stdin,
env=env)
runner.Run()
return runner.result
def __init__(self, args, timeout=None, stdin=None, env=None): self.args = [encoding.Encode(a, encoding='utf-8') for a in args] self.timeout = timeout self.stdin = stdin self.env = encoding.EncodeEnv(env, encoding='utf-8') self.thread = None self.p = None self.result = None self.exc_info = None
def ExecuteLegacyScriptAsync(cls,
script_name,
args,
interpreter=None,
match_strings=None,
timeout=None,
stdin=None,
env=None,
add_python_paths=True):
"""Execute the given legacy script asynchronously in another thread.
Same as ExecuteScript() except it does not wait for the process to return.
Instead it returns a context manager that will kill the process once the
scope exits. If match_strings is given, the current thread will block until
the process outputs lines of text that matches the strings in order, then
the context manager will be returned and the process is kept alive.
Args:
script_name: str, The script to run.
args: [str], The arguments for the script.
interpreter: str, An interpreter to use rather than trying to derive it
from the extension name.
match_strings: [str], The output strings to match before returning.
timeout: int, The number of seconds to wait before killing the process.
stdin: str, Optional input for the script.
env: {str: str}, Optional environment variables for the script.
add_python_paths: bool, Add PYTHONPATH=sys.path to env if True.
Returns:
A context manager that will kill the process once the scope exists.
"""
args = exec_utils.GetArgsForLegacyScript(script_name,
args,
interpreter=interpreter)
env = encoding.EncodeEnv(dict(env if env else os.environ.copy()))
if add_python_paths:
env = cls._AddPythonPathsToEnv(env)
# pylint: disable=protected-access
runner = exec_utils._ProcessRunner(args,
timeout=timeout,
stdin=stdin,
env=env)
runner.RunAsync(match_strings=match_strings)
# pylint: disable=protected-access
return exec_utils._ProcessContext(runner)
def testStart(self):
port = self.GetPort()
test_call_out = None
# We wrap the execution for error handling purposes. We do not want to
# interleave the emulator stdout with the grpc call's standard out, but if
# there is an error it useful to have both. When the test catches the
# TimeoutError it will print the emulator's stdout, so we need to print the
# grpc calling process's.
try:
with self.ExecuteScriptAsync(
'gcloud', ['alpha', 'emulators', 'start', '--emulators=pubsub',
'--proxy-port=' + port],
match_strings=['Logging pubsub to:',
'Executing:',
'routes configuration written to file:',
'proxy configuration written to file:',
('INFO: creating handler for service [pubsub] to '
'local port ['),
'com.google.cloudsdk.emulators.EmulatorProxy main',
'INFO: Starting server on port: {}'.format(port)],
timeout=30):
env_dict = enc.EncodeEnv(dict(os.environ))
test_call_env = dict(
env_dict,
PUBSUB_EMULATOR_HOST=str('localhost:{}'.format(port)),
GOOGLE_CLOUD_PROJECT=str('cloudsdktest'))
test_call_out_no, test_call_out = tempfile.mkstemp()
with proxy_util.RunEmulatorProxyClient(
log_file=test_call_out_no, env=test_call_env) as proc:
proc.wait()
output = files.ReadFileContents(test_call_out)
self.assertIn('Created with target localhost:' + port, output)
self.assertIn('topic created', output)
self.assertIn('topic deleted', output)
self.assertIn('it all checks out!', output)
self.assertIn('Terminated', output)
finally:
if test_call_out is not None:
log.status.Print(
('had exception, printing out subprocess '
'stdout and stderr\n{line}\n{content}\n{line}').format(
line='='*80,
content=files.ReadFileContents(test_call_out)))
def _AddPythonPathsToEnv(env):
"""Returns a copy of env with Python specific path vars added.
This preserves any environment specific test runner tweaks.
Args:
env: {str: str}, Optional environment variables for the script.
Returns:
A copy of env with Python specific path vars added.
"""
env_with_pythonpaths = encoding.EncodeEnv(dict(env if env else os.environ))
# sys.path was initialized from PYTHONPATH at startup so we don't have to
# check PYTHONPATH here. The result will be the original PYTHONPATH dirs
# plus and dirs inserted/appened by Python startup and test runner
# initialization.
encoding.SetEncodedValue(
env_with_pythonpaths, 'PYTHONPATH', os.pathsep.join(sys.path))
return env_with_pythonpaths
def ExecuteLegacyScript(cls,
script_name,
args,
interpreter=None,
timeout=None,
stdin=None,
env=None,
add_python_paths=True):
"""Execute the given legacy script with the given args.
Args:
script_name: str, The script to run.
args: [str], The arguments for the script.
interpreter: str, An interpreter to use rather than trying to derive it
from the extension name.
timeout: int, The number of seconds to wait before killing the process.
stdin: str, Optional input for the script.
env: {str: str}, Optional environment variables for the script.
add_python_paths: bool, Add PYTHONPATH=sys.path to env if True.
Returns:
An ExecutionResult object.
"""
args = exec_utils.GetArgsForLegacyScript(script_name,
args,
interpreter=interpreter)
env = encoding.EncodeEnv(dict(env if env else os.environ.copy()))
if add_python_paths:
env = cls._AddPythonPathsToEnv(env)
# pylint: disable=protected-access
runner = exec_utils._ProcessRunner(args,
timeout=timeout,
stdin=stdin,
env=env)
runner.Run()
return runner.result
def KillSubprocess(p):
"""Kills a subprocess using an OS specific method when python can't do it.
This also kills all processes rooted in this process.
Args:
p: the Popen or multiprocessing.Process object to kill
Raises:
RuntimeError: if it fails to kill the process
"""
# This allows us to kill a Popen object or a multiprocessing.Process object
code = None
if hasattr(p, 'returncode'):
code = p.returncode
elif hasattr(p, 'exitcode'):
code = p.exitcode
if code is not None:
# already dead
return
if platforms.OperatingSystem.Current(
) == platforms.OperatingSystem.WINDOWS:
# Consume stdout so it doesn't show in the shell
taskkill_process = subprocess.Popen(
['taskkill', '/F', '/T', '/PID',
six.text_type(p.pid)],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
(stdout, stderr) = taskkill_process.communicate()
if taskkill_process.returncode != 0 and _IsTaskKillError(stderr):
# Sometimes taskkill does things in the wrong order and the processes
# disappear before it gets a chance to kill it. This is exposed as an
# error even though it's the outcome we want.
raise RuntimeError(
'Failed to call taskkill on pid {0}\nstdout: {1}\nstderr: {2}'.
format(p.pid, stdout, stderr))
else:
# Create a mapping of ppid to pid for all processes, then kill all
# subprocesses from the main process down
# set env LANG for subprocess.Popen to be 'en_US.UTF-8'
new_env = encoding.EncodeEnv(dict(os.environ))
new_env['LANG'] = 'en_US.UTF-8'
get_pids_process = subprocess.Popen(
['ps', '-e', '-o', 'ppid=', '-o', 'pid='],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
env=new_env)
(stdout, stderr) = get_pids_process.communicate()
stdout = stdout.decode('utf-8')
if get_pids_process.returncode != 0:
raise RuntimeError(
'Failed to get subprocesses of process: {0}'.format(p.pid))
# Create the process map
pid_map = {}
for line in stdout.strip().split('\n'):
(ppid, pid) = re.match(r'\s*(\d+)\s+(\d+)', line).groups()
ppid = int(ppid)
pid = int(pid)
children = pid_map.get(ppid)
if not children:
pid_map[ppid] = [pid]
else:
children.append(pid)
# Expand all descendants of the main process
all_pids = [p.pid]
to_process = [p.pid]
while to_process:
current = to_process.pop()
children = pid_map.get(current)
if children:
to_process.extend(children)
all_pids.extend(children)
# Kill all the subprocesses we found
for pid in all_pids:
_KillPID(pid)
def Exec(args,
env=None,
no_exit=False,
out_func=None,
err_func=None,
in_str=None,
**extra_popen_kwargs):
"""Emulates the os.exec* set of commands, but uses subprocess.
This executes the given command, waits for it to finish, and then exits this
process with the exit code of the child process.
Args:
args: [str], The arguments to execute. The first argument is the command.
env: {str: str}, An optional environment for the child process.
no_exit: bool, True to just return the exit code of the child instead of
exiting.
out_func: str->None, a function to call with the stdout of the executed
process. This can be e.g. log.file_only_logger.debug or log.out.write.
err_func: str->None, a function to call with the stderr of the executed
process. This can be e.g. log.file_only_logger.debug or log.err.write.
in_str: bytes or str, input to send to the subprocess' stdin.
**extra_popen_kwargs: Any additional kwargs will be passed through directly
to subprocess.Popen
Returns:
int, The exit code of the child if no_exit is True, else this method does
not return.
Raises:
PermissionError: if user does not have execute permission for cloud sdk bin
files.
InvalidCommandError: if the command entered cannot be found.
"""
log.debug('Executing command: %s', args)
# We use subprocess instead of execv because windows does not support process
# replacement. The result of execv on windows is that a new processes is
# started and the original is killed. When running in a shell, the prompt
# returns as soon as the parent is killed even though the child is still
# running. subprocess waits for the new process to finish before returning.
env = encoding.EncodeEnv(_GetToolEnv(env=env))
process_holder = _ProcessHolder()
with _ReplaceSignal(signal.SIGTERM, process_holder.Handler):
with _ReplaceSignal(signal.SIGINT, process_holder.Handler):
if out_func:
extra_popen_kwargs['stdout'] = subprocess.PIPE
if err_func:
extra_popen_kwargs['stderr'] = subprocess.PIPE
if in_str:
extra_popen_kwargs['stdin'] = subprocess.PIPE
try:
if args and isinstance(args, list):
# On Python 2.x on Windows, the first arg can't be unicode. We encode
# encode it anyway because there is really nothing else we can do if
# that happens.
# https://bugs.python.org/issue19264
args = [encoding.Encode(a) for a in args]
p = subprocess.Popen(args, env=env, **extra_popen_kwargs)
except OSError as err:
if err.errno == errno.EACCES:
raise PermissionError(err.strerror)
elif err.errno == errno.ENOENT:
raise InvalidCommandError(args[0])
raise
process_holder.process = p
if process_holder.signum is not None:
# This covers the small possibility that process_holder handled a
# signal when the process was starting but not yet set to
# process_holder.process.
if p.poll() is None:
p.terminate()
if isinstance(in_str, six.text_type):
in_str = in_str.encode('utf-8')
stdout, stderr = list(
map(encoding.Decode, p.communicate(input=in_str)))
if out_func:
out_func(stdout)
if err_func:
err_func(stderr)
ret_val = p.returncode
if no_exit and process_holder.signum is None:
return ret_val
sys.exit(ret_val)
def MakeProcess(module_name,
package_root,
args=None,
cluster=None,
task_type=None,
index=None,
**extra_popen_args):
"""Make a Popen object that runs the module, with the correct env.
If task_type is primary instead replaces the current process with the
subprocess via execution_utils.Exec
Args:
module_name: str. Name of the module to run, e.g. trainer.task
package_root: str. Absolute path to the package root for the module.
used as CWD for the subprocess.
args: [str]. Additional user args. Any relative paths will not work.
cluster: dict. Cluster configuration dictionary. Suitable for passing to
tf.train.ClusterSpec.
task_type: str. Task type of this process. Only relevant if cluster is
specified.
index: int. Task index of this process.
**extra_popen_args: extra args passed to Popen. Used for testing.
Returns:
a subprocess.Popen object corresponding to the subprocesses or an int
corresponding to the return value of the subprocess
(if task_type is primary)
Raises:
RuntimeError: if there is no python executable on the user system
"""
if args is None:
args = []
exe_override = properties.VALUES.ml_engine.local_python.Get()
python_executable = exe_override or files.FindExecutableOnPath('python')
if not python_executable:
raise RuntimeError('No python interpreter found on local machine')
cmd = [python_executable, '-m', module_name] + args
config = {
'job': {
'job_name': module_name,
'args': args
},
'task': {
'type': task_type,
'index': index
} if cluster else {},
'cluster': cluster or {},
'environment': 'cloud'
}
log.info(('launching training process:\n'
'command: {cmd}\n config: {config}').format(cmd=' '.join(cmd),
config=json.dumps(
config,
indent=2,
sort_keys=True)))
env = os.environ.copy()
# the tf_config environment variable is used to pass the tensorflow
# configuration options to the training module. the module specific
# arguments are passed as command line arguments.
env['TF_CONFIG'] = json.dumps(config)
if task_type == _GetPrimaryNodeName():
return execution_utils.Exec(cmd,
env=env,
no_exit=True,
cwd=package_root,
**extra_popen_args)
else:
env = encoding.EncodeEnv(env)
task = subprocess.Popen(cmd,
env=env,
cwd=package_root,
**extra_popen_args)
atexit.register(execution_utils.KillSubprocess, task)
return task
def __Popen(self, *args, **kwargs):
"""Makes sure os.environ gets export to subprocess children."""
if 'env' not in kwargs:
kwargs['env'] = encoding.EncodeEnv(os.environ)
return self._REAL_POPEN(*args, **kwargs)
