feat: support CPU and memory limit for docker runner (#1486)

* feat: support CPU and memory limit for docker runner Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix typo Signed-off-by: Keming <kemingyang@tensorchord.ai> * add cpu-set, use cpu quota on UNIX Signed-off-by: Keming <kemingyang@tensorchord.ai> * panic if both cpus and cpu-set are set Signed-off-by: Keming <kemingyang@tensorchord.ai> --------- Signed-off-by: Keming <kemingyang@tensorchord.ai>
tensorchord · Feb 18, 2023 · 53ba50e · 53ba50e
1 parent b0143c3
commit 53ba50e
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 2 deletions.
diff --git a/pkg/app/run.go b/pkg/app/run.go
@@ -87,7 +87,7 @@ var CommandCreate = &cli.Command{
 			Value: 2048,
 		},
 		&cli.StringFlag{
-			Name:  "cpu",
+			Name:  "cpus",
 			Usage: "Request CPU resources (number of cores), such as 0.5, 1, 2",
 			Value: "",
 		},
@@ -134,7 +134,7 @@ func run(clicontext *cli.Context) error {
 		Image:           clicontext.String("image"),
 		Timeout:         clicontext.Duration("timeout"),
 		NumMem:          clicontext.String("memory"),
-		NumCPU:          clicontext.String("cpu"),
+		NumCPU:          clicontext.String("cpus"),
 		NumGPU:          clicontext.Int("gpu"),
 		ShmSize:         clicontext.Int("shm-size"),
 		EnvironmentName: name,

diff --git a/pkg/app/up.go b/pkg/app/up.go
@@ -90,6 +90,21 @@ var CommandUp = &cli.Command{
 			Usage: "Configure the shared memory size (megabyte)",
 			Value: 2048,
 		},
+		&cli.StringFlag{
+			Name:  "cpus",
+			Usage: "Request CPU resources (number of cores), such as 0.5, 1, 2",
+			Value: "",
+		},
+		&cli.StringFlag{
+			Name:  "cpu-set",
+			Usage: "Limit the specific CPUs or cores the environment can use, such as `0-3`, `1,3`",
+			Value: "",
+		},
+		&cli.StringFlag{
+			Name:  "memory",
+			Usage: "Request Memory, such as 512Mb, 2Gb",
+			Value: "",
+		},
 		&cli.BoolFlag{
 			Name:  "detach",
 			Usage: "Detach from the container",
@@ -205,7 +220,14 @@ func up(clicontext *cli.Context) error {
 		Timeout:         clicontext.Duration("timeout"),
 		SshdHost:        clicontext.String("host"),
 		ShmSize:         clicontext.Int("shm-size"),
+		NumCPU:          clicontext.String("cpus"),
+		NumMem:          clicontext.String("memory"),
+		CPUSet:          clicontext.String("cpu-set"),
 	}
+	if len(startOptions.NumCPU) > 0 && len(startOptions.CPUSet) > 0 {
+		return errors.New("`--cpus` and `--cpu-set` are mutually exclusive")
+	}
+
 	if c.Runner != types.RunnerTypeEnvdServer {
 		startOptions.EngineSource = envd.EngineSource{
 			DockerSource: &envd.DockerSource{

diff --git a/pkg/envd/docker.go b/pkg/envd/docker.go
@@ -19,6 +19,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"path/filepath"
+	"runtime"
 	"strconv"
 	"strings"
 	"time"
@@ -30,6 +31,7 @@ import (
 	"github.com/docker/docker/api/types/mount"
 	"github.com/docker/docker/client"
 	"github.com/docker/go-connections/nat"
+	dockerutils "github.com/docker/go-units"
 	"github.com/sirupsen/logrus"
 
 	envdconfig "github.com/tensorchord/envd/pkg/config"
@@ -348,6 +350,10 @@ func (e dockerEngine) StartEnvd(ctx context.Context, so StartOptions) (*StartRes
 		"tag":           so.Image,
 		"environment":   so.EnvironmentName,
 		"gpu":           so.NumGPU,
+		"shm":           so.ShmSize,
+		"cpu":           so.NumCPU,
+		"cpu-set":       so.CPUSet,
+		"memory":        so.NumMem,
 		"build-context": so.BuildContext,
 	})
 
@@ -459,6 +465,30 @@ func (e dockerEngine) StartEnvd(ctx context.Context, so StartOptions) (*StartRes
 	if so.ShmSize > 0 {
 		hostConfig.ShmSize = int64(so.ShmSize) * 1024 * 1024
 	}
+	// resource
+	if len(so.NumCPU) > 0 {
+		cpu, err := strconv.ParseFloat(so.NumCPU, 64)
+		if err != nil {
+			logger.Infof("parse `cpu` error: %v, ignore this argument", err)
+		} else if runtime.GOOS == "windows" {
+			hostConfig.NanoCPUs = int64(cpu * 10e9)
+		} else {
+			// refer to https://docs.docker.com/config/containers/resource_constraints/#configure-the-default-cfs-scheduler
+			// CPU quota and CPU period only work for UNIX platform
+			hostConfig.CPUQuota = int64(cpu * 10e5)
+		}
+	}
+	if len(so.CPUSet) > 0 {
+		hostConfig.CpusetCpus = so.CPUSet
+	}
+	if len(so.NumMem) > 0 {
+		mem, err := dockerutils.RAMInBytes(so.NumMem)
+		if err != nil {
+			logger.Infof("parse `memory` error: %v, ignore this argument", err)
+		} else {
+			hostConfig.Memory = mem
+		}
+	}
 
 	// Configure ssh port.
 	natPort := nat.Port(fmt.Sprintf("%d/tcp", envdconfig.SSHPortInContainer))

diff --git a/pkg/envd/types.go b/pkg/envd/types.go
@@ -40,6 +40,7 @@ type StartOptions struct {
 	BuildContext    string
 	NumGPU          int
 	NumCPU          string
+	CPUSet          string
 	NumMem          string
 	Timeout         time.Duration
 	ShmSize         int