Drop BigQuery ML slots at midnight

deploy_function.sh

#!/bin/bash

# Needs google cloud SDK installed and initated
# Deploy before scheduler

FUNC_NAME="delete_bqml_pipe_slots"
PROJECT_ID="my-project"
REGION="europe-west1"
# Service Account with bigquery.resourceAdmin role
SA_ID="bqml-pipe-slot-admin"

echo "Function name: ${FUNC_NAME}"
echo "Project ID: ${PROJECT_ID}"
echo "Region: ${REGION}"
echo "SA ID: ${SA_ID}"

echo "Deploy bqml pipe function"

gcloud functions deploy ${FUNC_NAME} \
--entry-point entry_point \
--runtime python39 \
--trigger-http \
--memory 128MB \
--timeout 180 \
--project ${PROJECT_ID} \
--region ${REGION} \
--service-account ${SA_ID}@${PROJECT_ID}.iam.gserviceaccount.com \

deploy_scheduler.sh

#!/bin/bash

# Needs google cloud SDK installed and initated
# Deploy after function

FUNC_NAME="delete_bqml_pipe_slots"
PROJECT_ID="my-project"
REGION="europe-west1"
# Service Account with cloudfunctions.invoker role
SA_ID="bqml-pipe-function-invoker"

echo "Function name: ${FUNC_NAME}"
echo "Project ID: ${PROJECT_ID}"
echo "Region: ${REGION}"
echo "SA ID: ${SA_ID}"

echo "Deploy bqml_pipe scheduler"

gcloud scheduler jobs create http daily_${FUNC_NAME} \
  --description "Call delete_bqml_pipe_slots function every midnight" \
  --schedule "0 0 * * *" \
  --time-zone "Europe/Berlin" \
  --uri "https://${REGION}-${PROJECT_ID}.cloudfunctions.net/delete_bqml_pipe_slots" \
  --http-method POST \
  --oidc-service-account-email ${SA_ID}@${PROJECT_ID}.iam.gserviceaccount.com \
  --message-body '{}'


echo "Test scheduler by running"
gcloud scheduler jobs run daily_${FUNC_NAME}

echo "List of schedulers"
gcloud scheduler jobs list

main.py

"""Main file that contains the delete_slots cloud function."""
import os
from time import sleep

import google.api_core.exceptions
from flask import Request
from google.cloud.bigquery_reservation_v1.services import reservation_service

PROJECT_ID = os.getenv("PROJECT_ID", "my-project")
REGION = os.getenv("REGION", "europe-west4")

SLOT_BASE = "bqml-matrix-factor-pipeline"
COMMITMENT_ID = os.getenv("COMMITMENT_ID", f"{SLOT_BASE}-commitment")
RESERVATION_ID = os.getenv("RESERVATION_ID", f"{SLOT_BASE}-reservation")
ASSIGNMENT_ID = os.getenv("ASSIGNMENT_ID", f"{SLOT_BASE}-assignment")

reservation_client = reservation_service.ReservationServiceClient()


def drop_slots(project_id: str, region: str, commitment_id: str, reservation_id: str, assignment_id: str,) -> None:
    """Drop all slots with the specified assignments and region.

    Args:
        project_id: Billing and assigned project
        region: Slot server region
        commitment_id: Name of commitment
        reservation_id: Name of reservation
        assignment_id: Name of assignment

    Returns:
        Returns nothing.
    """

    
    assignment_name = reservation_client.assignment_path(
        project=project_id, location=region, reservation=reservation_id, assignment=assignment_id
    )
    print("Assignment name: " + assignment_name)

    reservation_name = reservation_client.reservation_path(
        project=project_id, location=region, reservation=reservation_id
    )
    print("Reservation name: " + reservation_name)

    commitment_name = reservation_client.capacity_commitment_path(
        project=project_id, location=region, capacity_commitment=commitment_id
    )
    print("Commitment name: " + commitment_name)

    try:
        reservation_client.delete_assignment(name=assignment_name)
        print("Assignment deleted")
    except google.api_core.exceptions.NotFound as error:
        print(error.message)
    try:
        reservation_client.delete_reservation(name=reservation_name)
        print("Reservation deleted")
    except google.api_core.exceptions.NotFound as error:
        print(error.message)
    try:
        reservation_client.delete_capacity_commitment(name=commitment_name)
        print("Commitment deleted")
    except google.api_core.exceptions.NotFound as error:
        print(error.message)
    return


def entry_point(request: Request) -> str:
    """The entry point of the cloud function. Incoming requests will be sent to this function.
    Returns:
        String
    """
    try:
        drop_slots(project_id=PROJECT_ID, region=REGION, commitment_id=COMMITMENT_ID, reservation_id=RESERVATION_ID, assignment_id=ASSIGNMENT_ID)
    except google.api_core.exceptions.FailedPrecondition as error:
        print(error.message)
        # Slots can be deleted 60 seconds after purchase at the earliest
        print("Wait 62s and try one more time")
        sleep(62)
        drop_slots(project_id=PROJECT_ID, region=REGION, commitment_id=COMMITMENT_ID, reservation_id=RESERVATION_ID,
                   assignment_id=ASSIGNMENT_ID)

    return "Slots are dropped"

requirements.txt

# Function dependencies, for example:
# package>=version

google-cloud-bigquery-reservation==1.4.1; python_version >= "3.6"