add scheduler for backfill (#111)

Co-authored-by: James Kwon <[email protected]>
Comfy-Org · Jan 2, 2025 · 3bd5425 · 3bd5425
1 parent d51f850
commit 3bd5425
Show file tree

Hide file tree

Showing 7 changed files with 74 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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())
+	})
 }