feat: Support forward (#1014)

* feat: Support forward

Signed-off-by: Ce Gao <[email protected]>

* feat: Support port forwarding

Signed-off-by: Ce Gao <[email protected]>

Signed-off-by: Ce Gao <[email protected]>

examples/python-basic/build.envd

@@ -13,3 +13,4 @@ def build():
         }
     )
     runtime.environ(env={"ENVD_MODE": "DEV"})
+    config.jupyter()

go.mod

@@ -29,7 +29,7 @@ require (
 	github.com/sirupsen/logrus v1.9.0
 	github.com/spf13/viper v1.13.0
 	github.com/stretchr/testify v1.8.1
-	github.com/tensorchord/envd-server v0.0.5
+	github.com/tensorchord/envd-server v0.0.6
 	github.com/tonistiigi/units v0.0.0-20180711220420-6950e57a87ea
 	github.com/tonistiigi/vt100 v0.0.0-20210615222946-8066bb97264f
 	github.com/urfave/cli/v2 v2.20.3

go.sum

@@ -610,8 +610,8 @@ github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o
 github.com/subosito/gotenv v1.4.1 h1:jyEFiXpy21Wm81FBN71l9VoMMV8H8jG+qIK3GCpY6Qs=
 github.com/subosito/gotenv v1.4.1/go.mod h1:ayKnFf/c6rvx/2iiLrJUk1e6plDbT3edrFNGqEflhK0=
 github.com/syndtr/gocapability v0.0.0-20200815063812-42c35b437635/go.mod h1:hkRG7XYTFWNJGYcbNJQlaLq0fg1yr4J4t/NcTQtrfww=
-github.com/tensorchord/envd-server v0.0.5 h1:tbkMU79PTp8OONFwD4uGGQSpOY//gUouCWs1z5d+74s=
-github.com/tensorchord/envd-server v0.0.5/go.mod h1:V76OpMczrgBeu19K+rgFqSf4ZK3DXo/U82/0xel/jYg=
+github.com/tensorchord/envd-server v0.0.6 h1:U/FLcIDRSIEavtfuJYJzW/pZQ1yC8umN/UaNUG0aWPs=
+github.com/tensorchord/envd-server v0.0.6/go.mod h1:TveWA1l+UWI87y4kbNS4f6OaYJaAFb9XOhTB1hcBPDw=
 github.com/tmc/grpc-websocket-proxy v0.0.0-20190109142713-0ad062ec5ee5/go.mod h1:ncp9v5uamzpCO7NfCPTXjqaC+bZgJeR0sMTm6dMHP7U=
 github.com/tonistiigi/fsutil v0.0.0-20220115021204-b19f7f9cb274 h1:wbyZxD6IPFp0sl5uscMOJRsz5UKGFiNiD16e+MVfKZY=
 github.com/tonistiigi/fsutil v0.0.0-20220115021204-b19f7f9cb274/go.mod h1:oPAfvw32vlUJSjyDcQ3Bu0nb2ON2B+G0dtVN/SZNJiA=

pkg/app/create.go

@@ -15,6 +15,7 @@
 package app
 
 import (
+	"fmt"
 	"strings"
 	"time"
 
@@ -29,6 +30,7 @@ import (
 	"github.com/tensorchord/envd/pkg/ssh"
 	sshconfig "github.com/tensorchord/envd/pkg/ssh/config"
 	"github.com/tensorchord/envd/pkg/types"
+	"github.com/tensorchord/envd/pkg/util/netutil"
 )
 
 var CommandCreate = &cli.Command{
@@ -52,7 +54,7 @@ var CommandCreate = &cli.Command{
 		&cli.DurationFlag{
 			Name:  "timeout",
 			Usage: "Timeout of container creation",
-			Value: time.Second * 30,
+			Value: time.Second * 1800,
 		},
 		&cli.BoolFlag{
 			Name:  "detach",
@@ -132,6 +134,7 @@ func create(clicontext *cli.Context) error {
 
 	// TODO(gaocegege): Test why it fails.
 	if !clicontext.Bool("detach") {
+		outputChannel := make(chan error)
 		opt := ssh.DefaultOptions()
 		opt.PrivateKeyPath = clicontext.Path("private-key")
 		opt.Port = res.SSHPort
@@ -141,10 +144,40 @@ func create(clicontext *cli.Context) error {
 
 		sshClient, err := ssh.NewClient(opt)
 		if err != nil {
-			return errors.Wrap(err, "failed to create the ssh client")
+			outputChannel <- errors.Wrap(err, "failed to create the ssh client")
 		}
-		if err := sshClient.Attach(); err != nil {
-			return errors.Wrap(err, "failed to attach to the container")
+
+		ports := res.Ports
+
+		for _, p := range ports {
+			if p.Port == 2222 {
+				continue
+			}
+
+			// TODO(gaocegege): Use one remote port.
+			localPort, err := netutil.GetFreePort()
+			if err != nil {
+				return errors.Wrap(err, "failed to get a free port")
+			}
+			localAddress := fmt.Sprintf("%s:%d", "localhost", localPort)
+			remoteAddress := fmt.Sprintf("%s:%d", "localhost", p.Port)
+			logrus.Infof("service \"%s\" is listening at %s\n", p.Name, localAddress)
+			go func() {
+				if err := sshClient.LocalForward(localAddress, remoteAddress); err != nil {
+					outputChannel <- errors.Wrap(err, "failed to forward to local port")
+				}
+			}()
+		}
+
+		go func() {
+			if err := sshClient.Attach(); err != nil {
+				outputChannel <- errors.Wrap(err, "failed to attach to the container")
+			}
+			outputChannel <- nil
+		}()
+
+		if err := <-outputChannel; err != nil {
+			return err
 		}
 	}
 	return nil

pkg/envd/envdserver.go

@@ -156,14 +156,15 @@ func (e *envdServerEngine) StartEnvd(ctx context.Context, so StartOptions) (*Sta
 	}
 
 	if err := e.WaitUntilRunning(
-		ctx, resp.ID, so.Timeout); err != nil {
+		ctx, resp.Created.Name, so.Timeout); err != nil {
 		return nil, errors.Wrap(err, "failed to wait until the container is running")
 	}
 
 	result := &StartResult{
 		SSHPort: 2222,
 		Address: "",
-		Name:    resp.ID,
+		Name:    resp.Created.Name,
+		Ports:   resp.Created.Spec.Ports,
 	}
 	return result, nil
 }

pkg/envd/types.go

@@ -17,6 +17,8 @@ package envd
 import (
 	"time"
 
+	"github.com/tensorchord/envd-server/api/types"
+
 	"github.com/tensorchord/envd/pkg/lang/ir"
 )
 
@@ -62,4 +64,6 @@ type StartResult struct {
 	SSHPort int
 	Address string
 	Name    string
+
+	Ports []types.EnvironmentPort
 }

pkg/lang/ir/compile.go

@@ -25,6 +25,7 @@ import (
 	"github.com/moby/buildkit/client/llb"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/viper"
+	servertypes "github.com/tensorchord/envd-server/api/types"
 
 	"github.com/tensorchord/envd/pkg/config"
 	"github.com/tensorchord/envd/pkg/flag"
@@ -147,6 +148,39 @@ func (g Graph) Labels() (map[string]string, error) {
 	}
 	labels[types.RuntimeGraphCode] = code
 
+	ports := []servertypes.EnvironmentPort{}
+	ports = append(ports, servertypes.EnvironmentPort{
+		Name: "ssh",
+		Port: config.SSHPortInContainer,
+	})
+	if g.JupyterConfig != nil {
+		ports = append(ports, servertypes.EnvironmentPort{
+			Name: "jupyter",
+			Port: config.JupyterPortInContainer,
+		})
+	}
+	if g.RStudioServerConfig != nil {
+		ports = append(ports, servertypes.EnvironmentPort{
+			Name: "rstudio-server",
+			Port: config.RStudioServerPortInContainer,
+		})
+	}
+
+	if g.RuntimeExpose != nil && len(g.RuntimeExpose) > 0 {
+		for _, item := range g.RuntimeExpose {
+			ports = append(ports, servertypes.EnvironmentPort{
+				Name: item.ServiceName,
+				Port: int32(item.EnvdPort),
+			})
+		}
+	}
+
+	portsData, err := json.Marshal(ports)
+	if err != nil {
+		return labels, err
+	}
+	labels[types.ImageLabelPorts] = string(portsData)
+
 	return labels, nil
 }
 

pkg/lang/ir/supervisor.go

@@ -36,7 +36,7 @@ working-directory = "${%[3]s}"
 
 [environment]
 keep-env = true
-re-export = [ "PATH", "SHELL", "USER", "%[3]s" ]
+re-export = [ "PATH", "SHELL", "USER", "%[3]s", "ENVD_AUTHORIZED_KEYS_PATH", "ENVD_HOST_KEY" ]
 
 [restart]
 strategy = "on-failure"

pkg/ssh/ssh.go

@@ -42,6 +42,7 @@ import (
 type Client interface {
 	Attach() error
 	ExecWithOutput(cmd string) ([]byte, error)
+	LocalForward(localAddress, targetAddress string) error
 	Close() error
 }
 
@@ -282,6 +283,42 @@ func (c generalClient) Attach() error {
 	return errors.Wrap(err, "command failed")
 }
 
+func (c generalClient) LocalForward(localAddress, targetAddress string) error {
+	localListener, err := net.Listen("tcp", localAddress)
+	if err != nil {
+		return errors.Wrap(err, "net.Listen failed")
+	}
+
+	logrus.Debug("begin to forward " + localAddress + " to " + targetAddress)
+	for {
+		localCon, err := localListener.Accept()
+		if err != nil {
+			return errors.Wrap(err, "listen.Accept failed")
+		}
+
+		sshConn, err := c.cli.Dial("tcp", targetAddress)
+		if err != nil {
+			return errors.Wrap(err, "listen.Accept failed")
+		}
+
+		// Copy local.Reader to sshConn.Writer
+		go func() {
+			_, err = io.Copy(sshConn, localCon)
+			if err != nil {
+				logrus.Debugf("io.Copy failed: %v", err)
+			}
+		}()
+
+		// Copy sshConn.Reader to localCon.Writer
+		go func() {
+			_, err = io.Copy(localCon, sshConn)
+			if err != nil {
+				logrus.Debugf("io.Copy failed: %v", err)
+			}
+		}()
+	}
+}
+
 func isTerminal(r io.Reader) (int, bool) {
 	switch v := r.(type) {
 	case *os.File:

pkg/types/label.go

@@ -22,6 +22,7 @@ const (
 
 	ImageLabelVendor    = "ai.tensorchord.envd.vendor"
 	ImageLabelGPU       = "ai.tensorchord.envd.gpu"
+	ImageLabelPorts     = "ai.tensorchord.envd.ports"
 	ImageLabelAPT       = "ai.tensorchord.envd.apt.packages"
 	ImageLabelPyPI      = "ai.tensorchord.envd.pypi.commands"
 	ImageLabelR         = "ai.tensorchord.envd.r.packages"