polish TaskContainer interface

WDL/runtime/backend/cli_subprocess.py

@@ -11,7 +11,7 @@
 from contextlib import ExitStack
 from ..._util import PygtailLogger
 from ..._util import StructuredLogMessage as _
-from .. import config, _statusbar
+from .. import config
 from ..error import Terminated, DownloadFailed
 from ..task_container import TaskContainer
 
@@ -71,6 +71,7 @@ def _run(self, logger: logging.Logger, terminating: Callable[[], bool], command:
             cli_log_filename = os.path.join(self.host_dir, f"{self.cli_name}.log.txt")
             cli_log = cleanup.enter_context(open(cli_log_filename, "wb"))
             cli_logger = logger.getChild(self.cli_name)
+            poll_stderr = cleanup.enter_context(self.poll_stderr_context(logger))
             poll_cli_log = cleanup.enter_context(
                 PygtailLogger(
                     logger,
@@ -79,13 +80,6 @@ def _run(self, logger: logging.Logger, terminating: Callable[[], bool], command:
                     level=logging.INFO,
                 )
             )
-            poll_stderr = cleanup.enter_context(
-                PygtailLogger(
-                    logger,
-                    self.host_stderr_txt(),
-                    callback=self.stderr_callback,
-                )
-            )
 
             # prepare command
             with open(os.path.join(self.host_dir, "command"), "w") as outfile:
@@ -103,12 +97,7 @@ def _run(self, logger: logging.Logger, terminating: Callable[[], bool], command:
             logger.notice(  # pyre-ignore
                 _(f"{self.cli_name} run", pid=proc.pid, log=cli_log_filename)
             )
-            cleanup.enter_context(
-                _statusbar.task_running(
-                    self.runtime_values.get("cpu", 0),
-                    self.runtime_values.get("memory_reservation", 0),
-                )
-            )
+            cleanup.enter_context(self.task_running_context())
 
             # long-poll for completion
             exit_code = None

WDL/runtime/backend/docker_swarm.py

@@ -18,12 +18,9 @@
 from typing import List, Dict, Set, Optional, Any, Callable, Tuple, Iterable
 import docker
 from ... import Error
-from ..._util import (
-    chmod_R_plus,
-    PygtailLogger,
-)
+from ..._util import chmod_R_plus
 from ..._util import StructuredLogMessage as _
-from .. import config, _statusbar
+from .. import config
 from ..error import Interrupted, Terminated
 from ..task_container import TaskContainer
 
@@ -232,13 +229,7 @@ def _run(self, logger: logging.Logger, terminating: Callable[[], bool], command:
 
             # stream stderr into log
             with contextlib.ExitStack() as cleanup:
-                poll_stderr = cleanup.enter_context(
-                    PygtailLogger(
-                        logger,
-                        self.host_stderr_txt(),
-                        callback=self.stderr_callback,
-                    )
-                )
+                poll_stderr = cleanup.enter_context(self.poll_stderr_context(logger))
 
                 # poll for container exit
                 running_states = {"preparing", "running"}
@@ -277,12 +268,7 @@ def _run(self, logger: logging.Logger, terminating: Callable[[], bool], command:
                         # indicate actual container start in status bar
                         # 'preparing' is when docker is pulling and extracting the image, which can
                         # be a lengthy and somewhat intensive operation, so we count it as running.
-                        cleanup.enter_context(
-                            _statusbar.task_running(
-                                self.runtime_values.get("cpu", 0),
-                                self.runtime_values.get("memory_reservation", 0),
-                            )
-                        )
+                        cleanup.enter_context(self.task_running_context())
                         was_running = True
                     if "running" in self._observed_states:
                         poll_stderr()

WDL/runtime/task_container.py

@@ -1,21 +1,22 @@
 """
-Abstract interface for task container runtime + default Docker Swarm backend
+Abstract interface for task container runtime
 """
 import os
 import logging
 import shutil
 import threading
-from typing import Callable, Iterable, List, Set, Tuple, Type, Any, Dict, Optional
+from typing import Callable, Iterable, List, Set, Type, Any, Dict, Optional, ContextManager
 from abc import ABC, abstractmethod
 from contextlib import suppress
 from .. import Error
 from .._util import (
     TerminationSignalFlag,
     path_really_within,
     rmtree_atomic,
+    PygtailLogger,
 )
 from .._util import StructuredLogMessage as _
-from . import config
+from . import config, _statusbar
 from .error import OutputError, Terminated, CommandFailed
 
 
@@ -29,16 +30,16 @@ class TaskContainer(ABC):
     @classmethod
     def global_init(cls, cfg: config.Loader, logger: logging.Logger) -> None:
         """
-        Perform any necessary one-time initialization of the underlying container backend. Must be
+        Perform any necessary one-time initialization of the underlying container backend. To be
         invoked once per process prior to any instantiation of the class.
         """
         raise NotImplementedError()
 
     @classmethod
     def detect_resource_limits(cls, cfg: config.Loader, logger: logging.Logger) -> Dict[str, int]:
         """
-        Detect the maximum resources (cpu and mem_bytes) that the underlying container backend
-        would be able to provision.
+        Detect the maximum resources ("cpu" and "mem_bytes") that the underlying container backend
+        will be able to provision for any one task.
 
         If determining this is at all costly, then backend should memoize (thread-safely and
         perhaps front-loaded in global_init).
@@ -74,21 +75,25 @@ def detect_resource_limits(cls, cfg: config.Loader, logger: logging.Logger) -> D
     """
 
     input_path_map_rev: Dict[str, str]
+    """
+    Inverse of ``input_path_map`` (also maintained by ``add_paths``)
+    """
 
     try_counter: int
     """
     :type: int
 
     Counter for number of retries; starts at 1 on the first attempt. On subsequent attempts, the
-    names (on the host) of the working directory, stdout.txt, and stderr.txt will incorporate the
+    names (on the host) of the working directory, stdout.txt, and stderr.txt may incorporate the
     count, to ensure their uniqueness.
     """
 
     runtime_values: Dict[str, Any]
     """
-    Evaluted task runtime{} section. Typically the TaskContainer backend needs to honor cpu,
-    memory_limit, memory_reservation, docker. Resources must have already been fit to
-    get_resource_limits(). Retry logic (maxRetries, preemptible) is handled externally.
+    Evaluted task runtime{} section, to be populated externall before invoking run(). Typically the
+    TaskContainer backend needs to honor cpu, memory_limit, memory_reservation, docker. Resources
+    must have already been fit to get_resource_limits(). Retry logic (maxRetries, preemptible) is
+    handled externally.
     """
 
     stderr_callback: Optional[Callable[[str], None]]
@@ -359,6 +364,31 @@ def host_stderr_txt(self):
             self.host_dir, f"stderr{self.try_counter if self.try_counter > 1 else ''}.txt"
         )
 
+    def poll_stderr_context(self, logger: logging.Logger) -> ContextManager[Callable[[], None]]:
+        """
+        Implementation helper: open a context yielding a function to poll stderr.txt and log each
+        each line (to either logger or self.stderr_callback if set). _run() implementation should
+        call the function periodically while container is running, and close the context once
+        done/failed.
+        """
+        return PygtailLogger(
+            logger,
+            self.host_stderr_txt(),
+            callback=self.stderr_callback,
+        )
+
+    def task_running_context(self) -> ContextManager[None]:
+        """
+        Implementation helper: open a context which counts the task, and its CPU and memory
+        reservations, in the CLI status bar's "running" ticker. _run() implementation should open
+        this context once the container is truly running (not while e.g. still queued), and close
+        it once done/failed.
+        """
+        return _statusbar.task_running(
+            self.runtime_values.get("cpu", 0),
+            self.runtime_values.get("memory_reservation", 0),
+        )
+
 
 _backends: Dict[str, Type[TaskContainer]] = dict()
 _backends_lock: threading.Lock = threading.Lock()