pilot/util/container.py

0001 #!/usr/bin/env python
0002 # Licensed under the Apache License, Version 2.0 (the "License");
0003 # you may not use this file except in compliance with the License.
0004 # You may obtain a copy of the License at
0005 # http://www.apache.org/licenses/LICENSE-2.0
0006 #
0007 # Authors:
0008 # - Paul Nilsson, paul.nilsson@cern.ch, 2018-2021
0009
0010 import subprocess
0011 from os import environ, getcwd, setpgrp  #, getpgid  #setsid
0012 from sys import version_info
0013
0014 from pilot.common.errorcodes import ErrorCodes
0015
0016 import logging
0017 logger = logging.getLogger(__name__)
0018 errors = ErrorCodes()
0019
0020
0021 def is_python3():
0022     """
0023     Check if we are running on Python 3.
0024
0025     :return: boolean.
0026     """
0027
0028     return version_info >= (3, 0)
0029
0030
0031 def execute(executable, **kwargs):
0032     """
0033     Execute the command and its options in the provided executable list.
0034     The function also determines whether the command should be executed within a container.
0035
0036     :param executable: command to be executed (string or list).
0037     :param kwargs (timeout, usecontainer, returnproc):
0038     :return: exit code, stdout and stderr (or process if requested via returnproc argument)
0039     """
0040
0041     mute = kwargs.get('mute', False)
0042     mode = kwargs.get('mode', 'bash')
0043     cwd = kwargs.get('cwd', getcwd())
0044     stdout_name = kwargs.get('stdout', subprocess.PIPE)
0045     stderr_name = kwargs.get('stderr', subprocess.PIPE)
0046     usecontainer = kwargs.get('usecontainer', False)
0047     returnproc = kwargs.get('returnproc', False)
0048     # timeout = kwargs.get('timeout', None)  # Python 3
0049     job = kwargs.get('job')
0050
0051     # convert executable to string if it is a list
0052     if type(executable) is list:
0053         executable = ' '.join(executable)
0054
0055     # switch off pilot controlled containers for user defined containers
0056     if job and job.imagename != "" and "runcontainer" in executable:
0057         usecontainer = False
0058         job.usecontainer = usecontainer
0059
0060     # Import user specific code if necessary (in case the command should be executed in a container)
0061     # Note: the container.wrapper() function must at least be declared
0062     if usecontainer:
0063         executable, diagnostics = containerise_executable(executable, **kwargs)
0064         if not executable:
0065             return None if returnproc else -1, "", diagnostics
0066
0067     if not mute:
0068         executable_readable = executable
0069         executables = executable_readable.split(";")
0070         for sub_cmd in executables:
0071             if 'S3_SECRET_KEY=' in sub_cmd:
0072                 secret_key = sub_cmd.split('S3_SECRET_KEY=')[1]
0073                 secret_key = 'S3_SECRET_KEY=' + secret_key
0074                 executable_readable = executable_readable.replace(secret_key, 'S3_SECRET_KEY=********')
0075         logger.info('executing command: %s', executable_readable)
0076
0077     if mode == 'python':
0078         exe = ['/usr/bin/python'] + executable.split()
0079     else:
0080         exe = ['/bin/bash', '-c', executable]
0081
0082     # try: intercept exception such as OSError -> report e.g. error.RESOURCEUNAVAILABLE: "Resource temporarily unavailable"
0083     if is_python3():  # Python 3
0084         process = subprocess.Popen(exe,
0085                                    bufsize=-1,
0086                                    stdout=stdout_name,
0087                                    stderr=stderr_name,
0088                                    cwd=cwd,
0089                                    preexec_fn=setpgrp,
0090                                    encoding='utf-8',
0091                                    errors='replace')
0092     else:
0093         process = subprocess.Popen(exe,
0094                                    bufsize=-1,
0095                                    stdout=stdout_name,
0096                                    stderr=stderr_name,
0097                                    cwd=cwd,
0098                                    preexec_fn=setpgrp)
0099     if returnproc:
0100         return process
0101     else:
0102         stdout, stderr = process.communicate()
0103         exit_code = process.poll()
0104
0105         # remove any added \n
0106         if stdout and stdout.endswith('\n'):
0107             stdout = stdout[:-1]
0108
0109     return exit_code, stdout, stderr
0110
0111
0112 def containerise_executable(executable, **kwargs):
0113     """
0114     Wrap the containerisation command around the executable.
0115
0116     :param executable: command to be wrapper (string).
0117     :param kwargs: kwargs dictionary.
0118     :return: containerised executable (list).
0119     """
0120
0121     job = kwargs.get('job')
0122
0123     user = environ.get('PILOT_USER', 'generic').lower()  # TODO: replace with singleton
0124     container = __import__('pilot.user.%s.container' % user, globals(), locals(), [user], 0)  # Python 2/3
0125     if container:
0126         # should a container really be used?
0127         do_use_container = job.usecontainer if job else container.do_use_container(**kwargs)
0128         # overrule for event service
0129         if job and job.is_eventservice and do_use_container and environ.get('PILOT_ES_EXECUTOR_TYPE', 'generic') != 'raythena':
0130             logger.info('overruling container decision for event service grid job')
0131             do_use_container = False
0132
0133         if do_use_container:
0134             diagnostics = ""
0135             try:
0136                 executable = container.wrapper(executable, **kwargs)
0137             except Exception as exc:
0138                 diagnostics = 'failed to execute wrapper function: %s' % exc
0139                 logger.fatal(diagnostics)
0140             else:
0141                 if executable == "":
0142                     diagnostics = 'failed to prepare container command (error code should have been set)'
0143                     logger.fatal(diagnostics)
0144             if diagnostics != "":
0145                 return None, diagnostics
0146         else:
0147             logger.info('pilot user container module has decided to not use a container')
0148     else:
0149         logger.warning('container module could not be imported')
0150
0151     return executable, ""