feat(lang): Support RStudio server (#503)

* feat(lang): Support rstudio server Signed-off-by: Ce Gao <[email protected]> * fix: Use tensorchord Signed-off-by: Ce Gao <[email protected]>
tensorchord · Jul 5, 2022 · e443784 · e443784
1 parent 89eb6e8
commit e443784
Show file tree

Hide file tree

Showing 19 changed files with 161 additions and 47 deletions.
diff --git a/base-images/build.sh b/base-images/build.sh
@@ -31,13 +31,15 @@ docker buildx build \
     --pull --push --platform linux/x86_64,linux/arm64 \
     -t ${DOCKER_HUB_ORG}/python:3.8-ubuntu20.04 \
     -f python3.8-ubuntu20.04.Dockerfile .
+
+# TODO(gaocegege): Support linux/arm64
 docker buildx build \
     --build-arg ENVD_VERSION=${ENVD_VERSION} \
     --build-arg ENVD_SSH_IMAGE=ghcr.io/tensorchord/envd-ssh-from-scratch \
     --build-arg HTTP_PROXY=${HTTP_PROXY} \
     --build-arg HTTPS_PROXY=${HTTPS_PROXY} \
-    -t ${DOCKER_HUB_ORG}/r-base:4.2  \
-    --pull --push --platform linux/x86_64,linux/arm64 \
+    -t ${DOCKER_HUB_ORG}/r-base:4.2 \
+    --pull --push --platform linux/x86_64 \
     -f r4.2.Dockerfile .
 docker buildx build \
     --build-arg ENVD_VERSION=${ENVD_VERSION} \

diff --git a/base-images/r4.2.Dockerfile b/base-images/r4.2.Dockerfile
@@ -18,11 +18,25 @@ ENV DEBIAN_FRONTEND noninteractive
 ENV PATH="/usr/bin:${PATH}"
 
 RUN apt-get update && \
-    apt-get install apt-utils && \
     apt-get install -y --no-install-recommends --no-install-suggests --fix-missing \
-    bash-static libtinfo5 libncursesw5 \
+    apt-utils bash-static libtinfo5 libncursesw5 \
+    # rstudio dependencies
+    file libapparmor1 libclang-dev libcurl4-openssl-dev libedit2 libobjc4 libssl1.1 libssl-dev \
+    libpq5 psmisc procps python-setuptools pwgen lsb-release \
     # envd dependencies
     python3 curl openssh-client git tini sudo zsh vim \
     && rm -rf /var/lib/apt/lists/*
 
+RUN set -x && \
+    UNAME_M="$(uname -m)" && \
+    if [ "${UNAME_M}" = "x86_64" ]; then \
+      RSTUDIO_URL="https://download2.rstudio.org/server/bionic/amd64/rstudio-server-2022.02.3-492-amd64.deb"; \
+    elif [ "${UNAME_M}" = "aarch64" ]; then \
+      RSTUDIO_URL="https://rstudio.org/download/latest/latest/server/focal/rstudio-server-latest-arm64.deb"; \
+    fi && \
+    DOWNLOAD_FILE=rstudio-server.deb && \
+    wget "${RSTUDIO_URL}" -O ${DOWNLOAD_FILE} && \
+    dpkg -i "$DOWNLOAD_FILE" && \
+    rm ${DOWNLOAD_FILE} && rm -f /var/lib/rstudio-server/secure-cookie-key
+
 COPY --from=envd /usr/bin/envd-ssh /var/envd/bin/envd-ssh
diff --git a/examples/mnist/main.py b/examples/mnist/main.py
@@ -64,7 +64,7 @@
 """
 
 batch_size = 128
-epochs = 15
+epochs = 1
 
 model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
 

diff --git a/examples/r-basic/build.envd b/examples/r-basic/build.envd
@@ -4,4 +4,5 @@ def build():
             "remotes",
             "rlang",
         ])
+    config.rstudio_server()
     shell("zsh")
diff --git a/pkg/app/get_env.go b/pkg/app/get_env.go
@@ -19,6 +19,7 @@ import (
 	"io"
 	"os"
 	"strconv"
+	"strings"
 
 	"github.com/docker/docker/pkg/stringid"
 	"github.com/olekukonko/tablewriter"
@@ -55,7 +56,7 @@ func getEnvironment(clicontext *cli.Context) error {
 func renderEnvironments(envs []types.EnvdEnvironment, w io.Writer) {
 	table := tablewriter.NewWriter(w)
 	table.SetHeader([]string{
-		"Name", "jupyter", "SSH Target", "Context", "Image",
+		"Name", "Endpoint", "SSH Target", "Image",
 		"GPU", "CUDA", "CUDNN", "Status", "Container ID",
 	})
 
@@ -72,18 +73,28 @@ func renderEnvironments(envs []types.EnvdEnvironment, w io.Writer) {
 	table.SetNoWhiteSpace(true)
 
 	for _, env := range envs {
-		envRow := make([]string, 10)
+		envRow := make([]string, 9)
 		envRow[0] = env.Name
-		envRow[1] = env.JupyterAddr
+		envRow[1] = endpointOrNone(env)
 		envRow[2] = fmt.Sprintf("%s.envd", env.Name)
-		envRow[3] = stringOrNone(env.BuildContext)
-		envRow[4] = env.Container.Image
-		envRow[5] = strconv.FormatBool(env.GPU)
-		envRow[6] = stringOrNone(env.CUDA)
-		envRow[7] = stringOrNone(env.CUDNN)
-		envRow[8] = env.Status
-		envRow[9] = stringid.TruncateID(env.Container.ID)
+		envRow[3] = env.Container.Image
+		envRow[4] = strconv.FormatBool(env.GPU)
+		envRow[5] = stringOrNone(env.CUDA)
+		envRow[6] = stringOrNone(env.CUDNN)
+		envRow[7] = env.Status
+		envRow[8] = stringid.TruncateID(env.Container.ID)
 		table.Append(envRow)
 	}
 	table.Render()
 }
+
+func endpointOrNone(env types.EnvdEnvironment) string {
+	var res strings.Builder
+	if env.JupyterAddr != nil {
+		res.WriteString(fmt.Sprintf("jupyter: %s", *env.JupyterAddr))
+	}
+	if env.RStudioServerAddr != nil {
+		res.WriteString(fmt.Sprintf("rstudio: %s", *env.RStudioServerAddr))
+	}
+	return res.String()
+}
diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -16,9 +16,10 @@ package config
 type UpState string
 
 const (
-	PrivateKeyFile              = "id_rsa_envd"
-	PublicKeyFile               = "id_rsa_envd.pub"
-	ContainerAuthorizedKeysPath = "/var/envd/authorized_keys"
-	SSHPortInContainer          = 2222
-	JupyterPortInContainer      = 8888
+	PrivateKeyFile               = "id_rsa_envd"
+	PublicKeyFile                = "id_rsa_envd.pub"
+	ContainerAuthorizedKeysPath  = "/var/envd/authorized_keys"
+	SSHPortInContainer           = 2222
+	JupyterPortInContainer       = 8888
+	RStudioServerPortInContainer = 8787
 )
diff --git a/pkg/docker/docker.go b/pkg/docker/docker.go
@@ -92,7 +92,7 @@ func NewClient(ctx context.Context) (Client, error) {
 	if err != nil {
 		// Special note needed to give users
 		if strings.Contains(err.Error(), "permission denied") {
-			err = errors.New(`It seems that current user have no access to docker daemon, 
+			err = errors.New(`It seems that current user have no access to docker daemon,
 please visit https://docs.docker.com/engine/install/linux-postinstall/ for more info.`)
 		}
 		return nil, err
@@ -384,13 +384,30 @@ func (c generalClient) StartEnvd(ctx context.Context, tag, name, buildContext st
 		}
 		config.ExposedPorts[natPort] = struct{}{}
 	}
+	var rStudioPortInHost int
+	if g.RStudioServerConfig != nil {
+		var err error
+		rStudioPortInHost, err = netutil.GetFreePort()
+		if err != nil {
+			return "", "", errors.Wrap(err, "failed to get a free port")
+		}
+		natPort := nat.Port(fmt.Sprintf("%d/tcp", envdconfig.RStudioServerPortInContainer))
+		hostConfig.PortBindings[natPort] = []nat.PortBinding{
+			{
+				HostIP:   localhost,
+				HostPort: strconv.Itoa(rStudioPortInHost),
+			},
+		}
+		config.ExposedPorts[natPort] = struct{}{}
+	}
 
 	if gpuEnabled {
 		logger.Debug("GPU is enabled.")
 		hostConfig.DeviceRequests = deviceRequests(numGPUs)
 	}
 
-	config.Labels = labels(name, g.JupyterConfig, sshPortInHost, jupyterPortInHost)
+	config.Labels = labels(name, g,
+		sshPortInHost, jupyterPortInHost, rStudioPortInHost)
 
 	logger = logger.WithFields(logrus.Fields{
 		"entrypoint":  config.Entrypoint,

diff --git a/pkg/docker/label.go b/pkg/docker/label.go
@@ -24,15 +24,20 @@ import (
 	"github.com/tensorchord/envd/pkg/types"
 )
 
-func labels(name string, jupyterConfig *ir.JupyterConfig,
-	sshPortInHost, jupyterPortInHost int) map[string]string {
+func labels(name string, g ir.Graph,
+	sshPortInHost, jupyterPortInHost, rstudioServerPortInHost int) map[string]string {
 	res := make(map[string]string)
 	res[types.ContainerLabelName] = name
 	res[types.ContainerLabelSSHPort] = strconv.Itoa(sshPortInHost)
-	if jupyterConfig != nil {
+	if g.JupyterConfig != nil {
 		res[types.ContainerLabelJupyterAddr] =
 			fmt.Sprintf("http://localhost:%d", jupyterPortInHost)
 	}
+	if g.RStudioServerConfig != nil {
+		res[types.ContainerLabelRStudioServerAddr] =
+			fmt.Sprintf("http://localhost:%d", rstudioServerPortInHost)
+	}
+
 	return res
 }
 

diff --git a/pkg/lang/frontend/starlark/config/config.go b/pkg/lang/frontend/starlark/config/config.go
@@ -41,6 +41,7 @@ var Module = &starlarkstruct.Module{
 			ruleCondaChannel, ruleFuncCondaChannel),
 		"julia_pkg_server": starlark.NewBuiltin(
 			ruleJuliaPackageServer, ruleFuncJuliaPackageServer),
+		"rstudio_server": starlark.NewBuiltin(ruleRStudioServer, ruleFuncRStudioServer),
 	},
 }
 
@@ -167,6 +168,15 @@ func ruleFuncUbuntuAptSource(thread *starlark.Thread, _ *starlark.Builtin,
 	return starlark.None, nil
 }
 
+func ruleFuncRStudioServer(thread *starlark.Thread, _ *starlark.Builtin,
+	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	if err := ir.RStudioServer(); err != nil {
+		return nil, err
+	}
+
+	return starlark.None, nil
+}
+
 func ruleFuncCondaChannel(thread *starlark.Thread, _ *starlark.Builtin,
 	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var channel starlark.String

diff --git a/pkg/lang/frontend/starlark/config/const.go b/pkg/lang/frontend/starlark/config/const.go
@@ -22,4 +22,5 @@ const (
 	ruleCondaChannel       = "config.conda_channel"
 	ruleGPU                = "config.gpu"
 	ruleJuliaPackageServer = "config.julia_pkg_server"
+	ruleRStudioServer      = "config.rstudio_server"
 )
diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go
@@ -132,6 +132,10 @@ func (g Graph) ExposedPorts() (map[string]struct{}, error) {
 	if g.JupyterConfig != nil {
 		ports[fmt.Sprintf("%d/tcp", config.JupyterPortInContainer)] = struct{}{}
 	}
+	if g.RStudioServerConfig != nil {
+		ports[fmt.Sprintf("%d/tcp", config.RStudioServerPortInContainer)] = struct{}{}
+	}
+
 	return ports, nil
 }
 
@@ -144,21 +148,28 @@ func (g Graph) Entrypoint(buildContextDir string) ([]string, error) {
 	}
 
 	template := `set -e
-	/var/envd/bin/envd-ssh --authorized-keys %s --port %d --shell %s &
-	%s
-	wait -n`
+/var/envd/bin/envd-ssh --authorized-keys %s --port %d --shell %s &
+%s
+wait -n`
 
+	// Generate jupyter and rstudio server commands.
+	var customCmd strings.Builder
 	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
+		customCmd.WriteString(strings.Join(jupyterCmd, " "))
+		customCmd.WriteString("\n")
 	}
+	if g.RStudioServerConfig != nil {
+		workingDir := fmt.Sprintf("/home/envd/%s", fileutil.Base(buildContextDir))
+		rstudioCmd := g.generateRStudioCommand(workingDir)
+		customCmd.WriteString(strings.Join(rstudioCmd, " "))
+		customCmd.WriteString("\n")
+	}
+
 	cmd := fmt.Sprintf(template,
-		config.ContainerAuthorizedKeysPath, config.SSHPortInContainer, g.Shell, "")
+		config.ContainerAuthorizedKeysPath,
+		config.SSHPortInContainer, g.Shell, customCmd.String())
 	ep = append(ep, cmd)
 	return ep, nil
 }

diff --git a/pkg/lang/ir/editor.go b/pkg/lang/ir/editor.go
@@ -54,13 +54,20 @@ func (g Graph) compileVSCode() (*llb.State, error) {
 	return &layer, nil
 }
 
-func (g *Graph) compileJupyter() {
+func (g *Graph) compileJupyter() error {
 	if g.JupyterConfig != nil {
 		g.PyPIPackages = append(g.PyPIPackages, "jupyter")
+		switch g.Language.Name {
+		case "python":
+			return nil
+		default:
+			return errors.Newf("Jupyter is not supported in %s yet", g.Language.Name)
+		}
 	}
+	return nil
 }
 
-func (g Graph) generateJupyterCommand(notebookDir string) []string {
+func (g Graph) generateJupyterCommand(workingDir string) []string {
 	if g.JupyterConfig == nil {
 		return nil
 	}
@@ -75,7 +82,7 @@ func (g Graph) generateJupyterCommand(notebookDir string) []string {
 
 	cmd = append(cmd, []string{
 		"-m", "notebook",
-		"--ip", "0.0.0.0", "--notebook-dir", notebookDir,
+		"--ip", "0.0.0.0", "--notebook-dir", workingDir,
 	}...)
 
 	if g.JupyterConfig.Password != "" {
@@ -91,3 +98,16 @@ func (g Graph) generateJupyterCommand(notebookDir string) []string {
 	}
 	return cmd
 }
+
+func (g Graph) generateRStudioCommand(workingDir string) []string {
+	if g.RStudioServerConfig == nil {
+		return nil
+	}
+
+	return []string{
+		// TODO(gaocegege): Remove root permission here.
+		"sudo",
+		"/usr/lib/rstudio-server/bin/rserver",
+		// TODO(gaocegege): Support working dir.
+	}
+}
diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go
@@ -132,6 +132,11 @@ func Jupyter(pwd string, port int64) error {
 	return nil
 }
 
+func RStudioServer() error {
+	DefaultGraph.RStudioServerConfig = &RStudioServerConfig{}
+	return nil
+}
+
 func Run(commands []string) error {
 	// TODO(gaocegege): Support order-based exec.
 	DefaultGraph.Exec = commands

diff --git a/pkg/lang/ir/julia.go b/pkg/lang/ir/julia.go
@@ -24,7 +24,10 @@ import (
 )
 
 func (g Graph) compileJulia(aptStage llb.State) (llb.State, error) {
-	g.compileJupyter()
+	if err := g.compileJupyter(); err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to compile jupyter")
+	}
+
 	builtinSystemStage := aptStage
 
 	sshStage, err := g.copySSHKey(builtinSystemStage)

diff --git a/pkg/lang/ir/python.go b/pkg/lang/ir/python.go
@@ -28,7 +28,9 @@ func (g Graph) compilePython(aptStage llb.State) (llb.State, error) {
 	condaChanelStage := g.compileCondaChannel(aptStage)
 	pypiMirrorStage := g.compilePyPIIndex(condaChanelStage)
 
-	g.compileJupyter()
+	if err := g.compileJupyter(); err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to compile jupyter")
+	}
 	builtinSystemStage := pypiMirrorStage
 
 	sshStage, err := g.copySSHKey(builtinSystemStage)

diff --git a/pkg/lang/ir/r.go b/pkg/lang/ir/r.go
@@ -23,7 +23,9 @@ import (
 )
 
 func (g Graph) compileRLang(aptStage llb.State) (llb.State, error) {
-	g.compileJupyter()
+	if err := g.compileJupyter(); err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to compile jupyter")
+	}
 	builtinSystemStage := aptStage
 
 	sshStage, err := g.copySSHKey(builtinSystemStage)

diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go
@@ -52,11 +52,15 @@ type Graph struct {
 	*JupyterConfig
 	*GitConfig
 	*CondaConfig
+	*RStudioServerConfig
 
 	Writer      compileui.Writer
 	CachePrefix string
 }
 
+type RStudioServerConfig struct {
+}
+
 type Language struct {
 	Name    string
 	Version *string