<commit_before>#!/bin/bash

echo "~~~ :mag: Testing connection to $SUPERCLUSTER_LOGIN_HOST"

# Do the first check with BatchMode=yes (which won't prompt for password access)
SUPERCLUSTER_NAME=$(ssh -oBatchMode=yes "$SUPERCLUSTER_LOGIN_HOST" "echo \"\$(whoami)@\$(hostname)\"")
SSH_TEST_EXIT_STATUS=$?

# If the SSH connection test failed, we should just bail now
if [[ $SSH_TEST_EXIT_STATUS -ne 0 ]]; then
  echo "SSH connection test exited with $SSH_TEST_EXIT_STATUS"
  echo $SUPERCLUSTER_HOSTNAME
  exit $SSH_TEST_EXIT_STATUS
fi

echo "Connection to $SUPERCLUSTER_LOGIN_HOST is all good! ✅"

# Since we did that test, we can use the information we got back to construct a
# nice display name for prompts
SUPERCLUSTER_PROMPT="\033[90m${SUPERCLUSTER_NAME}:$\033[0m"

COMMAND="hostname && sleep 20"

# CHECKOUT_FOLDER=".buildkite/$BUILDKITE_ORGANIZATION_SLUG/$BUILDKITE_PIPELINE_SLUG"

SUPERCLUSTER_CHECKOUT_FOLDER=".buildkite/${BUILDKITE_JOB_ID}"

# The job name is suffixed with the current PID (so if a job is retried with te
# same $BUILDKITE_JOB_ID, when we check job statuses it won't return the
# previous jobs run status)
SUPERCLUSTER_JOB_NAME="buildkite-job-${BUILDKITE_JOB_ID}-$$"

SUPERCLUSTER_BOOTSTRAP_SCRIPT_NAME="buildkite-bootstrap-${BUILDKITE_JOB_ID}.sh"
SUPERCLUSTER_RUNNER_SCRIPT_NAME="buildkite-runner-${BUILDKITE_JOB_ID}.sh"

# The command script will run the desired command from Buildkite within the compute node
SUPERCLUSTER_COMMAND_SCRIPT_NAME="buildkite-exec-${BUILDKITE_JOB_ID}.sh"
SUPERCLUSTER_COMMAND_LOG_NAME="buildkite-output-${BUILDKITE_JOB_ID}.log"
SUPERCLUSTER_EXIT_STATUS_FILE="buildkite-exit-status-${BUILDKITE_JOB_ID}"

read -r -d '' SUPERCLUSTER_COMMAND_SCRIPT << EOM
#!/bin/bash
set -e
if [[ -f "${COMMAND}" ]]; then
  chmod +x "${COMMAND}"
  ./"${COMMAND}"
else
  ${COMMAND}
fi
EOM

# The runner script is what is submitted to the cluster. It will run the
# command script, stream the output, and store the exit status.
read -r -d '' SUPERCLUSTER_RUNNER_SCRIPT << EOM
#!/bin/bash
./${SUPERCLUSTER_COMMAND_SCRIPT_NAME} >"${SUPERCLUSTER_COMMAND_LOG_NAME}" 2>&1
echo \$? > "${SUPERCLUSTER_EXIT_STATUS_FILE}"
EOM

# The bootstrap script is what is run on the remote login node, that prepares
# it for running a buildkite job
cat > $SUPERCLUSTER_BOOTSTRAP_SCRIPT_NAME <<- EOM
echo 'Connection established to "${SUPERCLUSTER_LOGIN_HOST}"'

function run {
  echo -e "${SUPERCLUSTER_PROMPT} \$1"
  eval "\$1"
  EVAL_EXIT_STATUS=\$?

  if [[ \$EVAL_EXIT_STATUS -ne 0 ]]; then
    exit \$EVAL_EXIT_STATUS
  fi
}

echo '~~~ :package: Preparing repository on the login node'

run "rm -rf \"${SUPERCLUSTER_CHECKOUT_FOLDER}\""
run "mkdir -p \"${SUPERCLUSTER_CHECKOUT_FOLDER}\""

