Add resource usage measurement injection to CWL interpreter

src/toil/cwl/cwltoil.py

@@ -31,13 +31,21 @@
 import pprint
 import shutil
 import stat
+import subprocess
 import sys
 import textwrap
 import uuid
 
 # This is also in configargparse but MyPy doesn't know it
 from argparse import RawDescriptionHelpFormatter
-from collections.abc import Callable, Iterator, Mapping, MutableMapping, MutableSequence
+from collections.abc import (
+    Callable,
+    Generator,
+    Iterator,
+    Mapping,
+    MutableMapping,
+    MutableSequence,
+)
 from tempfile import NamedTemporaryFile, TemporaryFile, gettempdir
 from threading import Thread
 from typing import IO, Any, Literal, Optional, Protocol, TextIO, TypeVar, Union, cast
@@ -67,13 +75,16 @@
     fill_in_defaults,
     shortname,
 )
+from cwltool.docker import DockerCommandLineJob, PodmanCommandLineJob
+from cwltool.job import ContainerCommandLineJob
 from cwltool.secrets import SecretStore
 from cwltool.singularity import SingularityCommandLineJob
 from cwltool.software_requirements import (
     DependenciesConfiguration,
     get_container_from_software_requirements,
 )
 from cwltool.stdfsaccess import StdFsAccess, abspath
