overhaul how number of requested cores is handled

Author: jhrmnn

mona/dirtask.py

@@ -2,6 +2,7 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 import logging
+import os
 import subprocess
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -151,14 +152,19 @@ async def dir_task(exe: File, inputs: List[DirtaskInput]) -> Dict[str, File]:
     input_names = {
         str(inp if isinstance(inp, File) else inp[0]) for inp in [exe, *inputs]
     }
+    ncores = Session.active().running_task.storage.get('ncores')
     dirtask_tmpdir = DirtaskTmpdir(lambda p: p not in input_names)
     with dirtask_tmpdir as tmpdir:
         checkout_files(tmpdir, exe, inputs)
         out_path, err_path = tmpdir / 'STDOUT', tmpdir / 'STDERR'
         try:
             with out_path.open('w') as stdout, err_path.open('w') as stderr:
                 await run_process(
-                    str(tmpdir / exe.path), stdout=stdout, stderr=stderr, cwd=tmpdir
+                    str(tmpdir / exe.path),
+                    stdout=stdout,
+                    stderr=stderr,
+                    cwd=tmpdir,
+                    ncores=ncores,
                 )
         except subprocess.CalledProcessError as e:
             if dirtask_tmpdir.has_tmpdir_manager():

mona/plugins/parallel.py

@@ -105,11 +105,16 @@ async def _acquire(self, ncores: int) -> AsyncGenerator[None, None]:
 
     async def _run_coro(self, corofunc: Corofunc[_T], *args: Any, **kwargs: Any) -> _T:
         task = Session.active().running_task
-        n = cast(int, task.storage.get('ncores', 1))
+        n: Optional[int] = kwargs.get('ncores')
+        if n is not None:
+            if n == -1:
+                n = self._ncores
+            kwargs['ncores'] = n
+        else:
+            n = 1
         if n > self._available:
             log.debug(
-                f'Waiting for {n-self._available}/{n} '
-                f'unavailable cores for "{task}"'
+                f'Waiting for {n-self._available}/{n} unavailable cores for {task}'
             )
             waited = True
         else:

mona/runners.py

@@ -3,14 +3,15 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 import asyncio
 import logging
+import os
 import subprocess
 from typing import Any, Callable, Optional, Tuple, TypeVar, Union
 from typing_extensions import Protocol, runtime
 
 from .sessions import Session
 from .tasks import Corofunc
 
-__version__ = '0.1.0'
+__version__ = '0.1.1'
 __all__ = ['run_shell', 'run_process', 'run_thread']
 
 log = logging.getLogger(__name__)
@@ -33,14 +34,15 @@ def _scheduler() -> Optional[Scheduler]:
     return scheduler
 
 
-async def run_shell(cmd: str, **kwargs: Any) -> ProcessOutput:
+async def run_shell(cmd: str, ncores: int = None, **kwargs: Any) -> ProcessOutput:
     """Execute a command in a shell.
 
     Wrapper around :func:`asyncio.create_subprocess_shell` that handles errors
     and whose behavior can be modified by session plugins.
 
     :param str cmd: a shell command to be executed
-    :param kwargs: all keyword arguments are passed to
+    :param int ncores: number of cores that should be taken by the process
+    :param kwargs: all other keyword arguments are passed to
                    :func:`~asyncio.create_subprocess_shell`.
                    :data:`~subprocess.PIPE` is passed to `stdin` and `stdout`
                    keyword arguments by default.
@@ -56,14 +58,15 @@ async def run_shell(cmd: str, **kwargs: Any) -> ProcessOutput:
     return await _run_process(cmd, **kwargs)
 
 
-async def run_process(*args: str, **kwargs: Any) -> ProcessOutput:
+async def run_process(*args: str, ncores: int = None, **kwargs: Any) -> ProcessOutput:
     """Create a subprocess.
 
     Wrapper around :func:`asyncio.create_subprocess_exec` that handles errors
     and whose behavior can be modified by session plugins.
 
     :param str args: arguments of the subprocess
-    :param kwargs: all keyword arguments are passed to
+    :param int ncores: number of cores that should be taken by the process
+    :param kwargs: all other keyword arguments are passed to
                    :func:`~asyncio.create_subprocess_exec`.
                    :data:`~subprocess.PIPE` is passed to `stdin` and `stdout`
                    keyword arguments by default.
@@ -73,18 +76,22 @@ async def run_process(*args: str, **kwargs: Any) -> ProcessOutput:
     """
     scheduler = _scheduler()
     if scheduler:
-        return await scheduler(_run_process, args, **kwargs)
+        return await scheduler(_run_process, args, ncores=ncores, **kwargs)
     return await _run_process(args, **kwargs)
 
 
 async def _run_process(
     args: Union[str, Tuple[str, ...]],
     shell: bool = False,
     input: bytes = None,
+    ncores: int = None,
     **kwargs: Any,
 ) -> Union[bytes, Tuple[bytes, bytes]]:
     kwargs.setdefault('stdin', subprocess.PIPE)
     kwargs.setdefault('stdout', subprocess.PIPE)
+    kwargs.setdefault('env', os.environ.copy())
+    if ncores is not None:
+        kwargs['env']['MONA_NCORES'] = str(ncores)
     if shell:
         assert isinstance(args, str)
         proc = await asyncio.create_subprocess_shell(args, **kwargs)