DP bug fix: wait after starting rank 0 sim (#193)

* Wait after starting rank 0 sim

Signed-off-by: Ira <[email protected]>

* Start rank 0 similar to others

Signed-off-by: Ira <[email protected]>

---------

Signed-off-by: Ira <[email protected]>

Author: irar2

pkg/llm-d-inference-sim/simulator.go

@@ -155,8 +155,8 @@ func (s *VllmSimulator) Start(ctx context.Context) error {
 	}
 
 	// For Data Parallel, start data-parallel-size - 1 additional simulators
+	g, ctx := errgroup.WithContext(ctx)
 	if s.config.DPSize > 1 {
-		g, ctx := errgroup.WithContext(context.Background())
 		for i := 2; i <= s.config.DPSize; i++ {
 			newConfig, err := s.config.Copy()
 			if err != nil {
@@ -173,12 +173,15 @@ func (s *VllmSimulator) Start(ctx context.Context) error {
 				return newSim.startSim(ctx)
 			})
 		}
-		if err := g.Wait(); err != nil {
-			return err
-		}
 		s.logger = klog.LoggerWithValues(s.logger, "rank", 0)
 	}
-	return s.startSim(ctx)
+	g.Go(func() error {
+		return s.startSim(ctx)
+	})
+	if err := g.Wait(); err != nil {
+		return err
+	}
+	return nil
 }
 
 func (s *VllmSimulator) startSim(ctx context.Context) error {