fix(ray): fix deployment check

instill-ai · Oct 16, 2024 · 487db2d · 487db2d
1 parent d71b5ac
commit 487db2d
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/pkg/worker/workflow.go b/pkg/worker/workflow.go
@@ -177,14 +177,23 @@ func (w *worker) TriggerModelActivity(ctx context.Context, param *TriggerModelAc
 	// wait for model instance to come online to start processing the request
 	// temporary solution to not overcharge for credits
 	// TODO: design a better flow
-	for {
-		if state, _, numOfActiveReplica, err := w.ray.ModelReady(ctx, fmt.Sprintf("%s/%s/%s", param.OwnerType, param.OwnerUID, param.ModelID), param.ModelVersion.Version); err != nil {
+	state, _, numOfActiveReplica, err := w.ray.ModelReady(ctx, fmt.Sprintf("%s/%s/%s", param.OwnerType, param.OwnerUID, param.ModelID), param.ModelVersion.Version)
+	if err != nil {
+		return w.toApplicationError(err, param.ModelID, ModelActivityError)
+	}
+	for *state == modelpb.State_STATE_OFFLINE {
+		time.Sleep(time.Millisecond * 500)
+	}
+	for *state != modelpb.State_STATE_ACTIVE && numOfActiveReplica <= 0 {
+		logger.Debug(fmt.Sprintf("model upscale state: %v", state))
+		logger.Debug(fmt.Sprintf("model upscale numOfActiveReplica: %v", numOfActiveReplica))
+		if state, _, numOfActiveReplica, err = w.ray.ModelReady(ctx, fmt.Sprintf("%s/%s/%s", param.OwnerType, param.OwnerUID, param.ModelID), param.ModelVersion.Version); err != nil {
 			return w.toApplicationError(err, param.ModelID, ModelActivityError)
-		} else if *state == modelpb.State_STATE_ACTIVE && numOfActiveReplica > 0 {
-			break
 		} else if *state != modelpb.State_STATE_SCALING_UP && *state != modelpb.State_STATE_STARTING {
 			logger.Error(fmt.Sprintf("model upscale failed: current model state: %v", state), zap.Error(err))
 			return w.toApplicationError(fmt.Errorf("model upscale failed: current model state: %v", state), param.ModelID, ModelActivityError)
+		} else {
+			time.Sleep(time.Millisecond * 500)
 		}
 	}