Repository: tensorchord/envd (language: python)

Issue:
feat: Add envd-server runtime proposal
Ref #179 

Proposal preview: https://github.com/gaocegege/envd/blob/proposal/docs/proposals/20220603-kubernetes-vendor.md


Discussion / hints:
/cc @Xuanwo @hezhizhen @knight42 
Other options for file syncing:
- reverse sshfs https://github.com/lima-vm/sshocker
- ksync https://github.com/ksync/ksync
> One thing I'm concerning is that whether syncing is needed at MVP. Code should be on a PV-like thing I think and usually people works with git

My 2 cents is syncing might be needed ultimately, but might in some different ways. Let's say we are working with git, and push some new commits to a branch, I think it would be better if envd could pull the new commits automatically to simulate the local development experience.
> My 2 cents is syncing might be needed ultimately, but might in some different ways. Let's say we are working with git, and push some new commits to a branch, I think it would be better if envd could pull the new commits automatically to simulate the local development experience.

Thanks for the advice! Automatic push/pull looks magic to me. And it is complex. If the container is crashed, we also may lost the commits if we do not run push.
As discussed with some infra engineers interested in envd, port-forwarding may consume many API server CPUs. And tools like virtual kubelet does not support port forwarding.
> port-forwarding may consume many API server CPUs.

Would you mind elaborating? AFAIK if there is no much traffic, port-forwarding should not consume too much cpu resources, as it is simply a SPDY connection under the hood.

> virtual kubelet does not support port forwarding

Indeed. But what are we going to do to access the services inside the container without port-forwarding?
> Thanks for the advice! Automatic push/pull looks magic to me. And it is complex. If the container is crashed, we also may lost the commits if we do not run push.

Or, we can forget the sync things. We request users to develop on the remote container, instead of the host. The build.envd may look like this:

```diff
def build():
    base(os="ubuntu20.04", language="python3")
    install.vscode_extensions([
        "ms-python.python",
    ])
    #config.pip_index(url = "https://pypi.tuna.tsinghua.edu.cn/simple")
    install.python_packages([
        "tensorflow",
        "numpy",
    ])
    shell("zsh")
    config.jupyter(password="", port=8888)
+    config.working_dir(local=".", remote="https://github.com/tensorchord/envd.git")
```

envd mounts the local dir with docker runner and downloads the repo with Kubernetes runner.
> Would you mind elaborating? AFAIK if there is no much traffic, port-forwarding should not consume too much cpu resources, as it is simply a SPDY connection under the hood.

There should not be huge traffic by design. But algorithm engineers may use it to copy data:

```
scp <10G-file> container:~
```
> But what are we going to do to access the services inside the container without port-forwarding?

They may use service and ingress to achieve this. Thus we may need a mechanism to support customization here.

Maybe just like the design of the device plugin, we provide an interface to communicate between envd and a CLI shim. The shim does the critical logic like port forwarding. The envd just communicate with the shim and show information to users.

Port forwarding can be used in our default shim, while users can write their own shim to customize, e.g. using service and ingress.
> > But what are we going to do to access the services inside the container without port-forwarding?
> 
> They may use service and ingress to achieve this. Thus we may need a mechanism to support customization here.
> 
> Maybe just like the design of the device plugin, we provide an interface to communicate between envd and a CLI shim. The shim does the critical logic like port forwarding. The envd just communicate with the shim and show information to users.
> 
> Port forwarding can be used in our default shim, while users can write their own shim to customize, e.g. using service and ingress.

Device plugin is a not-so-good comparison. Let's say kubectl plugin mechanism. The shim maintained by users can be integrated into the envd. 
> +    config.working_dir(local=".", remote="https://github.com/tensorchord/envd.git")

My concern is that working dir seems to be a command line argument to me, otherwise it might prevent the reuse of `build.envd` in different working dir, just like we don't specify the build context in Dockerfile. Besides if we need to specify the repo address, should we need to specify the branch as well?

> algorithm engineers may use it to copy data:

Got it 👌 If we need to transfer such huge file via port-forwarding without rate limiting, the functionality of apiserver might be affected.

> while users can write their own shim to customize, e.g. using service and ingress.

I think it make sense 👍 

> My concern is that working dir seems to be a command line argument to me, otherwise it might prevent the reuse of build.envd in different working dir, just like we don't specify the build context in Dockerfile. Besides if we need to specify the repo address, should we need to specify the branch as well?

Sounds reasonable. It should be a runtime argument instead of build time.

@VoVAllen Do you have opinion on it?


Agree it better to be a runtime option. If using git, user should handle the sync related thing by himself (clone git repo and git pull/push). Therefore  `config.working_dir` might not be needed.

