Auto set GOMAXPROCS based on container limits in executor and operator

SeldonIO · Aug 6, 2021 · 701fa92 · 701fa92
1 parent 83717e8
commit 701fa92
Show file tree

Hide file tree

Showing 7 changed files with 37 additions and 93 deletions.
diff --git a/executor/cmd/executor/main.go b/executor/cmd/executor/main.go
@@ -34,6 +34,7 @@ import (
 	predictor2 "github.com/seldonio/seldon-core/executor/predictor"
 	"github.com/seldonio/seldon-core/executor/proto/tensorflow/serving"
 	v1 "github.com/seldonio/seldon-core/operator/apis/machinelearning.seldon.io/v1"
+	"go.uber.org/automaxprocs/maxprocs"
 	"go.uber.org/zap"
 	"google.golang.org/grpc/reflection"
 	zapf "sigs.k8s.io/controller-runtime/pkg/log/zap"
@@ -259,6 +260,13 @@ func main() {
 	setupLogger()
 	logger := logf.Log.WithName("entrypoint")
 
+	// Set runtime.GOMAXPROCS to respect container limits if the env var GOMAXPROCS is not set or is invalid, preventing CPU throttling.
+	undo, err := maxprocs.Set(maxprocs.Logger(logger.Info))
+	defer undo()
+	if err != nil {
+		logger.Error(err, "failed to set GOMAXPROCS")
+	}
+
 	// Set hostname
 	if *hostname == "" {
 		*hostname = os.Getenv("POD_NAME")

diff --git a/executor/cmd/proxy/main.go b/executor/cmd/proxy/main.go
@@ -2,13 +2,15 @@ package main
 
 import (
 	"flag"
+	"os"
+
 	"github.com/prometheus/common/log"
 	"github.com/seldonio/seldon-core/executor/api"
 	"github.com/seldonio/seldon-core/executor/api/kafka"
 	"github.com/seldonio/seldon-core/executor/api/rest"
 	"github.com/seldonio/seldon-core/executor/k8s"
 	predictor2 "github.com/seldonio/seldon-core/executor/predictor"
-	"os"
+	"go.uber.org/automaxprocs/maxprocs"
 	logf "sigs.k8s.io/controller-runtime/pkg/runtime/log"
 )
 
@@ -64,6 +66,13 @@ func main() {
 	logf.SetLogger(logf.ZapLogger(false))
 	logger := logf.Log.WithName("entrypoint")
 
+	// Set runtime.GOMAXPROCS to respect container limits if the env var GOMAXPROCS is not set or is invalid, preventing CPU throttling.
+	undo, err := maxprocs.Set(maxprocs.Logger(logger.Info))
+	defer undo()
+	if err != nil {
+		logger.Error(err, "failed to set GOMAXPROCS")
+	}
+
 	kafkaProxy := kafka.NewKafkaProxy(client, *modelName, *predictorName, *sdepName, *namespace, *broker, *hostname, int32(*httpPort), logger)
 
 	err = kafkaProxy.Consume()

diff --git a/executor/go.mod b/executor/go.mod
@@ -20,12 +20,12 @@ require (
 	github.com/tensorflow/tensorflow/tensorflow/go/core v0.0.0-00010101000000-000000000000
 	github.com/uber/jaeger-client-go v2.25.0+incompatible
 	github.com/uber/jaeger-lib v2.2.0+incompatible // indirect
+	go.uber.org/automaxprocs v1.4.0
 	go.uber.org/zap v1.16.0
 	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1
 	google.golang.org/grpc v1.32.0
 	gotest.tools v2.2.0+incompatible
 	k8s.io/api v0.18.8
-	k8s.io/apimachinery v0.18.8
 	sigs.k8s.io/controller-runtime v0.6.4
 )
 

diff --git a/executor/go.sum b/executor/go.sum
diff --git a/operator/go.mod b/operator/go.mod
@@ -11,6 +11,7 @@ require (
 	github.com/kedacore/keda v0.0.0-20200911122749-717aab81817f
 	github.com/onsi/ginkgo v1.15.0
 	github.com/onsi/gomega v1.10.4
+	go.uber.org/automaxprocs v1.4.0
 	go.uber.org/zap v1.15.0
 	gopkg.in/yaml.v2 v2.4.0
 	istio.io/api v0.0.0-20200513175333-ae3da0d240e3

diff --git a/operator/go.sum b/operator/go.sum
diff --git a/operator/main.go b/operator/main.go
@@ -32,6 +32,7 @@ import (
 	machinelearningv1alpha3 "github.com/seldonio/seldon-core/operator/apis/machinelearning.seldon.io/v1alpha3"
 	"github.com/seldonio/seldon-core/operator/controllers"
 	k8sutils "github.com/seldonio/seldon-core/operator/utils/k8s"
+	"go.uber.org/automaxprocs/maxprocs"
 	"go.uber.org/zap"
 	istio "istio.io/client-go/pkg/apis/networking/v1alpha3"
 	appsv1 "k8s.io/api/apps/v1"
@@ -135,6 +136,13 @@ func main() {
 
 	config := ctrl.GetConfigOrDie()
 
+	// Set runtime.GOMAXPROCS to respect container limits if the env var GOMAXPROCS is not set or is invalid, preventing CPU throttling.
+	undo, err := maxprocs.Set(maxprocs.Logger(setupLog.Info))
+	defer undo()
+	if err != nil {
+		setupLog.Error(err, "failed to set GOMAXPROCS")
+	}
+
 	//Override operator namespace from environment variable as the source of truth
 	operatorNamespace = utils.GetEnv("POD_NAMESPACE", operatorNamespace)