1- #! /usr/bin/env bash
2- set -Eeuo pipefail
1+ #! /bin/bash
32
4- [[ " ${DEBUG :- } " == " 1 " ]] && set -x
3+ set -x
54
6- CONDA_BIN=" /usr/local/bin"
7- JOB_AD=" ${_CONDOR_JOB_AD:- } "
8- JOB_IWD=" ${_CONDOR_JOB_IWD:- } "
9-
10- # #######################################
11- # Conda init
12- # #######################################
13-
14- if [[ -f " $CONDA_BIN /conda" ]]; then
15- eval " $( " $CONDA_BIN /conda" shell.bash hook) "
16- else
17- echo " Conda not found at $CONDA_BIN /conda"
5+ if [ " $EXTRA_CONDA_PACKAGES " ]; then
6+ echo " conda: EXTRA_CONDA_PACKAGES environment variable found. Installing."
7+ /opt/conda/bin/mamba install -y $EXTRA_CONDA_PACKAGES
188fi
199
20- # #######################################
21- # Conda / Pip extra packages
22- # #######################################
23-
24- if [[ -n " ${EXTRA_CONDA_PACKAGES:- } " ]]; then
25- echo " Installing EXTRA_CONDA_PACKAGES..."
26- " $CONDA_BIN /mamba" install -y ${EXTRA_CONDA_PACKAGES}
10+ if [ " $EXTRA_PIP_PACKAGES " ]; then
11+ echo " pip: EXTRA_PIP_PACKAGES environment variable found. Installing"
12+ /opt/conda/bin/pip install $EXTRA_PIP_PACKAGES
2713fi
2814
29- if [[ -n " ${EXTRA_PIP_PACKAGES:- } " ]]; then
30- echo " Installing EXTRA_PIP_PACKAGES..."
31- " $CONDA_BIN /pip" install ${EXTRA_PIP_PACKAGES}
32- fi
33-
34- # #######################################
35- # Exit early if running as sidecar
36- # #######################################
37-
38- if [[ -v COFFEA_CASA_SIDECAR ]]; then
39- exec " $@ "
40- fi
41-
42- # #######################################
43- # Wait for required ClassAd fields
44- # #######################################
45-
46- wait_for_key () {
47- local key=" $1 "
48- echo " Waiting for $key ..."
49- until grep -q " $key " " $JOB_AD " ; do
50- sleep 1
15+ # If there is defined COFFEA_CASA_SIDECAR env variable (inside hub values.yml),
16+ # then we use this container as a sidecar for notebook.
17+ if [[ ! -v COFFEA_CASA_SIDECAR ]]; then
18+ if [ " ${GITHUB_ACTIONS:- } " == " true" ]; then
19+ echo " CI mode, no need to test dask_HostPort info..."
20+ else
21+ # From chtc/dask-chtc: wait for the job ad to be updated with <service>_HostPort
22+ # This happens during the first update, usually a few seconds after the job starts
23+ echo " Waiting for dask_HostPort and nanny_HostPort information..."
24+ # Check if we are not in GH CI environment (otherwise image check will stuck forever)
25+ # docs: Always set to true when GitHub Actions is running the workflow.
26+ # You can use this variable to differentiate when tests are being run locally or by GitHub Actions.
27+ while true ; do
28+ if grep dask_HostPort " $_CONDOR_JOB_AD " ; then
29+ break
30+ fi
31+ sleep 1
5132 done
52- echo " $key detected."
53- }
54-
55- if [[ " ${GITHUB_ACTIONS:- } " != " true" ]]; then
56- wait_for_key " dask_HostPort"
57- wait_for_key " nanny_HostPort"
58- fi
59-
60- [[ -z " $JOB_IWD " ]] && { echo " Missing _CONDOR_JOB_IWD" ; exit 1; }
61-
62- # #######################################
63- # Secure file transfers
64- # #######################################
65-
66- copy_if_exists () {
67- local src=" $1 "
68- local dst=" $2 "
69- [[ -f " $src " ]] && mkdir -p " $( dirname " $dst " ) " && cp " $src " " $dst "
70- }
71-
72- copy_if_exists " $JOB_IWD /condor_token" \
73- " ${SEC_TOKEN_SYSTEM_DIRECTORY:- } /condor_token"
74- copy_if_exists " $JOB_IWD /access_token" " /tmp/.xcache/access_token"
75-
76- # #######################################
77- # Environment installation
78- # #######################################
79-
80- if [[ -f " $JOB_IWD /environment.yml" || -f " $JOB_IWD /environment.yaml" ]]; then
81- ENV_FILE=$( ls " $JOB_IWD " /environment.y* ml | head -1)
82- echo " Updating conda environment..."
83- " $CONDA_BIN /mamba" env update -n base -f " $ENV_FILE "
84- fi
85-
86- if [[ -f " $JOB_IWD /requirements.txt" ]]; then
87- echo " Installing pip requirements..."
88- " $CONDA_BIN /python" -m pip install -r " $JOB_IWD /requirements.txt"
89- fi
33+ echo " Got dask_HostPort, proceeding..."
34+ echo
35+
36+ while true ; do
37+ if grep nanny_HostPort " $_CONDOR_JOB_AD " ; then
38+ break
39+ fi
40+ sleep 1
41+ done
42+ echo " Got nanny_HostPort, proceeding..."
43+ echo
9044
91- # #######################################
92- # Parse ClassAd safely (single awk pass)
93- # #######################################
94-
95- if [[ -n " $JOB_AD " ]]; then
96- declare -A AD
97-
98- while IFS=' =' read -r key value; do
99- key=$( echo " $key " | xargs)
100- value=$( echo " $value " | tr -d ' "' | xargs)
101- AD[" $key " ]=" $value "
102- done < " $JOB_AD "
103-
104- HOST=$( echo " ${AD[StartdIpAddr]:- } " | sed ' s/<\([^:]*\).*/\1/' )
105- WORKER_TYPE=" ${AD[CoffeaCasaWorkerType]:- dask} "
106-
107- echo " Loaded ClassAd:"
108- cat " $JOB_AD " >&2
109-
110- if [[ " $WORKER_TYPE " == " dask" ]]; then
111- MEMORY_MB=" ${AD[MemoryProvisioned]} .00MB"
112-
113- exec python -m distributed.cli.dask_worker " ${AD[DaskSchedulerAddress]} " \
114- --name " ${AD[DaskWorkerName]} " \
115- --tls-ca-file " ${JOB_IWD} /ca.pem" \
116- --tls-cert " ${JOB_IWD} /hostcert.pem" \
117- --tls-key " ${JOB_IWD} /hostcert.pem" \
118- --nthreads " ${AD[DaskWorkerCores]} " \
119- --memory-limit " $MEMORY_MB " \
120- --nanny \
121- --death-timeout 60 \
122- --protocol tls \
123- --lifetime 7200 \
124- --listen-address " tls://0.0.0.0:${AD[dask_ContainerPort]} " \
125- --contact-address " tls://${HOST} :${AD[dask_HostPort]} "
126-
127- elif [[ " $WORKER_TYPE " == " taskvine" ]]; then
128- DISK_MB=$(( AD[DiskProvisioned] / 1024 ))
129-
130- exec vine_worker --ssl -dall \
131- --contact-hostport " $HOST :${AD[dask_HostPort]} " \
132- --transfer-port " ${AD[dask_ContainerPort]} " \
133- --cores " ${AD[CpusProvisioned]} " \
134- --memory " ${AD[MemoryProvisioned]} " \
135- --disk " $DISK_MB " \
136- --timeout 7200 \
137- " ${AD[DaskSchedulerAddress]%:* } " \
138- " ${AD[DaskSchedulerAddress]##*: } "
45+ if [ -z " $_CONDOR_JOB_IWD " ]; then
46+ echo " Error: something is wrong, $_CONDOR_JOB_IWD (path to the initial working directory the job was born with) was not defined!"
47+ exit 1
13948 fi
140- fi
141-
142- # #######################################
143- # Default fallback
144- # #######################################
145-
146- exec " $@ "
49+ fi
50+
51+ # Condor token securily transfered from scheduler
52+ if [[ -f " $_CONDOR_JOB_IWD /condor_token" ]]; then
53+ mkdir -p $SEC_TOKEN_SYSTEM_DIRECTORY && cp $_CONDOR_JOB_IWD /condor_token $SEC_TOKEN_SYSTEM_DIRECTORY /condor_token
54+ fi
55+
56+ if [[ -f " $_CONDOR_JOB_IWD /ceph.conf" ]]; then
57+ cp $_CONDOR_JOB_IWD /ceph.conf ${CEPH_DIR}
58+ fi
59+
60+ if [[ -f " $_CONDOR_JOB_IWD /access_token" ]]; then
61+ mkdir -p /tmp/.xcache && cp $_CONDOR_JOB_IWD /access_token /tmp/.xcache/access_token
62+ fi
63+
64+ if [[ -f " $_CONDOR_JOB_IWD /keyring" ]]; then
65+ cp $_CONDOR_JOB_IWD /keyring ${CEPH_DIR}
66+ fi
67+
68+ if [ -e " $_CONDOR_JOB_IWD /environment.yml" ]; then
69+ echo " Conda: environment.yml found. Installing packages."
70+ /opt/conda/bin/mamba env update -n base -f $_CONDOR_JOB_IWD /environment.yml
71+ elif [ -e " $_CONDOR_JOB_IWD /environment.yaml" ]; then
72+ echo " Conda: environment.yaml found. Installing packages."
73+ /opt/conda/bin/mamba env update -n base -f $_CONDOR_JOB_IWD /environment.yaml
74+ else
75+ echo " No environment.yml, conda will not install any package."
76+ fi
77+
78+ if [ -e " $_CONDOR_JOB_IWD /requirements.txt" ]; then
79+ echo " Pip: requirements.txt found. Installing packages."
80+ /opt/conda/bin/python -m pip install -r $_CONDOR_JOB_IWD /requirements.txt
81+ else
82+ echo " No requirements.txt, pip will not install any module."
83+ fi
84+
85+ # CA certificate securily transfered from scheduler
86+ if [[ -f " $_CONDOR_JOB_IWD /ca.pem" ]]; then
87+ PATH_CA_FILE=" $_CONDOR_JOB_IWD /ca.pem"
88+ fi
89+ # Hostcertn securily transfered from scheduler and
90+ # userkey here is simply concatenated in usercert
91+ if [[ -f " $_CONDOR_JOB_IWD /hostcert.pem" ]]; then
92+ FILE_CERT=" $_CONDOR_JOB_IWD /hostcert.pem"
93+ FILE_KEY=" $_CONDOR_JOB_IWD /hostcert.pem"
94+ fi
95+
96+ # HTCondor port, hostname and external IP ("must" variables)
97+ if [ ! -z " $_CONDOR_JOB_AD " ]; then
98+ # We make sure that we use proper configuration (dask worker name, ports, number of CPUs,
99+ # memory requested for worker, hostname of scheduler), parcing HTCondor Job AD file`
100+ # Nanny container port will be used later...
101+ PORT=` cat $_CONDOR_JOB_AD | grep dask_HostPort | tr -d ' "' | awk ' {print $NF;}' `
102+ NANNYPORT=` cat $_CONDOR_JOB_AD | grep nanny_HostPort | tr -d ' "' | awk ' {print $NF;}' `
103+ NANNYCONTAINER_PORT=` cat $_CONDOR_JOB_AD | grep nanny_ContainerPort | tr -d ' "' | awk ' {print $NF;}' `
104+ CONTAINER_PORT=` cat $_CONDOR_JOB_AD | grep dask_ContainerPort | tr -d ' "' | awk ' {print $NF;}' `
105+ # FIXME:
106+ # NANNY_PORT=`cat $_CONDOR_JOB_AD | grep nanny_HostPort | tr -d '"' | awk '{print $NF;}'`
107+ # NANNYCONTAINER_PORT=`cat $_CONDOR_JOB_AD | grep nanny_ContainerPort | tr -d '"' | awk '{print $NF;}'`
108+ # HOST=`cat $_CONDOR_JOB_AD | grep RemoteHost | tr -d '"' | tr '@' ' ' | awk '{print $NF;}'`
109+ HOST=` cat $_CONDOR_JOB_AD | grep StartdIpAddr | sed ' s/StartdIpAddr = \"<\([^:]*\).*/\1/' `
110+ NAME=` cat $_CONDOR_JOB_AD | grep " DaskWorkerName " | tr -d ' "' | awk ' {print $NF;}' `
111+ # Requirement: to add to Condor job decription "+DaskSchedulerAddress": '"tcp://129.93.183.34:8787"',
112+ EXTERNALIP_PORT=` cat $_CONDOR_JOB_AD | grep DaskSchedulerAddress | tr -d ' "' | awk ' {print $NF;}' `
113+ # From jthiltges
114+ # WORKER_LIFETIME=${WORKER_LIFETIME:-1 hour}
115+
116+ WORKER_TYPE=` cat $_CONDOR_JOB_AD | grep " CoffeaCasaWorkerType " | tr -d ' "' | awk ' {print $NF;}' `
117+
118+ CPUS_PROVISIONED=` cat $_CONDOR_JOB_AD | grep " CpusProvisioned " | tr -d ' "' | awk ' {print $NF;}' `
119+ MEMORY_PROVISIONED=` cat $_CONDOR_JOB_AD | grep " MemoryProvisioned " | tr -d ' "' | awk ' {print $NF;}' `
120+ DISK_PROVISIONED=` cat $_CONDOR_JOB_AD | grep " DiskProvisioned " | tr -d ' "' | awk ' {print $NF;}' `
121+
122+ echo " Print ClassAd:" 1>&2
123+ cat $_CONDOR_JOB_AD 1>&2
124+
125+ if [ -z " ${WORKER_TYPE} " ] || [ " ${WORKER_TYPE} " = dask ]
126+ then
127+ CPUS=` cat $_CONDOR_JOB_AD | grep " DaskWorkerCores " | tr -d ' "' | awk ' {print $NF;}' `
128+ MEMORY=` cat $_CONDOR_JOB_AD | grep " RequestMemory " | tr -d ' "' | awk ' {print $NF;}' `
129+ MEMORY_MB_FORMATTED=$MEMORY_PROVISIONED " .00MB"
130+
131+ # Dask worker command execurted in HTCondor pool.
132+ # Communication protocol: in Coffea-casa we use only secured communications (over TLS)
133+ # re-apply John's patch and add: --nanny-contact-address tls://$HOST:$NANNYPORT --nanny-port $NANNYCONTAINER_PORT \ \
134+ HTCONDOR_COMMAND=" python -m distributed.cli.dask_worker $EXTERNALIP_PORT \
135+ --name $NAME \
136+ --tls-ca-file $PATH_CA_FILE \
137+ --tls-cert $FILE_CERT \
138+ --tls-key $FILE_KEY \
139+ --nthreads $CPUS \
140+ --memory-limit $MEMORY_MB_FORMATTED \
141+ --nanny \
142+ --death-timeout 60 \
143+ --protocol tls \
144+ --lifetime 7200 \
145+ --listen-address tls://0.0.0.0:$CONTAINER_PORT \
146+ --contact-address tls://$HOST :$PORT "
147+ elif [ " ${WORKER_TYPE} " = taskvine ]
148+ then
149+ DISK_MB=$(( DISK_PROVISIONED/ 1024 ))
150+ MANAGER_HOST=${EXTERNALIP_PORT%:* }
151+ MANAGER_HOST=${MANAGER_HOST#* tls: \/\/ } # XXX: fix on vine_worker, it should parsing this
152+ MANAGER_PORT=${EXTERNALIP_PORT##*: }
153+
154+ HTCONDOR_COMMAND=" vine_worker --ssl -dall \
155+ --contact-hostport $HOST :$PORT \
156+ --transfer-port $CONTAINER_PORT \
157+ --cores $CPUS_PROVISIONED \
158+ --memory $MEMORY_PROVISIONED \
159+ --disk $DISK_MB \
160+ --timeout 7200 \
161+ $MANAGER_HOST $MANAGER_PORT "
162+ else
163+ HTCONDOR_COMMAND=" $@ "
164+ fi
165+
166+ # Debug print
167+ echo $HTCONDOR_COMMAND 1>&2
168+ exec $HTCONDOR_COMMAND
169+
170+ export LD_LIBRARY_PATH=" ${LD_LIBRARY_PATH} :/opt/combine/CombinedLimit/build/lib"
171+ fi
172+ else
173+ exec " $@ "
174+ fi
0 commit comments