Create a new Parameter for batching the backfill operation (#113)

* Parameterize backfill * Resolve merge conflict --------- Co-authored-by: James Kwon <[email protected]>
Comfy-Org · Jan 3, 2025 · c46dd20 · c46dd20
1 parent aa7a13c
commit c46dd20
Show file tree

Hide file tree

Showing 9 changed files with 167 additions and 115 deletions.
diff --git a/drip/api.gen.go b/drip/api.gen.go
diff --git a/infrastructure/modules/node-pack-extract-trigger/main.tf b/infrastructure/modules/node-pack-extract-trigger/main.tf
@@ -88,7 +88,8 @@ resource "google_cloud_scheduler_job" "backfill" {
 
   http_target {
     http_method = "POST"
-    uri         = "${var.registry_backend_url}/comfy-nodes/backfill"
+    uri         = "${var.registry_backend_url}/comfy-nodes/backfill?max_node=${var.backfill_job_max_node}"
+
 
     oidc_token {
       service_account_email = data.google_service_account.cloudbuild_service_account.email

diff --git a/infrastructure/modules/node-pack-extract-trigger/variable.tf b/infrastructure/modules/node-pack-extract-trigger/variable.tf
@@ -45,6 +45,12 @@ variable "backfill_job_schedule" {
   default     = "30 3 29 2 *"
 }
 
+variable "backfill_job_max_node" {
+  type        = number
+  default     = 10
+  description = "maximum number of nodes to be backfilled"
+}
+
 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/integration-tests/registry_integration_test.go b/integration-tests/registry_integration_test.go
@@ -900,6 +900,9 @@ func TestRegistryComfyNode(t *testing.T) {
 	nodeVersionToBeBackfill := []*drip.NodeVersion{
 		randomNodeVersion(1),
 		randomNodeVersion(2),
+		randomNodeVersion(3),
+		randomNodeVersion(4),
+		randomNodeVersion(5),
 	}
 	for _, nv := range nodeVersionToBeBackfill {
 		_, err = withMiddleware(authz, impl.PublishNodeVersion)(ctx, drip.PublishNodeVersionRequestObject{
@@ -1044,10 +1047,25 @@ func TestRegistryComfyNode(t *testing.T) {
 	})
 
 	t.Run("TriggerBackfill", func(t *testing.T) {
-		impl.mockPubsubService.On("PublishNodePack", mock.Anything, mock.Anything).Return(nil)
-		res, err := withMiddleware(authz, impl.ComfyNodesBackfill)(ctx, drip.ComfyNodesBackfillRequestObject{})
-		require.NoError(t, err, "should return created node version")
-		require.IsType(t, drip.ComfyNodesBackfill204Response{}, res)
-		impl.mockPubsubService.AssertNumberOfCalls(t, "PublishNodePack", len(nodeVersionToBeBackfill))
+		mockCalled := 0
+		t.Run("Unlimited", func(t *testing.T) {
+			impl.mockPubsubService.On("PublishNodePack", mock.Anything, mock.Anything).Return(nil)
+			res, err := withMiddleware(authz, impl.ComfyNodesBackfill)(ctx, drip.ComfyNodesBackfillRequestObject{})
+			require.NoError(t, err, "should return created node version")
+			require.IsType(t, drip.ComfyNodesBackfill204Response{}, res)
+			impl.mockPubsubService.AssertNumberOfCalls(t, "PublishNodePack", len(nodeVersionToBeBackfill)+mockCalled)
+			mockCalled += len(nodeVersionToBeBackfill)
+		})
+
+		t.Run("Limited", func(t *testing.T) {
+			limit := 2
+			impl.mockPubsubService.On("PublishNodePack", mock.Anything, mock.Anything).Return(nil)
+			res, err := withMiddleware(authz, impl.ComfyNodesBackfill)(ctx, drip.ComfyNodesBackfillRequestObject{Params: drip.ComfyNodesBackfillParams{MaxNode: &limit}})
+			require.NoError(t, err, "should return created node version")
+			require.IsType(t, drip.ComfyNodesBackfill204Response{}, res)
+			impl.mockPubsubService.AssertNumberOfCalls(t, "PublishNodePack", limit+mockCalled)
+			mockCalled += limit
+		})
 	})
+
 }
diff --git a/node-pack-extract/cloudbuild.yaml b/node-pack-extract/cloudbuild.yaml
@@ -27,8 +27,8 @@ steps:
       - -c
       - gcloud auth print-identity-token --audiences="$_REGISTRY_BACKEND_URL" | tee /workspace/token
 
-  - name: "gcr.io/cloud-builders/curl"
-    entrypoint: "bash"
+  - name: "curlimages/curl"
+    entrypoint: "sh"
     args:
       - -c
       - |

diff --git a/node-pack-extract/test/trigger_test.go b/node-pack-extract/test/trigger_test.go
@@ -81,7 +81,7 @@ func TestApply(t *testing.T) {
 		require.NoError(t, err)
 		h := j.GetHttpTarget()
 		require.NotNil(t, h)
-		assert.Contains(t, h.GetUri(), "/comfy-nodes/backfill")
+		assert.Equal(t, h.GetUri(), "https://stagingapi.comfy.org/comfy-nodes/backfill?max_node=10")
 		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())

diff --git a/openapi.yml b/openapi.yml
@@ -1870,6 +1870,13 @@ paths:
       operationId: ComfyNodesBackfill
       tags:
         - ComfyNodes
+      parameters:
+        - in: query
+          name: max_node
+          required: false
+          schema:
+            type: integer
+            default: 10
       responses:
         '204':
           description: Backfill triggered

diff --git a/server/implementation/registry.go b/server/implementation/registry.go
@@ -1063,7 +1063,7 @@ func (impl *DripStrictServerImplementation) GetComfyNode(ctx context.Context, re
 
 func (impl *DripStrictServerImplementation) ComfyNodesBackfill(ctx context.Context, request drip.ComfyNodesBackfillRequestObject) (drip.ComfyNodesBackfillResponseObject, error) {
 	log.Ctx(ctx).Info().Msg("ComfyNodesBackfill request received")
-	err := impl.RegistryService.TriggerComfyNodesBackfill(ctx, impl.Client)
+	err := impl.RegistryService.TriggerComfyNodesBackfill(ctx, impl.Client, request.Params.MaxNode)
 	if err != nil {
 		log.Ctx(ctx).Error().Msgf("Failed to trigger comfy nodes backfill w/ err: %v", err)
 		return drip.ComfyNodesBackfill500JSONResponse{Message: "Failed to trigger comfy nodes backfill", Error: err.Error()}, nil

diff --git a/services/registry/registry_svc.go b/services/registry/registry_svc.go
@@ -700,12 +700,15 @@ func (s *RegistryService) GetComfyNode(ctx context.Context, client *ent.Client,
 	return nv.Edges.ComfyNodes[0], nil
 }
 
-func (s *RegistryService) TriggerComfyNodesBackfill(ctx context.Context, client *ent.Client) error {
-	nvs, err := client.NodeVersion.
+func (s *RegistryService) TriggerComfyNodesBackfill(ctx context.Context, client *ent.Client, max *int) error {
+	q := client.NodeVersion.
 		Query().
 		WithStorageFile().
-		Where(nodeversion.Not(nodeversion.HasComfyNodes())).
-		All(ctx)
+		Where(nodeversion.Not(nodeversion.HasComfyNodes()))
+	if max != nil {
+		q.Limit(*max)
+	}
+	nvs, err := q.All(ctx)
 	if err != nil {
 		return fmt.Errorf("failed to query node versions: %w", err)
 	}