feat(backend): Make number of persistence worker goroutine configurab…

…le (#3904) * Make number of persistence worker configurable * address comments * address comments * address comments
kubeflow · Jun 30, 2020 · 8a2d11c · 8a2d11c
1 parent 50fd45f
commit 8a2d11c
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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