toil 6.1.0a1__py3-none-any.whl → 8.0.0__py3-none-any.whl

toil/provisioners/node.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 import datetime
 import logging
-import pipes
 import socket
 import subprocess
 import time
 from itertools import count
+from shlex import quote
+from typing import Any, Optional, Union
 
 from toil.lib.memoize import parse_iso_utc
 
@@ -29,15 +30,31 @@ logger = logging.getLogger(__name__)
 class Node:
     maxWaitTime = 7 * 60
 
-    def __init__(self, publicIP, privateIP, name, launchTime, nodeType, preemptible, tags=None, use_private_ip=None):
+    def __init__(
+        self,
+        publicIP: str,
+        privateIP: str,
+        name: str,
+        launchTime: Union[datetime.datetime, str],
+        nodeType: Optional[str],
+        preemptible: bool,
+        tags: Optional[dict[str, str]] = None,
+        use_private_ip: Optional[bool] = None,
+    ) -> None:
         self.publicIP = publicIP
         self.privateIP = privateIP
         if use_private_ip:
-            self.effectiveIP = self.privateIP #or self.publicIP?
+            self.effectiveIP = self.privateIP  # or self.publicIP?
         else:
             self.effectiveIP = self.publicIP or self.privateIP
         self.name = name
-        self.launchTime = launchTime
+        if isinstance(launchTime, datetime.datetime):
+            self.launchTime = launchTime
+        else:
+            try:
+                self.launchTime = parse_iso_utc(launchTime)
+            except ValueError:
+                self.launchTime = datetime.datetime.fromisoformat(launchTime)
         self.nodeType = nodeType
         self.preemptible = preemptible
         self.tags = tags
@@ -65,12 +82,12 @@ class Node:
         """
         if self.launchTime:
             now = datetime.datetime.utcnow()
-            delta = now - parse_iso_utc(self.launchTime)
+            delta = now - self.launchTime
             return 1 - delta.total_seconds() / 3600.0 % 1.0
         else:
             return 1
 
-    def waitForNode(self, role, keyName='core'):
+    def waitForNode(self, role: str, keyName: str = "core") -> None:
         self._waitForSSHPort()
         # wait here so docker commands can be used reliably afterwards
         self._waitForSSHKeys(keyName=keyName)
@@ -78,8 +95,8 @@ class Node:
         self._waitForAppliance(role=role, keyName=keyName)
 
     def copySshKeys(self, keyName):
-        """ Copy authorized_keys file to the core user from the keyName user."""
-        if keyName == 'core':
+        """Copy authorized_keys file to the core user from the keyName user."""
+        if keyName == "core":
             return  # No point.
 
         # Make sure that keys are there.
@@ -88,9 +105,17 @@ class Node:
         # copy keys to core user so that the ssh calls will work
         # - normal mechanism failed unless public key was in the google-ssh format
         # - even so, the key wasn't copied correctly to the core account
-        keyFile = '/home/%s/.ssh/authorized_keys' % keyName
-        self.sshInstance('/usr/bin/sudo', '/usr/bin/cp', keyFile, '/home/core/.ssh', user=keyName)
-        self.sshInstance('/usr/bin/sudo', '/usr/bin/chown', 'core', '/home/core/.ssh/authorized_keys', user=keyName)
+        keyFile = "/home/%s/.ssh/authorized_keys" % keyName
+        self.sshInstance(
+            "/usr/bin/sudo", "/usr/bin/cp", keyFile, "/home/core/.ssh", user=keyName
+        )
+        self.sshInstance(
+            "/usr/bin/sudo",
+            "/usr/bin/chown",
+            "core",
+            "/home/core/.ssh/authorized_keys",
+            user=keyName,
+        )
 
     def injectFile(self, fromFile, toFile, role):
         """
@@ -102,9 +127,13 @@ class Node:
                 self.coreRsync([fromFile, ":" + toFile], applianceName=role)
                 return True
             except Exception as e:
