Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-04-19 08:00:04

0001 import argparse
0002 import logging
0003 import os
0004 import re
0005 import shutil
0006 import subprocess
0007 import sys
0008 import tempfile
0009 
0010 import paramiko
0011 from pandaharvester.harvestercore.queue_config_mapper import QueueConfigMapper
0012 
0013 
0014 class TemporaryDirectory(object):
0015     def __enter__(self):
0016         self.name = tempfile.mkdtemp()
0017         return self.name
0018 
0019     def __exit__(self, exc_type, exc_value, traceback):
0020         shutil.rmtree(self.name)
0021 
0022 
0023 def make_ssh_connection(ssh_host, ssh_port, ssh_username, ssh_password, pass_phrase, private_key, jump_host, jump_port):
0024     # ssh
0025     sshClient = paramiko.SSHClient()
0026     sshClient.set_missing_host_key_policy(paramiko.AutoAddPolicy())
0027     if jump_host is None:
0028         # direct SSH
0029         sshClient.connect(ssh_host, ssh_port, username=ssh_username, password=ssh_password, passphrase=pass_phrase, key_filename=private_key)
0030     else:
0031         # via jump host
0032         sshClient.connect(jump_host, jump_port, username=ssh_username, password=ssh_password, passphrase=pass_phrase, key_filename=private_key)
0033         transport = sshClient.get_transport()
0034         dst_address = (ssh_host, ssh_port)
0035         src_address = (jump_host, jump_port)
0036         channel = transport.open_channel("direct-tcpip", dst_address, src_address)
0037         jumpClient = paramiko.SSHClient()
0038         jumpClient.set_missing_host_key_policy(paramiko.AutoAddPolicy())
0039         jumpClient.connect(ssh_host, ssh_port, username=ssh_username, password=ssh_password, passphrase=pass_phrase, key_filename=private_key, sock=channel)
0040     return sshClient
0041 
0042 
0043 def main():
0044     logging.basicConfig()
0045 
0046     parser = argparse.ArgumentParser()
0047     parser.add_argument(
0048         "--remoteDir", action="store", dest="remoteDir", default="harvester", help="directory on the remote target machine where harvester is installed"
0049     )
0050     parser.add_argument(
0051         "--remoteBuildDir",
0052         action="store",
0053         dest="remoteBuildDir",
0054         default="harvester_build",
0055         help="directory on the remote target machine where harvester is build",
0056     )
0057     parser.add_argument("--remotePythonSetup", action="store", dest="remotePythonSetup", default="", help="python setup on remote target machine")
0058     parser.add_argument("--queueName", action="store", dest="queueName", default=None, required=True, help="the name of queue where harvester is installed")
0059     parser.add_argument("--middleware", action="store", dest="middleware", default="rpc", help="middleware to access the remote target machine")
0060     options = parser.parse_args()
0061 
0062     # remove ~/ which doesn't work with sftp
0063     options.remoteDir = re.sub("^~/", "", options.remoteDir)
0064     options.remoteBuildDir = re.sub("^~/", "", options.remoteBuildDir)
0065 
0066     # get queue
0067     qcm = QueueConfigMapper()
0068     qcm.load_data()
0069     queueConfig = qcm.get_queue(options.queueName)
0070     if queueConfig is None:
0071         print(f"ERROR: queue={options.queueName} not found in panda_queueconfig.json")
0072         sys.exit(1)
0073 
0074     # get middleware
0075     if not hasattr(queueConfig, options.middleware):
0076         print(f"ERROR: middleware={options.middleware} is not defined for {options.queueName} in panda_queueconfig.json")
0077         sys.exit(1)
0078     middleware = getattr(queueConfig, options.middleware)
0079 
0080     # get ssh parameters
0081     sshHost = middleware["remoteHost"]
0082     try:
0083         sshPort = middleware["remotePort"]
0084     except Exception:
0085         sshPort = 22
0086     sshUserName = middleware["sshUserName"]
0087     try:
0088         sshPassword = middleware["sshPassword"]
0089     except Exception:
0090         sshPassword = None
0091 
0092     privateKey = None
0093     passPhrase = None
0094     if sshPassword is None:
0095         try:
0096             privateKey = middleware["privateKey"]
0097         except Exception:
0098             print(f"ERROR: set sshPassword or privateKey in middleware={options.middleware}")
0099             sys.exit(1)
0100         try:
0101             passPhrase = middleware["passPhrase"]
0102         except Exception:
0103             passPhrase = None
0104 
0105     try:
0106         jumpHost = middleware["jumpHost"]
0107     except Exception:
0108         jumpHost = None
0109     try:
0110         jumpPort = middleware["jumpPort"]
0111     except Exception:
0112         jumpPort = 22
0113 
0114     # ssh
0115     sshClient = make_ssh_connection(sshHost, sshPort, sshUserName, sshPassword, passPhrase, privateKey, jumpHost, jumpPort)
0116 
0117     # get remote python version
0118     exec_out = sshClient.exec_command(";".join([options.remotePythonSetup, """python -c 'import sys;print("{0}{1}".format(*(sys.version_info[:2])))' """]))
0119     remotePythonVer = exec_out[1].read().rstrip()
0120     sshClient.close()
0121     print(f"remote python version : {remotePythonVer}")
0122 
0123     # make tmp dir
0124     with TemporaryDirectory() as tmpDir:
0125         harvesterGit = "git+git://github.com/PanDAWMS/panda-harvester.git"
0126 
0127         # get all dependencies
0128         print("getting dependencies")
0129         p = subprocess.Popen("pip download -d {0} {1}; rm -rf {0}/*".format(tmpDir, harvesterGit), stdout=subprocess.PIPE, shell=True)
0130         stdout, stderr = p.communicate()
0131         packages = []
0132         for line in stdout.split("\n"):
0133             if line.startswith("Successfully downloaded"):
0134                 packages = line.split()[2:]
0135         packages.append(harvesterGit)
0136         packages.append("pip")
0137         packages.remove("pandaharvester")
0138 
0139         # download packages
0140         print(f"pip download to {tmpDir}")
0141         for package in packages:
0142             print(f"getting {package}")
0143             ret = subprocess.call(f"pip download --no-deps --python-version {remotePythonVer} -d {tmpDir} {package}", shell=True)
0144             if ret != 0:
0145                 print(f"ERROR: failed to download {package}")
0146                 sys.exit(1)
0147 
0148         # sftp
0149         sshClient = make_ssh_connection(sshHost, sshPort, sshUserName, sshPassword, passPhrase, privateKey, jumpHost, jumpPort)
0150         try:
0151             sshClient.exec_command("rm -rf {0}; mkdir -p {0}".format(options.remoteBuildDir))
0152         except Exception:
0153             pass
0154         sftp = sshClient.open_sftp()
0155         for name in os.listdir(tmpDir):
0156             path = os.path.join(tmpDir, name)
0157             if os.path.isdir(path):
0158                 continue
0159             remotePath = os.path.join(options.remoteBuildDir, name)
0160             print(f"copy {name} to {remotePath}")
0161             sftp.put(path, remotePath)
0162 
0163         # install
0164         print("install harvester")
0165         buildDir = options.remoteBuildDir
0166         if not buildDir.startswith("/"):
0167             buildDir = "~/" + buildDir
0168         exec_out = sshClient.exec_command(
0169             ";".join([options.remotePythonSetup, f"cd {options.remoteDir}", f"pip install pip pandaharvester --no-index --find-links {buildDir}"])
0170         )
0171         print(exec_out[1].read())
0172         print(exec_out[2].read())
0173         sshClient.close()
0174 
0175 
0176 if __name__ == "__main__":
0177     main()