Fix patch:
diff --git a/README.md b/README.md
index 8518a457f..6a986e325 100644
--- a/README.md
+++ b/README.md
@@ -187,17 +187,6 @@ $ cd envd-quick-start && envd up
 (envd) ➜  demo git:(master) ✗ # You are in the container-based environment!
 ```
 
-### Play with the environment
-
-You can run `ssh envd-quick-start.envd` to reconnect if you exit from the environment. Or you can execute `git` or `python` commands inside.
-
-```bash
-$ python demo.py
-[2 3 4]
-$ git fetch
-$
-```
-
 ### Set up Jupyter notebook
 
 Please edit the `build.envd` to enable jupyter notebook:
@@ -209,7 +198,7 @@ def build():
         "numpy",
     ])
     shell("zsh")
-    config.jupyter(password="", port=8888)
+    config.jupyter(password="")
 ```
 
 You can get the endpoint of the running Jupyter notebook via `envd get envs`.
@@ -218,7 +207,7 @@ You can get the endpoint of the running Jupyter notebook via `envd get envs`.
 $ envd up --detach
 $ envd get env
 NAME                    JUPYTER                 SSH TARGET              CONTEXT                                 IMAGE                   GPU     CUDA    CUDNN   STATUS          CONTAINER ID 
-envd-quick-start        http://localhost:8888   envd-quick-start.envd   /home/gaocegege/code/envd-quick-start   envd-quick-start:dev    false   <none>  <none>  Up 54 seconds   bd3f6a729e94
+envd-quick-start        http://localhost:42779   envd-quick-start.envd   /home/gaocegege/code/envd-quick-start   envd-quick-start:dev    false   <none>  <none>  Up 54 seconds   bd3f6a729e94
 ```
 
 ## Roadmap 🗂️
diff --git a/pkg/app/build.go b/pkg/app/build.go
index a4e578c93..a383c6f89 100644
--- a/pkg/app/build.go
+++ b/pkg/app/build.go
@@ -15,6 +15,7 @@
 package app
 
 import (
+	"fmt"
 	"path/filepath"
 
 	"github.com/cockroachdb/errors"
@@ -94,7 +95,7 @@ func build(clicontext *cli.Context) error {
 	tag := clicontext.String("tag")
 	if tag == "" {
 		logrus.Debug("tag not specified, using default")
-		tag = fileutil.Base(buildContext)
+		tag = fmt.Sprintf("%s:%s", fileutil.Base(buildContext), "dev")
 	}
 
 	logger := logrus.WithFields(logrus.Fields{
diff --git a/pkg/app/up.go b/pkg/app/up.go
index e8a89980a..a5034c1a9 100644
--- a/pkg/app/up.go
+++ b/pkg/app/up.go
@@ -146,6 +146,7 @@ func up(clicontext *cli.Context) error {
 	if err := builder.Build(clicontext.Context, clicontext.Path("public-key")); err != nil {
 		return errors.Wrap(err, "failed to build the image")
 	}
+	// Do not attach GPU if the flag is set.
 	gpuEnable := clicontext.Bool("no-gpu")
 	var gpu bool
 	if gpuEnable {
@@ -169,14 +170,14 @@ func up(clicontext *cli.Context) error {
 		}
 	}
 
-	sshPort, err := netutil.GetFreePort()
+	sshPortInHost, err := netutil.GetFreePort()
 	if err != nil {
 		return errors.Wrap(err, "failed to get a free port")
 	}
 	numGPUs := builder.NumGPUs()
 
 	containerID, containerIP, err := dockerClient.StartEnvd(clicontext.Context,
-		tag, ctr, buildContext, gpu, numGPUs, sshPort, *ir.DefaultGraph, clicontext.Duration("timeout"),
+		tag, ctr, buildContext, gpu, numGPUs, sshPortInHost, *ir.DefaultGraph, clicontext.Duration("timeout"),
 		clicontext.StringSlice("volume"))
 	if err != nil {
 		return errors.Wrap(err, "failed to start the envd environment")
@@ -185,14 +186,14 @@ func up(clicontext *cli.Context) error {
 
 	logrus.Debugf("Add entry %s to SSH config. at %s", buildContext, containerIP)
 	if err = sshconfig.AddEntry(
-		ctr, localhost, sshPort, clicontext.Path("private-key")); err != nil {
+		ctr, localhost, sshPortInHost, clicontext.Path("private-key")); err != nil {
 		logrus.Infof("failed to add entry %s to your SSH config file: %s", ctr, err)
 		return errors.Wrap(err, "failed to add entry to your SSH config file")
 	}
 
 	if !detach {
 		sshClient, err := ssh.NewClient(
-			localhost, "envd", sshPort, true, clicontext.Path("private-key"), "")
+			localhost, "envd", sshPortInHost, true, clicontext.Path("private-key"), "")
 		if err != nil {
 			return errors.Wrap(err, "failed to create the ssh client")
 		}
diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go
index ca0321d05..60a13e617 100644
--- a/pkg/builder/builder.go
+++ b/pkg/builder/builder.go
@@ -159,13 +159,22 @@ func (b generalBuilder) compile(ctx context.Context, pub string) (*llb.Definitio
 	return def, nil
 }
 
-func (b generalBuilder) labels(ctx context.Context) (string, error) {
+func (b generalBuilder) imageConfig(ctx context.Context) (string, error) {
 	labels, err := ir.Labels()
 	if err != nil {
 		return "", errors.Wrap(err, "failed to get labels")
 	}
+	ports, err := ir.ExposedPorts()
+	if err != nil {
+		return "", errors.Wrap(err, "failed to get expose ports")
+	}
 	labels[types.ImageLabelContext] = b.buildContextDir
-	data, err := ImageConfigStr(labels)
+
+	ep, err := ir.Entrypoint(b.buildContextDir)
+	if err != nil {
+		return "", errors.Wrap(err, "failed to get entrypoint")
+	}
+	data, err := ImageConfigStr(labels, ports, ep)
 	if err != nil {
 		return "", errors.Wrap(err, "failed to get image config")
 	}
@@ -173,7 +182,7 @@ func (b generalBuilder) labels(ctx context.Context) (string, error) {
 }
 
 func (b generalBuilder) build(ctx context.Context, def *llb.Definition, pw progresswriter.Writer) error {
-	labels, err := b.labels(ctx)
+	imageConfig, err := b.imageConfig(ctx)
 	if err != nil {
 		return errors.Wrap(err, "failed to get labels")
 	}
@@ -198,7 +207,7 @@ func (b generalBuilder) build(ctx context.Context, def *llb.Definition, pw progr
 						Attrs: map[string]string{
 							"name": b.tag,
 							// Ref https://github.com/r2d4/mockerfile/blob/140c6a912bbfdae220febe59ab535ef0acba0e1f/pkg/build/build.go#L65
-							"containerimage.config": labels,
+							"containerimage.config": imageConfig,
 						},
 						Output: func(map[string]string) (io.WriteCloser, error) {
 							return pipeW, nil
diff --git a/pkg/builder/util.go b/pkg/builder/util.go
index 937400e0c..46d790a24 100644
--- a/pkg/builder/util.go
+++ b/pkg/builder/util.go
@@ -29,13 +29,16 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
-func ImageConfigStr(labels map[string]string) (string, error) {
+func ImageConfigStr(labels map[string]string,
+	ports map[string]struct{}, entrypoint []string) (string, error) {
 	pl := platforms.Normalize(platforms.DefaultSpec())
 	img := v1.Image{
 		Config: v1.ImageConfig{
-			Labels:     labels,
-			WorkingDir: "/",
-			Env:        []string{"PATH=" + DefaultPathEnv(pl.OS)},
+			Labels:       labels,
+			WorkingDir:   "/",
+			Env:          []string{"PATH=" + DefaultPathEnv(pl.OS)},
+			ExposedPorts: ports,
+			Entrypoint:   entrypoint,
 		},
 		Architecture: pl.Architecture,
 		// Refer to https://github.com/tensorchord/envd/issues/269#issuecomment-1152944914
diff --git a/pkg/config/config.go b/pkg/config/config.go
index 5301dd06b..8e2174a0b 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -19,4 +19,6 @@ const (
 	PrivateKeyFile              = "id_rsa_envd"
 	PublicKeyFile               = "id_rsa_envd.pub"
 	ContainerAuthorizedKeysPath = "/var/envd/authorized_keys"
+	SSHPortInContainer          = 2222
+	JupyterPortInContainer      = 8888
 )
diff --git a/pkg/docker/docker.go b/pkg/docker/docker.go
index 614ca3c49..cfe45782f 100644
--- a/pkg/docker/docker.go
+++ b/pkg/docker/docker.go
@@ -35,8 +35,10 @@ import (
 	"github.com/moby/term"
 	"github.com/sirupsen/logrus"
 
+	envdconfig "github.com/tensorchord/envd/pkg/config"
 	"github.com/tensorchord/envd/pkg/lang/ir"
 	"github.com/tensorchord/envd/pkg/util/fileutil"
+	"github.com/tensorchord/envd/pkg/util/netutil"
 )
 
 const (
@@ -305,7 +307,7 @@ func (c generalClient) StartBuildkitd(ctx context.Context, tag, name, mirror str
 
 // StartEnvd creates the container for the given tag and container name.
 func (c generalClient) StartEnvd(ctx context.Context, tag, name, buildContext string,
-	gpuEnabled bool, numGPUs int, sshPort int, g ir.Graph, timeout time.Duration,
+	gpuEnabled bool, numGPUs int, sshPortInHost int, g ir.Graph, timeout time.Duration,
 	mountOptionsStr []string) (string, string, error) {
 	logger := logrus.WithFields(logrus.Fields{
 		"tag":           tag,
@@ -315,21 +317,13 @@ func (c generalClient) StartEnvd(ctx context.Context, tag, name, buildContext st
 		"build-context": buildContext,
 	})
 	config := &container.Config{
-		Image: tag,
-		User:  "envd",
-		Entrypoint: []string{
-			"tini",
-			"--",
-			"bash",
-			"-c",
-		},
+		Image:        tag,
+		User:         "envd",
 		ExposedPorts: nat.PortSet{},
 	}
 	base := fileutil.Base(buildContext)
 	base = filepath.Join("/home/envd", base)
 	config.WorkingDir = base
-	config.Entrypoint = append(config.Entrypoint,
-		entrypointSH(g, config.WorkingDir, sshPort))
 
 	mountOption := make([]mount.Mount, len(mountOptionsStr)+1)
 	for i, option := range mountOptionsStr {
@@ -365,22 +359,27 @@ func (c generalClient) StartEnvd(ctx context.Context, tag, name, buildContext st
 	}
 
 	// Configure ssh port.
-	natPort := nat.Port(fmt.Sprintf("%d/tcp", sshPort))
+	natPort := nat.Port(fmt.Sprintf("%d/tcp", envdconfig.SSHPortInContainer))
 	hostConfig.PortBindings[natPort] = []nat.PortBinding{
 		{
 			HostIP:   localhost,
-			HostPort: strconv.Itoa(sshPort),
+			HostPort: strconv.Itoa(sshPortInHost),
 		},
 	}
-	config.ExposedPorts[natPort] = struct{}{}
 
+	var jupyterPortInHost int
 	// TODO(gaocegege): Avoid specific logic to set the port.
 	if g.JupyterConfig != nil {
-		natPort := nat.Port(fmt.Sprintf("%d/tcp", g.JupyterConfig.Port))
+		var err error
+		jupyterPortInHost, err = netutil.GetFreePort()
+		if err != nil {
+			return "", "", errors.Wrap(err, "failed to get a free port")
+		}
+		natPort := nat.Port(fmt.Sprintf("%d/tcp", envdconfig.JupyterPortInContainer))
 		hostConfig.PortBindings[natPort] = []nat.PortBinding{
 			{
 				HostIP:   localhost,
-				HostPort: strconv.Itoa(int(g.JupyterConfig.Port)),
+				HostPort: strconv.Itoa(jupyterPortInHost),
 			},
 		}
 		config.ExposedPorts[natPort] = struct{}{}
@@ -391,7 +390,7 @@ func (c generalClient) StartEnvd(ctx context.Context, tag, name, buildContext st
 		hostConfig.DeviceRequests = deviceRequests(numGPUs)
 	}
 
-	config.Labels = labels(name, g.JupyterConfig, sshPort)
+	config.Labels = labels(name, g.JupyterConfig, sshPortInHost, jupyterPortInHost)
 
 	logger = logger.WithFields(logrus.Fields{
 		"entrypoint":  config.Entrypoint,
@@ -408,7 +407,8 @@ func (c generalClient) StartEnvd(ctx context.Context, tag, name, buildContext st
 		logger.Warnf("run with warnings: %s", w)
 	}
 
-	if err := c.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{}); err != nil {
+	if err := c.ContainerStart(
+		ctx, resp.ID, types.ContainerStartOptions{}); err != nil {
 		errCause := errors.UnwrapAll(err)
 		// Hack to check if the port is already allocated.
 		if strings.Contains(errCause.Error(), "port is already allocated") {
diff --git a/pkg/docker/entrypoint.go b/pkg/docker/entrypoint.go
deleted file mode 100644
index f897cdb35..000000000
--- a/pkg/docker/entrypoint.go
+++ /dev/null
@@ -1,42 +0,0 @@
-// Copyright 2022 The envd Authors
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package docker
-
-import (
-	"fmt"
-	"strings"
-
-	"github.com/tensorchord/envd/pkg/config"
-	"github.com/tensorchord/envd/pkg/editor/jupyter"
-	"github.com/tensorchord/envd/pkg/lang/ir"
-)
-
-const (
-	template = `set -e
-/var/envd/bin/envd-ssh --authorized-keys %s --port %d --shell %s &
-%s
-wait -n`
-)
-
-func entrypointSH(g ir.Graph, workingDir string, sshPort int) string {
-	if g.JupyterConfig != nil {
-		cmds := jupyter.GenerateCommand(g, workingDir)
-		return fmt.Sprintf(template,
-			config.ContainerAuthorizedKeysPath, sshPort, g.Shell,
-			strings.Join(cmds, " "))
-	}
-	return fmt.Sprintf(template,
-		config.ContainerAuthorizedKeysPath, sshPort, g.Shell, "")
-}
diff --git a/pkg/docker/label.go b/pkg/docker/label.go
index 5167772c2..4c07bf7af 100644
--- a/pkg/docker/label.go
+++ b/pkg/docker/label.go
@@ -24,12 +24,14 @@ import (
 	"github.com/tensorchord/envd/pkg/types"
 )
 
-func labels(name string, jupyterConfig *ir.JupyterConfig, sshPort int) map[string]string {
+func labels(name string, jupyterConfig *ir.JupyterConfig,
+	sshPortInHost, jupyterPortInHost int) map[string]string {
 	res := make(map[string]string)
 	res[types.ContainerLabelName] = name
-	res[types.ContainerLabelSSHPort] = strconv.Itoa(sshPort)
+	res[types.ContainerLabelSSHPort] = strconv.Itoa(sshPortInHost)
 	if jupyterConfig != nil {
-		res[types.ContainerLabelJupyterAddr] = fmt.Sprintf("http://localhost:%d", jupyterConfig.Port)
+		res[types.ContainerLabelJupyterAddr] =
+			fmt.Sprintf("http://localhost:%d", jupyterPortInHost)
 	}
 	return res
 }
diff --git a/pkg/editor/jupyter/util.go b/pkg/editor/jupyter/util.go
deleted file mode 100644
index 3aea62e52..000000000
--- a/pkg/editor/jupyter/util.go
+++ /dev/null
@@ -1,53 +0,0 @@
-// Copyright 2022 The envd Authors
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package jupyter
-
-import (
-	"strconv"
-
-	"github.com/tensorchord/envd/pkg/lang/ir"
-)
-
-func GenerateCommand(g ir.Graph, notebookDir string) []string {
-	if g.JupyterConfig == nil {
-		return nil
-	}
-
-	var cmd []string
-	// Use python in conda env.
-	if g.CondaEnabled() {
-		cmd = append(cmd, "/opt/conda/bin/python3")
-	} else {
-		cmd = append(cmd, "python3")
-	}
-
-	cmd = append(cmd, []string{
-		"-m", "notebook",
-		"--ip", "0.0.0.0", "--notebook-dir", notebookDir,
-	}...)
-
-	if g.JupyterConfig.Password != "" {
-		cmd = append(cmd, "--NotebookApp.password", g.JupyterConfig.Password,
-			"--NotebookApp.token", "''")
-	} else {
-		cmd = append(cmd, "--NotebookApp.password", "''",
-			"--NotebookApp.token", "''")
-	}
-	if g.JupyterConfig.Port != 0 {
-		p := strconv.Itoa(int(g.JupyterConfig.Port))
-		cmd = append(cmd, "--port", p)
-	}
-	return cmd
-}
diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go
index 5af5e5b79..4187a2906 100644
--- a/pkg/lang/ir/compile.go
+++ b/pkg/lang/ir/compile.go
@@ -17,14 +17,18 @@ package ir
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"os"
+	"strings"
 
 	"github.com/cockroachdb/errors"
 	"github.com/moby/buildkit/client/llb"
 	"github.com/sirupsen/logrus"
 
+	"github.com/tensorchord/envd/pkg/config"
 	"github.com/tensorchord/envd/pkg/progress/compileui"
 	"github.com/tensorchord/envd/pkg/types"
+	"github.com/tensorchord/envd/pkg/util/fileutil"
 )
 
 func NewGraph() *Graph {
@@ -81,6 +85,14 @@ func Labels() (map[string]string, error) {
 	return DefaultGraph.Labels()
 }
 
+func ExposedPorts() (map[string]struct{}, error) {
+	return DefaultGraph.ExposedPorts()
+}
+
+func Entrypoint(buildContextDir string) ([]string, error) {
+	return DefaultGraph.Entrypoint(buildContextDir)
+}
+
 func (g Graph) GPUEnabled() bool {
 	return g.CUDA != nil
 }
@@ -114,6 +126,43 @@ func (g Graph) Labels() (map[string]string, error) {
 	return labels, nil
 }
 
+func (g Graph) ExposedPorts() (map[string]struct{}, error) {
+	ports := make(map[string]struct{})
+	ports[fmt.Sprintf("%d/tcp", config.SSHPortInContainer)] = struct{}{}
+	if g.JupyterConfig != nil {
+		ports[fmt.Sprintf("%d/tcp", config.JupyterPortInContainer)] = struct{}{}
+	}
+	return ports, nil
+}
+
+func (g Graph) Entrypoint(buildContextDir string) ([]string, error) {
+	ep := []string{
+		"tini",
+		"--",
+		"bash",
+		"-c",
+	}
+
+	template := `set -e
+	/var/envd/bin/envd-ssh --authorized-keys %s --port %d --shell %s &
+	%s
+	wait -n`
+
+	if g.JupyterConfig != nil {
+		workingDir := fmt.Sprintf("/home/envd/%s", fileutil.Base(buildContextDir))
+		jupyterCmd := g.generateJupyterCommand(workingDir)
+		cmd := fmt.Sprintf(template,
+			config.ContainerAuthorizedKeysPath, config.SSHPortInContainer, g.Shell,
+			strings.Join(jupyterCmd, " "))
+		ep = append(ep, cmd)
+		return ep, nil
+	}
+	cmd := fmt.Sprintf(template,
+		config.ContainerAuthorizedKeysPath, config.SSHPortInContainer, g.Shell, "")
+	ep = append(ep, cmd)
+	return ep, nil
+}
+
 func (g Graph) Compile(uid, gid int) (llb.State, error) {
 	g.uid = uid
 
@@ -156,140 +205,3 @@ func (g Graph) Compile(uid, gid int) (llb.State, error) {
 	g.Writer.Finish()
 	return finalStage, nil
 }
-
-func (g Graph) compileJulia(aptStage llb.State) (llb.State, error) {
-	g.compileJupyter()
-	builtinSystemStage := aptStage
-
-	sshStage, err := g.copySSHKey(builtinSystemStage)
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to copy ssh keys")
-	}
-	diffSSHStage := llb.Diff(builtinSystemStage, sshStage, llb.WithCustomName("install ssh keys"))
-
-	shellStage, err := g.compileShell(builtinSystemStage)
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to compile shell")
-	}
-	diffShellStage := llb.Diff(builtinSystemStage, shellStage, llb.WithCustomName("install shell"))
-
-	systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage),
-		llb.WithCustomName("install system packages"))
-
-	juliaStage := llb.Diff(builtinSystemStage,
-		g.installJuliaPackages(builtinSystemStage), llb.WithCustomName("install julia packages"))
-
-	vscodeStage, err := g.compileVSCode()
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to get vscode plugins")
-	}
-
-	var merged llb.State
-	if vscodeStage != nil {
-		merged = llb.Merge([]llb.State{
-			builtinSystemStage, systemStage, diffShellStage,
-			diffSSHStage, juliaStage, *vscodeStage,
-		}, llb.WithCustomName("merging all components into one"))
-	} else {
-		merged = llb.Merge([]llb.State{
-			builtinSystemStage, systemStage, diffShellStage,
-			diffSSHStage, juliaStage,
-		}, llb.WithCustomName("merging all components into one"))
-	}
-	return merged, nil
-}
-
-func (g Graph) compileRLang(aptStage llb.State) (llb.State, error) {
-	g.compileJupyter()
-	builtinSystemStage := aptStage
-
-	sshStage, err := g.copySSHKey(builtinSystemStage)
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to copy ssh keys")
-	}
-	diffSSHStage := llb.Diff(builtinSystemStage, sshStage, llb.WithCustomName("install ssh keys"))
-
-	// Conda affects shell and python, thus we cannot do it in parallel.
-	shellStage, err := g.compileShell(builtinSystemStage)
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to compile shell")
-	}
-	diffShellStage := llb.Diff(builtinSystemStage, shellStage, llb.WithCustomName("install shell"))
-
-	systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage),
-		llb.WithCustomName("install system packages"))
-
-	// TODO(terrytangyuan): Support RStudio local server
-	rPackageInstallStage := llb.Diff(builtinSystemStage,
-		g.installRPackages(builtinSystemStage), llb.WithCustomName("install R packages"))
-
-	vscodeStage, err := g.compileVSCode()
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to get vscode plugins")
-	}
-
-	var merged llb.State
-	if vscodeStage != nil {
-		merged = llb.Merge([]llb.State{
-			builtinSystemStage, systemStage, diffShellStage,
-			diffSSHStage, rPackageInstallStage, *vscodeStage,
-		}, llb.WithCustomName("merging all components into one"))
-	} else {
-		merged = llb.Merge([]llb.State{
-			builtinSystemStage, systemStage, diffShellStage,
-			diffSSHStage, rPackageInstallStage,
-		}, llb.WithCustomName("merging all components into one"))
-	}
-	return merged, nil
-}
-
-func (g Graph) compilePython(aptStage llb.State) (llb.State, error) {
-	condaChanelStage := g.compileCondaChannel(aptStage)
-	pypiMirrorStage := g.compilePyPIIndex(condaChanelStage)
-
-	g.compileJupyter()
-	builtinSystemStage := pypiMirrorStage
-
-	sshStage, err := g.copySSHKey(builtinSystemStage)
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to copy ssh keys")
-	}
-	diffSSHStage := llb.Diff(builtinSystemStage, sshStage, llb.WithCustomName("install ssh keys"))
-
-	// Conda affects shell and python, thus we cannot do it in parallel.
-	shellStage, err := g.compileShell(builtinSystemStage)
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to compile shell")
-	}
-
-	condaEnvStage := g.setCondaENV(shellStage)
-
-	condaStage := llb.Diff(builtinSystemStage,
-		g.compileCondaPackages(condaEnvStage),
-		llb.WithCustomName("install conda packages"))
-
-	pypiStage := llb.Diff(condaEnvStage,
-		g.compilePyPIPackages(condaEnvStage),
-		llb.WithCustomName("install PyPI packages"))
-	systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage),
-		llb.WithCustomName("install system packages"))
-
-	vscodeStage, err := g.compileVSCode()
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to get vscode plugins")
-	}
-
-	var merged llb.State
-	if vscodeStage != nil {
-		merged = llb.Merge([]llb.State{
-			builtinSystemStage, systemStage, condaStage,
-			diffSSHStage, pypiStage, *vscodeStage,
-		}, llb.WithCustomName("merging all components into one"))
-	} else {
-		merged = llb.Merge([]llb.State{
-			builtinSystemStage, systemStage, condaStage,
-			diffSSHStage, pypiStage,
-		}, llb.WithCustomName("merging all components into one"))
-	}
-	return merged, nil
-}
diff --git a/pkg/lang/ir/editor.go b/pkg/lang/ir/editor.go
index 769d6da64..abef3e396 100644
--- a/pkg/lang/ir/editor.go
+++ b/pkg/lang/ir/editor.go
@@ -15,6 +15,8 @@
 package ir
 
 import (
+	"strconv"
+
 	"github.com/cockroachdb/errors"
 	"github.com/moby/buildkit/client/llb"
 
@@ -57,3 +59,35 @@ func (g *Graph) compileJupyter() {
 		g.PyPIPackages = append(g.PyPIPackages, "jupyter")
 	}
 }
+
+func (g Graph) generateJupyterCommand(notebookDir string) []string {
+	if g.JupyterConfig == nil {
+		return nil
+	}
+
+	var cmd []string
+	// Use python in conda env.
+	if g.CondaEnabled() {
+		cmd = append(cmd, "/opt/conda/bin/python3")
+	} else {
+		cmd = append(cmd, "python3")
+	}
+
+	cmd = append(cmd, []string{
+		"-m", "notebook",
+		"--ip", "0.0.0.0", "--notebook-dir", notebookDir,
+	}...)
+
+	if g.JupyterConfig.Password != "" {
+		cmd = append(cmd, "--NotebookApp.password", g.JupyterConfig.Password,
+			"--NotebookApp.token", "''")
+	} else {
+		cmd = append(cmd, "--NotebookApp.password", "''",
+			"--NotebookApp.token", "''")
+	}
+	if g.JupyterConfig.Port != 0 {
+		p := strconv.Itoa(int(g.JupyterConfig.Port))
+		cmd = append(cmd, "--port", p)
+	}
+	return cmd
+}
diff --git a/pkg/lang/ir/julia.go b/pkg/lang/ir/julia.go
index 62c4544c3..affd739c1 100644
--- a/pkg/lang/ir/julia.go
+++ b/pkg/lang/ir/julia.go
@@ -18,10 +18,53 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/cockroachdb/errors"
 	"github.com/moby/buildkit/client/llb"
 	"github.com/sirupsen/logrus"
 )
 
+func (g Graph) compileJulia(aptStage llb.State) (llb.State, error) {
+	g.compileJupyter()
+	builtinSystemStage := aptStage
+
+	sshStage, err := g.copySSHKey(builtinSystemStage)
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to copy ssh keys")
+	}
+	diffSSHStage := llb.Diff(builtinSystemStage, sshStage, llb.WithCustomName("install ssh keys"))
+
+	shellStage, err := g.compileShell(builtinSystemStage)
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to compile shell")
+	}
+	diffShellStage := llb.Diff(builtinSystemStage, shellStage, llb.WithCustomName("install shell"))
+
+	systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage),
+		llb.WithCustomName("install system packages"))
+
+	juliaStage := llb.Diff(builtinSystemStage,
+		g.installJuliaPackages(builtinSystemStage), llb.WithCustomName("install julia packages"))
+
+	vscodeStage, err := g.compileVSCode()
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to get vscode plugins")
+	}
+
+	var merged llb.State
+	if vscodeStage != nil {
+		merged = llb.Merge([]llb.State{
+			builtinSystemStage, systemStage, diffShellStage,
+			diffSSHStage, juliaStage, *vscodeStage,
+		}, llb.WithCustomName("merging all components into one"))
+	} else {
+		merged = llb.Merge([]llb.State{
+			builtinSystemStage, systemStage, diffShellStage,
+			diffSSHStage, juliaStage,
+		}, llb.WithCustomName("merging all components into one"))
+	}
+	return merged, nil
+}
+
 func (g Graph) installJuliaPackages(root llb.State) llb.State {
 	if len(g.JuliaPackages) == 0 {
 		return root
diff --git a/pkg/lang/ir/python.go b/pkg/lang/ir/python.go
index bd2f2570d..e644116c1 100644
--- a/pkg/lang/ir/python.go
+++ b/pkg/lang/ir/python.go
@@ -19,10 +19,62 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/cockroachdb/errors"
 	"github.com/moby/buildkit/client/llb"
 	"github.com/sirupsen/logrus"
 )
 
+func (g Graph) compilePython(aptStage llb.State) (llb.State, error) {
+	condaChanelStage := g.compileCondaChannel(aptStage)
+	pypiMirrorStage := g.compilePyPIIndex(condaChanelStage)
+
+	g.compileJupyter()
+	builtinSystemStage := pypiMirrorStage
+
+	sshStage, err := g.copySSHKey(builtinSystemStage)
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to copy ssh keys")
+	}
+	diffSSHStage := llb.Diff(builtinSystemStage, sshStage, llb.WithCustomName("install ssh keys"))
+
+	// Conda affects shell and python, thus we cannot do it in parallel.
+	shellStage, err := g.compileShell(builtinSystemStage)
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to compile shell")
+	}
+
+	condaEnvStage := g.setCondaENV(shellStage)
+
+	condaStage := llb.Diff(builtinSystemStage,
+		g.compileCondaPackages(condaEnvStage),
+		llb.WithCustomName("install conda packages"))
+
+	pypiStage := llb.Diff(condaEnvStage,
+		g.compilePyPIPackages(condaEnvStage),
+		llb.WithCustomName("install PyPI packages"))
+	systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage),
+		llb.WithCustomName("install system packages"))
+
+	vscodeStage, err := g.compileVSCode()
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to get vscode plugins")
+	}
+
+	var merged llb.State
+	if vscodeStage != nil {
+		merged = llb.Merge([]llb.State{
+			builtinSystemStage, systemStage, condaStage,
+			diffSSHStage, pypiStage, *vscodeStage,
+		}, llb.WithCustomName("merging all components into one"))
+	} else {
+		merged = llb.Merge([]llb.State{
+			builtinSystemStage, systemStage, condaStage,
+			diffSSHStage, pypiStage,
+		}, llb.WithCustomName("merging all components into one"))
+	}
+	return merged, nil
+}
+
 func (g Graph) compilePyPIPackages(root llb.State) llb.State {
 	if len(g.PyPIPackages) == 0 {
 		return root
diff --git a/pkg/lang/ir/r.go b/pkg/lang/ir/r.go
index d93bcc80d..597cbfc0a 100644
--- a/pkg/lang/ir/r.go
+++ b/pkg/lang/ir/r.go
@@ -18,9 +18,54 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/cockroachdb/errors"
 	"github.com/moby/buildkit/client/llb"
 )
 
+func (g Graph) compileRLang(aptStage llb.State) (llb.State, error) {
+	g.compileJupyter()
+	builtinSystemStage := aptStage
+
+	sshStage, err := g.copySSHKey(builtinSystemStage)
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to copy ssh keys")
+	}
+	diffSSHStage := llb.Diff(builtinSystemStage, sshStage, llb.WithCustomName("install ssh keys"))
+
+	// Conda affects shell and python, thus we cannot do it in parallel.
+	shellStage, err := g.compileShell(builtinSystemStage)
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to compile shell")
+	}
+	diffShellStage := llb.Diff(builtinSystemStage, shellStage, llb.WithCustomName("install shell"))
+
+	systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage),
+		llb.WithCustomName("install system packages"))
+
+	// TODO(terrytangyuan): Support RStudio local server
+	rPackageInstallStage := llb.Diff(builtinSystemStage,
+		g.installRPackages(builtinSystemStage), llb.WithCustomName("install R packages"))
+
+	vscodeStage, err := g.compileVSCode()
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to get vscode plugins")
+	}
+
+	var merged llb.State
+	if vscodeStage != nil {
+		merged = llb.Merge([]llb.State{
+			builtinSystemStage, systemStage, diffShellStage,
+			diffSSHStage, rPackageInstallStage, *vscodeStage,
+		}, llb.WithCustomName("merging all components into one"))
+	} else {
+		merged = llb.Merge([]llb.State{
+			builtinSystemStage, systemStage, diffShellStage,
+			diffSSHStage, rPackageInstallStage,
+		}, llb.WithCustomName("merging all components into one"))
+	}
+	return merged, nil
+}
+
 func (g Graph) installRPackages(root llb.State) llb.State {
 	if len(g.RPackages) == 0 {
 		return root


Test patch:
diff --git a/pkg/editor/jupyter/util_test.go b/pkg/lang/ir/editor_test.go
similarity index 86%
rename from pkg/editor/jupyter/util_test.go
rename to pkg/lang/ir/editor_test.go
index 1e8505b15..c379dc7f4 100644
--- a/pkg/editor/jupyter/util_test.go
+++ b/pkg/lang/ir/editor_test.go
@@ -12,23 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package jupyter
+package ir
 
 import (
 	"testing"
-
-	"github.com/tensorchord/envd/pkg/lang/ir"
 )
 
 func TestGenerateCommand(t *testing.T) {
 	testcases := []struct {
-		graph    ir.Graph
+		graph    Graph
 		dir      string
 		expected []string
 	}{
 		{
-			graph: ir.Graph{
-				JupyterConfig: &ir.JupyterConfig{
+			graph: Graph{
+				JupyterConfig: &JupyterConfig{
 					Password: "",
 					Port:     8888,
 				},
@@ -41,8 +39,8 @@ func TestGenerateCommand(t *testing.T) {
 			},
 		},
 		{
-			graph: ir.Graph{
-				JupyterConfig: &ir.JupyterConfig{
+			graph: Graph{
+				JupyterConfig: &JupyterConfig{
 					Password: "test",
 					Port:     8888,
 				},
@@ -55,13 +53,13 @@ func TestGenerateCommand(t *testing.T) {
 			},
 		},
 		{
-			graph:    ir.Graph{},
+			graph:    Graph{},
 			dir:      "test",
 			expected: []string{},
 		},
 	}
 	for _, tc := range testcases {
-		actual := GenerateCommand(tc.graph, tc.dir)
+		actual := tc.graph.generateJupyterCommand(tc.dir)
 		if !equal(actual, tc.expected) {
 			t.Errorf("failed to generate the command: expected %v, got %v", tc.expected, actual)
 		}