Repository: DataBiosphere/toil

Issue:
Docker in the Toil appliance container does not know about --gpus
@glennhickey says that you can't use the Docker client we ship in the Toil appliance with the newly-upstreamed `--gpus` option:

```
docker run -it --rm --gpus all  nvidia/cuda:11.4.3-runtime-ubuntu20.04 bash
unknown flag: --gpus
See 'docker run --help'.
```

We should update our Docker client to one that supports this option.

Related to #4199.

┆Issue is synchronized with this [Jira Story](https://ucsc-cgl.atlassian.net/browse/TOIL-1247)
┆Issue Number: TOIL-1247

Docker in the Toil appliance container does not know about --gpus
@glennhickey says that you can't use the Docker client we ship in the Toil appliance with the newly-upstreamed `--gpus` option:

```
docker run -it --rm --gpus all  nvidia/cuda:11.4.3-runtime-ubuntu20.04 bash
unknown flag: --gpus
See 'docker run --help'.
```

We should update our Docker client to one that supports this option.

Related to #4199.

┆Issue is synchronized with this [Jira Story](https://ucsc-cgl.atlassian.net/browse/TOIL-1247)
┆Issue Number: TOIL-1247

Toil's `cpu_count()` function should respect `os.sched_getaffinity()` CPU pinning
In https://github.com/ComparativeGenomicsToolkit/cactus/issues/820, someone is running Toil workflow leaders under an HPC batch system, and getting them scheduled on non-container-based slices of machines. In this case the Linux-specific `sched_setaffinity`/`sched_getaffinity` API is being used to try and isolate the workflows to particular sets of CPU cores.

The Toil leader tries to schedule on `toil.lib.threading.cpu_count()` cores, which is the number of cores actually in the host, but not the number of cores that are supposed to be being used for the workflow. `nproc` knows about `sched_getaffinity()` and returns the number of cores that we ought to be pinned to, but `cpu_count()` doesn't.

We should, on platforms where `os` has a `sched_getaffinity()`, cap `cpu_count()` at the number of cores that we are actually allowed to run on, since the user will probably not expect to also need to provide `--maxCores` and will want us to be as smart as `nproc`.

┆Issue is synchronized with this [Jira Story](https://ucsc-cgl.atlassian.net/browse/TOIL-1236)
┆friendlyId: TOIL-1236



Discussion:
According to https://github.com/ComparativeGenomicsToolkit/cactus/pull/844, Glenn wants this for the next Toil release.
According to https://github.com/ComparativeGenomicsToolkit/cactus/pull/844, Glenn wants this for the next Toil release.


Resolving patch:
diff --git a/docker/Dockerfile.py b/docker/Dockerfile.py
--- a/docker/Dockerfile.py
+++ b/docker/Dockerfile.py
@@ -53,6 +53,7 @@
                          'less',
                          'vim',
                          'git',
+                         'docker.io',
                          'time',
                          # Dependencies for Mesos which the deb doesn't actually list
                          'libsvn1',
@@ -83,7 +84,7 @@ def heredoc(s):
 motd = ''.join(l + '\\n\\\n' for l in motd.splitlines())
 
 print(heredoc('''
-    FROM ubuntu:20.04
+    FROM ubuntu:22.04
 
     ARG TARGETARCH
 
@@ -102,12 +103,13 @@ def heredoc(s):
 
     # Find a repo with a Mesos build.
     # See https://rpm.aventer.biz/README.txt
-    # A working snapshot is https://ipfs.io/ipfs/QmfTy9sXhHsgyWwosCJDfYR4fChTosA8HhoaMgmeJ5LSmS/
+    # A working snapshot is https://ipfs.io/ipfs/QmfTy9sXhHsgyWwosCJDfYR4fChTosA8HhoaMgmeJ5LSmS/ for https://rpm.aventer.biz/Ubuntu
+    # And one that works with https://rpm.aventer.biz/Ubuntu/focal (the new URL) is at https://ipfs.io/ipfs/Qmcrmx7T1YkEnyexMXdd7QjoBZxf7DMDrQ5ErUKi9mDRw6/
     # As archived with:
     # mkdir mesos-repo && cd mesos-repo
     # wget --recursive --restrict-file-names=windows -k --convert-links --no-parent --page-requisites https://rpm.aventer.biz/Ubuntu/ https://www.aventer.biz/assets/support_aventer.asc https://rpm.aventer.biz/README.txt
     # ipfs add -r .
-    RUN echo "deb https://rpm.aventer.biz/Ubuntu focal main" \
+    RUN echo "deb https://rpm.aventer.biz/Ubuntu/focal focal main" \
         > /etc/apt/sources.list.d/mesos.list \
         && curl https://www.aventer.biz/assets/support_aventer.asc | apt-key add -
 
@@ -119,10 +121,32 @@ def heredoc(s):
         rm -rf /var/lib/apt/lists/*
     
     # Install a particular old Debian Sid Singularity from somewhere.
-    ADD singularity-sources.tsv /etc/singularity/singularity-sources.tsv
-    RUN wget -q "$(cat /etc/singularity/singularity-sources.tsv | grep "^$TARGETARCH" | cut -f3)" && \
-        dpkg -i singularity-container_3*.deb && \
-        rm singularity-container_3*.deb && \
+    # It's 3.10, which is new enough to use cgroups2, but it needs a newer libc
+    # than Ubuntu 20.04 ships. So we need a 22.04+ base.
+    #
+    # But 22.04 ships squashfs-tools 4.4 or 4.5 or so, which is new enough that
+    # errors encountered during extraction produce a nonzero exit code, without
+    # a special option:
+    # <https://github.com/plougher/squashfs-tools/commit/1dd7f32e79b7600d379a4f26fb8d138ebdfc70be>.
+    # If unsquashfs thinks it is root, but it can't change UIDs and GIDs freely, it
+    # will continue but fail the whole command instead of returning success. It complains:
+    #
+    # set_attributes: failed to change uid and gids on /image/rootfs/etc/gshadow, because Invalid argument
+    #
+    # But inside a Kubernetes container we can be root but not actually be
+    # allowed to set UIDs and GIDs arbitrarily. Singularity can't handle this,
+    # and can't pass the flag to ignore these errors, and we can't wrap
+    # unsquashfs with a shell script because of how it gets mounted into the
+    # container under construction along with its libraries (see
+    # <https://github.com/apptainer/singularity/issues/6113#issuecomment-901897566>).
+    #
+    # So we need to make sure to install a downgraded squashfs first.
+    ADD extra-debs.tsv /etc/singularity/extra-debs.tsv
+    RUN wget -q "$(cat /etc/singularity/extra-debs.tsv | grep "^squashfs-tools.$TARGETARCH" | cut -f4)" && \
+        dpkg -i squashfs-tools_*.deb && \
+        wget -q "$(cat /etc/singularity/extra-debs.tsv | grep "^singularity-container.$TARGETARCH" | cut -f4)" && \
+        dpkg -i singularity-container_*.deb && \
+        rm singularity-container_*.deb && \
         sed -i 's!bind path = /etc/localtime!#bind path = /etc/localtime!g' /etc/singularity/singularity.conf && \
         mkdir -p /usr/local/libexec/toil && \
         mv /usr/bin/singularity /usr/local/libexec/toil/singularity-real \
@@ -135,9 +159,11 @@ def heredoc(s):
 
     ADD customDockerInit.sh /usr/bin/customDockerInit.sh
 
-    ADD singularity-wrapper.sh /usr/local/bin/singularity
+    # Wrap Singularity to use a Docker mirror instead of always Docker Hub
+    # We need to put it where the installed singularity expects singularity to actually be.
+    ADD singularity-wrapper.sh /usr/bin/singularity
 
-    RUN chmod 777 /usr/bin/waitForKey.sh && chmod 777 /usr/bin/customDockerInit.sh && chmod 777 /usr/local/bin/singularity
+    RUN chmod 777 /usr/bin/waitForKey.sh && chmod 777 /usr/bin/customDockerInit.sh && chmod 777 /usr/bin/singularity
 
     # The stock pip is too old and can't install from sdist with extras
     RUN curl -sS https://bootstrap.pypa.io/get-pip.py | {python}
@@ -153,12 +179,6 @@ def heredoc(s):
         && /home/s3am/bin/pip install s3am==2.0 \
         && ln -s /home/s3am/bin/s3am /usr/local/bin/
 
-    # Install statically linked version of docker client
-    RUN curl https://download.docker.com/linux/static/stable/$(if [ $TARGETARCH = amd64 ] ; then echo x86_64 ; else echo aarch64 ; fi)/docker-18.06.1-ce.tgz \
-        | tar -xvzf - --transform='s,[^/]*/,,g' -C /usr/local/bin/ \
-        && chmod u+x /usr/local/bin/docker \
-        && /usr/local/bin/docker -v
-
     # Fix for https://issues.apache.org/jira/browse/MESOS-3793
     ENV MESOS_LAUNCHER=posix
 
diff --git a/src/toil/lib/threading.py b/src/toil/lib/threading.py
--- a/src/toil/lib/threading.py
+++ b/src/toil/lib/threading.py
@@ -25,7 +25,7 @@
 import threading
 import traceback
 from contextlib import contextmanager
-from typing import Any, Dict, Iterator, Optional
+from typing import Any, Dict, Iterator, Optional, Union, cast
 
 import psutil  # type: ignore
 
@@ -84,7 +84,7 @@ def join(self, *args: Optional[float], **kwargs: Optional[float]) -> None:
             raise_(exc_type, exc_value, traceback)
 
 
-def cpu_count() -> Any:
+def cpu_count() -> int:
     """
     Get the rounded-up integer number of whole CPUs available.
 
@@ -105,42 +105,77 @@ def cpu_count() -> Any:
     cached = getattr(cpu_count, 'result', None)
     if cached is not None:
         # We already got a CPU count.
-        return cached
+        return cast(int, cached)
 
     # Get the fallback answer of all the CPUs on the machine
-    total_machine_size = psutil.cpu_count(logical=True)
+    total_machine_size = cast(int, psutil.cpu_count(logical=True))
 
     logger.debug('Total machine size: %d cores', total_machine_size)
 
-    try:
-        with open('/sys/fs/cgroup/cpu/cpu.cfs_quota_us') as stream:
-            # Read the quota
-            quota = int(stream.read())
-
-        logger.debug('CPU quota: %d', quota)
+    # cgroups may limit the size
+    cgroup_size: Union[float, int] = float('inf')
 
-        if quota == -1:
-            # Assume we can use the whole machine
-            return total_machine_size
-
-        with open('/sys/fs/cgroup/cpu/cpu.cfs_period_us') as stream:
-            # Read the period in which we are allowed to burn the quota
-            period = int(stream.read())
+    try:
+        # See if we can fetch these and use them
+        quota: Optional[int] = None
+        period: Optional[int] = None
+
+        # CGroups v1 keeps quota and period separate
+        CGROUP1_QUOTA_FILE = '/sys/fs/cgroup/cpu/cpu.cfs_quota_us'
+        CGROUP1_PERIOD_FILE = '/sys/fs/cgroup/cpu/cpu.cfs_period_us'
+        # CGroups v2 keeps both in one file, space-separated, quota first
+        CGROUP2_COMBINED_FILE = '/sys/fs/cgroup/cpu.max'
+
+        if os.path.exists(CGROUP1_QUOTA_FILE) and os.path.exists(CGROUP1_PERIOD_FILE):
+            logger.debug('CPU quota and period available from cgroups v1')
+            with open(CGROUP1_QUOTA_FILE) as stream:
+                # Read the quota
+                quota = int(stream.read())
+
+            with open(CGROUP1_PERIOD_FILE) as stream:
+                # Read the period in which we are allowed to burn the quota
+                period = int(stream.read())
+        elif os.path.exists(CGROUP2_COMBINED_FILE):
+            logger.debug('CPU quota and period available from cgroups v2')
+            with open(CGROUP2_COMBINED_FILE) as stream:
+                # Read the quota and the period together
+                quota, period = [int(part) for part in stream.read().split(' ')]
+        else:
+            logger.debug('CPU quota/period not available from cgroups v1 or cgroups v2')
 
-        logger.debug('CPU quota period: %d', period)
+        if quota is not None and period is not None:
+            # We got a quota and a period.
+            logger.debug('CPU quota: %d period: %d', quota, period)
 
-        # The thread count is how many multiples of a wall clcok period we can burn in that period.
-        cgroup_size = int(math.ceil(float(quota)/float(period)))
+            if quota == -1:
+                # But the quota can be -1 for unset.
+                # Assume we can use the whole machine.
+                return total_machine_size
 
-        logger.debug('Cgroup size in cores: %d', cgroup_size)
+            # The thread count is how many multiples of a wall clock period we
+            # can burn in that period.
+            cgroup_size = int(math.ceil(float(quota)/float(period)))
 
+            logger.debug('Control group size in cores: %d', cgroup_size)
     except:
         # We can't actually read these cgroup fields. Maybe we are a mac or something.
         logger.debug('Could not inspect cgroup: %s', traceback.format_exc())
-        cgroup_size = float('inf') # type: ignore
+
+    # CPU affinity may limit the size
+    affinity_size: Union[float, int] = float('inf')
+    if hasattr(os, 'sched_getaffinity'):
+        try:
+            logger.debug('CPU affinity available')
+            affinity_size = len(os.sched_getaffinity(0))
+            logger.debug('CPU affinity is restricted to %d cores', affinity_size)
+        except:
+             # We can't actually read this even though it exists.
+            logger.debug('Could not inspect scheduling affinity: %s', traceback.format_exc())
+    else:
+        logger.debug('CPU affinity not available')
 
     # Return the smaller of the actual thread count and the cgroup's limit, minimum 1.
-    result = max(1, min(cgroup_size, total_machine_size))
+    result = cast(int, max(1, min(min(affinity_size, cgroup_size), total_machine_size)))
     logger.debug('cpu_count: %s', str(result))
     # Make sure to remember it for the next call
     setattr(cpu_count, 'result', result)


Test patch: