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

toil/batchSystems/mesos/executor.py

@@ -51,14 +51,14 @@ class MesosExecutor(Executor):
         self.popenLock = threading.Lock()
         self.runningTasks = {}
         self.workerCleanupInfo = None
-        log.debug('Preparing system for resource download')
+        log.debug("Preparing system for resource download")
         Resource.prepareSystem()
         self.address = None
         self.id = None
         # Setting this value at this point will ensure that the toil workflow directory will go to
         # the mesos sandbox if the user hasn't specified --workDir on the command line.
-        if not os.getenv('TOIL_WORKDIR'):
-            os.environ['TOIL_WORKDIR'] = os.getcwd()
+        if not os.getenv("TOIL_WORKDIR"):
+            os.environ["TOIL_WORKDIR"] = os.getcwd()
 
     def registered(self, driver, executorInfo, frameworkInfo, agentInfo):
         """
@@ -66,11 +66,13 @@ class MesosExecutor(Executor):
         """
 
         # Get the ID we have been assigned, if we have it
-        self.id = executorInfo.executor_id.get('value', None)
+        self.id = executorInfo.executor_id.get("value", None)
 
         log.debug("Registered executor %s with framework", self.id)
         self.address = socket.gethostbyname(agentInfo.hostname)
-        nodeInfoThread = threading.Thread(target=self._sendFrameworkMessage, args=[driver], daemon=True)
+        nodeInfoThread = threading.Thread(
+            target=self._sendFrameworkMessage, args=[driver], daemon=True
+        )
         nodeInfoThread.start()
 
     def reregistered(self, driver, agentInfo):
@@ -99,12 +101,12 @@ class MesosExecutor(Executor):
             os.killpg(pgid, signal.SIGKILL)
 
     def shutdown(self, driver):
-        log.critical('Shutting down executor ...')
+        log.critical("Shutting down executor ...")
         for taskId in list(self.runningTasks.keys()):
             self.killTask(driver, taskId)
         Resource.cleanSystem()
         BatchSystemSupport.workerCleanup(self.workerCleanupInfo)
