run_dev_instance.sh 12.2 KB
Newer Older
1
2
3
4
5
6
#!/usr/bin/bash

# RUN DEV INSTANCE
# Spin up an Airflow instance for development purposes.
#
# This script is intended to help Airflow users develop DAGs.
Mforns's avatar
Mforns committed
7
# It should be run in a stat machine (i.e. stat1008.eqiad.wmnet),
8
9
10
# or any machine within the analytics VLAN that has the following:
#   - Access to the Hadoop cluster via Kerberos.
#   - The Airflow base Conda environment in /usr/lib/airflow (or a custom one).
11
#   - Spark3 configuration files.
Mforns's avatar
Mforns committed
12
#   - It's not a production machine.
13
14
15
16
17
18
19
20
#
# The script creates a Conda environment with Airflow in it. Then,
# initializes and configures Airflow, and finally launches its services.
# The user should pass an AIRFLOW_PROJECT argument, which indicates the project
# directory within the airflow-dags repository that is going to be executed.
# Finally, the user can access the Airflow UI via ssh tunnel.
# The user can also conda install necessary packages in the environment.
#
21
22
23
24
# For testing jobs that use Skein or jobs with Spark deploy-mode=cluster
# you need to run this script with a user that has a kerberos keytab
# (usually analytics-privatedata). The problem is that those users normally
# don't have a $HOME set, and Conda needs it to operate properly. Thus,
Mforns's avatar
Mforns committed
25
26
# whenever you run this script with i.e. `sudo -u analytics-privatedata`,
# you need to specify the parameter -m USER_HOME.
27
#
28
29
# Examples:
#   ./run_dev_instance.sh -h
Mforns's avatar
Mforns committed
30
31
32
33
34
#   ./run_dev_instance.sh analytics
#   ./run_dev_instance.sh -a ~/my_airflow_home analytics
#   ./run_dev_instance.sh -p 8081 analytics
#   ./run_dev_instance.sh -b ~/my_base_env_path analytics
#   ./run_dev_instance.sh -n my_dev_env_name analytics
Mforns's avatar
Mforns committed
35
#   sudo -u analytics-privatedata ./run_dev_instance.sh -m /tmp/my_home analytics
36
37


Mforns's avatar
Mforns committed
38
# Helper: Print usage text.
39
40
41
42
43
44
45
46
47
48
print_usage () {
    echo
    echo "Usage: ${0} [ OPTIONS ] AIRFLOW_PROJECT"
    echo
    echo "ARGUMENTS"
    echo "    AIRFLOW_PROJECT   Name of the project directory in airflow-dags"
    echo "                      repository to execute Airflow for."
    echo
    echo "OPTIONS"
    echo "    -h                Show this message and exit."
49
50
51
    echo "    -m USER_HOME      Absolute path to use as a user home folder."
    echo "                      Mandatory if user has no \$HOME defined."
    echo "                      Default: ${HOME}"
52
    echo "    -a AIRFLOW_HOME   Absolute path where to install Airflow."
53
    echo "                      Default: <USER_HOME>/airflow"
54
55
56
57
58
59
    echo "    -p SERVER_PORT    Port for the Airflow webserver."
    echo "                      Default: 8080"
    echo "    -b BASE_ENV_PATH  Absolute path to base Conda env to clone."
    echo "                      Default: /usr/lib/airflow"
    echo "    -n DEV_ENV_NAME   Name of the development Conda environment."
    echo "                      Default: airflow_development"
60
    echo "    -i                Install wmf_airflow_common into conda env"
61
62
}

