From 2ed89594875713940d89f81a4dff0144e576787d Mon Sep 17 00:00:00 2001
From: Junru Shao <junrushao1994@gmail.com>
Date: Mon, 4 Sep 2023 16:22:01 -0700
Subject: [PATCH] [CI] Allow Limit CPUs in Docker

This PR adds a new flag `--cpus` in `./docker/bash.sh`, which is passed
to docker command that allows limiting the number of CPU cores of a
docker container.

Related materials: https://docs.docker.com/config/containers/resource_constraints/#configure-the-default-cfs-scheduler
---
 docker/bash.sh | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/docker/bash.sh b/docker/bash.sh
index b04270d0a819..a5dec23f91ad 100755
--- a/docker/bash.sh
+++ b/docker/bash.sh
@@ -38,9 +38,11 @@ set -euo pipefail
 function show_usage() {
     cat <<EOF
 Usage: docker/bash.sh [-i|--interactive] [--net=host] [-t|--tty]
-         [--mount MOUNT_DIR] [--repo-mount-point REPO_MOUNT_POINT]
-         [--dry-run] [--name NAME]
-         <DOCKER_IMAGE_NAME> [--] [COMMAND]
+          [--cpus NUM_CPUS] [--mount MOUNT_DIR]
+          [--repo-mount-point REPO_MOUNT_POINT]
+          [--dry-run] [--name NAME]
+          <DOCKER_IMAGE_NAME> [--] [COMMAND]
+
 
 -h, --help
 
@@ -54,6 +56,10 @@ Usage: docker/bash.sh [-i|--interactive] [--net=host] [-t|--tty]
 
     Start the docker session with a pseudo terminal (tty).
 
+--cpus NUM_CPUS
+
+    Limit the number of CPU cores to be used.
+
 --net=host
 
     Expose servers run into the container to the host, passing the
@@ -184,6 +190,11 @@ while (( $# )); do
             eval $break_joined_flag
             ;;
 
+        --cpus)
+            DOCKER_FLAGS+=(--cpus "$2")
+            shift 2
+            ;;
+
         --net=host)
             USE_NET_HOST=true
             shift