Skip to content

Commit ab3fe74

Browse files
committed
Revert script back to see if it is a reason
1 parent 4138219 commit ab3fe74

1 file changed

Lines changed: 165 additions & 137 deletions

File tree

Lines changed: 165 additions & 137 deletions
Original file line numberDiff line numberDiff line change
@@ -1,146 +1,174 @@
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
188
fi
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
2713
fi
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

Comments
 (0)