Add NotifyContext and pass it to http request

to abort the request of the worker is killed, so worker won't be stuck on http request Signed-off-by: Anil Vishnoi <[email protected]>
instructlab · Apr 27, 2024 · fbc98fc · fbc98fc
1 parent 47e4179
commit fbc98fc
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/worker/cmd/generate.go b/worker/cmd/generate.go
@@ -132,7 +132,9 @@ var generateCmd = &cobra.Command{
 	Run: func(cmd *cobra.Command, args []string) {
 		logger := initLogger(Debug)
 		sugar := logger.Sugar()
-		ctx := cmd.Context()
+
+		ctx, cancel := signal.NotifyContext(cmd.Context(), syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGINT)
+		defer cancel()
 
 		sugar.Info("Starting generate worker")
 
@@ -916,7 +918,7 @@ func (w *Worker) datagenSvc(taxonomyFiles []string, outputDir string, numSamples
 			requestURL = strings.Replace(requestURL, "skill", "knowledge", -1)
 		}
 
-		request, err := http.NewRequest("POST", requestURL, bytes.NewBuffer(jsonData))
+		request, err := http.NewRequestWithContext(w.ctx, "POST", requestURL, bytes.NewBuffer(jsonData))
 		if err != nil {
 			return nil, fmt.Errorf("failed to create request: %w", err)
 		}