feat(v1): support copy from an image

Signed-off-by: Keming <kemingyang@tensorchord.ai>

envd/api/v1/io.py

@@ -30,12 +30,21 @@
 from typing import Optional
 
 
-def copy(host_path: str, envd_path: str):
+def copy(source: str, target: str, image: Optional[str]):
     """Copy from host path to container path (build time)
 
     Args:
-        host_path (str): source path in the host machine
-        envd_path (str): destination path in the envd container
+        source (str): source path in the host machine or in the ``image``
+        target (str): destination path in the envd container
+        image(Optional[str]): image name, if not specified, will use the host
+
+    Examples:
+    ```
+    # copy from host to container
+    io.copy(source='main.py', target='/home/envd/')
+    # copy from image to container
+    io.copy(source='/bin/micromamba', target='/usr/local/bin/micromamba', image='mambaorg/micromamba:1.0.0')
+    ```
     """
 
 

examples/python-basic/build.envd

@@ -6,7 +6,7 @@ def build():
             "via",
         ]
     )
-    io.copy(host_path="./build.envd", envd_path="/")
+    io.copy("./build.envd", "/")
     runtime.command(
         commands={
             "test": "ls /",

pkg/lang/frontend/starlark/v1/io/io.go

@@ -36,19 +36,16 @@ var Module = &starlarkstruct.Module{
 
 func ruleFuncCopy(thread *starlark.Thread, _ *starlark.Builtin,
 	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
-	var source, destination starlark.String
+	var source, destination, image string
 
 	if err := starlark.UnpackArgs(ruleCopy, args, kwargs,
-		"host_path?", &source, "envd_path?", &destination); err != nil {
+		"source", &source, "target", &destination, "image?", &image); err != nil {
 		return nil, err
 	}
 
-	sourceStr := source.GoString()
-	destinationStr := destination.GoString()
-
-	logger.Debugf("rule `%s` is invoked, src=%s, dest=%s\n",
-		ruleCopy, sourceStr, destinationStr)
-	ir.Copy(sourceStr, destinationStr)
+	logger.Debugf("rule `%s` is invoked, src=%s, dest=%s, image=%s\n",
+		ruleCopy, source, destination, image)
+	ir.Copy(source, destination, image)
 
 	return starlark.None, nil
 }

pkg/lang/ir/types.go

@@ -31,6 +31,7 @@ type RuntimeGraph struct {
 type CopyInfo struct {
 	Source      string
 	Destination string
+	Image       string
 }
 
 type MountInfo struct {

pkg/lang/ir/v1/interface.go

@@ -261,12 +261,13 @@ func CondaPackage(deps []string, channel []string, envFile string) error {
 	return nil
 }
 
-func Copy(src, dest string) {
+func Copy(src, dest, image string) {
 	g := DefaultGraph.(*generalGraph)
 
 	g.Copy = append(g.Copy, ir.CopyInfo{
 		Source:      src,
 		Destination: dest,
+		Image:       image,
 	})
 }
 

pkg/lang/ir/v1/system.go

@@ -172,8 +172,14 @@ func (g generalGraph) compileCopy(root llb.State) llb.State {
 	result := root
 	// Compose the copy command.
 	for _, c := range g.Copy {
+		var from llb.State
+		if c.Image == "" {
+			from = llb.Local(flag.FlagBuildContext)
+		} else {
+			from = llb.Image(c.Image)
+		}
 		result = result.File(llb.Copy(
-			llb.Local(flag.FlagBuildContext), c.Source, c.Destination,
+			from, c.Source, c.Destination,
 			&llb.CopyInfo{CreateDestPath: true},
 			llb.WithUIDGID(g.uid, g.gid)))
 	}