From e903e9019c3b7d9e5762d014440a287f966e5f92 Mon Sep 17 00:00:00 2001
From: AJ Schmidt <aschmidt@nvidia.com>
Date: Wed, 10 Jan 2024 16:16:01 -0500
Subject: [PATCH] Update initialization script

This PR removes event restrictions from non-NVIDIA orgs since they run exclusively in the cloud.

These changes have been tested with the following commands:

```sh
GITHUB_REPOSITORY_OWNER=rapidsai GITHUB_EVENT_NAME=pull_request ./scripts/helpers/initialize_runner.sh

GITHUB_REPOSITORY_OWNER=rapidsai GITHUB_EVENT_NAME=pull_request_target ./scripts/helpers/initialize_runner.sh

GITHUB_REPOSITORY_OWNER=rapidsai GITHUB_EVENT_NAME=push ./scripts/helpers/initialize_runner.sh

GITHUB_REPOSITORY_OWNER=NERSC GITHUB_EVENT_NAME=pull_request ./scripts/helpers/initialize_runner.sh

GITHUB_REPOSITORY_OWNER=NERSC GITHUB_EVENT_NAME=push ./scripts/helpers/initialize_runner.sh
```
---
 scripts/helpers/check_event_type.sh  | 11 ----------
 scripts/helpers/initialize_runner.sh | 33 ++++++++++++++++++++++++++++
 scripts/helpers/runner.sh            |  2 +-
 scripts/installers/runner.sh         |  2 +-
 4 files changed, 35 insertions(+), 13 deletions(-)
 delete mode 100755 scripts/helpers/check_event_type.sh
 create mode 100755 scripts/helpers/initialize_runner.sh

diff --git a/scripts/helpers/check_event_type.sh b/scripts/helpers/check_event_type.sh
deleted file mode 100755
index 7fe450b..0000000
--- a/scripts/helpers/check_event_type.sh
+++ /dev/null
@@ -1,11 +0,0 @@
-#!/bin/bash
-
-set -eu
-
-if [ "${GITHUB_EVENT_NAME}" = "pull_request" ] || [ "${GITHUB_EVENT_NAME}" = "pull_request_target" ]; then
-  echo "Workflows triggered by \"${GITHUB_EVENT_NAME}\" events are not allowed to run on self-hosted runners."
-  exit 1
-fi
-
-echo "Event \"${GITHUB_EVENT_NAME}\" allowed to run on self-hosted runners."
-exit 0
diff --git a/scripts/helpers/initialize_runner.sh b/scripts/helpers/initialize_runner.sh
new file mode 100755
index 0000000..26e2556
--- /dev/null
+++ b/scripts/helpers/initialize_runner.sh
@@ -0,0 +1,33 @@
+#!/bin/bash
+
+set -eu
+
+# these strings are case-sensitive and should match the case of the GH org
+CLOUD_ONLY_ORGS=(
+  "dask-contrib"
+  "dask"
+  "NERSC"
+  "numba"
+)
+
+for ORG in "${CLOUD_ONLY_ORGS[@]}"; do
+  if [ "${GITHUB_REPOSITORY_OWNER}" = "${ORG}" ] ; then
+    echo "Runner initialized"
+    exit 0
+  fi
+done
+
+BLOCKED_EVENTS=(
+  "pull_request"
+  "pull_request_target"
+)
+
+for EVENT in "${BLOCKED_EVENTS[@]}"; do
+  if [ "${GITHUB_EVENT_NAME}" = "${EVENT}" ] ; then
+    echo "Workflows triggered by \"${GITHUB_EVENT_NAME}\" events are not allowed to run on self-hosted runners."
+    exit 1
+  fi
+done
+
+echo "Runner initialized"
+exit 0
diff --git a/scripts/helpers/runner.sh b/scripts/helpers/runner.sh
index ef1411b..160a93b 100755
--- a/scripts/helpers/runner.sh
+++ b/scripts/helpers/runner.sh
@@ -8,7 +8,7 @@ if [ ! -f "/jitconfig" ]; then
 fi
 
 ACTIONS_RUNNER_INPUT_JITCONFIG="$(cat /jitconfig)"
-ACTIONS_RUNNER_HOOK_JOB_STARTED=/home/runner/.check_event_type.sh
+ACTIONS_RUNNER_HOOK_JOB_STARTED=/home/runner/.initialize_runner.sh
 PARALLEL_LEVEL=$(nproc)
 export ACTIONS_RUNNER_INPUT_JITCONFIG ACTIONS_RUNNER_HOOK_JOB_STARTED PARALLEL_LEVEL
 
diff --git a/scripts/installers/runner.sh b/scripts/installers/runner.sh
index 825a630..f8ea3bb 100755
--- a/scripts/installers/runner.sh
+++ b/scripts/installers/runner.sh
@@ -20,7 +20,7 @@ rm -rf ./actions-runner.tar.gz
 
 # Copy scripts and services
 sudo cp "${NV_HELPER_SCRIPTS}/runner.sh" /runner.sh
-sudo cp "${NV_HELPER_SCRIPTS}/check_event_type.sh" /home/runner/.check_event_type.sh
+sudo cp "${NV_HELPER_SCRIPTS}/initialize_runner.sh" /home/runner/.initialize_runner.sh
 
 _UID=$(id -u)
 _GID=$(id -g)