diff --git a/examples/python-basic/build.envd b/examples/python-basic/build.envd index d3459ad7b..a006b104a 100644 --- a/examples/python-basic/build.envd +++ b/examples/python-basic/build.envd @@ -13,3 +13,4 @@ def build(): } ) runtime.environ(env={"ENVD_MODE": "DEV"}) + config.jupyter() diff --git a/go.mod b/go.mod index fe0fe3c49..4cffb54b0 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 6a6ebf520..39161640e 100644 --- a/go.sum +++ b/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= diff --git a/pkg/app/create.go b/pkg/app/create.go index be8c99116..678a973bb 100644 --- a/pkg/app/create.go +++ b/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 diff --git a/pkg/envd/envdserver.go b/pkg/envd/envdserver.go index 12e1cb5cb..be281f659 100644 --- a/pkg/envd/envdserver.go +++ b/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 } diff --git a/pkg/envd/types.go b/pkg/envd/types.go index 887df97c1..9d2cbb50a 100644 --- a/pkg/envd/types.go +++ b/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 } diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go index 180882090..7b446b00b 100644 --- a/pkg/lang/ir/compile.go +++ b/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 } diff --git a/pkg/lang/ir/supervisor.go b/pkg/lang/ir/supervisor.go index 9e8de1564..3bf3ed332 100644 --- a/pkg/lang/ir/supervisor.go +++ b/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" diff --git a/pkg/ssh/ssh.go b/pkg/ssh/ssh.go index 737374f7a..d1edc402a 100644 --- a/pkg/ssh/ssh.go +++ b/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: diff --git a/pkg/types/label.go b/pkg/types/label.go index a693f0d79..061554a4d 100644 --- a/pkg/types/label.go +++ b/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"