Go-Slot-Scheduler

A simple web utility service to auto schedule BigQuery flex slots for an organization at certain hours of the day, in a serverless mode with Cloud Run, Tasks and Scheduler.

Deploy

• Clone repo in Cloud Shell of project assigned for BigQuery Reservation and change directory to go-slot-scheduler/
• Set GCloud parameters if not set, choose region of compute

export PROJECT_ID=$(gcloud config get-value project)
gcloud config set compute/region us-east4

Create a Task Queue

QUEUE_ID=commit-delete-queue
QUEUE_LOCATION=us-east4
gcloud tasks queues create $QUEUE_ID --location=$QUEUE_LOCATION

Create or Grant service account with Bigquery resource admin permission

• Use default compute service account for Cloud Run

export PROJECT_ID=$(gcloud config get-value project)
SERV_ACCT=`gcloud iam service-accounts list --format="value(email)" | grep [email protected]`
gcloud projects add-iam-policy-binding $PROJECT_ID \
--member="serviceAccount:${SERV_ACCT}" \
--role="roles/bigquery.resourceAdmin" \
--condition=None

OR

• (Recommended) For additional security or least privilege, create a custom service account and grant roles/bigquery.resourceAdmin, roles/cloudtasks.admin and roles/run.admin

gcloud iam service-accounts create slot-scheduler-sa \
    --description="go-slot-scheduler service account" \
    --display-name="Slot Scheduler"
SERV_ACCT=slot-scheduler-sa@$PROJECT_ID.iam.gserviceaccount.com

gcloud iam service-accounts add-iam-policy-binding \
  $(gcloud projects describe $PROJECT_ID --format="value(projectNumber)")[email protected] \
  --member="serviceAccount:${SERV_ACCT}" \
  --role="roles/iam.serviceAccountUser" \
  --condition=None

gcloud projects add-iam-policy-binding $PROJECT_ID \
--member="serviceAccount:${SERV_ACCT}" \
--role="roles/bigquery.resourceAdmin" \
--condition=None

gcloud projects add-iam-policy-binding $PROJECT_ID \
--member="serviceAccount:${SERV_ACCT}" \
--role="roles/run.admin" \
--condition=None

gcloud projects add-iam-policy-binding $PROJECT_ID \
--member="serviceAccount:${SERV_ACCT}" \
--role="roles/cloudtasks.admin" \
--condition=None

Deploy to CloudRun

• Ensure you have the right permissions to deploy to CloudRun, using Cloud Build from source and Artifact Registry
• Update environment variables for MAX_SLOTS for the organization or slot commitment
• Deploy service to Cloudrun

REGION=$(gcloud config get-value compute/region)
MAX_SLOTS=500

gcloud run deploy go-slot-scheduler --region ${REGION} --set-env-vars=MAX_SLOTS=${MAX_SLOTS},QUEUE_ID=${QUEUE_ID},QUEUE_LOCATION=${QUEUE_LOCATION} --no-allow-unauthenticated --service-account=$SERV_ACCT --source .

• Payload of http request in data.json

# if extra_slot is less than 100, scheduler will default to minimum slot of 100
{
    "extra_slot":100,
    "region":"us",
    "minutes": 180
}

# Get the Cloudrun service https endpoint
ENDPOINT=$(gcloud run services describe go-slot-scheduler --region $REGION --format 'value(status.url)')

# Call the service with sample data 
# Permission will be denied because it is an internal service. If you want to test use `--allow-unauthenticated` in cloud run deploy command
curl -d '@data.json' $ENDPOINT/add_capacity -H "Content-Type:application/json"

Set up schedule with Cloud Scheduler

# Schedule 100 extra slots at 6AM M-F, for 10 hours
# https://cloud.google.com/sdk/gcloud/reference/scheduler/jobs/create/http 
# Default timezone is UTC, but you can change with an extra command `--time-zone="est"`
gcloud scheduler jobs create http slot-schedule \
    --location=$QUEUE_LOCATION \
    --schedule="* 6 * * 1-5" \
    --headers="Content-Type=application/json" \
    --uri="${ENDPOINT}/add_capacity" \
    --message-body-from-file=data.json \
    --oidc-service-account-email=${SERV_ACCT}

Development

gcloud builds submit --pack image=[IMAGE] us-east1 
and 
gcloud run deploy go-slot-scheduler --image [IMAGE]

OR run on Docker locally

Future Work

• Adjust dedicated slot assignments for projects after capacity adjustment
• Add schedule frequency to environment and create job schedules

Credit

Go-slot-scheduler is inspired by bq-slot-scheduler in python written by Patrick Dunn for AppEngine