llvm/utils/remote-exec.py

   1 #!/usr/bin/env python
   2 # ===----------------------------------------------------------------------===##
   3 #
   4 # Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   5 # See https://llvm.org/LICENSE.txt for license information.
   6 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   7 #
   8 # ===----------------------------------------------------------------------===##
   9
  10 """
  11 Runs an executable on a remote host.
  12
  13 This is meant to be used as an executor when running the LLVM and the Libraries
  14 tests on a target.
  15 """
  16
  17 import argparse
  18 import os
  19 import posixpath
  20 import shlex
  21 import subprocess
  22 import sys
  23 import tarfile
  24 import tempfile
  25 import re
  26
  27
  28 def ssh(args, command):
  29     cmd = ["ssh", "-oBatchMode=yes"]
  30     if args.extra_ssh_args is not None:
  31         cmd.extend(shlex.split(args.extra_ssh_args))
  32     return cmd + [args.host, command]
  33
  34
  35 def scp(args, src, dst):
  36     cmd = ["scp", "-q", "-oBatchMode=yes"]
  37     if args.extra_scp_args is not None:
  38         cmd.extend(shlex.split(args.extra_scp_args))
  39     return cmd + [src, "{}:{}".format(args.host, dst)]
  40
  41
  42 def main():
  43     parser = argparse.ArgumentParser()
  44     parser.add_argument("--host", type=str, required=True)
  45     parser.add_argument("--execdir", type=str, required=False)
  46     parser.add_argument("--extra-ssh-args", type=str, required=False)
  47     parser.add_argument("--extra-scp-args", type=str, required=False)
  48     parser.add_argument("--codesign_identity", type=str, required=False, default=None)
  49     parser.add_argument("--env", type=str, nargs="*", required=False, default=dict())
  50
  51     # Note: The default value is for the backward compatibility with a hack in
  52     # libcxx test suite.
  53     # If an argument is a file that ends in `.tmp.exe`, assume it is the name
  54     # of an executable generated by a test file. We call these test-executables
  55     # below. This allows us to do custom processing like codesigning test-executables
  56     # and changing their path when running on the remote host. It's also possible
  57     # for there to be no such executable, for example in the case of a .sh.cpp
  58     # test.
  59     parser.add_argument(
  60         "--exec-pattern",
  61         type=str,
  62         required=False,
  63         default=".*",
  64         help="The name regex pattern of the executables generated by \
  65                               a test file. Specifying it allows us to do custom \
  66                               processing like codesigning test-executables \
  67                               and changing their path when running on \
  68                               the remote host. It's also possible for there \
  69                               to be no such executable, for example in \
  70                               the case of a .sh.cpp test.",
  71     )
  72
  73     parser.add_argument("command", nargs=argparse.ONE_OR_MORE)
  74     args = parser.parse_args()
  75     commandLine = args.command
  76
  77     execdir = args.execdir
  78     if execdir == ".":
  79         # Retrieve the exec directory from the command line.
  80         execdir, _ = os.path.split(commandLine[0])
  81         if execdir == "":
  82             # Get the current directory in that case.
  83             execdir = os.getcwd()
  84     arcname = os.path.basename(execdir) if execdir else None
  85
  86     # Create a temporary directory where the test will be run.
  87     # That is effectively the value of %T on the remote host.
  88     tmp = subprocess.check_output(
  89         ssh(args, "mktemp -d /tmp/llvm.XXXXXXXXXX"), universal_newlines=True
  90     ).strip()
  91
  92     isExecutable = lambda exe: re.match(args.exec_pattern, exe) and os.path.exists(exe)
  93     pathOnRemote = lambda file: posixpath.join(tmp, os.path.basename(file))
  94
  95     remoteCommands = []
  96
  97     try:
  98         # Do any necessary codesigning of test-executables found in the command line.
  99         if args.codesign_identity:
 100             for exe in filter(isExecutable, commandLine):
 101                 subprocess.check_call(
 102                     ["xcrun", "codesign", "-f", "-s", args.codesign_identity, exe],
 103                     env={},
 104                 )
 105
 106         # tar up the execution directory (which contains everything that's needed
 107         # to run the test), and copy the tarball over to the remote host.
 108         if execdir:
 109             try:
 110                 tmpTar = tempfile.NamedTemporaryFile(suffix=".tar", delete=False)
 111                 with tarfile.open(fileobj=tmpTar, mode="w") as tarball:
 112                     tarball.add(execdir, arcname=arcname)
 113
 114                 # Make sure we close the file before we scp it, because accessing
 115                 # the temporary file while still open doesn't work on Windows.
 116                 tmpTar.close()
 117                 remoteTarball = pathOnRemote(tmpTar.name)
 118                 subprocess.check_call(scp(args, tmpTar.name, remoteTarball))
 119             finally:
 120                 # Make sure we close the file in case an exception happens before
 121                 # we've closed it above -- otherwise close() is idempotent.
 122                 tmpTar.close()
 123                 os.remove(tmpTar.name)
 124
 125             # Untar the dependencies in the temporary directory and remove the tarball.
 126             remoteCommands.extend(
 127                 [
 128                     "tar -xf {} -C {} --strip-components 1".format(remoteTarball, tmp),
 129                     "rm {}".format(remoteTarball),
 130                 ]
 131             )
 132         else:
 133             # Copy only the files, which are specified in the command line.
 134             # Copy them to remote host one by one.
 135             for x in commandLine:
 136                 _, f = os.path.split(x)
 137                 subprocess.check_call(scp(args, x, pathOnRemote(f)))
 138
 139         # Make sure all executables in the remote command line have 'execute'
 140         # permissions on the remote host. The host that compiled the test-executable
 141         # might not have a notion of 'executable' permissions.
 142         for exe in filter(isExecutable, commandLine):
 143             remoteCommands.append("chmod +x {}".format(pathOnRemote(exe)))
 144
 145         # Execute the command through SSH in the temporary directory, with the
 146         # correct environment. We tweak the command line to run it on the remote
 147         # host by transforming the path of test-executables to their path in the
 148         # temporary directory on the remote host.
 149         for i, x in enumerate(commandLine):
 150             if isExecutable(x):
 151                 commandLine[i] = pathOnRemote(x)
 152         remoteCommands.append("cd {}".format(tmp))
 153         if args.env:
 154             remoteCommands.append("export {}".format(" ".join(args.env)))
 155         remoteCommands.append(subprocess.list2cmdline(commandLine))
 156
 157         # Finally, SSH to the remote host and execute all the commands.
 158         rc = subprocess.call(ssh(args, " && ".join(remoteCommands)))
 159         return rc
 160
 161     finally:
 162         # Make sure the temporary directory is removed when we're done.
 163         subprocess.check_call(ssh(args, "rm -r {}".format(tmp)))
 164
 165
 166 if __name__ == "__main__":
 167     exit(main())