From ea6992c80b99153f9a01c40bb83568b142ddb0df Mon Sep 17 00:00:00 2001
From: Robert Bartel <robert.bartel@noaa.gov>
Date: Fri, 12 Jul 2024 09:37:47 -0400
Subject: [PATCH] Update worker entrypoints for local data.

Updating main entrypoint scripts for ngen and calibration worker images
for local data handling.
---
 docker/main/ngen/ngen_cal_entrypoint.sh | 8 ++++++++
 docker/main/ngen/ngen_entrypoint.sh     | 8 ++++++++
 2 files changed, 16 insertions(+)

diff --git a/docker/main/ngen/ngen_cal_entrypoint.sh b/docker/main/ngen/ngen_cal_entrypoint.sh
index 2483f9855..1a843f27f 100755
--- a/docker/main/ngen/ngen_cal_entrypoint.sh
+++ b/docker/main/ngen/ngen_cal_entrypoint.sh
@@ -39,6 +39,10 @@ while [ ${#} -gt 0 ]; do
             declare -x CALIBRATION_CONFIG_BASENAME="${2:?}"
             shift
             ;;
+        --primary-workers)
+            declare -x PRIMARY_WORKERS="${2:?}"
+            shift
+            ;;
     esac
     shift
 done
@@ -101,6 +105,10 @@ start_calibration() {
     return ${NGEN_RETURN}
 }
 
+# Run make_data_local Python functions to make necessary data local
+# Called for every worker, but Python code will make sure only one worker per node makes a call that has effect
+py_funcs make_data_local ${WORKER_INDEX:-0} ${PRIMARY_WORKERS:-0}
+
 # We can allow worker index to not be supplied when executing serially
 if [ "${WORKER_INDEX:-0}" = "0" ]; then
     if [ "$(whoami)" = "${MPI_USER:?MPI user not defined}" ]; then
diff --git a/docker/main/ngen/ngen_entrypoint.sh b/docker/main/ngen/ngen_entrypoint.sh
index c22b86770..7bba15f80 100755
--- a/docker/main/ngen/ngen_entrypoint.sh
+++ b/docker/main/ngen/ngen_entrypoint.sh
@@ -35,6 +35,10 @@ while [ ${#} -gt 0 ]; do
             declare -x WORKER_INDEX="${2:?}"
             shift
             ;;
+        --primary-workers)
+            declare -x PRIMARY_WORKERS="${2:?}"
+            shift
+            ;;
     esac
     shift
 done
@@ -60,6 +64,10 @@ if [ ! -e /dmod/datasets/linked_job_output ]; then
     ln -s ${JOB_OUTPUT_WRITE_DIR} /dmod/datasets/linked_job_output
 fi
 
+# Run make_data_local Python functions to make necessary data local
+# Called for every worker, but Python code will make sure only one worker per node makes a call that has effect
+py_funcs make_data_local ${WORKER_INDEX:-0} ${PRIMARY_WORKERS:-0}
+
 # We can allow worker index to not be supplied when executing serially
 if [ "${WORKER_INDEX:-0}" = "0" ]; then
     if [ "$(whoami)" = "${MPI_USER:?MPI user not defined}" ]; then