From bc0c4ac3e0b334638151aa14785d6d464caaf22d Mon Sep 17 00:00:00 2001
From: Piero Toffanin <pt@masseranolabs.com>
Date: Mon, 4 Dec 2023 23:40:23 -0500
Subject: [PATCH] Add --worker-cpus option

---
 docker-compose.worker-cpu.yml |  4 ++++
 webodm.sh                     | 32 ++++++++++++++++++++++----------
 2 files changed, 26 insertions(+), 10 deletions(-)
 create mode 100644 docker-compose.worker-cpu.yml

diff --git a/docker-compose.worker-cpu.yml b/docker-compose.worker-cpu.yml
new file mode 100644
index 00000000..b002e0e1
--- /dev/null
+++ b/docker-compose.worker-cpu.yml
@@ -0,0 +1,4 @@
+version: '2.2'
+services:
+  worker:
+    cpus: ${WO_WORKER_CPUS}
\ No newline at end of file
diff --git a/webodm.sh b/webodm.sh
index 0c794bf7..55faeecb 100755
--- a/webodm.sh
+++ b/webodm.sh
@@ -56,13 +56,13 @@ case $key in
     shift # past argument
     shift # past value
     ;;
-	--media-dir)
+    --media-dir)
     WO_MEDIA_DIR=$(realpath "$2")
     export WO_MEDIA_DIR
     shift # past argument
     shift # past value
     ;;
-	--db-dir)
+    --db-dir)
     WO_DB_DIR=$(realpath "$2")
     export WO_DB_DIR
     shift # past argument
@@ -72,19 +72,19 @@ case $key in
     export WO_SSL=YES
     shift # past argument
     ;;
-	--ssl-key)
+    --ssl-key)
     WO_SSL_KEY=$(realpath "$2")
     export WO_SSL_KEY
     shift # past argument
     shift # past value
     ;;
-	--ssl-cert)
-	WO_SSL_CERT=$(realpath "$2")
-	export WO_SSL_CERT
+    --ssl-cert)
+    WO_SSL_CERT=$(realpath "$2")
+    export WO_SSL_CERT
     shift # past argument
     shift # past value
     ;;
-	--ssl-insecure-port-redirect)
+    --ssl-insecure-port-redirect)
     export WO_SSL_INSECURE_PORT_REDIRECT="$2"
     shift # past argument
     shift # past value
@@ -103,11 +103,11 @@ case $key in
     dev_mode=true
     shift # past argument
     ;;
-	--gpu)
+    --gpu)
     gpu=true
     shift # past argument
     ;;
-	--broker)
+    --broker)
     export WO_BROKER="$2"
     shift # past argument
     shift # past value
@@ -136,11 +136,17 @@ case $key in
     shift # past argument
     shift # past value
     ;;    
-	--worker-memory)
+    --worker-memory)
     WO_WORKER_MEMORY="$2"
     export WO_WORKER_MEMORY
     shift # past argument
     shift # past value
+    ;;	
+    --worker-cpus)
+    WO_WORKER_CPUS="$2"
+    export WO_WORKER_CPUS
+    shift # past argument
+    shift # past value
     ;;
     *)    # unknown option
     POSITIONAL+=("$1") # save it in an array for later
@@ -185,6 +191,7 @@ usage(){
   echo "	--gpu	Use GPU NodeODM nodes (Linux only) (default: disabled)"
   echo "	--settings	Path to a settings.py file to enable modifications of system settings (default: None)"
   echo "	--worker-memory	Maximum amount of memory allocated for the worker process (default: unlimited)"
+  echo "	--worker-cpus	Maximum number of CPUS allocated for the worker process (default: all)"
   
   exit
 }
@@ -371,6 +378,7 @@ start(){
 	echo "Default Nodes: $WO_DEFAULT_NODES"
 	echo "Settings: $WO_SETTINGS"
 	echo "Worker memory limit: $WO_WORKER_MEMORY"
+	echo "Worker cpus limit: $WO_WORKER_CPUS"
 	echo "================================"
 	echo "Make sure to issue a $0 down if you decide to change the environment."
 	echo ""
@@ -445,6 +453,10 @@ start(){
 		command+=" -f docker-compose.worker-memory.yml"
 	fi
 
+	if [ ! -z "$WO_WORKER_CPUS" ]; then
+		command+=" -f docker-compose.worker-cpu.yml"
+	fi
+
 	command="$command up"
 
 	if [[ $detached = true ]]; then