-        log.critical('... executor shut down.')
+        log.critical("... executor shut down.")
 
     def error(self, driver, message):
         """
@@ -123,13 +125,15 @@ class MesosExecutor(Executor):
                 message = Expando(address=self.address)
                 psutil.cpu_percent()
             else:
-                message.nodeInfo = dict(coresUsed=float(psutil.cpu_percent()) * .01,
-                                        memoryUsed=float(psutil.virtual_memory().percent) * .01,
-                                        coresTotal=cpu_count(),
-                                        memoryTotal=psutil.virtual_memory().total,
-                                        workers=len(self.runningTasks))
+                message.nodeInfo = dict(
+                    coresUsed=float(psutil.cpu_percent()) * 0.01,
+                    memoryUsed=float(psutil.virtual_memory().percent) * 0.01,
+                    coresTotal=cpu_count(),
+                    memoryTotal=psutil.virtual_memory().total,
+                    workers=len(self.runningTasks),
+                )
             log.debug("Send framework message: %s", message)
-            driver.sendFrameworkMessage(encode_data(repr(message).encode('utf-8')))
+            driver.sendFrameworkMessage(encode_data(repr(message).encode("utf-8")))
             # Prevent workers launched together from repeatedly hitting the leader at the same time
             time.sleep(random.randint(45, 75))
 
@@ -144,16 +148,21 @@ class MesosExecutor(Executor):
 
             log.debug("Running task %s", task.task_id.value)
             startTime = time.time()
-            sendUpdate(task, 'TASK_RUNNING', wallTime=0)
+            sendUpdate(task, "TASK_RUNNING", wallTime=0)
 
             # try to unpickle the task
             try:
                 taskData = pickle.loads(decode_data(task.data))
             except:
                 exc_info = sys.exc_info()
-                log.error('Exception while unpickling task: ', exc_info=exc_info)
+                log.error("Exception while unpickling task: ", exc_info=exc_info)
                 exc_type, exc_value, exc_trace = exc_info
-                sendUpdate(task, 'TASK_FAILED', wallTime=0, msg=''.join(traceback.format_exception_only(exc_type, exc_value)))
+                sendUpdate(
+                    task,
+                    "TASK_FAILED",
+                    wallTime=0,
+                    msg="".join(traceback.format_exception_only(exc_type, exc_value)),
+                )
                 return
 
             # This is where task.data is first invoked. Using this position to setup cleanupInfo
@@ -170,23 +179,27 @@ class MesosExecutor(Executor):
                     exitStatus = process.wait()
                     wallTime = time.time() - startTime
                     if 0 == exitStatus:
-                        sendUpdate(task, 'TASK_FINISHED', wallTime)
+                        sendUpdate(task, "TASK_FINISHED", wallTime)
                     elif -9 == exitStatus:
-                        sendUpdate(task, 'TASK_KILLED', wallTime)
+                        sendUpdate(task, "TASK_KILLED", wallTime)
                     else:
-                        sendUpdate(task, 'TASK_FAILED', wallTime, msg=str(exitStatus))
+                        sendUpdate(task, "TASK_FAILED", wallTime, msg=str(exitStatus))
                 finally:
                     del self.runningTasks[task.task_id.value]
             except:
                 wallTime = time.time() - startTime
                 exc_info = sys.exc_info()
-                log.error('Exception while running task:', exc_info=exc_info)
+                log.error("Exception while running task:", exc_info=exc_info)
                 exc_type, exc_value, exc_trace = exc_info
-                sendUpdate(task, 'TASK_FAILED', wallTime=wallTime, msg=''.join(traceback.format_exception_only(exc_type, exc_value)))
+                sendUpdate(
+                    task,
+                    "TASK_FAILED",
+                    wallTime=wallTime,
+                    msg="".join(traceback.format_exception_only(exc_type, exc_value)),
+                )
 
             wallTime = time.time() - startTime
-            sendUpdate(task, 'TASK_FINISHED', wallTime)
-
+            sendUpdate(task, "TASK_FINISHED", wallTime)
 
         def runJob(job):
             """
