From d9e4e76d02d481edc237d3f9632769b2587b9a33 Mon Sep 17 00:00:00 2001
From: Keming <kemingyang@tensorchord.ai>
Date: Mon, 14 Aug 2023 13:07:06 +0800
Subject: [PATCH] feat(v1): support copy from an image

Signed-off-by: Keming <kemingyang@tensorchord.ai>
---
 envd/api/v1/io.py                      | 15 ++++++++++++---
 examples/python-basic/build.envd       |  2 +-
 pkg/lang/frontend/starlark/v1/io/io.go | 13 +++++--------
 pkg/lang/ir/types.go                   |  1 +
 pkg/lang/ir/v1/interface.go            |  3 ++-
 pkg/lang/ir/v1/system.go               |  8 +++++++-
 6 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/envd/api/v1/io.py b/envd/api/v1/io.py
index d33744364..55c44d65c 100644
--- a/envd/api/v1/io.py
+++ b/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')
+    ```
     """
 
 
diff --git a/examples/python-basic/build.envd b/examples/python-basic/build.envd
index a006b104a..5b008a357 100644
--- a/examples/python-basic/build.envd
+++ b/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 /",
diff --git a/pkg/lang/frontend/starlark/v1/io/io.go b/pkg/lang/frontend/starlark/v1/io/io.go
index 2c1b7ac58..5cc50af75 100644
--- a/pkg/lang/frontend/starlark/v1/io/io.go
+++ b/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
 }
diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go
index a92597cf8..f2a8f21a8 100644
--- a/pkg/lang/ir/types.go
+++ b/pkg/lang/ir/types.go
@@ -31,6 +31,7 @@ type RuntimeGraph struct {
 type CopyInfo struct {
 	Source      string
 	Destination string
+	Image       string
 }
 
 type MountInfo struct {
diff --git a/pkg/lang/ir/v1/interface.go b/pkg/lang/ir/v1/interface.go
index d69bf9ddb..c08bdf3b3 100644
--- a/pkg/lang/ir/v1/interface.go
+++ b/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,
 	})
 }
 
diff --git a/pkg/lang/ir/v1/system.go b/pkg/lang/ir/v1/system.go
index b555a5da2..0f5a02f91 100644
--- a/pkg/lang/ir/v1/system.go
+++ b/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)))
 	}