From 7d15da431fee8402339d35092a20178a54c6d631 Mon Sep 17 00:00:00 2001
From: Jinjing Zhou <VoVAllen@users.noreply.github.com>
Date: Sun, 8 May 2022 16:47:47 +0800
Subject: [PATCH] support runtime mount directory (#118)

Co-authored-by: Jinjing.Zhou <allenzhou@tensorchord.ai>
---
 cmd/midi/up.go       |  7 ++++++-
 pkg/docker/docker.go | 37 ++++++++++++++++++++++++++++---------
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/cmd/midi/up.go b/cmd/midi/up.go
index e31aa6cf8..46c2c98ad 100644
--- a/cmd/midi/up.go
+++ b/cmd/midi/up.go
@@ -34,6 +34,11 @@ var CommandUp = &cli.Command{
 	Aliases: []string{"u"},
 	Usage:   "build and run the MIDI environment",
 	Flags: []cli.Flag{
+		&cli.StringSliceFlag{
+			Name:    "volume",
+			Usage:   "Mount host directory into container",
+			Aliases: []string{"v"},
+		},
 		&cli.StringFlag{
 			Name:    "tag",
 			Usage:   "Name and optionally a tag in the 'name:tag' format",
@@ -95,7 +100,7 @@ func up(clicontext *cli.Context) error {
 		return err
 	}
 	containerID, containerIP, err := dockerClient.StartMIDI(
-		clicontext.Context, tag, "midi", gpu, *ir.DefaultGraph, clicontext.Duration("timeout"))
+		clicontext.Context, tag, "midi", gpu, *ir.DefaultGraph, clicontext.Duration("timeout"), clicontext.StringSlice("volume"))
 	if err != nil {
 		return err
 	}
diff --git a/pkg/docker/docker.go b/pkg/docker/docker.go
index 0876111ff..769852fd4 100644
--- a/pkg/docker/docker.go
+++ b/pkg/docker/docker.go
@@ -20,6 +20,7 @@ import (
 	"io"
 	"os"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/cockroachdb/errors"
@@ -43,7 +44,7 @@ type Client interface {
 	Load(ctx context.Context, r io.ReadCloser, quiet bool) error
 	// Start creates the container for the given tag and container name.
 	StartMIDI(ctx context.Context, tag, name string,
-		gpuEnabled bool, g ir.Graph, timeout time.Duration) (string, string, error)
+		gpuEnabled bool, g ir.Graph, timeout time.Duration, mountOptionsStr []string) (string, string, error)
 	StartBuildkitd(ctx context.Context, tag, name string) (string, error)
 	IsRunning(ctx context.Context, name string) (bool, error)
 	IsCreated(ctx context.Context, name string) (bool, error)
@@ -158,7 +159,7 @@ func (g generalClient) StartBuildkitd(ctx context.Context,
 
 // Start creates the container for the given tag and container name.
 func (c generalClient) StartMIDI(ctx context.Context, tag, name string,
-	gpuEnabled bool, g ir.Graph, timeout time.Duration) (string, string, error) {
+	gpuEnabled bool, g ir.Graph, timeout time.Duration, mountOptionsStr []string) (string, string, error) {
 	logger := logrus.WithFields(logrus.Fields{
 		"tag":       tag,
 		"container": name,
@@ -183,19 +184,37 @@ func (c generalClient) StartMIDI(ctx context.Context, tag, name string,
 	base = fmt.Sprintf("/root/%s", base)
 	config.WorkingDir = base
 
+	mountOption := make([]mount.Mount, len(mountOptionsStr)+1)
+	for i, option := range mountOptionsStr {
+		mStr := strings.Split(option, ":")
+		if len(mStr) != 2 {
+			return "", "", errors.Wrap(err, fmt.Sprintf("Invalid mount options %s", option))
+		}
+
+		logger.WithFields(logrus.Fields{
+			"mount-path":     mStr[0],
+			"container-path": mStr[1],
+		}).Debug("setting up container working directory")
+		mountOption[i] = mount.Mount{
+			Type:   mount.TypeBind,
+			Source: mStr[0],
+			Target: mStr[1],
+		}
+	}
+	mountOption[len(mountOptionsStr)] = mount.Mount{
+		Type:   mount.TypeBind,
+		Source: path,
+		Target: base,
+	}
+
 	logger.WithFields(logrus.Fields{
 		"mount-path":  path,
 		"working-dir": base,
 	}).Debug("setting up container working directory")
+
 	hostConfig := &container.HostConfig{
 		PortBindings: nat.PortMap{},
-		Mounts: []mount.Mount{
-			{
-				Type:   mount.TypeBind,
-				Source: path,
-				Target: base,
-			},
-		},
+		Mounts:       mountOption,
 	}
 	// TODO(gaocegege): Avoid specific logic to set the port.
 	if g.JupyterConfig != nil {