Repository: numaproj/numaflow (go)

Issue:
Jetstream Isbsvc Replicas
# Summary

What change needs making?

I've noticed in the documentation and in the source code that if the number of replicas of the jetsream isbsvc is <3, it will still create 3 pods and replicas.  I would like like to have a single isbsvc pod with no extra replicas.  Are there any fundamental limitations that would prevent this?

# Use Cases

When would you use this?

I am running on an embedded device and trying to cut down on overhead as much as possible.  The device is frequently powered on and off as well.  On startup, we sometimes have issues with the isbsvc pods struggling to communicate with each other (e.g. establish a raft leader - Some of these issues have been addressed in recent jetstream releases, but we still have issues sometimes).  I think that life would be a lot easier if we just had a single pod for the buffer rather than 3.  Since we are running on an embedded device, I don't see any advantage to having 3 replicas and we are paying extra costs for it.

---
<!-- Issue Author: Don't delete this message to encourage other users to support your issue! -->
**Message from the maintainers**:

If you wish to see this enhancement implemented please add a 👍 reaction to this issue! We often sort issues this way to know what to prioritize.


Pull request description:
feat: Jetstream support for replica of 1 Fixes #944
Fixes #944

This allows for Jetstream event buses to have a replica set 1 instead of the minimum required 3. Modifications taken mostly from the same change done to Argo Events.

Tested this locally by modifying the numaflow install to set the bufferConfig to have replica of 1 for each (instead of the default 3) and then set my ISBS to explicitly have a replica of 1. Then verified I could run the example simple pipeline on the website. Verified in the ISBS that if replicas are not set that it defaults to 3

Patch:
diff --git a/docs/core-concepts/inter-step-buffer-service.md b/docs/core-concepts/inter-step-buffer-service.md
index 2f4bf966..5af8f602 100644
--- a/docs/core-concepts/inter-step-buffer-service.md
+++ b/docs/core-concepts/inter-step-buffer-service.md
@@ -46,7 +46,7 @@ The version `latest` in the ConfigMap should only be used for testing purpose. I
 
 ### Replicas
 
-An optional property `spec.jetstream.replicas` (defaults to 3) can be specified, which gives the total number of nodes. An odd number 3 or 5 is suggested. If the given number < 3, 3 will be used.
+An optional property `spec.jetstream.replicas` (defaults to 3) can be specified, which gives the total number of nodes. 
 
 ### Persistence
 