@@ -196,16 +209,17 @@ class MesosExecutor(Executor):
             """
             if job.userScript:
                 job.userScript.register()
-            log.debug("Invoking command: '%s'", job.command)
+            command = job.command
+            log.debug("Invoking command: '%s'", command)
             # Construct the job's environment
             jobEnv = dict(os.environ, **job.environment)
-            log.debug('Using environment variables: %s', jobEnv.keys())
+            log.debug("Using environment variables: %s", jobEnv.keys())
             with self.popenLock:
-                return subprocess.Popen(job.command,
-                                        preexec_fn=lambda: os.setpgrp(),
-                                        shell=True, env=jobEnv)
+                return subprocess.Popen(
+                    command, preexec_fn=lambda: os.setpgrp(), shell=True, env=jobEnv
+                )
 
-        def sendUpdate(task, taskState, wallTime, msg=''):
+        def sendUpdate(task, taskState, wallTime, msg=""):
             update = addict.Dict()
             update.task_id.value = task.task_id.value
             if self.id is not None:
@@ -216,7 +230,7 @@ class MesosExecutor(Executor):
 
             # Add wallTime as a label.
             labels = addict.Dict()
-            labels.labels = [{'key': 'wallTime', 'value': str(wallTime)}]
+            labels.labels = [{"key": "wallTime", "value": str(wallTime)}]
             update.labels = labels
 
             driver.sendStatusUpdate(update)
@@ -238,34 +252,48 @@ def main():
     if not os.environ.get("MESOS_AGENT_ENDPOINT"):
         # Some Mesos setups in our tests somehow lack this variable. Provide a
         # fake one to maybe convince the executor driver to work.
-        os.environ["MESOS_AGENT_ENDPOINT"] = os.environ.get("MESOS_SLAVE_ENDPOINT", "127.0.0.1:5051")
-        log.warning("Had to fake MESOS_AGENT_ENDPOINT as %s" % os.environ["MESOS_AGENT_ENDPOINT"])
+        os.environ["MESOS_AGENT_ENDPOINT"] = os.environ.get(
+            "MESOS_SLAVE_ENDPOINT", "127.0.0.1:5051"
+        )
+        log.warning(
+            "Had to fake MESOS_AGENT_ENDPOINT as %s"
+            % os.environ["MESOS_AGENT_ENDPOINT"]
+        )
 
     # must be set manually to enable toggling of the mesos log level for debugging jenkins
     # may be useful: https://github.com/DataBiosphere/toil/pull/2338#discussion_r223854931
     if False:
         try:
-            urlopen("http://%s/logging/toggle?level=1&duration=15mins" % os.environ["MESOS_AGENT_ENDPOINT"]).read()
+            urlopen(
+                "http://%s/logging/toggle?level=1&duration=15mins"
+                % os.environ["MESOS_AGENT_ENDPOINT"]
+            ).read()
             log.debug("Toggled agent log level")
         except Exception:
             log.debug("Failed to toggle agent log level")
 
     # Parse the agent state
-    agent_state = json.loads(urlopen("http://%s/state" % os.environ["MESOS_AGENT_ENDPOINT"]).read())
-    if 'completed_frameworks' in agent_state:
+    agent_state = json.loads(
+        urlopen("http://%s/state" % os.environ["MESOS_AGENT_ENDPOINT"]).read()
+    )
+    if "completed_frameworks" in agent_state:
         # Drop the completed frameworks which grow over time
-        del agent_state['completed_frameworks']
+        del agent_state["completed_frameworks"]
     log.debug("Agent state: %s", str(agent_state))
 
     log.debug("Virtual memory info in executor: %s" % repr(psutil.virtual_memory()))
 
-    if os.path.exists('/sys/fs/cgroup/memory'):
+    if os.path.exists("/sys/fs/cgroup/memory"):
         # Mesos can limit memory with a cgroup, so we should report on that.
-        for (dirpath, dirnames, filenames) in os.walk('/sys/fs/cgroup/memory', followlinks=True):
+        for dirpath, dirnames, filenames in os.walk(
+            "/sys/fs/cgroup/memory", followlinks=True
+        ):
             for filename in filenames:
-                if 'limit_in_bytes' not in filename:
+                if "limit_in_bytes" not in filename:
                     continue
-                log.debug('cgroup memory info from %s:' % os.path.join(dirpath, filename))
+                log.debug(
+                    "cgroup memory info from %s:" % os.path.join(dirpath, filename)
+                )
                 try:
                     for line in open(os.path.join(dirpath, filename)):
                         log.debug(line.rstrip())
@@ -274,14 +302,13 @@ def main():
 
     # Mesos can also impose rlimit limits, including on things that really
     # ought to not be limited, like virtual address space size.
-    log.debug('DATA rlimit: %s', str(resource.getrlimit(resource.RLIMIT_DATA)))
-    log.debug('STACK rlimit: %s', str(resource.getrlimit(resource.RLIMIT_STACK)))
-    log.debug('RSS rlimit: %s', str(resource.getrlimit(resource.RLIMIT_RSS)))
-    log.debug('AS rlimit: %s', str(resource.getrlimit(resource.RLIMIT_AS)))
-
+    log.debug("DATA rlimit: %s", str(resource.getrlimit(resource.RLIMIT_DATA)))
+    log.debug("STACK rlimit: %s", str(resource.getrlimit(resource.RLIMIT_STACK)))
+    log.debug("RSS rlimit: %s", str(resource.getrlimit(resource.RLIMIT_RSS)))
+    log.debug("AS rlimit: %s", str(resource.getrlimit(resource.RLIMIT_AS)))
 
     executor = MesosExecutor()
-    log.debug('Made executor')
+    log.debug("Made executor")
     driver = MesosExecutorDriver(executor, use_addict=True)
 
     old_on_event = driver.on_event
@@ -295,13 +322,15 @@ def main():
 
     driver.on_event = patched_on_event
 
-    log.debug('Made driver')
+    log.debug("Made driver")
     driver.start()
-    log.debug('Started driver')
+    log.debug("Started driver")
     driver_result = driver.join()
-    log.debug('Joined driver')
+    log.debug("Joined driver")
 
     # Tolerate a None in addition to the code the docs suggest we should receive from join()
-    exit_value = 0 if (driver_result is None or driver_result == 'DRIVER_STOPPED') else 1
+    exit_value = (
+        0 if (driver_result is None or driver_result == "DRIVER_STOPPED") else 1
+    )
     assert len(executor.runningTasks) == 0
     sys.exit(exit_value)

toil/batchSystems/mesos/test/__init__.py

@@ -17,16 +17,18 @@ log = logging.getLogger(__name__)
 class MesosTestSupport:
     """Mixin for test cases that need a running Mesos master and agent on the local host."""
 
-    @retry(intervals=[1, 1, 2, 4, 8, 16, 32, 64, 128],
-           log_message=(log.info, 'Checking if Mesos is ready...'))
+    @retry(
+        intervals=[1, 1, 2, 4, 8, 16, 32, 64, 128],
+        log_message=(log.info, "Checking if Mesos is ready..."),
+    )
     def wait_for_master(self):
-        with closing(urlopen('http://127.0.0.1:5050/version')) as content:
+        with closing(urlopen("http://127.0.0.1:5050/version")) as content:
             content.read()
 
     def _startMesos(self, numCores=None):
         if numCores is None:
             numCores = cpu_count()
-        shutil.rmtree('/tmp/mesos', ignore_errors=True)
+        shutil.rmtree("/tmp/mesos", ignore_errors=True)
         self.master = self.MesosMasterThread(numCores)
         self.master.start()
         self.agent = self.MesosAgentThread(numCores)
@@ -35,7 +37,7 @@ class MesosTestSupport:
         # Bad Things will happen if the master is not yet ready when Toil tries to use it.
         self.wait_for_master()
 
-        log.info('Mesos is ready! Running test.')
+        log.info("Mesos is ready! Running test.")
 
     def _stopProcess(self, process, timeout=10) -> None:
         """Gracefully stop a process on a timeout, given the Popen object for the process."""
@@ -47,7 +49,7 @@ class MesosTestSupport:
             waited += 1
         if process.poll() is None:
             # It didn't shut down gracefully
-            log.warning('Forcibly killing child which ignored SIGTERM')
+            log.warning("Forcibly killing child which ignored SIGTERM")
             process.kill()
 
     def _stopMesos(self):
@@ -71,7 +73,7 @@ class MesosTestSupport:
 
         def tryRun(self):
             self.popen.wait()
-            log.info('Exiting %s', self.__class__.__name__)
+            log.info("Exiting %s", self.__class__.__name__)
 
         def findMesosBinary(self, names):
             if isinstance(names, str):
@@ -86,7 +88,7 @@ class MesosTestSupport:
                         # Special case for users of PyCharm on OS X. This is where Homebrew installs
                         # it. It's hard to set PATH for PyCharm (or any GUI app) on OS X so let's
                         # make it easy for those poor souls.
-                        return which(name, path='/usr/local/sbin')
+                        return which(name, path="/usr/local/sbin")
                     except StopIteration:
                         pass
 
@@ -94,18 +96,22 @@ class MesosTestSupport:
             if len(names) == 1:
                 sought = "binary '%s'" % names[0]
             else:
-                sought = 'any binary in %s' % str(names)
+                sought = "any binary in %s" % str(names)
 
-            raise RuntimeError("Cannot find %s. Make sure Mesos is installed "
-                                "and it's 'bin' directory is present on the PATH." % sought)
+            raise RuntimeError(
+                "Cannot find %s. Make sure Mesos is installed "
+                "and it's 'bin' directory is present on the PATH." % sought
+            )
 
     class MesosMasterThread(MesosThread):
         def mesosCommand(self):
-            return [self.findMesosBinary('mesos-master'),
-                    '--registry=in_memory',
-                    '--ip=127.0.0.1',
-                    '--port=5050',
-                    '--allocation_interval=500ms']
+            return [
+                self.findMesosBinary("mesos-master"),
+                "--registry=in_memory",
+                "--ip=127.0.0.1",
+                "--port=5050",
+                "--allocation_interval=500ms",
+            ]
 
     class MesosAgentThread(MesosThread):
         def mesosCommand(self):
@@ -114,10 +120,12 @@ class MesosTestSupport:
             # We also make sure to point it explicitly at the right temp work directory, and
             # to disable systemd support because we have to be root to make systemd make us
             # things and we probably aren't when testing.
-            return [self.findMesosBinary(['mesos-agent']),
-                    '--ip=127.0.0.1',
-                    '--master=127.0.0.1:5050',
-                    '--attributes=preemptible:False',
-                    '--resources=cpus(*):%i' % self.numCores,
-                    '--work_dir=/tmp/mesos',
-                    '--no-systemd_enable_support']
+            return [
+                self.findMesosBinary(["mesos-agent"]),
+                "--ip=127.0.0.1",
+                "--master=127.0.0.1:5050",
+                "--attributes=preemptible:False",
+                "--resources=cpus(*):%i" % self.numCores,
+                "--work_dir=/tmp/mesos",
+                "--no-systemd_enable_support",
+            ]

toil/batchSystems/options.py

@@ -12,22 +12,19 @@
 # See the License for the specific language governing permissions and
 
 import logging
-import sys
 from argparse import ArgumentParser, _ArgumentGroup
-from typing import Any, Callable, List, Optional, TypeVar, Union
+from typing import Any, Callable, Optional, Protocol, TypeVar, Union
 
-if sys.version_info >= (3, 8):
-    from typing import Protocol
-else:
-    from typing_extensions import Protocol
-
-from toil.batchSystems.registry import (DEFAULT_BATCH_SYSTEM,
-                                        get_batch_system,
-                                        get_batch_systems)
+from toil.batchSystems.registry import (
+    DEFAULT_BATCH_SYSTEM,
+    get_batch_system,
+    get_batch_systems,
+)
 from toil.lib.threading import cpu_count
 
 logger = logging.getLogger(__name__)
 
+
 class OptionSetter(Protocol):
     """
     Protocol for the setOption function we get to let us set up CLI options for
