feat(context): Get ssh hostname from context, instead of hard-coded s…

…tring (#1020) * fix: Use envd-server in the context Signed-off-by: Ce Gao <cegao@tensorchord.ai> * feat: Add ssh address discovery Signed-off-by: Ce Gao <cegao@tensorchord.ai> * fix: Remove unused struct Signed-off-by: Ce Gao <cegao@tensorchord.ai> Signed-off-by: Ce Gao <cegao@tensorchord.ai>
tensorchord · Oct 13, 2022 · 0fc4d8d · 0fc4d8d
1 parent 51af450
commit 0fc4d8d
Show file tree

Hide file tree

Showing 8 changed files with 118 additions and 15 deletions.
diff --git a/docs/proposals/20220603-kubernetes-vendor.md b/docs/proposals/20220603-kubernetes-vendor.md
@@ -63,15 +63,15 @@ Users may use `envd` to build the image, and use it on Kubernetes. Thus they nee
 The end-to-end process will be:
 
 ```
-$ envd context create --name test --builder-name test --use --builder kube-pod --runner server --runner-addr http://localhost:2222
+$ envd context create --name test --builder-name test --use --builder kube-pod --runner server --runner-addr http://localhost:8080
 $ envd login
 $ envd build
 $ envd push
 $ envd run --env test --image test
 
 or 
 
-$ envd context create --name test --builder-name test --use --builder kube-pod --runner server --runner-addr http://localhost:2222
+$ envd context create --name test --builder-name test --use --builder kube-pod --runner server --runner-addr http://localhost:8080
 $ envd login
 $ envd up
 ```

diff --git a/pkg/app/create.go b/pkg/app/create.go
@@ -91,9 +91,14 @@ func create(clicontext *cli.Context) error {
  logrus.Debugf("container %s is running", res.Name)
 
  logrus.Debugf("add entry %s to SSH config.", res.Name)
+ hostname, err := c.GetSSHHostname()
+ if err != nil {
+ return errors.Wrap(err, "failed to get the ssh hostname")
+ }
+
  eo := sshconfig.EntryOptions{
  Name: res.Name,
- IFace: localhost,
+ IFace: hostname,
  Port: res.SSHPort,
  PrivateKeyPath: clicontext.Path("private-key"),
  EnableHostKeyCheck: false,
@@ -111,6 +116,8 @@ func create(clicontext *cli.Context) error {
  opt.Port = res.SSHPort
  opt.AgentForwarding = false
  opt.User = res.Name
+ opt.Server = hostname
+
  sshClient, err := ssh.NewClient(opt)
  if err != nil {
  return errors.Wrap(err, "failed to create the ssh client")

diff --git a/pkg/app/up.go b/pkg/app/up.go
@@ -30,10 +30,6 @@ import (
  "github.com/tensorchord/envd/pkg/types"
 )
 
-const (
- localhost = "127.0.0.1"
-)
-
 var CommandUp = &cli.Command{
  Name: "up",
  Category: CategoryBasic,
@@ -120,9 +116,13 @@ var CommandUp = &cli.Command{
 }
 
 func up(clicontext *cli.Context) error {
+ c, err := home.GetManager().ContextGetCurrent()
+ if err != nil {
+ return errors.Wrap(err, "failed to get the current context")
+ }
  buildOpt, err := ParseBuildOpt(clicontext)
  if err != nil {
- return err
+ return errors.Wrap(err, "failed to parse the build options")
  }
 
  ctr := filepath.Base(buildOpt.BuildContextDir)
@@ -197,9 +197,14 @@ func up(clicontext *cli.Context) error {
  logrus.Debugf("container %s is running", res.Name)
 
  logrus.Debugf("add entry %s to SSH config.", ctr)
+ hostname, err := c.GetSSHHostname()
+ if err != nil {
+ return errors.Wrap(err, "failed to get the ssh hostname")
+ }
+
  eo := sshconfig.EntryOptions{
  Name: ctr,
- IFace: localhost,
+ IFace: hostname,
  Port: res.SSHPort,
  PrivateKeyPath: clicontext.Path("private-key"),
  EnableHostKeyCheck: false,
@@ -218,6 +223,8 @@ func up(clicontext *cli.Context) error {
  if err != nil {
  return errors.Wrap(err, "failed to create the ssh client")
  }
+ opt.Server = hostname
+
  if err := sshClient.Attach(); err != nil {
  return errors.Wrap(err, "failed to attach to the container")
  }

diff --git a/pkg/envd/factory.go b/pkg/envd/factory.go
@@ -36,12 +36,24 @@ func New(ctx context.Context, opt Options) (Engine, error) {
  if opt.Context.Runner == types.RunnerTypeEnvdServer {
  ac, err := home.GetManager().AuthGetCurrent()
  if err != nil {
- return nil, err
+ return nil, errors.Wrap(err, "failed to get the auth information")
  }
 
- cli, err := envdclient.NewClientWithOpts(envdclient.FromEnv)
+ // Get the runner host.
+ opts := []envdclient.Opt{
+ envdclient.WithTLSClientConfigFromEnv(),
+ }
+ if opt.Context.RunnerAddress != nil {
+ opts = append(opts,
+ envdclient.WithHost(*opt.Context.RunnerAddress))
+ } else {
+ opts = append(opts,
+ envdclient.WithHostFromEnv())
+ }
+
+ cli, err := envdclient.NewClientWithOpts(opts...)
  if err != nil {
- return nil, err
+ return nil, errors.Wrap(err, "failed to create the envd-server client")
  }
  return &envdServerEngine{
  Client: cli,
@@ -51,7 +63,7 @@ func New(ctx context.Context, opt Options) (Engine, error) {
  cli, err := client.NewClientWithOpts(
  client.FromEnv, client.WithAPIVersionNegotiation())
  if err != nil {
- return nil, err
+ return nil, errors.Wrap(err, "failed to create the docker client")
  }
  return &dockerEngine{
  Client: cli,

diff --git a/pkg/ssh/ssh.go b/pkg/ssh/ssh.go
@@ -57,7 +57,6 @@ type Options struct {
 
 func DefaultOptions() Options {
  return Options{
- Server: "localhost",
  User: "envd",
  Auth: true,
  PrivateKeyPwd: "",

diff --git a/pkg/types/envd.go b/pkg/types/envd.go
@@ -22,6 +22,7 @@ import (
  "github.com/docker/docker/api/types"
  "github.com/moby/buildkit/util/system"
 
+ "github.com/tensorchord/envd/pkg/util/netutil"
  "github.com/tensorchord/envd/pkg/version"
 )
 
@@ -278,3 +279,16 @@ func parsePyPICommands(lst string) ([]string, error) {
  err := json.Unmarshal([]byte(lst), &pkgs)
  return pkgs, err
 }
+
+func (c Context) GetSSHHostname() (string, error) {
+ if c.RunnerAddress == nil {
+ return "localhost", nil
+ }
+
+ // TODO(gaocegege): Check ENVD_SERVER_HOST.
+ hostname, err := netutil.GetHost(*c.RunnerAddress)
+ if err != nil {
+ return "", err
+ }
+ return hostname, nil
+}
diff --git a/pkg/util/netutil/netutil.go b/pkg/util/netutil/netutil.go
@@ -14,8 +14,13 @@
 
 package netutil
 
-import "net"
+import (
+ "fmt"
+ "net"
+ "net/url"
+)
 
+// GetFreePort returns an available port in the host.
 func GetFreePort() (int, error) {
  l, err := net.Listen("tcp", ":0")
  if err != nil {
@@ -24,3 +29,16 @@ func GetFreePort() (int, error) {
  defer l.Close()
  return l.Addr().(*net.TCPAddr).Port, nil
 }
+
+// GetHost get the IP address from the address.
+func GetHost(addr string) (string, error) {
+ if u, err := url.Parse(addr); err != nil {
+ return "", err
+ } else {
+ h := u.Hostname()
+ if h == "" {
+ return "", fmt.Errorf("failed to get the hostname from %s", addr)
+ }
+ return h, nil
+ }
+}
diff --git a/pkg/util/netutil/netutil_test.go b/pkg/util/netutil/netutil_test.go
@@ -25,3 +25,49 @@ func TestGetFreePort(t *testing.T) {
  assert.NoError(t, err)
  assert.NotEqual(t, port, 0)
 }
+
+func TestGetHost(t *testing.T) {
+ tcs := []struct {
+ host string
+ expected string
+ err bool
+ }{
+ {
+ host: "https://localhost:8080",
+ expected: "localhost",
+ err: false,
+ },
+ {
+ host: "localhost:8080",
+ expected: "",
+ err: true,
+ },
+ {
+ host: "http://localhost:8080",
+ expected: "localhost",
+ err: false,
+ },
+ {
+ host: "http://1.1.1.1:8080",
+ expected: "1.1.1.1",
+ err: false,
+ },
+ }
+ for _, tc := range tcs {
+ host, err := GetHost(tc.host)
+ if tc.err == true {
+ if err == nil {
+ t.Error("expect to get the error, but got nil")
+ }
+ continue
+ }
+ if tc.err == false {
+ if err != nil {
+ t.Errorf("unexpected error %v", err)
+ }
+ if host != tc.expected {
+ t.Errorf("expected %s, got %s", tc.expected, host)
+ }
+ }
+ }
+}