Some neat bash tricks

job_submission.bash

#! /bin/bash
set -euxo pipefail

job_file="cluster_jobs/${1}.py"
if [[ ! -f "${job_file}" ]]; then
	echo "${job_file} not found"
	exit 1
fi

RAY_ADDRESS=${RAY_ADDRESS:-""}
PKG="distributed-ml"
command="pip uninstall ${PKG} && pip install . && python ${job_file}"
if [[ -z "${RAY_ADDRESS}" ]]; then
	echo "RAY_ADDRESS is not defined, assuming local cluster"
	RAY_ADDRESS="http://127.0.0.1:8265"
	echo "Refreshing ${PKG} install"
	if [[ -z $(pip list | grep -F "${PKG}") ]]; then
		echo "Package ${PKG} not installed, installing"
		pip install -e .
	else
		echo "Package ${PKG} already installed"
	fi
	# For local development, keep the package installed
	command="python ${job_file}"
fi

# Workaround Ray not syncing files in .gitignore
mv .gitignore .gitignore.bak
function finish {
	mv .gitignore.bak .gitignore
}
trap finish EXIT

job_id=$(python-petname)
env_values=$(python3 -c 'from dotenv import dotenv_values; import json; print(json.dumps(dotenv_values(".job.env")))')

# TODO: I'm not sure why this check is here
if [ "$(jq 'has("JOB_ID")' <<<$env_values)" == "true" ]; then
	echo "JOB_ID already defined in .job.env"
	exit 1
fi

env_values=$(echo "${env_values}" | jq '.JOB_ID = "'${job_id}'"')
output=$(ray job submit \
	--working-dir . \
	--no-wait \
	--submission-id ${job_id} \
	--runtime-env-json="{\"env_vars\": ${env_values}}" \
	-- bash -c "${command}")

# output_no_escape_codes=$(echo "${output}" | sed 's/\x1B\[[0-9;]\{1,\}[A-Za-z]//g')
# job_id=$(echo "${output_no_escape_codes}" | grep 'ray job logs' | awk '{print $4}')
# echo "${job_id}"

mkdir -p temp_job_scripts
stop_script="temp_job_scripts/${job_id}_kill.bash"
logs_script="temp_job_scripts/${job_id}_logs.bash"
tail_logs_script="temp_job_scripts/${job_id}_tail.bash"
status_script="temp_job_scripts/${job_id}_status.bash"

prefix="#! /usr/bin/env bash\n env RAY_ADDRESS=${RAY_ADDRESS} ray job"
echo -e "${prefix} logs ${job_id}" >>"${logs_script}"
echo -e "${prefix} stop ${job_id}" >>"${stop_script}"
echo -e "${prefix} status ${job_id}" >>"${status_script}"
echo -e "#! /usr/bin/env bash\n env RAY_ADDRESS=${RAY_ADDRESS} SUBMISSION_ID=${job_id} ./dev_scripts/tail_job_logs.py" >>"${tail_logs_script}"

chmod +x "${logs_script}"
chmod +x "${stop_script}"
chmod +x "${status_script}"
chmod +x "${tail_logs_script}"

echo "job_id: ${job_id}"

# TODO: Might be broken right now
TAIL=${2:-""}
if [[ "${TAIL}" == "tail" ]]; then
	"${tail_logs_script}"
fi