feat(run): add log message for successful detachment from SSH target

Signed-off-by: zhangwei <[email protected]>

Author: rainbend

pkg/app/run.go

@@ -29,21 +29,18 @@ import (
 	"github.com/tensorchord/envd/pkg/app/telemetry"
 	"github.com/tensorchord/envd/pkg/envd"
 	"github.com/tensorchord/envd/pkg/home"
-	"github.com/tensorchord/envd/pkg/ssh"
 	sshconfig "github.com/tensorchord/envd/pkg/ssh/config"
 	"github.com/tensorchord/envd/pkg/syncthing"
 	"github.com/tensorchord/envd/pkg/types"
 	"github.com/tensorchord/envd/pkg/util/fileutil"
-	"github.com/tensorchord/envd/pkg/util/netutil"
 )
 
 var CommandCreate = &cli.Command{
-	Name:        "run",
-	Category:    CategoryBasic,
-	Aliases:     []string{"c"},
-	Usage:       "Run the envd environment from the existing image",
-	Hidden:      false,
-	Description: `run is only supported in envd-server runner currently`,
+	Name:     "run",
+	Category: CategoryBasic,
+	Aliases:  []string{"c"},
+	Usage:    "Run the envd environment from the existing image",
+	Hidden:   false,
 	Flags: []cli.Flag{
 		&cli.StringFlag{
 			Name:        "image",
@@ -105,7 +102,7 @@ var CommandCreate = &cli.Command{
 		},
 		&cli.BoolFlag{
 			Name:  "sync",
-			Usage: "Sync the local directory with the remote container",
+			Usage: "Sync the local directory with the remote container (only supported in envd-server runner currently)",
 			Value: false,
 		},
 		&cli.StringSliceFlag{
@@ -132,9 +129,9 @@ func run(clicontext *cli.Context) error {
 		return err
 	}
 
-	name := clicontext.String("name")
-	if name == "" {
-		name = strings.ToLower(randomdata.SillyName())
+	environmentName := clicontext.String("name")
+	if environmentName == "" {
+		environmentName = strings.ToLower(randomdata.SillyName())
 	}
 	opt := envd.StartOptions{
 		SshdHost:        clicontext.String("host"),
@@ -144,7 +141,7 @@ func run(clicontext *cli.Context) error {
 		NumCPU:          clicontext.String("cpus"),
 		NumGPU:          clicontext.Int("gpu"),
 		ShmSize:         clicontext.Int("shm-size"),
-		EnvironmentName: name,
+		EnvironmentName: environmentName,
 	}
 	switch c.Runner {
 	case types.RunnerTypeEnvdServer:
@@ -171,6 +168,8 @@ func run(clicontext *cli.Context) error {
 		return err
 	}
 
+	ctr := filepath.Base(opt.BuildContext)
+	detach := clicontext.Bool("detach")
 	logger := logrus.WithFields(logrus.Fields{
 		"cmd":          "run",
 		"StartOptions": opt,
@@ -185,105 +184,83 @@ func run(clicontext *cli.Context) error {
 		return errors.Wrap(err, "failed to get the ssh hostname")
 	}
 
-	ac, err := home.GetManager().AuthGetCurrent()
-	if err != nil {
-		return errors.Wrap(err, "failed to get the auth information")
-	}
-	username, err := sshname.Username(ac.Name, res.Name)
-	if err != nil {
-		return errors.Wrap(err, "failed to get the username")
+	var eo sshconfig.EntryOptions
+	switch c.Runner {
+	case types.RunnerTypeEnvdServer:
+		ac, err := home.GetManager().AuthGetCurrent()
+		if err != nil {
+			return errors.Wrap(err, "failed to get the auth information")
+		}
+		username, err := sshname.Username(ac.Name, res.Name)
+		if err != nil {
+			return errors.Wrap(err, "failed to get the username")
+		}
+		eo = sshconfig.EntryOptions{
+			Name:               res.Name,
+			IFace:              hostname,
+			Port:               res.SSHPort,
+			PrivateKeyPath:     clicontext.Path("private-key"),
+			EnableHostKeyCheck: false,
+			EnableAgentForward: false,
+			User:               username,
+		}
+	case types.RunnerTypeDocker:
+		eo, err = engine.GenerateSSHConfig(ctr, hostname,
+			clicontext.Path("private-key"), res)
+		if err != nil {
+			return errors.Wrap(err, "failed to get the ssh entry")
+		}
 	}
 
-	eo := sshconfig.EntryOptions{
-		Name:               res.Name,
-		IFace:              hostname,
-		Port:               res.SSHPort,
-		PrivateKeyPath:     clicontext.Path("private-key"),
-		EnableHostKeyCheck: false,
-		EnableAgentForward: false,
-		User:               username,
-	}
 	if err = sshconfig.AddEntry(eo); err != nil {
 		logger.WithError(err).
 			Infof("failed to add entry %s to your SSH config file", res.Name)
 		return errors.Wrap(err, "failed to add entry to your SSH config file")
 	}
 
-	// TODO(gaocegege): Test why it fails.
-	if !clicontext.Bool("detach") {
+	if !detach {
 		outputChannel := make(chan error)
-		opt := ssh.DefaultOptions()
-		opt.PrivateKeyPath = clicontext.Path("private-key")
-		opt.Port = res.SSHPort
-		opt.AgentForwarding = false
-		opt.User = username
-		opt.Server = hostname
-
-		sshClient, err := ssh.NewClient(opt)
-		if err != nil {
-			outputChannel <- errors.Wrap(err, "failed to create the ssh client")
-		}
-
-		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)
-			logger.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")
+		if c.Runner == types.RunnerTypeEnvdServer {
+			if clicontext.Bool("sync") {
+				go func() {
+					if err = engine.LocalForward(hostname, clicontext.Path("private-key"), res, syncthing.DefaultRemoteAPIAddress, syncthing.DefaultRemoteAPIAddress); err != nil {
+						outputChannel <- errors.Wrap(err, "failed to forward to remote api port")
+					}
+				}()
+
+				go func() {
+					syncthingRemoteAddr := fmt.Sprintf("127.0.0.1:%s", syncthing.ParsePortFromAddress(syncthing.DefaultRemoteDeviceAddress))
+					if err = engine.LocalForward(hostname, clicontext.Path("private-key"), res, syncthingRemoteAddr, syncthingRemoteAddr); err != nil {
+						outputChannel <- errors.Wrap(err, "failed to forward to remote port")
+					}
+				}()
+
+				go func() {
+					syncthingLocalAddr := fmt.Sprintf("127.0.0.1:%s", syncthing.ParsePortFromAddress(syncthing.DefaultLocalDeviceAddress))
+					if err = engine.RemoteForward(hostname, clicontext.Path("private-key"), res, syncthingLocalAddr, syncthingLocalAddr); err != nil {
+						outputChannel <- errors.Wrap(err, "failed to forward to local port")
+					}
+				}()
+
+				localSyncthing, _, err := startSyncthing(res.Name)
+				if err != nil {
+					return errors.Wrap(err, "failed to start syncthing")
 				}
-			}()
-		}
-
-		if clicontext.Bool("sync") {
-			go func() {
-				if err := sshClient.LocalForward(syncthing.DefaultRemoteAPIAddress, syncthing.DefaultRemoteAPIAddress); err != nil {
-					outputChannel <- errors.Wrap(err, "failed to forward to remote api port")
-				}
-			}()
-
-			go func() {
-				syncthingRemoteAddr := fmt.Sprintf("127.0.0.1:%s", syncthing.ParsePortFromAddress(syncthing.DefaultRemoteDeviceAddress))
-				if err := sshClient.LocalForward(syncthingRemoteAddr, syncthingRemoteAddr); err != nil {
-					outputChannel <- errors.Wrap(err, "failed to forward to remote port")
-				}
-			}()
-
-			go func() {
-				syncthingLocalAddr := fmt.Sprintf("127.0.0.1:%s", syncthing.ParsePortFromAddress(syncthing.DefaultLocalDeviceAddress))
-				if err := sshClient.RemoteForward(syncthingLocalAddr, syncthingLocalAddr); err != nil {
-					outputChannel <- errors.Wrap(err, "failed to forward to local port")
-				}
-			}()
-
-			localSyncthing, _, err := startSyncthing(res.Name)
-			if err != nil {
-				return errors.Wrap(err, "failed to start syncthing")
+				defer localSyncthing.StopLocalSyncthing()
 			}
-			defer localSyncthing.StopLocalSyncthing()
-
 		}
 
 		go func() {
-			// TODO(gaocegege): Avoid the hard code.
-			if err := sshClient.Attach(); err != nil {
-				outputChannel <- errors.Wrap(err, "failed to attach to the container")
+			if err = engine.Attach(ctr, hostname,
+				clicontext.Path("private-key"), res, nil); err != nil {
+				outputChannel <- errors.Wrap(err, "failed to attach to the ssh target")
 			}
+			logrus.Infof("Detached successfully. You can attach to the container with command `ssh %s.envd`\n",
+				environmentName)
 			outputChannel <- nil
 		}()
 
-		if err := <-outputChannel; err != nil {
+		if err = <-outputChannel; err != nil {
 			return err
 		}
 	}

pkg/envd/docker.go

@@ -328,24 +328,58 @@ func (e dockerEngine) GenerateSSHConfig(name, iface, privateKeyPath string,
 	return eo, nil
 }
 
-func (e dockerEngine) Attach(name, iface, privateKeyPath string,
-	startResult *StartResult, g ir.Graph) error {
+func (e dockerEngine) newSSHClient(server string, port int, privateKeyPath string) (ssh.Client, error) {
 	opt := ssh.DefaultOptions()
-	opt.Server = iface
+	opt.Server = server
 	opt.PrivateKeyPath = privateKeyPath
-	opt.Port = startResult.SSHPort
+	opt.Port = port
 	sshClient, err := ssh.NewClient(opt)
 	if err != nil {
-		return errors.Wrap(err, "failed to create the ssh client")
+		return nil, errors.Wrap(err, "failed to create the ssh client")
 	}
-	opt.Server = iface
+	return sshClient, nil
+}
 
-	if err := sshClient.Attach(); err != nil {
+func (e dockerEngine) Attach(name, iface, privateKeyPath string,
+	startResult *StartResult, g ir.Graph) error {
+	sshClient, err := e.newSSHClient(iface, startResult.SSHPort, privateKeyPath)
+	if err != nil {
+		return err
+	}
+	defer sshClient.Close()
+
+	if err = sshClient.Attach(); err != nil {
 		return errors.Wrap(err, "failed to attach to the container")
 	}
 	return nil
 }
 
+func (e dockerEngine) LocalForward(iface, privateKeyPath string, startResult *StartResult, localAddress, targetAddress string) error {
+	sshClient, err := e.newSSHClient(iface, startResult.SSHPort, privateKeyPath)
+	if err != nil {
+		return err
+	}
+	defer sshClient.Close()
+
+	if err = sshClient.LocalForward(localAddress, targetAddress); err != nil {
+		return errors.Wrap(err, "failed to forward to local port")
+	}
+	return nil
+}
+
+func (e dockerEngine) RemoteForward(iface, privateKeyPath string, startResult *StartResult, localAddress, targetAddress string) error {
+	sshClient, err := e.newSSHClient(iface, startResult.SSHPort, privateKeyPath)
+	if err != nil {
+		return err
+	}
+	defer sshClient.Close()
+
+	if err = sshClient.RemoteForward(localAddress, targetAddress); err != nil {
+		return errors.Wrap(err, "failed to forward to remote port")
+	}
+	return nil
+}
+
 // StartEnvd creates the container for the given tag and container name.
 func (e dockerEngine) StartEnvd(ctx context.Context, so StartOptions) (*StartResult, error) {
 	logger := logrus.WithFields(logrus.Fields{

pkg/envd/engine.go

@@ -57,6 +57,8 @@ type SSHClient interface {
 		startResult *StartResult) (sshconfig.EntryOptions, error)
 	Attach(name, iface, privateKeyPath string,
 		startResult *StartResult, g ir.Graph) error
+	LocalForward(iface, privateKeyPath string, startResult *StartResult, localAddress, targetAddress string) error
+	RemoteForward(iface, privateKeyPath string, startResult *StartResult, localAddress, targetAddress string) error
 }
 
 type ImageClient interface {

pkg/envd/envdserver.go

@@ -152,24 +152,32 @@ func (e envdServerEngine) GenerateSSHConfig(name, iface, privateKeyPath string,
 	return eo, nil
 }
 
+func (e envdServerEngine) newSSHClient(username, server string, port int, privateKeyPath string) (ssh.Client, error) {
+	opt := ssh.DefaultOptions()
+	opt.PrivateKeyPath = privateKeyPath
+	opt.Port = port
+	opt.AgentForwarding = false
+	opt.User = username
+	opt.Server = server
+	sshClient, err := ssh.NewClient(opt)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to create the ssh client")
+	}
+	return sshClient, nil
+}
+
 func (e envdServerEngine) Attach(name, iface, privateKeyPath string, startResult *StartResult, g ir.Graph) error {
 	username, err := sshname.Username(e.Loginname, startResult.Name)
 	if err != nil {
 		return errors.Wrap(err, "failed to get the username")
 	}
 
 	outputChannel := make(chan error)
-	opt := ssh.DefaultOptions()
-	opt.PrivateKeyPath = privateKeyPath
-	opt.Port = startResult.SSHPort
-	opt.AgentForwarding = false
-	opt.User = username
-	opt.Server = iface
-
-	sshClient, err := ssh.NewClient(opt)
+	sshClient, err := e.newSSHClient(username, iface, startResult.SSHPort, privateKeyPath)
 	if err != nil {
-		outputChannel <- errors.Wrap(err, "failed to create the ssh client")
+		outputChannel <- err
 	}
+	defer sshClient.Close()
 
 	ports := startResult.Ports
 
@@ -207,6 +215,41 @@ func (e envdServerEngine) Attach(name, iface, privateKeyPath string, startResult
 	return nil
 }
 
+func (e *envdServerEngine) LocalForward(iface, privateKeyPath string, startResult *StartResult, localAddress, targetAddress string) error {
+	username, err := sshname.Username(e.Loginname, startResult.Name)
+	if err != nil {
+		return errors.Wrap(err, "failed to get the username")
+	}
+	sshClient, err := e.newSSHClient(username, iface, startResult.SSHPort, privateKeyPath)
+	if err != nil {
+		return err
+	}
+	defer sshClient.Close()
+
+	if err = sshClient.LocalForward(localAddress, targetAddress); err != nil {
+		return errors.Wrap(err, "failed to forward to local port")
+	}
+	return nil
+}
+
+func (e *envdServerEngine) RemoteForward(iface, privateKeyPath string, startResult *StartResult, localAddress, targetAddress string) error {
+	username, err := sshname.Username(e.Loginname, startResult.Name)
+	if err != nil {
+		return errors.Wrap(err, "failed to get the username")
+	}
+
+	sshClient, err := e.newSSHClient(username, iface, startResult.SSHPort, privateKeyPath)
+	if err != nil {
+		return err
+	}
+	defer sshClient.Close()
+
+	if err = sshClient.RemoteForward(localAddress, targetAddress); err != nil {
+		return errors.Wrap(err, "failed to forward to remote port")
+	}
+	return nil
+}
+
 func (e *envdServerEngine) ListEnvRuntimeGraph(ctx context.Context, env string) (*ir.RuntimeGraph, error) {
 	resp, err := e.EnvironmentGet(ctx, env)
 	if err != nil {