-                logger.debug("Rsync to new node failed, trying again. Error message: %s" % e)
+                logger.debug(
+                    "Rsync to new node failed, trying again. Error message: %s" % e
+                )
                 time.sleep(10 * retry)
-        raise RuntimeError(f"Failed to inject file {fromFile} to {role} with ip {self.effectiveIP}")
+        raise RuntimeError(
+            f"Failed to inject file {fromFile} to {role} with ip {self.effectiveIP}"
+        )
 
     def extractFile(self, fromFile, toFile, role):
         """
@@ -116,74 +145,111 @@ class Node:
                 self.coreRsync([":" + fromFile, toFile], applianceName=role)
                 return True
             except Exception as e:
-                logger.debug("Rsync from new node failed, trying again. Error message: %s" % e)
+                logger.debug(
+                    "Rsync from new node failed, trying again. Error message: %s" % e
+                )
                 time.sleep(10 * retry)
-        raise RuntimeError(f"Failed to extract file {fromFile} from {role} with ip {self.effectiveIP}")
+        raise RuntimeError(
+            f"Failed to extract file {fromFile} from {role} with ip {self.effectiveIP}"
+        )
 
-    def _waitForSSHKeys(self, keyName='core'):
+    def _waitForSSHKeys(self, keyName="core"):
         # the propagation of public ssh keys vs. opening the SSH port is racey, so this method blocks until
         # the keys are propagated and the instance can be SSH into
         start_time = time.time()
         last_error = None
         while True:
             if time.time() - start_time > self.maxWaitTime:
-                raise RuntimeError(f"Key propagation failed on machine with ip {self.effectiveIP}." +
-                                   ("\n\nMake sure that your public key is attached to your account and you are using "
-                                    "the correct private key. If you are using a key with a passphrase, be sure to "
-                                    "set up ssh-agent. For details, refer to "
-                                    "https://toil.readthedocs.io/en/latest/running/cloud/cloud.html."
-                                    if last_error and 'Permission denied' in last_error else ""))
+                raise RuntimeError(
+                    f"Key propagation failed on machine with ip {self.effectiveIP}."
+                    + (
+                        "\n\nMake sure that your public key is attached to your account and you are using "
+                        "the correct private key. If you are using a key with a passphrase, be sure to "
+                        "set up ssh-agent. For details, refer to "
+                        "https://toil.readthedocs.io/en/latest/running/cloud/cloud.html."
+                        if last_error and "Permission denied" in last_error
+                        else ""
+                    )
+                )
             try:
-                logger.info('Attempting to establish SSH connection to %s@%s...', keyName, self.effectiveIP)
-                self.sshInstance('ps', sshOptions=['-oBatchMode=yes'], user=keyName)
+                logger.info(
+                    "Attempting to establish SSH connection to %s@%s...",
+                    keyName,
+                    self.effectiveIP,
+                )
+                self.sshInstance("ps", sshOptions=["-oBatchMode=yes"], user=keyName)
             except RuntimeError as err:
                 last_error = str(err)
-                logger.info('Connection rejected, waiting for public SSH key to be propagated. Trying again in 10s.')
+                logger.info(
+                    "Connection rejected, waiting for public SSH key to be propagated. Trying again in 10s."
+                )
                 time.sleep(10)
             else:
-                logger.info('...SSH connection established.')
+                logger.info("...SSH connection established.")
                 return
 
-    def _waitForDockerDaemon(self, keyName='core'):
-        logger.info('Waiting for docker on %s to start...', self.effectiveIP)
+    def _waitForDockerDaemon(self, keyName="core"):
+        logger.info("Waiting for docker on %s to start...", self.effectiveIP)
         sleepTime = 10
         startTime = time.time()
         while True:
             if time.time() - startTime > self.maxWaitTime:
-                raise RuntimeError("Docker daemon failed to start on machine with ip %s" % self.effectiveIP)
+                raise RuntimeError(
+                    "Docker daemon failed to start on machine with ip %s"
+                    % self.effectiveIP
+                )
             try:
-                output = self.sshInstance('/usr/bin/ps', 'auxww', sshOptions=['-oBatchMode=yes'], user=keyName)
-                if b'dockerd' in output:
+                output = self.sshInstance(
+                    "/usr/bin/ps", "auxww", sshOptions=["-oBatchMode=yes"], user=keyName
+                )
+                if b"dockerd" in output:
                     # docker daemon has started
-                    logger.info('Docker daemon running')
+                    logger.info("Docker daemon running")
                     break
                 else:
-                    logger.info('... Still waiting for docker daemon, trying in %s sec...' % sleepTime)
+                    logger.info(
+                        "... Still waiting for docker daemon, trying in %s sec..."
+                        % sleepTime
+                    )
                     time.sleep(sleepTime)
             except RuntimeError:
                 logger.info("Wait for docker daemon failed ssh, trying again.")
                 sleepTime += 20
 
-    def _waitForAppliance(self, role, keyName='core'):
-        logger.info('Waiting for %s Toil appliance to start...', role)
+    def _waitForAppliance(self, role, keyName="core"):
+        logger.info("Waiting for %s Toil appliance to start...", role)
         sleepTime = 20
         startTime = time.time()
         while True:
             if time.time() - startTime > self.maxWaitTime:
-                raise RuntimeError("Appliance failed to start on machine with IP: " + self.effectiveIP +
-                                   "\nCheck if TOIL_APPLIANCE_SELF is set correctly and the container exists.")
+                raise RuntimeError(
+                    "Appliance failed to start on machine with IP: "
+                    + self.effectiveIP
+                    + "\nCheck if TOIL_APPLIANCE_SELF is set correctly and the container exists."
+                )
             try:
-                output = self.sshInstance('/usr/bin/docker', 'ps', sshOptions=['-oBatchMode=yes'], user=keyName)
-
-                role = bytes(role, encoding='utf-8') if type(role) != type(output) else role
+                output = self.sshInstance(
+                    "/usr/bin/docker",
+                    "ps",
+                    sshOptions=["-oBatchMode=yes"],
+                    user=keyName,
+                )
+
+                role = (
+                    bytes(role, encoding="utf-8")
+                    if type(role) != type(output)
+                    else role
+                )
 
                 if role in output:
-                    logger.info('...Toil appliance started')
+                    logger.info("...Toil appliance started")
                     break
                 else:
-                    logger.info('...Still waiting for appliance, trying again in %s sec...' % sleepTime)
-                    logger.debug(f'Role: {role}\n'
-                                 f'Output: {output}\n\n')
+                    logger.info(
+                        "...Still waiting for appliance, trying again in %s sec..."
+                        % sleepTime
+                    )
+                    logger.debug(f"Role: {role}\n" f"Output: {output}\n\n")
                     time.sleep(sleepTime)
             except RuntimeError:
                 # ignore exceptions, keep trying
@@ -197,13 +263,13 @@ class Node:
         :return: the number of unsuccessful attempts to connect to the port before a the first
         success
         """
