refact: io.mount => runtime.mount (#861)

Signed-off-by: Keming <kemingyang@tensorchord.ai> Signed-off-by: Keming <kemingyang@tensorchord.ai>
tensorchord · Sep 9, 2022 · a0fbaa0 · a0fbaa0
1 parent 8178844
commit a0fbaa0
Show file tree

Hide file tree

Showing 9 changed files with 76 additions and 78 deletions.
diff --git a/envd/api/io/__init__.py b/envd/api/io/__init__.py
@@ -29,12 +29,3 @@ def copy(host_path: str, envd_path: str):
         host_path (str): source path in the host machine
         envd_path (str): destination path in the envd container
     """
-
-
-def mount(host_path: str, envd_path: str):
-    """Mount from host path to container path (runtime)
-
-    Args:
-        host_path (str): source path in the host machine
-        envd_path (str): destination path in the envd container
-    """
diff --git a/envd/api/runtime/__init__.py b/envd/api/runtime/__init__.py
@@ -84,3 +84,12 @@ def environ(env: Dict[str, str]):
     runtime.environ(env={"ENVD_MODE": "DEV"})
     ```
     """
+
+
+def mount(host_path: str, envd_path: str):
+    """Mount from host path to container path (runtime)
+
+    Args:
+        host_path (str): source path in the host machine
+        envd_path (str): destination path in the envd container
+    """
diff --git a/examples/dgl/build.envd b/examples/dgl/build.envd
@@ -8,8 +8,8 @@ def build():
     # Select the shell environment you like
     shell("zsh")
 
-    # io.mount(host_path="~/.envd/data/dgl", envd_path="~/.dgl")
-    io.mount(host_path=data.envd("dgl"), envd_path=data.path.dgl)
+    # runtime.mount(host_path="~/.envd/data/dgl", envd_path="~/.dgl")
+    runtime.mount(host_path=data.envd("dgl"), envd_path=data.path.dgl)
 
 def build_gpu():
     # Use ubuntu20.04 as base image and install python    
@@ -26,4 +26,4 @@ def build_gpu():
     # Select the shell environment you like
     shell("zsh")
 
-    io.mount(host_path=data.envd("dgl"), envd_path=data.path.dgl)
+    runtime.mount(host_path=data.envd("dgl"), envd_path=data.path.dgl)
diff --git a/pkg/home/data.go b/pkg/home/data.go
@@ -33,7 +33,7 @@ func (m *generalManager) InitDataDir(name string) (string, error) {
 		return newDataDir, nil
 	}
 
-	err := os.Mkdir(newDataDir, 0777) // Avoid UID/GID issues
+	err := os.MkdirAll(newDataDir, 0777) // Avoid UID/GID issues
 	if err != nil {
 		return "", errors.Wrap(err, "failed to create data dir")
 	}

diff --git a/pkg/lang/frontend/starlark/data/rule.go b/pkg/lang/frontend/starlark/data/rule.go
@@ -50,5 +50,5 @@ func ruleValueEnvdManagedDataSource(thread *starlark.Thread, _ *starlark.Builtin
 	logger.Debugf("rule `%s` is invoked, name=%s",
 		ruleEnvdManagedDataSource, name)
 
-	return NewDataSourceValue(envddata.NewEnvdManagedDataSource(name.String())), nil
+	return NewDataSourceValue(envddata.NewEnvdManagedDataSource(name.GoString())), nil
 }
diff --git a/pkg/lang/frontend/starlark/io/const.go b/pkg/lang/frontend/starlark/io/const.go
@@ -15,6 +15,5 @@
 package io
 
 const (
-	ruleCopy  = "io.copy"
-	ruleMount = "io.mount"
+	ruleCopy = "io.copy"
 )
diff --git a/pkg/lang/frontend/starlark/io/io.go b/pkg/lang/frontend/starlark/io/io.go
@@ -15,17 +15,10 @@
 package io
 
 import (
-	"os/user"
-	"path/filepath"
-	"strings"
-
-	"github.com/cockroachdb/errors"
-
 	"github.com/sirupsen/logrus"
 	"go.starlark.net/starlark"
 	"go.starlark.net/starlarkstruct"
 
-	"github.com/tensorchord/envd/pkg/lang/frontend/starlark/data"
 	"github.com/tensorchord/envd/pkg/lang/ir"
 )
 
@@ -36,64 +29,10 @@ var (
 var Module = &starlarkstruct.Module{
 	Name: "io",
 	Members: starlark.StringDict{
-		"copy":  starlark.NewBuiltin(ruleCopy, ruleFuncCopy),
-		"mount": starlark.NewBuiltin(ruleMount, ruleFuncMount),
+		"copy": starlark.NewBuiltin(ruleCopy, ruleFuncCopy),
 	},
 }
 
-func ruleFuncMount(thread *starlark.Thread, _ *starlark.Builtin,
-	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
-	var source starlark.Value
-	var destination starlark.String
-
-	if err := starlark.UnpackArgs(ruleMount, args, kwargs,
-		"host_path?", &source, "envd_path?", &destination); err != nil {
-		return nil, err
-	}
-
-	var sourceStr string
-	var err error
-
-	if v, ok := source.(*data.DataSourceValue); ok {
-		err = v.Init()
-		if err != nil {
-			return starlark.None, err
-		}
-		sourceStr, err = v.GetHostDir()
-		if err != nil {
-			return starlark.None, err
-		}
-	} else if vs, ok := source.(starlark.String); ok {
-		sourceStr = vs.GoString()
-	} else {
-		return starlark.None, errors.New("invalid data source")
-	}
-
-	destinationStr := destination.GoString()
-
-	logger.Debugf("rule `%s` is invoked, src=%s, dest=%s",
-		ruleMount, sourceStr, destinationStr)
-
-	// Expand source directory based on host user
-	usr, _ := user.Current()
-	dir := usr.HomeDir
-	if sourceStr == "~" {
-		sourceStr = dir
-	} else if strings.HasPrefix(sourceStr, "~/") {
-		sourceStr = filepath.Join(dir, sourceStr[2:])
-	}
-	// Expand dest directory based on container user envd
-	dir = "/home/envd/"
-	if destinationStr == "~" {
-		destinationStr = dir
-	} else if strings.HasPrefix(destinationStr, "~/") {
-		destinationStr = filepath.Join(dir, destinationStr[2:])
-	}
-	ir.Mount(sourceStr, destinationStr)
-
-	return starlark.None, nil
-}
-
 func ruleFuncCopy(thread *starlark.Thread, _ *starlark.Builtin,
 	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var source, destination starlark.String

diff --git a/pkg/lang/frontend/starlark/runtime/const.go b/pkg/lang/frontend/starlark/runtime/const.go
@@ -19,4 +19,5 @@ const (
 	ruleExpose  = "runtime.expose"
 	ruleDaemon  = "runtime.daemon"
 	ruleEnviron = "runtime.environ"
+	ruleMount   = "runtime.mount"
 )
diff --git a/pkg/lang/frontend/starlark/runtime/runtime.go b/pkg/lang/frontend/starlark/runtime/runtime.go
@@ -15,11 +15,16 @@
 package runtime
 
 import (
+	"os/user"
+	"path/filepath"
+	"strings"
+
 	"github.com/cockroachdb/errors"
 	"github.com/sirupsen/logrus"
 	"go.starlark.net/starlark"
 	"go.starlark.net/starlarkstruct"
 
+	"github.com/tensorchord/envd/pkg/lang/frontend/starlark/data"
 	"github.com/tensorchord/envd/pkg/lang/ir"
 )
 
@@ -34,6 +39,7 @@ var Module = &starlarkstruct.Module{
 		"daemon":  starlark.NewBuiltin(ruleDaemon, ruleFuncDaemon),
 		"expose":  starlark.NewBuiltin(ruleExpose, ruleFuncExpose),
 		"environ": starlark.NewBuiltin(ruleEnviron, ruleFuncEnviron),
+		"mount":   starlark.NewBuiltin(ruleMount, ruleFuncMount),
 	},
 }
 
@@ -139,3 +145,56 @@ func ruleFuncEnviron(thread *starlark.Thread, _ *starlark.Builtin,
 	ir.RuntimeEnviron(envMap)
 	return starlark.None, nil
 }
+
+func ruleFuncMount(thread *starlark.Thread, _ *starlark.Builtin,
+	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var source starlark.Value
+	var destination starlark.String
+
+	if err := starlark.UnpackArgs(ruleMount, args, kwargs,
+		"host_path?", &source, "envd_path?", &destination); err != nil {
+		return nil, err
+	}
+
+	var sourceStr string
+	var err error
+
+	if v, ok := source.(*data.DataSourceValue); ok {
+		err = v.Init()
+		if err != nil {
+			return starlark.None, err
+		}
+		sourceStr, err = v.GetHostDir()
+		if err != nil {
+			return starlark.None, err
+		}
+	} else if vs, ok := source.(starlark.String); ok {
+		sourceStr = vs.GoString()
+	} else {
+		return starlark.None, errors.New("invalid data source")
+	}
+
+	destinationStr := destination.GoString()
+
+	logger.Debugf("rule `%s` is invoked, src=%s, dest=%s",
+		ruleMount, sourceStr, destinationStr)
+
+	// Expand source directory based on host user
+	usr, _ := user.Current()
+	dir := usr.HomeDir
+	if sourceStr == "~" {
+		sourceStr = dir
+	} else if strings.HasPrefix(sourceStr, "~/") {
+		sourceStr = filepath.Join(dir, sourceStr[2:])
+	}
+	// Expand dest directory based on container user envd
+	dir = "/home/envd/"
+	if destinationStr == "~" {
+		destinationStr = dir
+	} else if strings.HasPrefix(destinationStr, "~/") {
+		destinationStr = filepath.Join(dir, destinationStr[2:])
+	}
+	ir.Mount(sourceStr, destinationStr)
+
+	return starlark.None, nil
+}