diff --git a/cmd/gateway/main.go b/cmd/gateway/main.go
index 0be3d7cdd..727d87fa9 100644
--- a/cmd/gateway/main.go
+++ b/cmd/gateway/main.go
@@ -2,9 +2,6 @@ package main
 
 import (
 	"log"
-	"os"
-	"os/signal"
-	"syscall"
 
 	"github.com/beam-cloud/beta9/pkg/gateway"
 )
@@ -15,12 +12,6 @@ func main() {
 		log.Fatalf("Error creating gateway service: %+v\n", err)
 	}
 
-	stop := make(chan os.Signal, 1)
-	signal.Notify(stop, syscall.SIGTERM)
-
 	gw.Start()
-
-	<-stop
-
 	log.Println("Gateway stopped")
 }
diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go
index de6963247..1c44a9670 100644
--- a/pkg/scheduler/scheduler.go
+++ b/pkg/scheduler/scheduler.go
@@ -197,6 +197,14 @@ func (s *Scheduler) getControllers(request *types.ContainerRequest) ([]WorkerPoo
 
 func (s *Scheduler) StartProcessingRequests() {
 	for {
+		select {
+		case <-s.ctx.Done():
+			// Context has been cancelled
+			return
+		default:
+			// Continue processing requests
+		}
+
 		if s.requestBacklog.Len() == 0 {
 			time.Sleep(requestProcessingInterval)
 			continue
diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go
index 87343206e..f8e2ba57b 100644
--- a/pkg/scheduler/scheduler_test.go
+++ b/pkg/scheduler/scheduler_test.go
@@ -59,6 +59,7 @@ func NewSchedulerForTest() (*Scheduler, error) {
 	}
 
 	return &Scheduler{
+		ctx:               context.Background(),
 		eventBus:          eventBus,
 		workerRepo:        workerRepo,
 		workerPoolManager: workerPoolManager,