feat: add shm size support (#1746)

add shm size support Signed-off-by: lxb1226 <1944303766@qq.com>
tensorchord · Aug 24, 2023 · a6613c7 · a6613c7
1 parent 1aa4a8f
commit a6613c7
Show file tree

Hide file tree

Showing 16 changed files with 104 additions and 1 deletion.
diff --git a/envd/api/v0/config.py b/envd/api/v0/config.py
@@ -116,6 +116,19 @@ def gpu(count: int):
     """
 
 
+def shm_size(size: int):
+    """Configure the shared memory size (megabyte) of docker containers
+
+    Example usage:
+    ```python
+    config.shm_size(size=1024)
+    ```
+
+    Args:
+        size (int): the shared memory size (megabyte) of docker containers
+    """
+
+
 def cran_mirror(url: str):
     """Configure the mirror URL, default is https://cran.rstudio.com
 

diff --git a/envd/api/v1/config.py b/envd/api/v1/config.py
@@ -122,6 +122,19 @@ def gpu(count: int):
     """
 
 
+def shm_size(size: int):
+    """Configure the shared memory size (megabyte) of docker containers
+
+    Example usage:
+    ```python
+    config.shm_size(size=1024)
+    ```
+
+    Args:
+        size (int): the shared memory size (megabyte) of docker containers
+    """
+
+
 def cran_mirror(url: str):
     """Configure the mirror URL, default is https://cran.rstudio.com
 

diff --git a/pkg/app/up.go b/pkg/app/up.go
@@ -216,6 +216,12 @@ func up(clicontext *cli.Context) error {
 		numGPU = 1
 	}
 
+	shmSize := builder.ShmSize()
+	isSetShmSize := clicontext.IsSet("shm-size")
+	if shmSize == 0 || isSetShmSize {
+		shmSize = clicontext.Int("shm-size")
+	}
+
 	opt := envd.Options{
 		Context: c,
 	}
@@ -238,7 +244,7 @@ func up(clicontext *cli.Context) error {
 		Forced:          clicontext.Bool("force"),
 		Timeout:         clicontext.Duration("timeout"),
 		SshdHost:        clicontext.String("host"),
-		ShmSize:         clicontext.Int("shm-size"),
+		ShmSize:         shmSize,
 		NumCPU:          clicontext.String("cpus"),
 		NumMem:          clicontext.String("memory"),
 		CPUSet:          clicontext.String("cpu-set"),

diff --git a/pkg/builder/build.go b/pkg/builder/build.go
@@ -121,6 +121,10 @@ func (b generalBuilder) NumGPUs() int {
 	return b.graph.GetNumGPUs()
 }
 
+func (b generalBuilder) ShmSize() int {
+	return b.graph.GetShmSize()
+}
+
 func (b generalBuilder) Build(ctx context.Context, force bool) error {
 	if !force && !b.checkIfNeedBuild(ctx) {
 		return nil

diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go
@@ -29,5 +29,6 @@ type Builder interface {
 	Compile(ctx context.Context) (*llb.Definition, error)
 	GPUEnabled() bool
 	NumGPUs() int
+	ShmSize() int
 	GetGraph() ir.Graph
 }
diff --git a/pkg/lang/frontend/starlark/v0/config/config.go b/pkg/lang/frontend/starlark/v0/config/config.go
@@ -45,9 +45,28 @@ var Module = &starlarkstruct.Module{
 		"rstudio_server": starlark.NewBuiltin(ruleRStudioServer, ruleFuncRStudioServer),
 		"entrypoint":     starlark.NewBuiltin(ruleEntrypoint, ruleFuncEntrypoint),
 		"repo":           starlark.NewBuiltin(ruleRepo, ruleFuncRepo),
+		"shm_size":       starlark.NewBuiltin(ruleShmSize, ruleFuncShmSize),
 	},
 }
 
+func ruleFuncShmSize(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var shmSize starlark.Int
+
+	if err := starlark.UnpackArgs(ruleShmSize, args, kwargs, "size", &shmSize); err != nil {
+		return nil, err
+	}
+
+	shmSizeInt, ok := shmSize.Int64()
+	if ok {
+		ir.ShmSize(int(shmSizeInt))
+		logger.Debugf("Using %d shm size", int(shmSizeInt))
+	} else {
+		logger.Debugf("Failed to convert shm size to int64")
+	}
+
+	return starlark.None, nil
+}
+
 func ruleFuncGPU(thread *starlark.Thread, _ *starlark.Builtin,
 	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var numGPUs starlark.Int

diff --git a/pkg/lang/frontend/starlark/v0/config/const.go b/pkg/lang/frontend/starlark/v0/config/const.go
@@ -25,4 +25,5 @@ const (
 	ruleRStudioServer      = "config.rstudio_server"
 	ruleEntrypoint         = "config.entrypoint"
 	ruleRepo               = "config.repo"
+	ruleShmSize            = "config.shm_size"
 )
diff --git a/pkg/lang/frontend/starlark/v1/config/config.go b/pkg/lang/frontend/starlark/v1/config/config.go
@@ -46,9 +46,29 @@ var Module = &starlarkstruct.Module{
 		"entrypoint":     starlark.NewBuiltin(ruleEntrypoint, ruleFuncEntrypoint),
 		"repo":           starlark.NewBuiltin(ruleRepo, ruleFuncRepo),
 		"owner":          starlark.NewBuiltin(ruleOwner, ruleFuncOwner),
+		"shm_size":       starlark.NewBuiltin(ruleShmSize, ruleFuncShmSize),
 	},
 }
 
+func ruleFuncShmSize(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var shmSize starlark.Int
+
+	if err := starlark.UnpackArgs(ruleShmSize, args, kwargs, "size", &shmSize); err != nil {
+		return nil, err
+	}
+
+	shmSizeInt, ok := shmSize.Int64()
+
+	if ok {
+		ir.ShmSize(int(shmSizeInt))
+		logger.Debugf("Using %d shm size", int(shmSizeInt))
+	} else {
+		logger.Debugf("Failed to convert shm size to int64")
+	}
+
+	return starlark.None, nil
+}
+
 func ruleFuncGPU(thread *starlark.Thread, _ *starlark.Builtin,
 	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var numGPUs starlark.Int

diff --git a/pkg/lang/frontend/starlark/v1/config/const.go b/pkg/lang/frontend/starlark/v1/config/const.go
@@ -26,4 +26,5 @@ const (
 	ruleEntrypoint         = "config.entrypoint"
 	ruleRepo               = "config.repo"
 	ruleOwner              = "config.owner"
+	ruleShmSize            = "config.shm_size"
 )
diff --git a/pkg/lang/ir/graph.go b/pkg/lang/ir/graph.go
@@ -43,6 +43,7 @@ type graphVisitor interface {
 	GetDepsFiles(deps []string) []string
 	GPUEnabled() bool
 	GetNumGPUs() int
+	GetShmSize() int
 	Labels() (map[string]string, error)
 	ExposedPorts() (map[string]struct{}, error)
 	GetEntrypoint(buildContextDir string) ([]string, error)

diff --git a/pkg/lang/ir/v0/compile.go b/pkg/lang/ir/v0/compile.go
@@ -56,6 +56,7 @@ func NewGraph() ir.Graph {
 		CUDA:    nil,
 		CUDNN:   CUDNNVersionDefault,
 		NumGPUs: 0,
+		ShmSize: 0,
 
 		PyPIPackages:    [][]string{},
 		RPackages:       []string{},
@@ -73,6 +74,10 @@ func NewGraph() ir.Graph {
 
 var DefaultGraph = NewGraph()
 
+func (g generalGraph) GetShmSize() int {
+	return g.ShmSize
+}
+
 func (g generalGraph) GetNumGPUs() int {
 	return g.NumGPUs
 }

diff --git a/pkg/lang/ir/v0/interface.go b/pkg/lang/ir/v0/interface.go
@@ -75,6 +75,12 @@ func SystemPackage(deps []string) {
 	g.SystemPackages = append(g.SystemPackages, deps...)
 }
 
+func ShmSize(shmSize int) {
+	g := DefaultGraph.(*generalGraph)
+
+	g.ShmSize = shmSize
+}
+
 func GPU(numGPUs int) {
 	g := DefaultGraph.(*generalGraph)
 

diff --git a/pkg/lang/ir/v0/types.go b/pkg/lang/ir/v0/types.go
@@ -39,6 +39,7 @@ type generalGraph struct {
 	CUDA    *string
 	CUDNN   string
 	NumGPUs int
+	ShmSize int
 
 	UbuntuAPTSource    *string
 	CRANMirrorURL      *string

diff --git a/pkg/lang/ir/v1/compile.go b/pkg/lang/ir/v1/compile.go
@@ -49,6 +49,7 @@ func NewGraph() ir.Graph {
 		CUDA:              nil,
 		CUDNN:             CUDNNVersionDefault,
 		NumGPUs:           0,
+		ShmSize:           0,
 		EnvdSyntaxVersion: "v1",
 
 		PyPIPackages:    [][]string{},
@@ -74,6 +75,10 @@ func (g generalGraph) GetHTTP() []ir.HTTPInfo {
 	return g.HTTP
 }
 
+func (g generalGraph) GetShmSize() int {
+	return g.ShmSize
+}
+
 func (g generalGraph) GetNumGPUs() int {
 	return g.NumGPUs
 }

diff --git a/pkg/lang/ir/v1/interface.go b/pkg/lang/ir/v1/interface.go
@@ -121,6 +121,12 @@ func SystemPackage(deps []string) {
 	g.SystemPackages = append(g.SystemPackages, deps...)
 }
 
+func ShmSize(shmSize int) {
+	g := DefaultGraph.(*generalGraph)
+
+	g.ShmSize = shmSize
+}
+
 func GPU(numGPUs int) {
 	g := DefaultGraph.(*generalGraph)
 

diff --git a/pkg/lang/ir/v1/types.go b/pkg/lang/ir/v1/types.go
@@ -40,6 +40,7 @@ type generalGraph struct {
 	CUDA    *string
 	CUDNN   string
 	NumGPUs int
+	ShmSize int
 
 	UbuntuAPTSource    *string
 	CRANMirrorURL      *string