-        logger.debug('Waiting for ssh port on %s to open...', self.effectiveIP)
+        logger.debug("Waiting for ssh port on %s to open...", self.effectiveIP)
         for i in count():
             s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
             try:
                 s.settimeout(a_short_time)
                 s.connect((self.effectiveIP, 22))
-                logger.debug('...ssh port open')
+                logger.debug("...ssh port open")
                 return i
             except OSError:
                 pass
@@ -217,7 +283,7 @@ class Node:
             interactive SSHing. The default value is False. Input=string is passed as
             input to the Popen call.
         """
-        kwargs['appliance'] = True
+        kwargs["appliance"] = True
         return self.coreSSH(*args, **kwargs)
 
     def sshInstance(self, *args, **kwargs):
@@ -225,7 +291,7 @@ class Node:
         Run a command on the instance.
         Returns the binary output of the command.
         """
-        kwargs['collectStdout'] = True
+        kwargs["collectStdout"] = True
         return self.coreSSH(*args, **kwargs)
 
     def coreSSH(self, *args, **kwargs):
@@ -241,64 +307,74 @@ class Node:
         :param bytes input: UTF-8 encoded input bytes to send to the command
 
         """
-        commandTokens = ['ssh', '-tt']
-        if not kwargs.pop('strict', False):
-            kwargs['sshOptions'] = ['-oUserKnownHostsFile=/dev/null', '-oStrictHostKeyChecking=no'] + kwargs.get(
-                'sshOptions', [])
-        sshOptions = kwargs.pop('sshOptions', None)
+        commandTokens = ["ssh", "-tt"]
+        if not kwargs.pop("strict", False):
+            kwargs["sshOptions"] = [
+                "-oUserKnownHostsFile=/dev/null",
+                "-oStrictHostKeyChecking=no",
+            ] + kwargs.get("sshOptions", [])
+        sshOptions = kwargs.pop("sshOptions", None)
         # Forward ports:
         # 5050 for Mesos dashboard (although to talk to agents you will need a proxy)
-        commandTokens.extend(['-L', '5050:localhost:5050'])
+        commandTokens.extend(["-L", "5050:localhost:5050"])
         if sshOptions:
             # add specified options to ssh command
             assert isinstance(sshOptions, list)
             commandTokens.extend(sshOptions)
         # specify host
-        user = kwargs.pop('user', 'core')  # CHANGED: Is this needed?
-        commandTokens.append(f'{user}@{str(self.effectiveIP)}')
+        user = kwargs.pop("user", "core")  # CHANGED: Is this needed?
+        commandTokens.append(f"{user}@{str(self.effectiveIP)}")
 
-        inputString = kwargs.pop('input', None)
+        inputString = kwargs.pop("input", None)
         if inputString is not None:
-            kwargs['stdin'] = subprocess.PIPE
+            kwargs["stdin"] = subprocess.PIPE
 
-        if kwargs.pop('collectStdout', None):
-            kwargs['stdout'] = subprocess.PIPE
-        kwargs['stderr'] = subprocess.PIPE
+        if kwargs.pop("collectStdout", None):
+            kwargs["stdout"] = subprocess.PIPE
+        kwargs["stderr"] = subprocess.PIPE
 
-        tty = kwargs.pop('tty', None)
-        if kwargs.pop('appliance', None):
-            ttyFlag = '-t' if tty else ''
-            commandTokens += ['docker', 'exec', '-i', ttyFlag, 'toil_leader']
+        tty = kwargs.pop("tty", None)
+        if kwargs.pop("appliance", None):
+            ttyFlag = "-t" if tty else ""
+            commandTokens += ["docker", "exec", "-i", ttyFlag, "toil_leader"]
 
-        logger.debug('Node %s: %s', self.effectiveIP, ' '.join(args))
-        args = list(map(pipes.quote, args))
+        logger.debug("Node %s: %s", self.effectiveIP, " ".join(args))
+        args = list(map(quote, args))
         commandTokens += args
-        logger.debug('Full command %s', ' '.join(commandTokens))
+        logger.debug("Full command %s", " ".join(commandTokens))
         process = subprocess.Popen(commandTokens, **kwargs)
         stdout, stderr = process.communicate(input=inputString)
         # at this point the process has already exited, no need for a timeout
         exit_code = process.returncode
         # ssh has been throwing random 255 errors - why?
         if exit_code != 0:
-            logger.info('Executing the command "%s" on the appliance returned a non-zero '
-                        'exit code %s with stdout %s and stderr %s'
-                        % (' '.join(args), exit_code, stdout, stderr))
-            raise RuntimeError('Executing the command "%s" on the appliance returned a non-zero '
-                               'exit code %s with stdout %s and stderr %s'
-                               % (' '.join(args), exit_code, stdout, stderr))
+            logger.info(
+                'Executing the command "%s" on the appliance returned a non-zero '
+                "exit code %s with stdout %s and stderr %s"
+                % (" ".join(args), exit_code, stdout, stderr)
+            )
+            raise RuntimeError(
+                'Executing the command "%s" on the appliance returned a non-zero '
+                "exit code %s with stdout %s and stderr %s"
+                % (" ".join(args), exit_code, stdout, stderr)
+            )
         return stdout
 
-    def coreRsync(self, args, applianceName='toil_leader', **kwargs):
-        remoteRsync = "docker exec -i %s rsync -v" % applianceName  # Access rsync inside appliance
+    def coreRsync(
+        self, args: list[str], applianceName: str = "toil_leader", **kwargs: Any
+    ) -> int:
+        remoteRsync = (
+            "docker exec -i %s rsync -v" % applianceName
+        )  # Access rsync inside appliance
         parsedArgs = []
         sshCommand = "ssh"
-        if not kwargs.pop('strict', False):
+        if not kwargs.pop("strict", False):
             sshCommand = "ssh -oUserKnownHostsFile=/dev/null -oStrictHostKeyChecking=no"
         hostInserted = False
         # Insert remote host address
         for i in args:
             if i.startswith(":") and not hostInserted:
-                user = kwargs.pop('user', 'core')  # CHANGED: Is this needed?
+                user = kwargs.pop("user", "core")  # CHANGED: Is this needed?
                 i = (f"{user}@{self.effectiveIP}") + i
                 hostInserted = True
             elif i.startswith(":") and hostInserted:
@@ -306,7 +382,7 @@ class Node:
             parsedArgs.append(i)
         if not hostInserted:
             raise ValueError("No remote host found in argument list")
-        command = ['rsync', '-e', sshCommand, '--rsync-path', remoteRsync]
+        command = ["rsync", "-e", sshCommand, "--rsync-path", remoteRsync]
         logger.debug("Running %r.", command + parsedArgs)
 
         return subprocess.check_call(command + parsedArgs)

toil/realtimeLogger.py

@@ -20,7 +20,7 @@ import os.path
 import socketserver as SocketServer
 import threading
 from types import TracebackType
-from typing import TYPE_CHECKING, Any, Optional, Type
+from typing import TYPE_CHECKING, Any, Optional
 
 from toil.lib.misc import get_public_ip
 from toil.statsAndLogging import set_log_level
@@ -49,7 +49,7 @@ class LoggingDatagramHandler(SocketServer.BaseRequestHandler):
 
         try:
             # Parse it as JSON
-            message_attrs = json.loads(data.decode('utf-8'))
+            message_attrs = json.loads(data.decode("utf-8"))
             # Fluff it up into a proper logging record
             record = logging.makeLogRecord(message_attrs)
             if isinstance(record.args, list):
@@ -81,7 +81,7 @@ class JSONDatagramHandler(logging.handlers.DatagramHandler):
 
     def makePickle(self, record: logging.LogRecord) -> bytes:
         """Actually, encode the record as bare JSON instead."""
-        return json.dumps(record.__dict__).encode('utf-8')
+        return json.dumps(record.__dict__).encode("utf-8")
 
 
 class RealtimeLoggerMetaclass(type):
@@ -113,7 +113,7 @@ class RealtimeLogger(metaclass=RealtimeLoggerMetaclass):
     envPrefix = "TOIL_RT_LOGGING_"
 
     # Avoid duplicating the default level everywhere
-    defaultLevel = 'INFO'
+    defaultLevel = "INFO"
 
     # State maintained on server and client
 
@@ -131,19 +131,24 @@ class RealtimeLogger(metaclass=RealtimeLoggerMetaclass):
     logger = None
 
     @classmethod
-    def _startLeader(cls, batchSystem: 'AbstractBatchSystem', level: str = defaultLevel) -> None:
+    def _startLeader(
+        cls, batchSystem: "AbstractBatchSystem", level: str = defaultLevel
+    ) -> None:
         with cls.lock:
             if cls.initialized == 0:
                 cls.initialized += 1
                 if level:
-                    logger.info('Starting real-time logging.')
+                    logger.info("Starting real-time logging.")
                     # Start up the logging server
                     cls.loggingServer = SocketServer.ThreadingUDPServer(
-                            server_address=('0.0.0.0', 0),
-                            RequestHandlerClass=LoggingDatagramHandler)
+                        server_address=("0.0.0.0", 0),
+                        RequestHandlerClass=LoggingDatagramHandler,
+                    )
 
                     # Set up a thread to do all the serving in the background and exit when we do
-                    cls.serverThread = threading.Thread(target=cls.loggingServer.serve_forever)
+                    cls.serverThread = threading.Thread(
+                        target=cls.loggingServer.serve_forever
+                    )
                     cls.serverThread.daemon = True
                     cls.serverThread.start()
 
@@ -156,28 +161,30 @@ class RealtimeLogger(metaclass=RealtimeLoggerMetaclass):
                         os.environ[name] = value
                         batchSystem.setEnv(name)
 
-                    _setEnv('ADDRESS', '%s:%i' % (ip, port))
-                    _setEnv('LEVEL', level)
+                    _setEnv("ADDRESS", "%s:%i" % (ip, port))
+                    _setEnv("LEVEL", level)
                 else:
-                    logger.debug('Real-time logging disabled')
+                    logger.debug("Real-time logging disabled")
             else:
                 if level:
-                    logger.warning('Ignoring nested request to start real-time logging')
+                    logger.warning("Ignoring nested request to start real-time logging")
 
     @classmethod
     def _stopLeader(cls) -> None:
         """Stop the server on the leader."""
         with cls.lock:
             if cls.initialized == 0:
-                raise RuntimeError("Can't stop the server on the leader as the leader was never initialized.")
+                raise RuntimeError(
+                    "Can't stop the server on the leader as the leader was never initialized."
+                )
             cls.initialized -= 1
             if cls.initialized == 0:
                 if cls.loggingServer:
-                    logger.info('Stopping real-time logging server.')
+                    logger.info("Stopping real-time logging server.")
                     cls.loggingServer.shutdown()
                     cls.loggingServer = None
                 if cls.serverThread:
-                    logger.info('Joining real-time logging server thread.')
+                    logger.info("Joining real-time logging server thread.")
                     cls.serverThread.join()
                     cls.serverThread = None
                 for k in list(os.environ.keys()):
@@ -198,9 +205,9 @@ class RealtimeLogger(metaclass=RealtimeLoggerMetaclass):
         if cls.logger is None:
             with cls.lock:
                 if cls.logger is None:
-                    cls.logger = logging.getLogger('toil-rt')
+                    cls.logger = logging.getLogger("toil-rt")
                     try:
-                        level = os.environ[cls.envPrefix + 'LEVEL']
+                        level = os.environ[cls.envPrefix + "LEVEL"]
                     except KeyError:
                         # There is no server running on the leader, so suppress most log messages
                         # and skip the UDP stuff.
@@ -209,16 +216,16 @@ class RealtimeLogger(metaclass=RealtimeLoggerMetaclass):
                         # Adopt the logging level set on the leader.
                         set_log_level(level, cls.logger)
                         try:
-                            address = os.environ[cls.envPrefix + 'ADDRESS']
+                            address = os.environ[cls.envPrefix + "ADDRESS"]
                         except KeyError:
                             pass
                         else:
                             # We know where to send messages to, so send them.
-                            host, port = address.split(':')
+                            host, port = address.split(":")
                             cls.logger.addHandler(JSONDatagramHandler(host, int(port)))
         return cls.logger
 
-    def __init__(self, batchSystem: 'AbstractBatchSystem', level: str = defaultLevel):
+    def __init__(self, batchSystem: "AbstractBatchSystem", level: str = defaultLevel):
         """
         Create a context manager that starts up the UDP server.
 
@@ -237,5 +244,10 @@ class RealtimeLogger(metaclass=RealtimeLoggerMetaclass):
         RealtimeLogger._startLeader(self.__batchSystem, level=self.__level)
 
     # noinspection PyUnusedLocal
-    def __exit__(self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]) -> None:
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
         RealtimeLogger._stopLeader()