conda-spark-env.sh

# If /etc/hadoop/conf exists, use it as HADOOP_CONF_DIR
if [ -z "${HADOOP_CONF_DIR}" -a -e "/etc/hadoop/conf" ]; then
  export HADOOP_CONF_DIR=/etc/hadoop/conf
fi

# If hadoop command is execultable, then use it to add Hadoop jars to Spark's runtime classpath.
# See: https://spark.apache.org/docs/2.4.4/hadoop-provided.html
if [ -x "$(command -v hadoop)" ]; then
    SPARK_DIST_CLASSPATH=$(hadoop classpath)
fi

# If /usr/lib/hadoop/native exists, use Hadoop native libs from there
if [ -z "${LD_LIBRARY_PATH}" -a -e /usr/lib/hadoop/lib/native ]; then
    export LD_LIBRARY_PATH=/usr/lib/hadoop/lib/native
fi

# Set default python to python3
if [ -z "${PYSPARK_PYTHON}" -a -n "$(command -v python3)" ]; then
    # We want PYSPARK_PYTHON to default to the versioned python
    # executable on the node where spark is being launched.
    # E.g. on Stretch we want python3.5 and on Buster we want python3.7.
    # This will cause the versions of python that is used on driver
    # and on workers to be the same.
    # https://phabricator.wikimedia.org/T229347#5439259
    export PYSPARK_PYTHON="$(realpath $(command -v python3))"
fi


if [ -z "${PYSPARK_DRIVER_PYTHON}" ]; then
    # If we are using the pyspark shell, assume we want to use ipython if it is available.
    : ${PYSPARK_DRIVER_PREFER_IPYTHON='true'}
    if [[ "${0}" == *pyspark* && "${PYSPARK_DRIVER_PREFER_IPYTHON}" == 'true' && -n "$(command -v ipython3)" ]]; then
        export PYSPARK_DRIVER_PYTHON="$(realpath $(command -v ipython3))"
    else
        export PYSPARK_DRIVER_PYTHON="${PYSPARK_PYTHON}"
    fi
fi

# conda customization.  If PYSPARK_PYTHON happens to be in a user's conda environment, it
# will not exist on the remote yarn executors. We need to either: set it to a python we
# know exists on the remote yarn executors, OR conditionally pack the user's conda environment
# and ship it to HDFS so that executors can use it too.

# Search the CLI opts to find the master option, if it is given.
# This is used to automate the modification of PYSPARK_PYTHON if
# a conda environment is active.  It needs to be set to something
# that will work on remote executors.
spark_master=''
for ((i = 1; i <= $#; i++ )); do
    arg="${!i}"
    if [[ "${arg}" == --master ]]; then
        master_index=$((i+1))
        spark_master="${!master_index}"
    elif [[ "${arg}" == --master=* ]]; then
        spark_master="$(echo ${arg} | cut -f2 -d=)"
    fi
done






# conda_base_env_path is expected to exist on spark drivers and executors.
# If we are running in yarn, and conda is active but we aren't using the conda_base_env_path
# for PYSPARK_PYTHON, then PYSPARK_PYTHON needs adjusted to work on remote executors.
# Either it needs to be set explicitly to conda_base_env_path's python, OR we should
# pack and ship the active conda environment to the exectuors and use it.
: ${conda_base_env_path:=/usr/lib/anaconda-wmf}
if [[ "${spark_master}" == yarn* && -n "${CONDA_PREFIX}" && "${CONDA_PREFIX}" != "${conda_base_env_path}" ]]; then

    if [ "${SPARK_SHIP_CONDA_ENV}" == 'true' ]; then
        # Conda pack the active conda environment.
        conda_env_name="$(basename ${CONDA_PREFIX})"
        conda_env_tgz="${HOME}/conda_${conda_env_name}.tar.gz"

        if [ ! -f "${conda_env_tgz}" ]; then
            echo "Packing conda env ${conda_env_name} at ${conda_env_tgz}."
            conda pack --force  --ignore-editable-packages -o ${conda_env_tgz}
        else
            echo "Your active conda environment ${conda_env_name} is packed at ${conda_env_tgz}."
            echo "If you have recently installed new packages into your conda env, delete "
            echo "${conda_env_tgz} and it will be re-packed for you."
        fi

        # TODO: PYSPARK_SUBMIT_ARGS doesn't seem to work via pyspark cli.
        export PYSPARK_SUBMIT_ARGS="--archives ${conda_env_tgz}#conda_${conda_env_name}"
        export PYSPARK_PYTHON="conda_${conda_env_name}/bin/$(basename $(realpath ${CONDA_PREFIX}/bin/python3))"

    else
        # Else use conda_base_env_path on the remote executors
        export PYSPARK_PYTHON="$(realpath ${conda_base_env_path}/bin/python3)"

    # elif [ "${CONDA_PREFIX}" != '/usr/lib/anaconda']
    #     echo "WARNING: PYSPARK_PYTHON=${PYSPARK_PYTHON} which may not be available on remote Spark "
    #     echo "executors (e.g. in YARN). It looks like you are using a non-stacked conda environment."
    fi

# Else There is no conda environment active, so assume we are using system python.  If so,
# and if SPARK_HOME/pythonX.X exists, insert it into the front of PYTHONPATH
# So any provided packages override system installed ones.
# TODO: stop doing this after https://phabricator.wikimedia.org/T275786
else
    python_version_path=${SPARK_HOME}/$(${PYSPARK_PYTHON} -c 'import sys; print("python{}.{}".format(sys.version_info.major, sys.version_info.minor))')
    if [ -d "${python_version_path}" ]; then
        echo "EXPORTING PYTHONPATH=${python_version_path}:${PYTHONPATH}"
        export PYTHONPATH="${python_version_path}:${PYTHONPATH}"
    fi
fi

test -n "${PYTHONPATH}" && echo "PYTHONPATH=${PYTHONPATH}"
echo "PYSPARK_DRIVER_PYTHON=$PYSPARK_DRIVER_PYTHON"
echo "PYSPARK_PYTHON=$PYSPARK_PYTHON"
test -n "${PYSPARK_SUBMIT_ARGS}" && echo "PYSPARK_SUBMIT_ARGS=$PYSPARK_SUBMIT_ARGS"