@@ -36,19 +33,22 @@ class OptionSetter(Protocol):
     Actual functionality is defined in the Config class.
     """
 
-    OptionType = TypeVar('OptionType')
+    OptionType = TypeVar("OptionType")
+
     def __call__(
         self,
         option_name: str,
         parsing_function: Optional[Callable[[Any], OptionType]] = None,
         check_function: Optional[Callable[[OptionType], Union[None, bool]]] = None,
         default: Optional[OptionType] = None,
-        env: Optional[List[str]] = None,
-        old_names: Optional[List[str]] = None
-    ) -> bool:
-        ...
+        env: Optional[list[str]] = None,
+        old_names: Optional[list[str]] = None,
+    ) -> bool: ...
 
-def set_batchsystem_options(batch_system: Optional[str], set_option: OptionSetter) -> None:
+
+def set_batchsystem_options(
+    batch_system: Optional[str], set_option: OptionSetter
+) -> None:
     """
     Call set_option for all the options for the given named batch system, or
     all batch systems if no name is provided.
@@ -76,6 +76,7 @@ def set_batchsystem_options(batch_system: Optional[str], set_option: OptionSette
     set_option("manualMemArgs")
     set_option("run_local_jobs_on_workers")
     set_option("statePollingWait")
+    set_option("state_polling_timeout")
     set_option("batch_logs_dir")
 
 
@@ -109,11 +110,11 @@ def add_all_batchsystem_options(parser: Union[ArgumentParser, _ArgumentGroup]) -
     parser.add_argument(
         "--maxJobs",
         dest="max_jobs",
-        default=SYS_MAX_SIZE, # This is *basically* unlimited and saves a lot of Optional[]
+        default=SYS_MAX_SIZE,  # This is *basically* unlimited and saves a lot of Optional[]
         type=lambda x: int(x) or SYS_MAX_SIZE,
         help="Specifies the maximum number of jobs to submit to the "
-             "backing scheduler at once. Not supported on Mesos or "
-             "AWS Batch. Use 0 for unlimited. Defaults to unlimited.",
+        "backing scheduler at once. Not supported on Mesos or "
+        "AWS Batch. Use 0 for unlimited. Defaults to unlimited.",
     )
     parser.add_argument(
         "--maxLocalJobs",
@@ -121,8 +122,8 @@ def add_all_batchsystem_options(parser: Union[ArgumentParser, _ArgumentGroup]) -
         default=None,
         type=lambda x: int(x) or 0,
         help=f"Specifies the maximum number of housekeeping jobs to "
-             f"run sumultaneously on the local system. Use 0 for "
-             f"unlimited. Defaults to the number of local cores ({cpu_count()}).",
+        f"run sumultaneously on the local system. Use 0 for "
+        f"unlimited. Defaults to the number of local cores ({cpu_count()}).",
     )
     parser.add_argument(
         "--manualMemArgs",
@@ -161,8 +162,16 @@ def add_all_batchsystem_options(parser: Union[ArgumentParser, _ArgumentGroup]) -
         type=int,
         default=None,
         help="Time, in seconds, to wait before doing a scheduler query for job state.  "
-             "Return cached results if within the waiting period. Only works for grid "
-             "engine batch systems such as gridengine, htcondor, torque, slurm, and lsf."
+        "Return cached results if within the waiting period. Only works for grid "
+        "engine batch systems such as gridengine, htcondor, torque, slurm, and lsf.",
+    )
+    parser.add_argument(
+        "--statePollingTimeout",
+        dest="state_polling_timeout",
+        type=int,
+        default=1200,
+        help="Time, in seconds, to retry against a broken scheduler. Only works for grid "
+        "engine batch systems such as gridengine, htcondor, torque, slurm, and lsf.",
     )
     parser.add_argument(
         "--batchLogsDir",
@@ -170,10 +179,19 @@ def add_all_batchsystem_options(parser: Union[ArgumentParser, _ArgumentGroup]) -
         default=None,
         env_var="TOIL_BATCH_LOGS_DIR",
         help="Directory to tell the backing batch system to log into. Should be available "
-             "on both the leader and the workers, if the backing batch system writes logs "
-             "to the worker machines' filesystems, as many HPC schedulers do. If unset, "
-             "the Toil work directory will be used. Only works for grid engine batch "
-             "systems such as gridengine, htcondor, torque, slurm, and lsf."
+        "on both the leader and the workers, if the backing batch system writes logs "
+        "to the worker machines' filesystems, as many HPC schedulers do. If unset, "
+        "the Toil work directory will be used. Only works for grid engine batch "
+        "systems such as gridengine, htcondor, torque, slurm, and lsf.",
+    )
+
+    parser.add_argument(
+        "--memoryIsProduct",
+        dest="memory_is_product",
+        default=False,
+        action="store_true",
+        help="If the batch system understands requested memory as a product of the requested memory and the number"
+        "of cores, set this flag to properly allocate memory.",
     )
 
     for name in get_batch_systems():
@@ -185,5 +203,5 @@ def add_all_batchsystem_options(parser: Union[ArgumentParser, _ArgumentGroup]) -
             # Skip anything we can't import
             continue
         # Ask the batch system to create its options in the parser
-        logger.debug('Add options for %s batch system', name)
+        logger.debug("Add options for %s batch system", name)
         batch_system_type.add_options(parser)