diff --git a/pkg/apis/numaflow/v1alpha1/jetstream_buffer_service.go b/pkg/apis/numaflow/v1alpha1/jetstream_buffer_service.go
index 2bcd235f..57871bec 100644
--- a/pkg/apis/numaflow/v1alpha1/jetstream_buffer_service.go
+++ b/pkg/apis/numaflow/v1alpha1/jetstream_buffer_service.go
@@ -79,7 +79,7 @@ func (j JetStreamBufferService) GetReplicas() int {
 	if j.Replicas == nil {
 		return 3
 	}
-	if *j.Replicas < 3 {
+	if *j.Replicas == 2 {
 		return 3
 	}
 	return int(*j.Replicas)
diff --git a/pkg/isbsvc/jetstream_service.go b/pkg/isbsvc/jetstream_service.go
index 4c2fd0ca..10f64fac 100644
--- a/pkg/isbsvc/jetstream_service.go
+++ b/pkg/isbsvc/jetstream_service.go
@@ -97,7 +97,7 @@ func (jss *jetStreamSvc) CreateBuffersAndBuckets(ctx context.Context, buffers, b
 				TTL:          time.Hour * 24 * 30, // 30 days
 				MaxBytes:     0,
 				Storage:      nats.FileStorage,
-				Replicas:     3,
+				Replicas:     v.GetInt("stream.replicas"),
 			}); err != nil {
 				return fmt.Errorf("failed to create side inputs KV %q, %w", kvName, err)
 			}
diff --git a/pkg/reconciler/isbsvc/installer/assets/jetstream/nats-cluster.conf b/pkg/reconciler/isbsvc/installer/assets/jetstream/nats-cluster.conf
new file mode 100644
index 00000000..c9bb4186
--- /dev/null
+++ b/pkg/reconciler/isbsvc/installer/assets/jetstream/nats-cluster.conf
@@ -0,0 +1,44 @@
+port: {{.ClientPort}}
+pid_file: "/var/run/nats/nats.pid"
+
+max_payload: {{.MaxPayload}}
+
+###############
+#             #
+# Monitoring  #
+#             #
+###############
+http: {{.MonitorPort}}
+server_name: $POD_NAME
+###################################
+#                                 #
+# NATS JetStream                  #
+#                                 #
+###################################
+jetstream {
+  {{.EncryptionSettings}}
+  store_dir: "/data/jetstream/store"
+  max_memory_store: {{.MaxMemoryStore}}
+  max_file_store: {{.MaxFileStore}}
+}
+
+###################################
+#                                 #
+# NATS Cluster                    #
+#                                 #
+###################################
+cluster {
+  port: {{.ClusterPort}}
+  name: {{.ClusterName}}
+  routes: [{{.Routes}}]
+  cluster_advertise: $CLUSTER_ADVERTISE
+  connect_retries: 120
+  {{.TLSConfig}}
+}
+lame_duck_duration: 120s
+##################
+#                #
+# Authorization  #
+#                #
+##################
+include ./auth.conf
\ No newline at end of file
diff --git a/pkg/reconciler/isbsvc/installer/assets/jetstream/nats.conf b/pkg/reconciler/isbsvc/installer/assets/jetstream/nats.conf
index c9bb4186..b57a6022 100644
--- a/pkg/reconciler/isbsvc/installer/assets/jetstream/nats.conf
+++ b/pkg/reconciler/isbsvc/installer/assets/jetstream/nats.conf
@@ -22,19 +22,6 @@ jetstream {
   max_file_store: {{.MaxFileStore}}
 }
 
-###################################
-#                                 #
-# NATS Cluster                    #
-#                                 #
-###################################
-cluster {
-  port: {{.ClusterPort}}
-  name: {{.ClusterName}}
-  routes: [{{.Routes}}]
-  cluster_advertise: $CLUSTER_ADVERTISE
-  connect_retries: 120
-  {{.TLSConfig}}
-}
 lame_duck_duration: 120s
 ##################
 #                #
diff --git a/pkg/reconciler/isbsvc/installer/jetstream.go b/pkg/reconciler/isbsvc/installer/jetstream.go
index d1326823..798d75da 100644
--- a/pkg/reconciler/isbsvc/installer/jetstream.go
+++ b/pkg/reconciler/isbsvc/installer/jetstream.go
@@ -390,9 +390,6 @@ func (r *jetStreamInstaller) createConfigMap(ctx context.Context) error {
 	svcName := generateJetStreamServiceName(r.isbs)
 	ssName := generateJetStreamStatefulSetName(r.isbs)
 	replicas := r.isbs.Spec.JetStream.GetReplicas()
-	if replicas < 3 {
-		replicas = 3
-	}
 	routes := []string{}
 	for j := 0; j < replicas; j++ {
 		routes = append(routes, fmt.Sprintf("nats://%s-%s.%s.%s.svc:%s", ssName, strconv.Itoa(j), svcName, r.isbs.Namespace, strconv.Itoa(int(clusterPort))))
@@ -422,7 +419,13 @@ func (r *jetStreamInstaller) createConfigMap(ctx context.Context) error {
 			return fmt.Errorf("failed to merge customized jetstream settings, %w", err)
 		}
 	}
-	confTpl := template.Must(template.ParseFS(jetStremAssets, "assets/jetstream/nats.conf"))
+	var confTpl *template.Template
+	if replicas > 2 {
+		confTpl = template.Must(template.ParseFS(jetStremAssets, "assets/jetstream/nats-cluster.conf"))
+	} else {
+		confTpl = template.Must(template.ParseFS(jetStremAssets, "assets/jetstream/nats.conf"))
+	}
+
 	var confTplOutput bytes.Buffer
 	if err := confTpl.Execute(&confTplOutput, struct {
 		ClusterName        string
diff --git a/pkg/reconciler/isbsvc/validate.go b/pkg/reconciler/isbsvc/validate.go
index 98759b01..31e65b83 100644
--- a/pkg/reconciler/isbsvc/validate.go
+++ b/pkg/reconciler/isbsvc/validate.go
@@ -47,8 +47,8 @@ func ValidateInterStepBufferService(isbs *dfv1.InterStepBufferService) error {
 		if x.Version == "" {
 			return fmt.Errorf(`invalid spec: "spec.jetstream.version" is not defined`)
 		}
-		if x.Replicas != nil && *(x.Replicas) < 3 {
-			return fmt.Errorf(`invalid spec: min value for "spec.jetstream.replicas" is 3`)
+		if x.Replicas != nil && (*x.Replicas == 2 || *x.Replicas <= 0) {
+			return fmt.Errorf(`invalid spec: min value for "spec.jetstream.replicas" is 1 and can't be 2`)
 		}
 	}
 	return nil


Test patch:
diff --git a/pkg/apis/numaflow/v1alpha1/jetstream_buffer_service_test.go b/pkg/apis/numaflow/v1alpha1/jetstream_buffer_service_test.go
index 9a4b75fa..becf05d8 100644
--- a/pkg/apis/numaflow/v1alpha1/jetstream_buffer_service_test.go
+++ b/pkg/apis/numaflow/v1alpha1/jetstream_buffer_service_test.go
@@ -125,6 +125,9 @@ func Test_JSBufferGetReplicas(t *testing.T) {
 	five := int32(5)
 	s.Replicas = &five
 	assert.Equal(t, 5, s.GetReplicas())
+	one := int32(1)
+	s.Replicas = &one
+	assert.Equal(t, 1, s.GetReplicas())
 	two := int32(2)
 	s.Replicas = &two
 	assert.Equal(t, 3, s.GetReplicas())
diff --git a/test/api-e2e/api_test.go b/test/api-e2e/api_test.go
index 616d9501..0d6acbb6 100644
--- a/test/api-e2e/api_test.go
+++ b/test/api-e2e/api_test.go
@@ -71,6 +71,43 @@ func (s *APISuite) TestISBSVC() {
 	stopPortForward()
 }
 
+func (s *APISuite) TestISBSVCReplica1() {
+	var err error
+	numaflowServerPodName := s.GetNumaflowServerPodName()
+	if numaflowServerPodName == "" {
+		panic("failed to find the nuamflow-server pod")
+	}
+	stopPortForward := s.StartPortForward(numaflowServerPodName, 8443)
+
+	var testISBSVC v1alpha1.InterStepBufferService
+	err = json.Unmarshal(testISBSVCReplica1Spec, &testISBSVC)
+	assert.NoError(s.T(), err)
+	createISBSVCBody := HTTPExpect(s.T(), "https://localhost:8443").POST(fmt.Sprintf("/api/v1/namespaces/%s/isb-services", Namespace)).WithJSON(testISBSVC).
+		Expect().
+		Status(200).Body().Raw()
+	var createISBSVCSuccessExpect = `{"data":null}`
+	assert.Contains(s.T(), createISBSVCBody, createISBSVCSuccessExpect)
+
+	listISBSVCBody := HTTPExpect(s.T(), "https://localhost:8443").GET(fmt.Sprintf("/api/v1/namespaces/%s/isb-services", Namespace)).
+		Expect().
+		Status(200).Body().Raw()
+	assert.Contains(s.T(), listISBSVCBody, testISBSVCReplica1Name)
+
+	getISBSVCBody := HTTPExpect(s.T(), "https://localhost:8443").GET(fmt.Sprintf("/api/v1/namespaces/%s/isb-services/%s", Namespace, testISBSVCReplica1Name)).
+		Expect().
+		Status(200).Body().Raw()
+	assert.Contains(s.T(), getISBSVCBody, fmt.Sprintf(`"name":"%s"`, testISBSVCReplica1Name))
+	assert.Contains(s.T(), getISBSVCBody, `"status":"healthy"`)
+
+	deleteISBSVC := HTTPExpect(s.T(), "https://localhost:8443").DELETE(fmt.Sprintf("/api/v1/namespaces/%s/isb-services/%s", Namespace, testISBSVCReplica1Name)).
+		Expect().
+		Status(200).Body().Raw()
+	var deleteISBSVCSuccessExpect = `{"data":null}`
+	assert.Contains(s.T(), deleteISBSVC, deleteISBSVCSuccessExpect)
+
+	stopPortForward()
+}
+
 func (s *APISuite) TestPipeline0() {
 	var err error
 	numaflowServerPodName := s.GetNumaflowServerPodName()
diff --git a/test/api-e2e/testdata.go b/test/api-e2e/testdata.go
index d3f1c1b0..653e596d 100644
--- a/test/api-e2e/testdata.go
+++ b/test/api-e2e/testdata.go
@@ -119,5 +119,24 @@ var (
         }
     }
 }
+`)
+	testISBSVCReplica1Name = "test-isbsvc-replica-1"
+	testISBSVCReplica1Spec = []byte(`
+{
+    "apiVersion": "numaflow.numaproj.io/v1alpha1",
+    "kind": "InterStepBufferService",
+    "metadata": {
+        "name": "test-isbsvc"
+    },
+    "spec": {
+        "jetstream": {
+            "persistence": {
+                "volumeSize": "3Gi"
+            },
+            "replicas": 1,
+            "version": "latest"
+        }
+    }
+}
 `)
 )
diff --git a/test/diamond-e2e/diamond_test.go b/test/diamond-e2e/diamond_test.go
index b0dc0f4d..2e405185 100644
--- a/test/diamond-e2e/diamond_test.go
+++ b/test/diamond-e2e/diamond_test.go
@@ -19,8 +19,11 @@ limitations under the License.
 package e2e
 
 import (
+	"context"
 	"fmt"
+	"strconv"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/suite"
 
@@ -31,7 +34,62 @@ type DiamondSuite struct {
 	E2ESuite
 }
 
-func (s *DiamondSuite) TestJoinSinkVertex() {
+func (s *DiamondSuite) TestJoinOnReducePipeline() {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
+	defer cancel()
+	w := s.Given().Pipeline("@testdata/join-on-reduce-pipeline.yaml").
+		When().
+		CreatePipelineAndWait()
+	defer w.DeletePipelineAndWait()
+	pipelineName := "join-on-reduce"
+
+	// wait for all the pods to come up
+	w.Expect().VertexPodsRunning()
+
+	done := make(chan struct{})
+	go func() {
+		// publish messages to source vertex, with event time starting from 60000
+		startTime := 60000
+		for i := 0; true; i++ {
+			select {
+			case <-ctx.Done():
+				return
+			case <-done:
+				return
+			default:
+				eventTime := strconv.Itoa(startTime + i*1000)
+				w.SendMessageTo(pipelineName, "in", NewHttpPostRequest().WithBody([]byte("1")).WithHeader("X-Numaflow-Event-Time", eventTime)).
+					SendMessageTo(pipelineName, "in", NewHttpPostRequest().WithBody([]byte("2")).WithHeader("X-Numaflow-Event-Time", eventTime)).
+					SendMessageTo(pipelineName, "in", NewHttpPostRequest().WithBody([]byte("3")).WithHeader("X-Numaflow-Event-Time", eventTime))
+			}
+		}
+	}()
+	// todo: this only tests for one occurrence: ideally should verify all
+	w.Expect().
+		SinkContains("sink", "40"). // per 10-second window: (10 * 2) * 2 atoi vertices
+		SinkContains("sink", "80")  // per 10-second window: 10 * (1 + 3) * 2 atoi vertices
+	done <- struct{}{}
+}
+
+func (s *DiamondSuite) TestJoinOnMapPipeline() {
+	w := s.Given().Pipeline("@testdata/join-on-map-pipeline.yaml").
+		When().
+		CreatePipelineAndWait()
+	defer w.DeletePipelineAndWait()
+	pipelineName := "join-on-map"
+
+	// wait for all the pods to come up
+	w.Expect().VertexPodsRunning()
+
+	w.SendMessageTo(pipelineName, "in-0", NewHttpPostRequest().WithBody([]byte("1")))
+	w.SendMessageTo(pipelineName, "in-1", NewHttpPostRequest().WithBody([]byte("2")))
+
+	w.Expect().
+		SinkContains("sink", "1").
+		SinkContains("sink", "2")
+}
+
+func (s *DiamondSuite) TestJoinOnSinkVertex() {
 	w := s.Given().Pipeline("@testdata/join-on-sink.yaml").
 		When().
 		CreatePipelineAndWait()
diff --git a/test/diamond-e2e/testdata/join-on-map-pipeline.yaml b/test/diamond-e2e/testdata/join-on-map-pipeline.yaml
new file mode 100644
index 00000000..a8d9ac5e
--- /dev/null
+++ b/test/diamond-e2e/testdata/join-on-map-pipeline.yaml
@@ -0,0 +1,40 @@
+apiVersion: numaflow.numaproj.io/v1alpha1
+kind: Pipeline
+metadata:
+  name: join-on-map
+spec:
+  limits:
+    readBatchSize: 50
+  vertices:
+    - name: in-0
+      scale:
+        min: 1
+      source:
+        http: {}
+    - name: in-1
+      scale:
+        min: 1
+      source:
+        http: {}
+    - name: cat
+      partitions: 2
+      scale:
+        min: 1
+      udf:
+        builtin:
+          name: cat
+    - name: sink
+      scale:
+        min: 1
+      sink:
+        udsink:
+          container:
+            # A redis sink for e2e testing, see https://github.com/numaproj/numaflow-sinks/tree/main/redis-e2e-test-sink
+            image: quay.io/numaio/numaflow-sink/redis-e2e-test-sink:v0.5.0
+  edges:
+    - from: in-0
+      to: cat
+    - from: in-1
+      to: cat
+    - from: cat
+      to: sink
diff --git a/test/reduce-e2e/testdata/join-on-reduce-pipeline.yaml b/test/diamond-e2e/testdata/join-on-reduce-pipeline.yaml
similarity index 100%
rename from test/reduce-e2e/testdata/join-on-reduce-pipeline.yaml
rename to test/diamond-e2e/testdata/join-on-reduce-pipeline.yaml
diff --git a/test/reduce-e2e/reduce_test.go b/test/reduce-e2e/reduce_test.go
index 45102716..8c8a0179 100644
--- a/test/reduce-e2e/reduce_test.go
+++ b/test/reduce-e2e/reduce_test.go
@@ -23,8 +23,9 @@ import (
 	"testing"
 	"time"
 
-	. "github.com/numaproj/numaflow/test/fixtures"
 	"github.com/stretchr/testify/suite"
+
+	. "github.com/numaproj/numaflow/test/fixtures"
 )
 
 type ReduceSuite struct {
@@ -147,87 +148,6 @@ func (r *ReduceSuite) TestComplexReducePipelineKeyedNonKeyed() {
 	done <- struct{}{}
 }
 
-func (r *ReduceSuite) TestJoinOnReducePipeline() {
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
-	defer cancel()
-	w := r.Given().Pipeline("@testdata/join-on-reduce-pipeline.yaml").
-		When().
-		CreatePipelineAndWait()
-	defer w.DeletePipelineAndWait()
-	pipelineName := "join-on-reduce"
-
-	// wait for all the pods to come up
-	w.Expect().VertexPodsRunning()
-
-	done := make(chan struct{})
-	go func() {
-		// publish messages to source vertex, with event time starting from 60000
-		startTime := 60000
-		for i := 0; true; i++ {
-			select {
-			case <-ctx.Done():
-				return
-			case <-done:
-				return
-			default:
-				eventTime := strconv.Itoa(startTime + i*1000)
-				w.SendMessageTo(pipelineName, "in", NewHttpPostRequest().WithBody([]byte("1")).WithHeader("X-Numaflow-Event-Time", eventTime)).
-					SendMessageTo(pipelineName, "in", NewHttpPostRequest().WithBody([]byte("2")).WithHeader("X-Numaflow-Event-Time", eventTime)).
-					SendMessageTo(pipelineName, "in", NewHttpPostRequest().WithBody([]byte("3")).WithHeader("X-Numaflow-Event-Time", eventTime))
-			}
-		}
-	}()
-
-	// todo: this only tests for one occurrence: ideally should verify all
-	w.Expect().
-		SinkContains("sink", "40"). // per 10 second window: (10 * 2) * 2 atoi vertices
-		SinkContains("sink", "80")  // per 10 second window: 10 * (1 + 3) * 2 atoi vertices
-	done <- struct{}{}
-}
-
-func (r *ReduceSuite) TestJoinOnMapPipeline() {
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
-	defer cancel()
-	w := r.Given().Pipeline("@testdata/join-on-map-pipeline.yaml").
-		When().
-		CreatePipelineAndWait()
-	defer w.DeletePipelineAndWait()
-	pipelineName := "join-on-map"
-
-	// wait for all the pods to come up
-	w.Expect().VertexPodsRunning()
-
-	done := make(chan struct{})
-	go func() {
-		// publish messages to source vertex, with event time starting from 60000
-		startTime := 60000
-		for i := 0; true; i++ {
-			select {
-			case <-ctx.Done():
-				return
-			case <-done:
-				return
-			default:
-				eventTime := strconv.Itoa(startTime + i*1000)
-				w.SendMessageTo(pipelineName, "in-0", NewHttpPostRequest().WithBody([]byte("1")).WithHeader("X-Numaflow-Event-Time", eventTime)).
-					SendMessageTo(pipelineName, "in-0", NewHttpPostRequest().WithBody([]byte("2")).WithHeader("X-Numaflow-Event-Time", eventTime)).
-					SendMessageTo(pipelineName, "in-0", NewHttpPostRequest().WithBody([]byte("3")).WithHeader("X-Numaflow-Event-Time", eventTime))
-				w.SendMessageTo(pipelineName, "in-1", NewHttpPostRequest().WithBody([]byte("5")).WithHeader("X-Numaflow-Event-Time", eventTime)).
-					SendMessageTo(pipelineName, "in-1", NewHttpPostRequest().WithBody([]byte("6")).WithHeader("X-Numaflow-Event-Time", eventTime)).
-					SendMessageTo(pipelineName, "in-1", NewHttpPostRequest().WithBody([]byte("7")).WithHeader("X-Numaflow-Event-Time", eventTime))
-
-			}
-		}
-	}()
-
-	// todo: this only tests for one occurrence: consider verifying that all match? (but test will take longer)
-	w.Expect().
-		SinkContains("sink", "80"). // per 10 second window: 10 * (2 + 6) = 80
-		SinkContains("sink", "160") // per 10 second window: 10 * (1 + 3 + 5 + 7) = 160
-
-	done <- struct{}{}
-}
-
 func (r *ReduceSuite) TestSimpleReducePipelineFailOverUsingWAL() {
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
 	defer cancel()
diff --git a/test/reduce-e2e/testdata/join-on-map-pipeline.yaml b/test/reduce-e2e/testdata/join-on-map-pipeline.yaml
deleted file mode 100644
index be8fe1ae..00000000
--- a/test/reduce-e2e/testdata/join-on-map-pipeline.yaml
+++ /dev/null
@@ -1,58 +0,0 @@
-apiVersion: numaflow.numaproj.io/v1alpha1
-kind: Pipeline
-metadata:
-  name: join-on-map
-spec:
-  limits:
-    readBatchSize: 50
-  vertices:
-    - name: in-0
-      scale:
-        min: 1
-      source:
-        http: {}
-    - name: in-1
-      scale:
-        min: 1
-      source:
-        http: {}
-    - name: atoi
-      partitions: 2
-      scale:
-        min: 1
-      udf:
-        container:
-          # Tell the input number is even or odd, see https://github.com/numaproj/numaflow-go/tree/main/pkg/mapper/examples/even_odd
-          image: quay.io/numaio/numaflow-go/map-even-odd:v0.5.0
-    - name: compute-sum
-      partitions: 2
-      udf:
-        container:
-          # Compute the sum, see https://github.com/numaproj/numaflow-go/tree/main/pkg/reducer/examples/sum
-          image: quay.io/numaio/numaflow-go/reduce-sum:v0.5.0
-        groupBy:
-          window:
-            fixed:
-              length: 10s
-          keyed: true
-          storage:
-            persistentVolumeClaim:
-              volumeSize: 10Gi
-              accessMode: ReadWriteOnce
-    - name: sink
-      scale:
-        min: 1
-      sink:
-        udsink:
-          container:
-            # A redis sink for e2e testing, see https://github.com/numaproj/numaflow-sinks/tree/main/redis-e2e-test-sink
-            image: quay.io/numaio/numaflow-sink/redis-e2e-test-sink:v0.5.0
-  edges:
-    - from: in-0
-      to: atoi
-    - from: in-1
-      to: atoi
-    - from: atoi
-      to: compute-sum
-    - from: compute-sum
-      to: sink
diff --git a/test/sdks-e2e/sdks_test.go b/test/sdks-e2e/sdks_test.go
index cb1bbca0..9b36250c 100644
--- a/test/sdks-e2e/sdks_test.go
+++ b/test/sdks-e2e/sdks_test.go
@@ -22,6 +22,7 @@ import (
 	"context"
 	"fmt"
 	"strconv"
+	"sync"
 	"testing"
 	"time"
 
@@ -123,17 +124,22 @@ func (s *SDKsSuite) TestReduceSDK() {
 	done <- struct{}{}
 }
 
-// TODO fix it when we have python sdk changes.
-func (s *SDKsSuite) TestSourceTransformerPython() {
-	s.testSourceTransformer("python")
-}
-
-func (s *SDKsSuite) TestSourceTransformerJava() {
-	s.testSourceTransformer("java")
-}
-
-func (s *SDKsSuite) TestSourceTransformerGo() {
-	s.testSourceTransformer("go")
+func (s *SDKsSuite) TestSourceTransformer() {
+	var wg sync.WaitGroup
+	wg.Add(3)
+	go func() {
+		defer wg.Done()
+		s.testSourceTransformer("python")
+	}()
+	go func() {
+		defer wg.Done()
+		s.testSourceTransformer("java")
+	}()
+	go func() {
+		defer wg.Done()
+		s.testSourceTransformer("go")
+	}()
+	wg.Wait()
 }
 
 func (s *SDKsSuite) testSourceTransformer(lang string) {