export GIT_TERMINAL_PROMPT=0

run "cd \"${SUPERCLUSTER_CHECKOUT_FOLDER}\""
run "git clone -v -- \"${BUILDKITE_REPO}\" ."

echo '~~~ :page_facing_up: Creating runner script'

cat > "$SUPERCLUSTER_RUNNER_SCRIPT_NAME" <<- BEOM
${SUPERCLUSTER_RUNNER_SCRIPT}
BEOM

run "cat \"${SUPERCLUSTER_RUNNER_SCRIPT_NAME}\""
run "chmod +x \"$SUPERCLUSTER_RUNNER_SCRIPT_NAME\""

echo '~~~ :page_facing_up: Creating command script'

cat > "$SUPERCLUSTER_COMMAND_SCRIPT_NAME" <<- BEOM
${SUPERCLUSTER_COMMAND_SCRIPT}
BEOM

run "cat \"${SUPERCLUSTER_COMMAND_SCRIPT_NAME}\""
run "chmod +x \"$SUPERCLUSTER_COMMAND_SCRIPT_NAME\""

echo '~~~ :desktop_computer: Submitting the job to the cluster'
run "sbatch --workdir=\"\$(pwd)\" --job-name=\"${SUPERCLUSTER_JOB_NAME}\" \"${SUPERCLUSTER_RUNNER_SCRIPT_NAME}\""
EOM

echo '~~~ Bootstrapping job on the supercluster'

ssh "$SUPERCLUSTER_LOGIN_HOST" "bash -s" < "$SUPERCLUSTER_BOOTSTRAP_SCRIPT_NAME"
SSH_BOOTSTRAP_EXIT_STATUS=$?

# If the bootstrapping process failed, bad things must have happened...
if [[ $SSH_BOOTSTRAP_EXIT_STATUS -ne 0 ]]; then
  exit $SSH_BOOTSTRAP_EXIT_STATUS
fi

echo '~~~ :hourglass: Waiting for supercluster job to start'
<commit_msg>More tweaks to the headings<commit_after>#!/bin/bash

echo "~~~ :mag: Testing connection to $SUPERCLUSTER_LOGIN_HOST"

# Do the first check with BatchMode=yes (which won't prompt for password access)
SUPERCLUSTER_NAME=$(ssh -oBatchMode=yes "$SUPERCLUSTER_LOGIN_HOST" "echo \"\$(whoami)@\$(hostname)\"")
SSH_TEST_EXIT_STATUS=$?

# If the SSH connection test failed, we should just bail now
if [[ $SSH_TEST_EXIT_STATUS -ne 0 ]]; then
  echo "SSH connection test exited with $SSH_TEST_EXIT_STATUS"
  echo $SUPERCLUSTER_HOSTNAME
  exit $SSH_TEST_EXIT_STATUS
fi

echo "Connection to $SUPERCLUSTER_LOGIN_HOST is all good! ✅"

# Since we did that test, we can use the information we got back to construct a
# nice display name for prompts
SUPERCLUSTER_PROMPT="\033[90m${SUPERCLUSTER_NAME}:$\033[0m"

COMMAND="hostname && sleep 20"

# CHECKOUT_FOLDER=".buildkite/$BUILDKITE_ORGANIZATION_SLUG/$BUILDKITE_PIPELINE_SLUG"

SUPERCLUSTER_CHECKOUT_FOLDER=".buildkite/${BUILDKITE_JOB_ID}"

# The job name is suffixed with the current PID (so if a job is retried with te
# same $BUILDKITE_JOB_ID, when we check job statuses it won't return the
# previous jobs run status)
SUPERCLUSTER_JOB_NAME="buildkite-job-${BUILDKITE_JOB_ID}-$$"

SUPERCLUSTER_BOOTSTRAP_SCRIPT_NAME="buildkite-bootstrap-${BUILDKITE_JOB_ID}.sh"
SUPERCLUSTER_RUNNER_SCRIPT_NAME="buildkite-runner-${BUILDKITE_JOB_ID}.sh"

