feat(CLI): Add dockerhub mirror flag (#242)

Signed-off-by: Ce Gao <cegao@tensorchord.ai>
tensorchord · Jun 8, 2022 · 289cb07 · 289cb07
1 parent 1a50052
commit 289cb07
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 6 deletions.
diff --git a/cmd/envd/bootstrap.go b/cmd/envd/bootstrap.go
@@ -37,6 +37,11 @@ var CommandBootstrap = &cli.Command{
  Usage: "Add envd autocompletions",
  Value: true,
  },
+ &cli.StringFlag{
+ Name: "dockerhub-mirror",
+ Usage: "Dockerhub mirror to use",
+ Aliases: []string{"m"},
+ },
  },
 
  Action: bootstrap,
@@ -64,7 +69,7 @@ func bootstrap(clicontext *cli.Context) error {
 
  if buildkit {
  logrus.Debug("bootstrap the buildkitd container")
- bkClient, err := buildkitd.NewClient(clicontext.Context)
+ bkClient, err := buildkitd.NewClient(clicontext.Context, clicontext.String("dockerhub-mirror"))
  if err != nil {
  return errors.Wrap(err, "failed to create buildkit client")
  }

diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go
@@ -67,7 +67,7 @@ func New(ctx context.Context,
  }),
  }
 
- cli, err := buildkitd.NewClient(ctx)
+ cli, err := buildkitd.NewClient(ctx, "")
  if err != nil {
  return nil, errors.Wrap(err, "failed to create buildkit client")
  }

diff --git a/pkg/buildkitd/buildkitd.go b/pkg/buildkitd/buildkitd.go
@@ -47,15 +47,17 @@ type Client interface {
 type generalClient struct {
  containerName string
  image string
+ mirror string
 
  *client.Client
  logger *logrus.Entry
 }
 
-func NewClient(ctx context.Context) (Client, error) {
+func NewClient(ctx context.Context, mirror string) (Client, error) {
  c := &generalClient{
  containerName: viper.GetString(flag.FlagBuildkitdContainer),
  image: viper.GetString(flag.FlagBuildkitdImage),
+ mirror: mirror,
  }
  c.logger = logrus.WithFields(logrus.Fields{
  "container": c.containerName,
@@ -103,7 +105,7 @@ func (c *generalClient) maybeStart(ctx context.Context,
 
  if !created {
  c.logger.Debug("container not created, creating...")
- if _, err := dockerClient.StartBuildkitd(ctx, c.image, c.containerName); err != nil {
+ if _, err := dockerClient.StartBuildkitd(ctx, c.image, c.containerName, c.mirror); err != nil {
  return "", err
  }
  if err := dockerClient.WaitUntilRunning(ctx, c.containerName, runningTimeout); err != nil {

diff --git a/pkg/docker/docker.go b/pkg/docker/docker.go
@@ -48,7 +48,7 @@ type Client interface {
  // Start creates the container for the given tag and container name.
  StartEnvd(ctx context.Context, tag, name, buildContext string,
  gpuEnabled bool, g ir.Graph, timeout time.Duration, mountOptionsStr []string) (string, string, error)
- StartBuildkitd(ctx context.Context, tag, name string) (string, error)
+ StartBuildkitd(ctx context.Context, tag, name, mirror string) (string, error)
 
  IsRunning(ctx context.Context, name string) (bool, error)
  IsCreated(ctx context.Context, name string) (bool, error)
@@ -225,10 +225,11 @@ func (c generalClient) Destroy(ctx context.Context, name string) (string, error)
 }
 
 func (g generalClient) StartBuildkitd(ctx context.Context,
- tag, name string) (string, error) {
+ tag, name, mirror string) (string, error) {
  logger := logrus.WithFields(logrus.Fields{
  "tag": tag,
  "container": name,
+ "mirror": mirror,
  })
  logger.Debug("starting buildkitd")
  if _, _, err := g.ImageInspectWithRaw(ctx, tag); err != nil {
@@ -252,6 +253,15 @@ func (g generalClient) StartBuildkitd(ctx context.Context,
  config := &container.Config{
  Image: tag,
  }
+ if mirror != "" {
+ cfg := fmt.Sprintf(`
+[registry."docker.io"]
+ mirrors = ["%s"]`, mirror)
+ config.Entrypoint = []string{
+ "/bin/sh", "-c", fmt.Sprintf("mkdir /etc/buildkit && echo '%s' > /etc/buildkit/buildkitd.toml && buildkitd", cfg),
+ }
+ logger.Debugf("setting buildkit config: %s", cfg)
+ }
  hostConfig := &container.HostConfig{
  Privileged: true,
  }