From 8a2d11c96a9625ce4a407fd21cfdbdfd8ad0186e Mon Sep 17 00:00:00 2001
From: frozeNinK <wenquan.xing@getcruise.com>
Date: Mon, 29 Jun 2020 21:37:58 -0700
Subject: [PATCH] feat(backend): Make number of persistence worker goroutine
 configurable (#3904)

* Make number of persistence worker configurable

* address comments

* address comments

* address comments
---
 backend/Dockerfile.persistenceagent                          | 5 ++++-
 backend/src/agent/persistence/main.go                        | 5 ++++-
 .../pipeline/ml-pipeline-persistenceagent-deployment.yaml    | 2 ++
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/backend/Dockerfile.persistenceagent b/backend/Dockerfile.persistenceagent
index 5dd03cf1e28..6921c677f0e 100644
--- a/backend/Dockerfile.persistenceagent
+++ b/backend/Dockerfile.persistenceagent
@@ -20,4 +20,7 @@ ENV NAMESPACE ""
 # Set Workflow TTL to 1 day. The way to use a different value for a particular Kubeflow Pipelines deployment is demonstrated in manifests/kustomize/base/pipeline/ml-pipeline-persistenceagent-deployment.yaml
 ENV TTL_SECONDS_AFTER_WORKFLOW_FINISH 86400
 
-CMD persistence_agent --logtostderr=true --namespace=${NAMESPACE} --ttlSecondsAfterWorkflowFinish=${TTL_SECONDS_AFTER_WORKFLOW_FINISH}
+# NUM_WORKERS indicates now many worker goroutines
+ENV NUM_WORKERS 2
+
+CMD persistence_agent --logtostderr=true --namespace=${NAMESPACE} --ttlSecondsAfterWorkflowFinish=${TTL_SECONDS_AFTER_WORKFLOW_FINISH} --numWorker ${NUM_WORKERS}
diff --git a/backend/src/agent/persistence/main.go b/backend/src/agent/persistence/main.go
index 21293674a70..9b0a77db05e 100644
--- a/backend/src/agent/persistence/main.go
+++ b/backend/src/agent/persistence/main.go
@@ -42,6 +42,7 @@ var (
 	mlPipelineServiceGRPCPort     string
 	namespace                     string
 	ttlSecondsAfterWorkflowFinish int64
+	numWorker                     int
 )
 
 const (
@@ -55,6 +56,7 @@ const (
 	mlPipelineAPIServerGRPCPortFlagName   = "mlPipelineServiceGRPCPort"
 	namespaceFlagName                     = "namespace"
 	ttlSecondsAfterWorkflowFinishFlagName = "ttlSecondsAfterWorkflowFinish"
+	numWorkerName                         = "numWorker"
 )
 
 func main() {
@@ -108,7 +110,7 @@ func main() {
 	go swfInformerFactory.Start(stopCh)
 	go workflowInformerFactory.Start(stopCh)
 
-	if err = controller.Run(2, stopCh); err != nil {
+	if err = controller.Run(numWorker, stopCh); err != nil {
 		log.Fatalf("Error running controller: %s", err.Error())
 	}
 }
@@ -125,4 +127,5 @@ func init() {
 		"/apis/v1beta1", "The base path for the ML pipeline API server.")
 	flag.StringVar(&namespace, namespaceFlagName, "", "The namespace name used for Kubernetes informers to obtain the listers.")
 	flag.Int64Var(&ttlSecondsAfterWorkflowFinish, ttlSecondsAfterWorkflowFinishFlagName, 604800 /* 7 days */, "The TTL for Argo workflow to persist after workflow finish.")
+	flag.IntVar(&numWorker, numWorkerName, 2, "Number of worker for sync job.")
 }
diff --git a/manifests/kustomize/base/pipeline/ml-pipeline-persistenceagent-deployment.yaml b/manifests/kustomize/base/pipeline/ml-pipeline-persistenceagent-deployment.yaml
index 2d40ac6cd70..58fe734c911 100644
--- a/manifests/kustomize/base/pipeline/ml-pipeline-persistenceagent-deployment.yaml
+++ b/manifests/kustomize/base/pipeline/ml-pipeline-persistenceagent-deployment.yaml
@@ -21,6 +21,8 @@ spec:
                 fieldPath: metadata.namespace
           - name: TTL_SECONDS_AFTER_WORKFLOW_FINISH
             value: "86400"
+          - name: NUM_WORKERS
+            value: "2"
         image: gcr.io/ml-pipeline/persistenceagent:dummy
         imagePullPolicy: IfNotPresent
         name: ml-pipeline-persistenceagent