Mforns's avatar
Mforns committed
63
64
65
66
# Helper: Clean up Airflow services and conda environment.
cleanup () {
    echo
    echo "Shutting down Airflow instance..."
67
    if [ -n "${webserver_pid}" ]; then
Mforns's avatar
Mforns committed
68
69
70
        # Kill the webserver and all its child processes.
        # Not using -9 on purpose (just send regular SIGTERM), since -9
        # doesn't let Airflow propagate the signal to the child processes.
71
        kill "${webserver_pid}"
Mforns's avatar
Mforns committed
72
    fi
73
    if [ -n "${scheduler_pid}" ]; then
Mforns's avatar
Mforns committed
74
75
76
        # Kill the scheduler and all its child processes.
        # Not using -9 on purpose (just send regular SIGTERM), since -9
        # doesn't let Airflow propagate the signal to the child processes.
77
        kill "${scheduler_pid}"
Mforns's avatar
Mforns committed
78
    fi
79
    # Deactivate the conda environment.
Mforns's avatar
Mforns committed
80
    conda deactivate > "/dev/null" 2>&1
81
82
83
84
    # Set $HOME to its previous value, if any.
    if [ -n "${previous_home}" ]; then
        export HOME="${previous_home}"
    fi
Mforns's avatar
Mforns committed
85
86
87
88
89
90
    echo "Airflow instance shut down successfully!"
    echo
    exit 2
}

# Helper: Print error message and exit.
91
92
93
94
error () {
    echo
    echo "Error: ${1}"
    echo
Mforns's avatar
Mforns committed
95
    cleanup
96
97
}

Mforns's avatar
Mforns committed
98
# Helper: Check if the last command failed, and error accordingly.
99
100
101
102
103
104
105
106
107
108
109
assert () {
    if [ ${?} != 0 ]; then
        error "${@}"
    fi
}

# Get this script's base path.
# https://stackoverflow.com/questions/4774054/reliable-way-for-a-bash-script-to-get-the-full-path-to-itself
script_dir="$(cd -- "$(dirname "${0}")" > "/dev/null" 2>&1 ; pwd -P)"

# Parse optional arguments.
110
while getopts "hm:a:p:b:n:i" options; do
111
112
113
114
115
    case "${options}" in
        h)
            print_usage
            exit 0
            ;;
