From 3bd5425d40d7d0c9759c9ef40f8a5cb4e9254c3e Mon Sep 17 00:00:00 2001 From: James Kwon <96548424+james03160927@users.noreply.github.com> Date: Thu, 2 Jan 2025 16:12:52 -0500 Subject: [PATCH] add scheduler for backfill (#111) Co-authored-by: James Kwon <96548424+hongil0316@users.noreply.github.com> --- .../examples/simple-node-pack-extract/main.tf | 6 +++++ .../modules/node-pack-extract-trigger/main.tf | 19 ++++++++++++++ .../node-pack-extract-trigger/output.tf | 4 +++ .../node-pack-extract-trigger/variable.tf | 12 +++++++++ node-pack-extract/test/go.mod | 11 +++++--- node-pack-extract/test/go.sum | 2 ++ node-pack-extract/test/trigger_test.go | 25 ++++++++++++++++++- 7 files changed, 74 insertions(+), 5 deletions(-) diff --git a/infrastructure/examples/simple-node-pack-extract/main.tf b/infrastructure/examples/simple-node-pack-extract/main.tf index e443ecd..05e80e2 100644 --- a/infrastructure/examples/simple-node-pack-extract/main.tf +++ b/infrastructure/examples/simple-node-pack-extract/main.tf @@ -58,6 +58,12 @@ output "topic_id" { output "bucket_notification_id" { value = module.node_pack_extract_trigger.bucket_notification_id } +output "backfill_scheduler_id" { + value = module.node_pack_extract_trigger.backfill_scheduler_id +} output "bucket_name" { value = google_storage_bucket.bucket.name } +output "service_account" { + value = google_service_account.service_account.email +} diff --git a/infrastructure/modules/node-pack-extract-trigger/main.tf b/infrastructure/modules/node-pack-extract-trigger/main.tf index 1b522e2..9144ef3 100644 --- a/infrastructure/modules/node-pack-extract-trigger/main.tf +++ b/infrastructure/modules/node-pack-extract-trigger/main.tf @@ -77,3 +77,22 @@ resource "google_cloudbuild_trigger" "trigger" { _REGISTRY_BACKEND_URL = var.registry_backend_url } } + + +resource "google_cloud_scheduler_job" "backfill" { + name = var.backfill_job_name + description = "A job to trigger comfy node pack backfill." + + schedule = var.backfill_job_schedule + paused = true + + http_target { + http_method = "POST" + uri = "${var.registry_backend_url}/comfy-nodes/backfill" + + oidc_token { + service_account_email = data.google_service_account.cloudbuild_service_account.email + audience = var.registry_backend_url + } + } +} diff --git a/infrastructure/modules/node-pack-extract-trigger/output.tf b/infrastructure/modules/node-pack-extract-trigger/output.tf index 32cbceb..fadf328 100644 --- a/infrastructure/modules/node-pack-extract-trigger/output.tf +++ b/infrastructure/modules/node-pack-extract-trigger/output.tf @@ -9,3 +9,7 @@ output "bucket_notification_id" { output "trigger_id" { value = google_cloudbuild_trigger.trigger.id } + +output "backfill_scheduler_id" { + value = google_cloud_scheduler_job.backfill.id +} diff --git a/infrastructure/modules/node-pack-extract-trigger/variable.tf b/infrastructure/modules/node-pack-extract-trigger/variable.tf index 41f67e8..d271009 100644 --- a/infrastructure/modules/node-pack-extract-trigger/variable.tf +++ b/infrastructure/modules/node-pack-extract-trigger/variable.tf @@ -33,6 +33,18 @@ variable "trigger_name" { default = "comfy-node-pack-extract" } +variable "backfill_job_name" { + type = string + description = "Cloud build trigger name" + default = "comfy-node-pack-backfill" +} + +variable "backfill_job_schedule" { + type = string + description = "Cron schedule for backfill job in UTC. Default is 3:30am on the 29th of february UTC." + default = "30 3 29 2 *" +} + variable "git_repo_uri" { type = string description = "Connected git repo containing the cloud build pipeline. See https://cloud.google.com/build/docs/repositories" diff --git a/node-pack-extract/test/go.mod b/node-pack-extract/test/go.mod index 8b55b7d..7650aef 100644 --- a/node-pack-extract/test/go.mod +++ b/node-pack-extract/test/go.mod @@ -2,19 +2,23 @@ module test go 1.23.2 -require github.com/gruntwork-io/terratest v0.48.1 +require ( + cloud.google.com/go/cloudbuild v1.19.0 + cloud.google.com/go/scheduler v1.11.2 + cloud.google.com/go/storage v1.47.0 + github.com/gruntwork-io/terratest v0.48.1 + github.com/stretchr/testify v1.9.0 +) require ( cel.dev/expr v0.16.1 // indirect cloud.google.com/go v0.116.0 // indirect cloud.google.com/go/auth v0.10.2 // indirect cloud.google.com/go/auth/oauth2adapt v0.2.5 // indirect - cloud.google.com/go/cloudbuild v1.19.0 // indirect cloud.google.com/go/compute/metadata v0.5.2 // indirect cloud.google.com/go/iam v1.2.2 // indirect cloud.google.com/go/longrunning v0.6.2 // indirect cloud.google.com/go/monitoring v1.21.2 // indirect - cloud.google.com/go/storage v1.47.0 // indirect filippo.io/edwards25519 v1.1.0 // indirect github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.24.1 // indirect github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.48.1 // indirect @@ -127,7 +131,6 @@ require ( github.com/russross/blackfriday/v2 v2.1.0 // indirect github.com/sirupsen/logrus v1.9.3 // indirect github.com/spf13/pflag v1.0.5 // indirect - github.com/stretchr/testify v1.9.0 // indirect github.com/tmccombs/hcl2json v0.6.4 // indirect github.com/ulikunitz/xz v0.5.10 // indirect github.com/urfave/cli v1.22.16 // indirect diff --git a/node-pack-extract/test/go.sum b/node-pack-extract/test/go.sum index f9e1a5a..462abb8 100644 --- a/node-pack-extract/test/go.sum +++ b/node-pack-extract/test/go.sum @@ -19,6 +19,8 @@ cloud.google.com/go/longrunning v0.6.2 h1:xjDfh1pQcWPEvnfjZmwjKQEcHnpz6lHjfy7Fo0 cloud.google.com/go/longrunning v0.6.2/go.mod h1:k/vIs83RN4bE3YCswdXC5PFfWVILjm3hpEUlSko4PiI= cloud.google.com/go/monitoring v1.21.2 h1:FChwVtClH19E7pJ+e0xUhJPGksctZNVOk2UhMmblmdU= cloud.google.com/go/monitoring v1.21.2/go.mod h1:hS3pXvaG8KgWTSz+dAdyzPrGUYmi2Q+WFX8g2hqVEZU= +cloud.google.com/go/scheduler v1.11.2 h1:PfkvJP1qKu9NvFB65Ja/s918bPZWMBcYkg35Ljdw1Oc= +cloud.google.com/go/scheduler v1.11.2/go.mod h1:GZSv76T+KTssX2I9WukIYQuQRf7jk1WI+LOcIEHUUHk= cloud.google.com/go/storage v1.47.0 h1:ajqgt30fnOMmLfWfu1PWcb+V9Dxz6n+9WKjdNg5R4HM= cloud.google.com/go/storage v1.47.0/go.mod h1:Ks0vP374w0PW6jOUameJbapbQKXqkjGd/OJRp2fb9IQ= cloud.google.com/go/trace v1.11.2 h1:4ZmaBdL8Ng/ajrgKqY5jfvzqMXbrDcBsUGXOT9aqTtI= diff --git a/node-pack-extract/test/trigger_test.go b/node-pack-extract/test/trigger_test.go index 55b1ae3..07f8032 100644 --- a/node-pack-extract/test/trigger_test.go +++ b/node-pack-extract/test/trigger_test.go @@ -2,11 +2,15 @@ package test import ( "context" + "net/http" "strings" "testing" "cloud.google.com/go/cloudbuild/apiv1/v2/cloudbuildpb" + cloudscheduler "cloud.google.com/go/scheduler/apiv1" + "cloud.google.com/go/scheduler/apiv1/schedulerpb" "cloud.google.com/go/storage" + "github.com/gruntwork-io/terratest/modules/environment" "github.com/gruntwork-io/terratest/modules/gcp" "github.com/gruntwork-io/terratest/modules/random" @@ -33,8 +37,10 @@ func TestApply(t *testing.T) { topicID := terraform.Output(t, terraformOptions, "topic_id") triggerID := terraform.Output(t, terraformOptions, "trigger_id") - bucketName := terraform.Output(t, terraformOptions, "bucket_name") bucketNotificationID := terraform.Output(t, terraformOptions, "bucket_notification_id") + schedulerID := terraform.Output(t, terraformOptions, "backfill_scheduler_id") + bucketName := terraform.Output(t, terraformOptions, "bucket_name") + serviceAccount := terraform.Output(t, terraformOptions, "service_account") t.Run("CheckCloudBuildTrigger", func(t *testing.T) { gcb := gcp.NewCloudBuildService(t) @@ -63,4 +69,21 @@ func TestApply(t *testing.T) { } assert.True(t, found) }) + + t.Run("CheckScheduler", func(t *testing.T) { + client, err := cloudscheduler.NewCloudSchedulerClient(context.Background()) + require.NoError(t, err) + defer client.Close() + + j, err := client.GetJob(context.Background(), &schedulerpb.GetJobRequest{ + Name: schedulerID, + }) + require.NoError(t, err) + h := j.GetHttpTarget() + require.NotNil(t, h) + assert.Contains(t, h.GetUri(), "/comfy-nodes/backfill") + assert.Equal(t, http.MethodPost, h.GetHttpMethod().String()) + assert.Equal(t, "https://stagingapi.comfy.org", h.GetOidcToken().GetAudience()) + assert.Equal(t, serviceAccount, h.GetOidcToken().GetServiceAccountEmail()) + }) }