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,20 +75,24 @@
     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,
     downloadHttpFile,
     get_listing,
     normalizeFilesDirs,
     visit_class,
+    OutputCallbackType,
 )
 from cwl_utils.types import (
     CWLDirectoryType,
@@ -102,6 +114,13 @@
 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 (
+    INJECTED_MESSAGE_DIR,
+    add_injections,
+    command_line_to_shell_script,
+    handle_injection_messages_from,
+    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 +1161,134 @@ 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 ToilDockerCommandLineJob(DockerCommandLineJob):
+    """Container job that collects resource stats from injected in-container code."""
+
+    def _execute(
+        self,
+        runtime: list[str],
+        env: MutableMapping[str, str],
+        runtimeContext: cwltool.context.RuntimeContext,
+        monitor_function: Callable[["subprocess.Popen[str]"], None] | None = None,
+    ) -> None:
+        # 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.
+
+        # First we need the command that's actually going to run. To be fully
+        # spec-compliant, we need to make sure Docker ENTRYPOINT keeps working.
+        # So we need to fetch it out, put it in the script, and remember to
+        # override it with a shell later. See
+        # <https://github.com/common-workflow-language/cwltool/blob/1bf74499ca1c4a5f98e7cffb0ad4aa89aa98cb9e/cwltool/schemas/v1.1/CommandLineTool.yml#L785-L792>
+
+        # We know the selected container image is the last element in runtime.
+        # Inspect it to get its entrypoint list
+        image_config = json.loads(
+            subprocess.check_output(
+                [
+                    self.docker_exec,
+                    "inspect",
+                    "-f",
+                    "{{json .Config}}",
+                    runtime[-1],
+                ]
+            )
+        )
+        # Null entrypoints are allowed; Docker doesn't seem to put literal
+        # nulls but who knows how it might change.
+        entrypoint: list[str] = image_config.get("Entrypoint") or []
+
+        # We also need to ensure that if the CWL baseCommand is empty, the
+        # image command list runs.
+        default_command: list[str] = image_config.get("Cmd") or []
+
+        # Make a command string that does the expected user work
+        script = command_line_to_shell_script(entrypoint + (self.command_line or default_command))
+
+        # Our injected code also checks to make sure files mounted on Docker
+        # for Mac are intact, so we need to tell it about the file
+        # mounts we are going to use.
+        file_mounts = self._file_mounts_from_pathmapper()
+
+        # Since we're not allowed to drop stuff in the working directory, we
+        # put the resource usage info in the temp directory, which cwltool
+        # mounts, as divined by Anthropic Claude. See
+        # <https://github.com/common-workflow-language/cwltool/blob/1bf74499ca1c4a5f98e7cffb0ad4aa89aa98cb9e/cwltool/docker.py#L332-L334>
+        script = add_injections(script, file_mounts, os.path.join(self.CONTAINER_TMPDIR, INJECTED_MESSAGE_DIR))
+        self.command_line = shell_script_to_command_line(script)
+
+        # Adjust the runtime to clear the entrypoint.
+        # See <https://github.com/moby/moby/issues/23498>.
+        # Stick the argument and value in before the last element.
+        runtime[-1:-1] = ["--entrypoint", ""]
+
+        # By the time the base class _execute() returns, the temp directory
+        # holding the resource usage information has already been deleted. So
+        # we need to hook into one of the functions it calls, from somewhere
+        # where we have access to our tmpdir (like here).
+        def message_handler_output_callback(
+            output_callbacks: OutputCallbackType | None,
+            outputs: CWLObjectType | None,
+            processStatus: str,
+        ) -> None:
+            """
+            Output callback that processes resource usage messages.
+
+            Needs to be partial'd into a callback chain.
+            """
+            handle_injection_messages_from(os.path.join(self.tmpdir, INJECTED_MESSAGE_DIR))
+            if output_callbacks:
+                output_callbacks(outputs, processStatus)
+        # Hook into the output callback chain
+        self.output_callback = functools.partial(message_handler_output_callback, self.output_callback)
+
+        # Run the container
+        super()._execute(runtime, env, runtimeContext, monitor_function)
+        # By the time we get here, the temp directory where we hide the
+        # resource usage info has already been deleted.
+
+    def _file_mounts_from_pathmapper(self) -> 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 self.pathmapper.files():
+            ent = self.pathmapper.mapper(location)
+            if ent.type == "File" and not ent.resolved.startswith("_:"):
+                file_mounts.append((ent.resolved, ent.target))
+        return file_mounts
+
 
 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 make_job_runner(
+        self, runtimeContext: cwltool.context.RuntimeContext
+    ) -> type[cwltool.job.JobBase]:
+        """
+        Use a job class that monitors container resource usage if needed.
+
+        We only act on container systems where the container isn't a descendant
+        process of us.
+        """
+        parent_class = super().make_job_runner(runtimeContext)
+        if parent_class is DockerCommandLineJob:
+            # There's only one daemon-containerized container system: it's
+            # Docker itself (and not any of the subclasses of this),
+            return ToilDockerCommandLineJob
+        return parent_class
 
     def _initialworkdir(
         self, j: cwltool.job.JobBase | None, builder: cwltool.builder.Builder
@@ -1186,7 +1330,7 @@ def _initialworkdir(
             # Notice that we have downloaded our inputs. Explain which files
             # those are here and what the task will expect to call them.
             self._toil_job.files_downloaded_hook(host_and_job_paths)
-
+    
 
 class ToilExpressionTool(ToilTool, cwltool.command_line_tool.ExpressionTool):
     """Subclass the cwltool expression tool to provide the custom ToilPathMapper."""

src/toil/jobStores/aws/jobStore.py

@@ -195,6 +195,11 @@ def initialize(self, config: Config) -> None:
 
         Create bucket, raise if it already exists.
         Set options from config.
+
+        Note that polling for bucket existence only happens once; there's not a
+        good way to guard against races to create the bucket that also handles
+        the possibility that a request to AWS could return an error but still
+        create the bucket.
         """
         logger.debug(
             f"Instantiating {self.__class__} for region {self.region} with bucket: '{self.bucket_name}'"

src/toil/lib/aws/s3.py

@@ -90,6 +90,10 @@ def create_s3_bucket(
     Create an AWS S3 bucket, using the given Boto3 S3 session, with the
     given name, in the given region.
 
+    If the bucket already exists, uses it. (It is difficult to implement
+    retries if this function is not idempotent and AWS can successfully create
+    the bucket but produce an error on the client.)
+
     Supports the us-east-1 region, where bucket creation is special.
 
     *ALL* S3 bucket creation should use this function.
@@ -107,13 +111,18 @@ def create_s3_bucket(
         False.
     """
     logger.info("Creating bucket '%s' in region %s.", bucket_name, region)
-    if region == "us-east-1":  # see https://github.com/boto/boto3/issues/125
-        bucket = s3_resource.create_bucket(Bucket=bucket_name)
-    else:
-        bucket = s3_resource.create_bucket(
-            Bucket=bucket_name,
-            CreateBucketConfiguration={"LocationConstraint": region},
-        )
+    try:
+        if region == "us-east-1":  # see https://github.com/boto/boto3/issues/125
+            bucket = s3_resource.create_bucket(Bucket=bucket_name)
+        else:
+            bucket = s3_resource.create_bucket(
+                Bucket=bucket_name,
+                CreateBucketConfiguration={"LocationConstraint": region},
+            )
+    except s3_resource.meta.client.exceptions.BucketAlreadyOwnedByYou:
+        # We could have created the bucket, then failed, then retried. So use
+        # it.
+        bucket = s3_resource.Bucket(bucket_name)
     # wait until the bucket exists before adding tags
     bucket.wait_until_exists()