116
117
118
119
120
121
        m)
            user_home="${OPTARG}"
            if [[ "${user_home}" != /* ]]; then
                error "USER_HOME ${user_home} is not an absolute path."
            fi
            ;;
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
        a)
            airflow_home="${OPTARG}"
            if [[ "${airflow_home}" != /* ]]; then
                error "AIRFLOW_HOME ${airflow_home} is not an absolute path."
            fi
            ;;
        p)
            server_port="${OPTARG}"
            if [[ ! "${server_port}" =~ ^[0-9]+$ ]] ; then
                error "SERVER_PORT ${server_port} is invalid."
            fi
            ;;
        b)
            base_env_path="${OPTARG}"
            if [[ "${base_env_path}" != /* ]]; then
                error "BASE_ENV_PATH ${base_env_path} is not an absolute path."
            fi
            ;;
        n)
            dev_env_name="${OPTARG}"
            if [[ ! "${dev_env_name}" =~ ^[a-zA-Z][a-zA-Z0-9_\-]+$ ]] ; then
                error "DEV_ENV_NAME ${dev_env_name} is invalid."
            fi
            ;;
146
147
148
        i)
            install_wmf_airflow_common=true
            ;;
149
150
151
152
153
154
155
        *)
            print_usage
            error "Invalid arguments."
            ;;
    esac
done

156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
# Make sure $HOME is set/defaulted correctly.
if [ -n "${user_home}" ]; then
    previous_home="${HOME}"
    export HOME="${user_home}"
fi
if [ "${HOME}" == "/nonexistent" ]; then
    print_usage
    error "USER_HOME should be specified when the user has no home folder."
fi

# Set defaults for other optional arguments.
if [ -z "${airflow_home}" ]; then
    airflow_home="${HOME}/airflow"
fi
if [ -z "${server_port}" ]; then
    server_port="8080"
fi
if [ -z "${base_env_path}" ]; then
    base_env_path="/usr/lib/airflow"
fi
if [ -z "${dev_env_name}" ]; then
    dev_env_name="airflow_development"
fi
179
180
181
if [ -z "${install_wmf_airflow_common}" ]; then
    install_wmf_airflow_common=false
fi
182

Mforns's avatar
Mforns committed
183
184
# Parse mandatory arguments.
shift $((OPTIND - 1)) # Shift away optional arguments.
185
186
187
188
if [ "${#}" != 1 ]; then
    print_usage
    error "Missing or too many mandatory arguments."
fi
Mforns's avatar
Mforns committed
189
airflow_project="${1}"
190
191
192
193
194
if ! test -d "${script_dir}/${airflow_project}"; then
    error "AIRFLOW_PROJECT ${airflow_project} not in airflow-dags repository."
fi

# Get Kerberos credentials.
Mforns's avatar
Mforns committed
195
196
197
keytab="/etc/security/keytabs/${USER}/${USER}.keytab"
if test -f "${keytab}"; then
    echo "Using keytab = ${keytab}"
198
    klist_prefix="kerberos-run-command ${USER}"
Mforns's avatar
Mforns committed
199
200
201
202
else
    echo "Keytab not found at ${keytab}. Omitting."
    unset keytab
fi
203
204
205
206
207
208
209
210
if ${klist_prefix} klist > "/dev/null" 2>&1; then
    principal=$(klist | head -n2 | tail -n1 | cut -d' ' -f3)
    cache_path=$(klist | head -n1 | cut -d' ' -f3 | cut -d':' -f2)
    echo "Using principal = ${principal}"
    echo "Using ccache = ${cache_path}"
else
    error "No Kerberos credentials found, please kinit or use kerberos-run-command."
fi
211
212
213
214
215
216
217
218
219
220
221
222

# SPIN UP AIRFLOW INSTANCE:
echo
echo "Spinning up Airflow instance:"

# Create airflow home if necessary.
echo "    Preparing AIRFLOW_HOME ${airflow_home}..."
if ! test -d "${airflow_home}"; then
    mkdir -p "${airflow_home}"
    assert "Can not create AIRFLOW_HOME ${airflow_home} directory."
fi
export AIRFLOW_HOME="${airflow_home}"
223
# We need script_dir on PYTHONPATH to import wmf_airflow_common, the airflow_project's config, etc.
224
export PYTHONPATH="${script_dir}"
225
226
227
cd "${airflow_home}"
assert "Can not cd into AIRFLOW_HOME ${airflow_home} directory."

Mforns's avatar
Mforns committed
228
229
230
# Create conda home and condarc file if necessary.
conda_home="${HOME}/.conda"
condarc_path="${conda_home}/condarc"
231
echo "    Preparing condarc file ${condarc_path}..."
Mforns's avatar
Mforns committed
232
233
234
if ! test -d "${conda_home}"; then
    mkdir -p "${conda_home}" > "/dev/null" 2>&1
    assert "Can not create ${conda_home} directory."
235
fi
Mforns's avatar
Mforns committed
236
237
echo -e "pkgs_dirs:\n  - ${conda_home}/pkgs\n" > "${condarc_path}"
assert "Can not create condarc file ${condarc_path}."
238
239
240
241
242
243
244

# Create and activate Conda environment.
echo "    Creating and activating Conda environment..."
echo "        Conda logs: ${airflow_home}/conda.log"
conda_execs_script="/usr/lib/airflow/etc/profile.d/conda.sh"
source "$conda_execs_script"
assert "Can not add Conda executables to PATH ($conda_execs_script failed)."
245
webproxy="http://webproxy.eqiad.wmnet:8080"
246
247
248
249
# Set http(s)_proxy to allow for package downloads.
export http_proxy="${webproxy}"
export https_proxy="${webproxy}"

Mforns's avatar
Mforns committed
250
if ! test -d "${conda_home}/envs/${dev_env_name}"; then
251
252
253
    conda create --clone "${base_env_path}" --name "${dev_env_name}" > "${airflow_home}/conda.log" 2>&1
    assert "Can not create Conda environment ${dev_env_name} using ${base_env_path}."
fi
254

255
256
conda activate "${dev_env_name}" >> "${airflow_home}/conda.log" 2>&1
assert "Can not activate Conda environment ${dev_env_name}."
257
258
259
260
261
262
263
264
265
266
267
268

if [ "${install_wmf_airflow_common}" = true ] ; then
    echo "    Installing wmf_airflow_dags into conda ${conda_home}/envs/${dev_env_name}..."
    cd "${script_dir}"
    pip install .
    assert "Can not install ${script_dir} in conda environment ${dev_env_name}."
    cd "${airflow_home}"
fi

unset http_proxy
unset https_proxy

269
270
# Set the following env vars to pick the "wmf" defaults for the dag config.
export AIRFLOW_ENVIRONMENT_NAME="dev_wmf"
Mforns's avatar
Mforns committed
271
export AIRFLOW_INSTANCE_NAME="${dev_env_name}_${airflow_project}_$(whoami)"
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286

# Initialize Airflow.
echo "    Initializing Airflow..."
echo "        DB init logs: ${airflow_home}/dbinit.log"
airflow db init > "${airflow_home}/dbinit.log" 2>&1
assert "Can not initialize Airflow DB."

# Configure Airflow.
echo "    Configuring Airflow..."
sed -i "s|dags_folder =.*|dags_folder = ${script_dir}/${airflow_project}/dags|g" "airflow.cfg"
sed -i "s|plugins_folder =.*|plugins_folder = ${script_dir}/${airflow_project}/plugins|g" "airflow.cfg"
sed -i "s|load_examples =.*|load_examples = False|g" "airflow.cfg"
sed -i "s|security =.*|security = kerberos|g" "airflow.cfg"
sed -i "s|principal =.*|principal = ${principal}|g" "airflow.cfg"
sed -i "s|ccache =.*|ccache = ${cache_path}|g" "airflow.cfg"
287
288
289
# Set or unset the keytab in config, depending on whether it is defined.
if [ -z ${keytab} ]; then
    sed -i "s|keytab =.*|keytab =|g" "airflow.cfg"
Mforns's avatar
Mforns committed
290
else
291
    sed -i "s|keytab =.*|keytab = ${keytab}|g" "airflow.cfg"
Mforns's avatar
Mforns committed
292
fi
293
sed -i "s|# AUTH_ROLE_PUBLIC = 'Public'|AUTH_ROLE_PUBLIC = 'Admin'|g" "webserver_config.py"
294
295
296
sed -i "s|backend_kwargs =.*|backend_kwargs = {\"connections_file_path\": \"${script_dir}/dev_instance/connections.yaml\"}|g" "airflow.cfg"
sed -i "0,/^backend =/ s/^backend =.*/backend = airflow.secrets.local_filesystem.LocalFilesystemBackend/" "airflow.cfg"
sed -i "s|expose_config =.*|expose_config = True|g" "airflow.cfg"
297
298
299
300
301

# Launching Airlfow webserver and scheduler
echo "    Launching Airflow webserver..."
airflow webserver -p "${server_port}" > "${airflow_home}/webserver.log" 2>&1 &
assert "Can not launch Airflow webserver."
Mforns's avatar
Mforns committed
302
webserver_pid="$!"
303
304
305
echo "    Launching Airflow scheduler..."
airflow scheduler > "${airflow_home}/scheduler.log" 2>&1 &
assert "Can not launch Airflow scheduler."
Mforns's avatar
Mforns committed
306
scheduler_pid="$!"
307
308
309
310
311
312
313
314
315
316
317
318
319
echo "Airflow instance spun up successfully!"

# Print help.
echo
echo "    Access Airflow logs at:"
echo "        Webserver logs: ${airflow_home}/webserver.log"
echo "        Scheduler logs: ${airflow_home}/scheduler.log"
echo
echo "    Access Airflow UI by setting up an ssh tunnel with:"
echo "        ssh ${HOSTNAME}.eqiad.wmnet -L ${server_port}:${HOSTNAME}.eqiad.wmnet:${server_port}"
echo "        and visiting http://localhost:${server_port}/home on your browser."
echo
echo "    Install Python packages in the Conda environment with:"
320
321
echo "        export http_proxy=\"${webproxy}\""
echo "        export https_proxy=\"${webproxy}\""
322
323
324
325
326
327
328
echo "        source \"${conda_execs_script}\""
echo "        conda activate \"${dev_env_name}\""
echo "        conda install <YOUR PACKAGE>"
echo "        conda deactivate"
echo
echo "Hit Ctrl+C to shut down Airflow instance."

329
330
# EXECUTE CLEANUP WHENEVER THE SCRIPT ERRORS OR EXITS.
trap "cleanup" ERR EXIT
331
332
333

# Wait so that the user can Ctrl+C to shut down the instance.
sleep infinity