# The command script will run the desired command from Buildkite within the compute node
SUPERCLUSTER_COMMAND_SCRIPT_NAME="buildkite-exec-${BUILDKITE_JOB_ID}.sh"
SUPERCLUSTER_COMMAND_LOG_NAME="buildkite-output-${BUILDKITE_JOB_ID}.log"
SUPERCLUSTER_EXIT_STATUS_FILE="buildkite-exit-status-${BUILDKITE_JOB_ID}"

read -r -d '' SUPERCLUSTER_COMMAND_SCRIPT << EOM
#!/bin/bash
set -e
if [[ -f "${COMMAND}" ]]; then
  chmod +x "${COMMAND}"
  ./"${COMMAND}"
else
  ${COMMAND}
fi
EOM

# The runner script is what is submitted to the cluster. It will run the
# command script, stream the output, and store the exit status.
read -r -d '' SUPERCLUSTER_RUNNER_SCRIPT << EOM
#!/bin/bash
./${SUPERCLUSTER_COMMAND_SCRIPT_NAME} >"${SUPERCLUSTER_COMMAND_LOG_NAME}" 2>&1
echo \$? > "${SUPERCLUSTER_EXIT_STATUS_FILE}"
EOM

# The bootstrap script is what is run on the remote login node, that prepares
# it for running a buildkite job
cat > $SUPERCLUSTER_BOOTSTRAP_SCRIPT_NAME <<- EOM
echo 'Connection established to "${SUPERCLUSTER_LOGIN_HOST}"'

function run {
  echo -e "${SUPERCLUSTER_PROMPT} \$1"
  eval "\$1"
  EVAL_EXIT_STATUS=\$?

  if [[ \$EVAL_EXIT_STATUS -ne 0 ]]; then
    exit \$EVAL_EXIT_STATUS
  fi
}

echo '~~~ :package: Preparing repository on the login node'

run "rm -rf \"${SUPERCLUSTER_CHECKOUT_FOLDER}\""
run "mkdir -p \"${SUPERCLUSTER_CHECKOUT_FOLDER}\""

export GIT_TERMINAL_PROMPT=0

run "cd \"${SUPERCLUSTER_CHECKOUT_FOLDER}\""
run "git clone -v -- \"${BUILDKITE_REPO}\" ."

cat > "$SUPERCLUSTER_RUNNER_SCRIPT_NAME" <<- BEOM
${SUPERCLUSTER_RUNNER_SCRIPT}
BEOM

chmod +x "$SUPERCLUSTER_RUNNER_SCRIPT_NAME"

cat > "$SUPERCLUSTER_COMMAND_SCRIPT_NAME" <<- BEOM
${SUPERCLUSTER_COMMAND_SCRIPT}
BEOM

chmod +x "$SUPERCLUSTER_COMMAND_SCRIPT_NAME"

echo '~~~ :floppy_disk: Submitting the job to the cluster'
run "sbatch --workdir=\"\$(pwd)\" --job-name=\"${SUPERCLUSTER_JOB_NAME}\" \"${SUPERCLUSTER_RUNNER_SCRIPT_NAME}\""
EOM

echo '~~~ :wrench: Connecting to the login node'

ssh "$SUPERCLUSTER_LOGIN_HOST" "bash -s" < "$SUPERCLUSTER_BOOTSTRAP_SCRIPT_NAME"
SSH_BOOTSTRAP_EXIT_STATUS=$?

# If the bootstrapping process failed, bad things must have happened...
if [[ $SSH_BOOTSTRAP_EXIT_STATUS -ne 0 ]]; then
  exit $SSH_BOOTSTRAP_EXIT_STATUS
fi

echo '~~~ :hourglass: Waiting for super cluster job to start'

{'fraction_non_alphanumeric': 0.12218989685268447, 'fraction_numerical': 0.005025125628140704, 'mean_word_length': 5.919487648673376, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1415800', 'n_tokens_mistral': 3004, 'n_tokens_neox': 2921, 'n_words': 694}