+from cwltool.udocker import UDockerCommandLineJob
 from cwltool.utils import (
     adjustDirObjs,
     aslist,
@@ -102,6 +113,12 @@
 from toil.common import Config, Toil, addOptions, InconsistentConfigurationError
 from toil.cwl import check_cwltool_version
 from toil.lib.directory import DirectoryContents, decode_directory, encode_directory
+from toil.lib.interpreter import (
+    add_injections,
+    command_line_to_shell_script,
+    handle_injection_messages_from_outdir,
+    shell_script_to_command_line,
+)
 from toil.lib.misc import call_command
 from toil.lib.trs import resolve_workflow
 from toil.provisioners.clusterScaler import JobTooBigError
@@ -1142,9 +1159,104 @@ def __str__(self) -> str:
         """Return string representation of this tool type."""
         return f'{self.__class__.__name__}({repr(getattr(self, "tool", {}).get("id", "???"))})'
 
+# We need to know which cwltool container implementations run the container
+# under the calling process's child tree. Unlisted container implementations
+# are assumed to run it elsewhere through a deamon.
+#
+# We have to list it this way around (and not list the implementations that use
+# a daemon) because some of these extend DockerCommandLineJob.
+CHILD_PROCESS_CONTAINER_JOBS = (PodmanCommandLineJob, SingularityCommandLineJob, UDockerCommandLineJob)
 
 class ToilCommandLineTool(ToilTool, cwltool.command_line_tool.CommandLineTool):
-    """Subclass the cwltool command line tool to provide the custom ToilPathMapper."""
+    """Subclass the cwltool command line tool to provide the custom ToilPathMapper
+    and add the monitoring code to the job's container command line."""
+
+    def _uses_container(self, runtimeContext: cwltool.context.RuntimeContext) -> bool:
+        """
+        Returns True if this tool will be run inside a container.
+        """
+        if not runtimeContext.use_container:
+            return False
+        docker_req, _ = self.get_requirement("DockerRequirement")
+        if docker_req is not None:
+            return True
+        if runtimeContext.find_default_container is not None:
+            return runtimeContext.find_default_container(self) is not None
+        return runtimeContext.default_container is not None
+
+    # TODO: Why is this on ToilCommandLineTool?
+    @staticmethod
+    def _file_mounts_from_pathmapper(
+        job: ContainerCommandLineJob,
+    ) -> list[tuple[str, str]]:
+        """
+        Get all the container mounts that will be used for a containerized job.
+
+        :returns: a list of (host path, container path) tuples, one per mount.
+        """
+        file_mounts: list[tuple[str, str]] = []
+        for location in job.pathmapper.files():
+            ent = job.pathmapper.mapper(location)
+            if ent.type == "File" and not ent.resolved.startswith("_:"):
+                file_mounts.append((ent.resolved, ent.target))
+        return file_mounts
+
+    def job(
+        self,
+        job_order: CWLObjectType,
+        output_callbacks: Callable[[CWLObjectType | None, str], None],
+        runtimeContext: cwltool.context.RuntimeContext,
+    ) -> Generator[
+        cwltool.job.JobBase | cwltool.command_line_tool.CallbackJob, None, None
+    ]:
+        """
+        Override the job method to inject resource (cpu & memory) monitoring code into
+        the job's container command line.
+        """
+        for job in super().job(job_order, output_callbacks, runtimeContext):
+            if (
+                isinstance(job, ContainerCommandLineJob) and
+                not isinstance(job, CHILD_PROCESS_CONTAINER_JOBS) and 
+                self._uses_container(
+                    runtimeContext
+                )
+            ):
+                # This job is going to run in a container that makes it not be
+                # a descendant of our process. So we need to inject code to
+                # count its resource usage.
+
+                # That code also checks to make sure files mounted on Docker
+                # for Mac are intact, so we need to tell it ablut the file
+                # mounts we are going to use.
+                file_mounts = self._file_mounts_from_pathmapper(job)
+                script = command_line_to_shell_script(job.command_line)
+                script = add_injections(script, file_mounts)
+                job.command_line = shell_script_to_command_line(script)
+            yield job
+
+    def collect_output_ports(
+        self,
+        ports: CommentedSeq | set[CWLObjectType],
+        builder: cwltool.builder.Builder,
+        outdir: str,
+        rcode: int,
+        compute_checksum: bool = True,
+        jobname: str = "",
+        readers: MutableMapping[str, CWLObjectType] | None = None,
+    ) -> cwltool.command_line_tool.OutputPortsType:
+        """
+        Hook output collection to also collect resource usage statistics.
+        """
+        handle_injection_messages_from_outdir(outdir)
+        return super().collect_output_ports(
+            ports,
+            builder,
+            outdir,
+            rcode,
+            compute_checksum,
+            jobname,
+            readers,
+        )        
 
     def _initialworkdir(
         self, j: cwltool.job.JobBase | None, builder: cwltool.builder.Builder

src/toil/lib/interpreter.py

@@ -0,0 +1,238 @@
+"""
+Shared code for workflow language interpreters implemented with Toil.
+"""
+
+# Copyright (C) 2026 Regents of the University of California
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import glob
+import logging
+import os
+import platform
+import shlex
+import textwrap
+from typing import Iterable
+
+from toil.lib.resources import ResourceMonitor
+
+logger = logging.getLogger(__name__)
+
+
+####
+# Shell command injection system
+# 
+# When a workflow steps runs inside a container like Docker, where the contained
+# process is not a descendant of Toil, the Toil worker process on the host cannot
+# see how much CPU and RAM that step actually used.
+# 
+# This system allows Toil to inject code into the container that will colklect
+# and export resource usage information back to Toil.
+#
+# It also allows Toil to do other checks on the container system.
+###
+
+# Code injected into the container communicates back to the rest of Toil
+# through files in this directory.
+INJECTED_MESSAGE_DIR = ".toil_runtime"
+
+
+# We mostly want to work with shell script strings, but CWL works with command
+# argument lists, so we use these functions to convert back and forth.
+
+def command_line_to_shell_script(command_line: list[str]) -> str:
+    """
+    Extract or synthesize the inner shell script from a cwltool argv list.
+
+    We don't want to disturb user CWL command line binding lists that
+    explicitly ask for things like ["bash", "-c"].
+    """
+    return " ".join(shlex.quote(arg) for arg in command_line) # this is the shell script string
+
+
+def shell_script_to_command_line(script: str) -> list[str]:
+    """
+    Wrap a shell script as a list of arguments for launchign a process.
+
+    The resulting command required Bash to be available.
+    """
+    return ["bash", "-c", script]
+
+# Main function
+
+def add_injections(
+    command_string: str,
+    file_mounts: Iterable[tuple[str, str]],
+    message_dir: str = INJECTED_MESSAGE_DIR,
+) -> str:
+    """
+    Add resource usage monitoring and file mount checking code to a command.
+
+    The command is expected to be about to run in a container, under a
+    container system that does not itself attribute resource usage to the
+    calling Toil process (such as Docker, which uses a daemon).
+
+    :param command_string: shell command or script to modify
+    :param file_mounts: collection of (host path, container path) tuples for
+        files mounted into the container. Code will be added to require that
+        the container sees the complete file that the host sees.
+    :param message_dir: directory relative to the working directory that the
+        command should record resource usage to
+
+    :returns: shell command string (possibly containing multiple commands) that
+        runs the original command and reports resource usage.
+
+    """
+
+    parts = []
+    # We're running on Docker or another platform where Toil isn't an ancestor
+    # process, so we need to monitor CPU usage and so on from inside the
+    # container, since it won't be attributed to Toil child processes in the
+    # leader's self-monitoring.
+
+    # TODO: Mount this script from a file Toil installs instead or something
+    # instead of injecting it in every command line, which makes it show up in
+    # logs.
+    script = textwrap.dedent(
+        """\
+        function _toil_resource_monitor () {
+            # Turn off error checking and echo in here
+            set +ex
+            MESSAGE_DIR="${1}"
+            mkdir -p "${MESSAGE_DIR}"
+
+            function sample_cpu_usec() {
+                if [[ -f  /sys/fs/cgroup/cpu.stat ]] ; then
+                    awk '{ if ($1 == "usage_usec") {print $2} }' /sys/fs/cgroup/cpu.stat
+                elif [[ -f /sys/fs/cgroup/cpuacct/cpuacct.stat ]] ; then
+                    echo $(( $(head -n 1 /sys/fs/cgroup/cpuacct/cpuacct.stat | cut -f2 -d' ') * 10000 ))
+                fi
+            }
+
+            function sample_memory_bytes() {
+                if [[ -f /sys/fs/cgroup/memory.stat ]] ; then
+                    awk '{ if ($1 == "anon") { print $2 } }' /sys/fs/cgroup/memory.stat
+                elif [[ -f /sys/fs/cgroup/memory/memory.stat ]] ; then
+                    awk '{ if ($1 == "total_rss") { print $2 } }' /sys/fs/cgroup/memory/memory.stat
+                fi
+            }
+
+            while true ; do
+                printf "CPU\\t" >> ${MESSAGE_DIR}/resources.tsv
+                sample_cpu_usec >> ${MESSAGE_DIR}/resources.tsv
+                printf "Memory\\t" >> ${MESSAGE_DIR}/resources.tsv
+                sample_memory_bytes >> ${MESSAGE_DIR}/resources.tsv
+                sleep 1
+            done
+        }
+        """
+    )
+    parts.append(script)
+    # Launch in a subshell so that it doesn't interfere with Bash "wait" in the main shell
+    parts.append(f"(_toil_resource_monitor {message_dir} &)")
+
+    if platform.system() == "Darwin":
+        # With gRPC FUSE file sharing, files immediately downloaded before
+        # being mounted may appear as size 0 in the container due to a race
+        # condition. Check for this and produce an approperiate error.
+
+        script = textwrap.dedent(
+            """\
+            function _toil_check_size () {
+                TARGET_FILE="${1}"
+                GOT_SIZE="$(stat -c %s "${TARGET_FILE}")"
+                EXPECTED_SIZE="${2}"
+                if [[ "${GOT_SIZE}" != "${EXPECTED_SIZE}" ]] ; then
+                    echo >&2 "Toil Error:"
+                    echo >&2 "File size visible in container for ${TARGET_FILE} is size ${GOT_SIZE} but should be size ${EXPECTED_SIZE}"
+                    echo >&2 "Are you using gRPC FUSE file sharing in Docker Desktop?"
+                    echo >&2 "It doesn't work: see <https://github.com/DataBiosphere/toil/issues/4542>."
+                    exit 1
+                fi
+            }
+        """
+        )
+        parts.append(script)
+        for host_path, job_path in file_mounts:
+            expected_size = os.path.getsize(host_path)
+            if expected_size != 0:
+                parts.append(f'_toil_check_size "{job_path}" {expected_size}')
+
+    parts.append(command_string)
+
+    return "\n".join(parts)
+
+# Helper functions to parse resource usage output
+
+def handle_message_file(file_path: str) -> None:
+    """
+    Handle a message file received from injected code from :meth:`add_injections()`.
+
+    Records the described usage as if it had occurred within a child process,
+    by talking to the :class:`toil.lib.resources.ResourceMonitor` system.
+
+    :param file_path: the host-side path of the message file.
+    """
+    if os.path.basename(file_path) == "resources.tsv":
+        # This is a TSV of resource usage info.
+        first_cpu_usec: int | None = None
+        last_cpu_usec: int | None = None
+        max_memory_bytes: int | None = None
+
+        for line in open(file_path):
+            if not line.endswith("\n"):
+                # Skip partial lines
+                continue
+            # For each full line we got
+            parts = line.strip().split("\t")
+            if len(parts) != 2:
+                # Skip odd-shaped lines
+                continue
+            if parts[0] == "CPU":
+                # Parse CPU usage
+                cpu_usec = int(parts[1])
+                # Update summary stats
+                if first_cpu_usec is None:
+                    first_cpu_usec = cpu_usec
+                last_cpu_usec = cpu_usec
+            elif parts[0] == "Memory":
+                # Parse memory usage
+                memory_bytes = int(parts[1])
+                # Update summary stats
+                if max_memory_bytes is None or max_memory_bytes < memory_bytes:
+                    max_memory_bytes = memory_bytes
+
+        if max_memory_bytes is not None:
+            logger.info(
+                "Container used at about %s bytes of memory at peak",
+                max_memory_bytes,
+            )
+            # Treat it as if used by a child process
+            ResourceMonitor.record_extra_memory(max_memory_bytes // 1024)
+        if last_cpu_usec is not None:
+            assert first_cpu_usec is not None
+            cpu_seconds = (last_cpu_usec - first_cpu_usec) / 1000000
+            logger.info("Container used about %s seconds of CPU time", cpu_seconds)
+            # Treat it as if used by a child process
+            ResourceMonitor.record_extra_cpu(cpu_seconds)
+
+def handle_injection_messages_from_outdir(outdir: str) -> None:
+    """
+    Handle any message files in the job outdir.
+
+    Files would have been left by injected code from :meth:`add_injections()`.
+    """
+    message_dir = os.path.join(outdir, INJECTED_MESSAGE_DIR)
+    for file_path in glob.glob(os.path.join(message_dir, "*")):
+        if os.path.isfile(file_path):
+            handle_message_file(file_path)