feat(docker): Add entrypoint and ports in image config (#533)

* feat: Add entrypoint to image config

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Remove jupyter port

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

README.md

@@ -187,17 +187,6 @@ $ cd envd-quick-start && envd up
 (envd) ➜ demo git:(master) ✗ # You are in the container-based environment!
 ```
 
-### Play with the environment
-
-You can run `ssh envd-quick-start.envd` to reconnect if you exit from the environment. Or you can execute `git` or `python` commands inside.
-
-```bash
-$ python demo.py
-[2 3 4]
-$ git fetch
-$
-```
-
 ### Set up Jupyter notebook
 
 Please edit the `build.envd` to enable jupyter notebook:
@@ -209,7 +198,7 @@ def build():
  "numpy",
  ])
  shell("zsh")
- config.jupyter(password="", port=8888)
+ config.jupyter(password="")
 ```
 
 You can get the endpoint of the running Jupyter notebook via `envd get envs`.
@@ -218,7 +207,7 @@ You can get the endpoint of the running Jupyter notebook via `envd get envs`.
 $ envd up --detach
 $ envd get env
 NAME JUPYTER SSH TARGET CONTEXT IMAGE GPU CUDA CUDNN STATUS CONTAINER ID 
-envd-quick-start http://localhost:8888 envd-quick-start.envd /home/gaocegege/code/envd-quick-start envd-quick-start:dev false <none> <none> Up 54 seconds bd3f6a729e94
+envd-quick-start http://localhost:42779 envd-quick-start.envd /home/gaocegege/code/envd-quick-start envd-quick-start:dev false <none> <none> Up 54 seconds bd3f6a729e94
 ```
 
 ## Roadmap 🗂️

pkg/app/build.go

@@ -15,6 +15,7 @@
 package app
 
 import (
+ "fmt"
  "path/filepath"
 
  "github.com/cockroachdb/errors"
@@ -94,7 +95,7 @@ func build(clicontext *cli.Context) error {
  tag := clicontext.String("tag")
  if tag == "" {
  logrus.Debug("tag not specified, using default")
- tag = fileutil.Base(buildContext)
+ tag = fmt.Sprintf("%s:%s", fileutil.Base(buildContext), "dev")
  }
 
  logger := logrus.WithFields(logrus.Fields{

pkg/app/up.go

@@ -146,6 +146,7 @@ func up(clicontext *cli.Context) error {
  if err := builder.Build(clicontext.Context, clicontext.Path("public-key")); err != nil {
  return errors.Wrap(err, "failed to build the image")
  }
+ // Do not attach GPU if the flag is set.
  gpuEnable := clicontext.Bool("no-gpu")
  var gpu bool
  if gpuEnable {
@@ -169,14 +170,14 @@ func up(clicontext *cli.Context) error {
  }
  }
 
- sshPort, err := netutil.GetFreePort()
+ sshPortInHost, err := netutil.GetFreePort()
  if err != nil {
  return errors.Wrap(err, "failed to get a free port")
  }
  numGPUs := builder.NumGPUs()
 
  containerID, containerIP, err := dockerClient.StartEnvd(clicontext.Context,
- tag, ctr, buildContext, gpu, numGPUs, sshPort, *ir.DefaultGraph, clicontext.Duration("timeout"),
+ tag, ctr, buildContext, gpu, numGPUs, sshPortInHost, *ir.DefaultGraph, clicontext.Duration("timeout"),
  clicontext.StringSlice("volume"))
  if err != nil {
  return errors.Wrap(err, "failed to start the envd environment")
@@ -185,14 +186,14 @@ func up(clicontext *cli.Context) error {
 
  logrus.Debugf("Add entry %s to SSH config. at %s", buildContext, containerIP)
  if err = sshconfig.AddEntry(
- ctr, localhost, sshPort, clicontext.Path("private-key")); err != nil {
+ ctr, localhost, sshPortInHost, clicontext.Path("private-key")); err != nil {
  logrus.Infof("failed to add entry %s to your SSH config file: %s", ctr, err)
  return errors.Wrap(err, "failed to add entry to your SSH config file")
  }
 
  if !detach {
  sshClient, err := ssh.NewClient(
- localhost, "envd", sshPort, true, clicontext.Path("private-key"), "")
+ localhost, "envd", sshPortInHost, true, clicontext.Path("private-key"), "")
  if err != nil {
  return errors.Wrap(err, "failed to create the ssh client")
  }

pkg/builder/builder.go

@@ -159,21 +159,30 @@ func (b generalBuilder) compile(ctx context.Context, pub string) (*llb.Definitio
  return def, nil
 }
 
-func (b generalBuilder) labels(ctx context.Context) (string, error) {
+func (b generalBuilder) imageConfig(ctx context.Context) (string, error) {
  labels, err := ir.Labels()
  if err != nil {
  return "", errors.Wrap(err, "failed to get labels")
  }
+ ports, err := ir.ExposedPorts()
+ if err != nil {
+ return "", errors.Wrap(err, "failed to get expose ports")
+ }
  labels[types.ImageLabelContext] = b.buildContextDir
- data, err := ImageConfigStr(labels)
+
+ ep, err := ir.Entrypoint(b.buildContextDir)
+ if err != nil {
+ return "", errors.Wrap(err, "failed to get entrypoint")
+ }
+ data, err := ImageConfigStr(labels, ports, ep)
  if err != nil {
  return "", errors.Wrap(err, "failed to get image config")
  }
  return data, nil
 }
 
 func (b generalBuilder) build(ctx context.Context, def *llb.Definition, pw progresswriter.Writer) error {
- labels, err := b.labels(ctx)
+ imageConfig, err := b.imageConfig(ctx)
  if err != nil {
  return errors.Wrap(err, "failed to get labels")
  }
@@ -198,7 +207,7 @@ func (b generalBuilder) build(ctx context.Context, def *llb.Definition, pw progr
  Attrs: map[string]string{
  "name": b.tag,
  // Ref https://github.com/r2d4/mockerfile/blob/140c6a912bbfdae220febe59ab535ef0acba0e1f/pkg/build/build.go#L65
- "containerimage.config": labels,
+ "containerimage.config": imageConfig,
  },
  Output: func(map[string]string) (io.WriteCloser, error) {
  return pipeW, nil

pkg/builder/util.go

@@ -29,13 +29,16 @@ import (
  "github.com/sirupsen/logrus"
 )
 
-func ImageConfigStr(labels map[string]string) (string, error) {
+func ImageConfigStr(labels map[string]string,
+ ports map[string]struct{}, entrypoint []string) (string, error) {
  pl := platforms.Normalize(platforms.DefaultSpec())
  img := v1.Image{
  Config: v1.ImageConfig{
- Labels: labels,
- WorkingDir: "/",
- Env: []string{"PATH=" + DefaultPathEnv(pl.OS)},
+ Labels: labels,
+ WorkingDir: "/",
+ Env: []string{"PATH=" + DefaultPathEnv(pl.OS)},
+ ExposedPorts: ports,
+ Entrypoint: entrypoint,
  },
  Architecture: pl.Architecture,
  // Refer to https://github.com/tensorchord/envd/issues/269#issuecomment-1152944914

pkg/config/config.go

@@ -19,4 +19,6 @@ const (
  PrivateKeyFile = "id_rsa_envd"
  PublicKeyFile = "id_rsa_envd.pub"
  ContainerAuthorizedKeysPath = "/var/envd/authorized_keys"
+ SSHPortInContainer = 2222
+ JupyterPortInContainer = 8888
 )

pkg/docker/docker.go

@@ -35,8 +35,10 @@ import (
  "github.com/moby/term"
  "github.com/sirupsen/logrus"
 
+ envdconfig "github.com/tensorchord/envd/pkg/config"
  "github.com/tensorchord/envd/pkg/lang/ir"
  "github.com/tensorchord/envd/pkg/util/fileutil"
+ "github.com/tensorchord/envd/pkg/util/netutil"
 )
 
 const (
@@ -305,7 +307,7 @@ func (c generalClient) StartBuildkitd(ctx context.Context, tag, name, mirror str
 
 // StartEnvd creates the container for the given tag and container name.
 func (c generalClient) StartEnvd(ctx context.Context, tag, name, buildContext string,
- gpuEnabled bool, numGPUs int, sshPort int, g ir.Graph, timeout time.Duration,
+ gpuEnabled bool, numGPUs int, sshPortInHost int, g ir.Graph, timeout time.Duration,
  mountOptionsStr []string) (string, string, error) {
  logger := logrus.WithFields(logrus.Fields{
  "tag": tag,
@@ -315,21 +317,13 @@ func (c generalClient) StartEnvd(ctx context.Context, tag, name, buildContext st
  "build-context": buildContext,
  })
  config := &container.Config{
- Image: tag,
- User: "envd",
- Entrypoint: []string{
- "tini",
- "--",
- "bash",
- "-c",
- },
+ Image: tag,
+ User: "envd",
  ExposedPorts: nat.PortSet{},
  }
  base := fileutil.Base(buildContext)
  base = filepath.Join("/home/envd", base)
  config.WorkingDir = base
- config.Entrypoint = append(config.Entrypoint,
- entrypointSH(g, config.WorkingDir, sshPort))
 
  mountOption := make([]mount.Mount, len(mountOptionsStr)+1)
  for i, option := range mountOptionsStr {
@@ -365,22 +359,27 @@ func (c generalClient) StartEnvd(ctx context.Context, tag, name, buildContext st
  }
 
  // Configure ssh port.
- natPort := nat.Port(fmt.Sprintf("%d/tcp", sshPort))
+ natPort := nat.Port(fmt.Sprintf("%d/tcp", envdconfig.SSHPortInContainer))
  hostConfig.PortBindings[natPort] = []nat.PortBinding{
  {
  HostIP: localhost,
- HostPort: strconv.Itoa(sshPort),
+ HostPort: strconv.Itoa(sshPortInHost),
  },
  }
- config.ExposedPorts[natPort] = struct{}{}
 
+ var jupyterPortInHost int
  // TODO(gaocegege): Avoid specific logic to set the port.
  if g.JupyterConfig != nil {
- natPort := nat.Port(fmt.Sprintf("%d/tcp", g.JupyterConfig.Port))
+ var err error
+ jupyterPortInHost, err = netutil.GetFreePort()
+ if err != nil {
+ return "", "", errors.Wrap(err, "failed to get a free port")
+ }
+ natPort := nat.Port(fmt.Sprintf("%d/tcp", envdconfig.JupyterPortInContainer))
  hostConfig.PortBindings[natPort] = []nat.PortBinding{
  {
  HostIP: localhost,
- HostPort: strconv.Itoa(int(g.JupyterConfig.Port)),
+ HostPort: strconv.Itoa(jupyterPortInHost),
  },
  }
  config.ExposedPorts[natPort] = struct{}{}
@@ -391,7 +390,7 @@ func (c generalClient) StartEnvd(ctx context.Context, tag, name, buildContext st
  hostConfig.DeviceRequests = deviceRequests(numGPUs)
  }
 
- config.Labels = labels(name, g.JupyterConfig, sshPort)
+ config.Labels = labels(name, g.JupyterConfig, sshPortInHost, jupyterPortInHost)
 
  logger = logger.WithFields(logrus.Fields{
  "entrypoint": config.Entrypoint,
@@ -408,7 +407,8 @@ func (c generalClient) StartEnvd(ctx context.Context, tag, name, buildContext st
  logger.Warnf("run with warnings: %s", w)
  }
 
- if err := c.ContainerStart(ctx, resp.ID, types.ContainerStartOptions{}); err != nil {
+ if err := c.ContainerStart(
+ ctx, resp.ID, types.ContainerStartOptions{}); err != nil {
  errCause := errors.UnwrapAll(err)
  // Hack to check if the port is already allocated.
  if strings.Contains(errCause.Error(), "port is already allocated") {

pkg/docker/label.go

@@ -24,12 +24,14 @@ import (
  "github.com/tensorchord/envd/pkg/types"
 )
 
-func labels(name string, jupyterConfig *ir.JupyterConfig, sshPort int) map[string]string {
+func labels(name string, jupyterConfig *ir.JupyterConfig,
+ sshPortInHost, jupyterPortInHost int) map[string]string {
  res := make(map[string]string)
  res[types.ContainerLabelName] = name
- res[types.ContainerLabelSSHPort] = strconv.Itoa(sshPort)
+ res[types.ContainerLabelSSHPort] = strconv.Itoa(sshPortInHost)
  if jupyterConfig != nil {
- res[types.ContainerLabelJupyterAddr] = fmt.Sprintf("http://localhost:%d", jupyterConfig.Port)
+ res[types.ContainerLabelJupyterAddr] =
+ fmt.Sprintf("http://localhost:%d", jupyterPortInHost)
  }
  return res
 }