feat: support registry mirror (#96)

Signed-off-by: Keming <kemingyang@tensorchord.ai>

mdz/pkg/cmd/server.go

@@ -7,8 +7,10 @@ import (
 )
 
 var (
-	serverVerbose         bool
-	serverPollingInterval time.Duration = 3 * time.Second
+	serverVerbose                 bool
+	serverPollingInterval         time.Duration = 3 * time.Second
+	serverRegistryMirrorName      string
+	serverRegistryMirrorEndpoints []string
 )
 
 // serverCmd represents the server command

mdz/pkg/cmd/server_join.go

@@ -28,6 +28,10 @@ func init() {
 
 	// Cobra supports local flags which will only run when this command
 	// is called directly, e.g.:
+	serverJoinCmd.Flags().StringVarP(&serverRegistryMirrorName, "mirror-name", "",
+		"", "Mirror name of the registry")
+	serverJoinCmd.Flags().StringArrayVarP(&serverRegistryMirrorEndpoints, "mirror-endpoints", "",
+		[]string{}, "Mirror endpoints of the registry")
 }
 
 func commandServerJoin(cmd *cobra.Command, args []string) error {
@@ -36,6 +40,10 @@ func commandServerJoin(cmd *cobra.Command, args []string) error {
 		OutputStream:  cmd.ErrOrStderr(),
 		RetryInternal: serverPollingInterval,
 		ServerIP:      args[0],
+		Mirror: server.Mirror{
+			Name:      serverRegistryMirrorName,
+			Endpoints: serverRegistryMirrorEndpoints,
+		},
 	})
 	if err != nil {
 		cmd.PrintErrf("Failed to join the cluster: %s\n", errors.Cause(err))

mdz/pkg/cmd/server_start.go

@@ -48,6 +48,10 @@ func init() {
 	serverStartCmd.Flags().MarkHidden("version")
 	serverStartCmd.Flags().BoolVarP(&serverStartWithGPU, "force-gpu", "g",
 		false, "Start the server with GPU support (ignore the GPU detection)")
+	serverStartCmd.Flags().StringVarP(&serverRegistryMirrorName, "mirror-name", "",
+		"", "Mirror name of the registry")
+	serverStartCmd.Flags().StringArrayVarP(&serverRegistryMirrorEndpoints, "mirror-endpoints", "",
+		[]string{}, "Mirror endpoints of the registry")
 }
 
 func commandServerStart(cmd *cobra.Command, args []string) error {
@@ -64,6 +68,10 @@ func commandServerStart(cmd *cobra.Command, args []string) error {
 		Domain:        domain,
 		Version:       serverStartVersion,
 		ForceGPU:      serverStartWithGPU,
+		Mirror: server.Mirror{
+			Name:      serverRegistryMirrorName,
+			Endpoints: serverRegistryMirrorEndpoints,
+		},
 	})
 	if err != nil {
 		cmd.PrintErrf("Failed to start the server: %s\n", errors.Cause(err))

mdz/pkg/server/engine.go

@@ -14,13 +14,23 @@ type Options struct {
 	Verbose       bool
 	OutputStream  io.Writer
 	Runtime       Runtime
+	Mirror        Mirror
 	RetryInternal time.Duration
 	ServerIP      string
 	Domain        *string
 	Version       string
 	ForceGPU      bool
 }
 
+type Mirror struct {
+	Name      string
+	Endpoints []string
+}
+
+func (m *Mirror) Configured() bool {
+	return m.Name != "" && len(m.Endpoints) > 0
+}
+
 type Runtime string
 
 var (
@@ -38,6 +48,9 @@ type Result struct {
 }
 
 func NewStart(o Options) (*Engine, error) {
+	if o.Verbose {
+		fmt.Fprintf(o.OutputStream, "Starting the server with config: %+v\n", o)
+	}
 	var engine *Engine
 	switch o.Runtime {
 	case RuntimeDocker:
@@ -54,6 +67,9 @@ func NewStart(o Options) (*Engine, error) {
 			options: o,
 			Steps: []Step{
 				// Install k3s and related tools.
+				&k3sPrepare{
+					options: o,
+				},
 				&k3sInstallStep{
 					options: o,
 				},

mdz/pkg/server/k3s_prepare.go

@@ -0,0 +1,66 @@
+package server
+
+import (
+	_ "embed"
+	"fmt"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"syscall"
+	"text/template"
+)
+
+//go:embed registries.yaml
+var registriesContent string
+
+const mirrorPath = "/etc/rancher/k3s"
+const mirrorFile = "registries.yaml"
+
+// k3sPrepare install everything required by k3s.
+type k3sPrepare struct {
+	options Options
+}
+
+func (s *k3sPrepare) Run() error {
+	if !s.options.Mirror.Configured() {
+		return nil
+	}
+	fmt.Fprintf(s.options.OutputStream, "🚧 Configure the mirror...\n")
+
+	tmpl, err := template.New("registries").Parse(registriesContent)
+	if err != nil {
+		panic(err)
+	}
+	buf := strings.Builder{}
+	err = tmpl.Execute(&buf, s.options.Mirror)
+	if err != nil {
+		panic(err)
+	}
+
+	cmd := exec.Command("/bin/sh", "-c", fmt.Sprintf(
+		"sudo mkdir -p %s && sudo tee %s > /dev/null << EOF\n%s\nEOF",
+		mirrorPath,
+		filepath.Join(mirrorPath, mirrorFile),
+		buf.String(),
+	))
+	cmd.SysProcAttr = &syscall.SysProcAttr{
+		Pdeathsig: syscall.SIGKILL,
+	}
+	if s.options.Verbose {
+		cmd.Stderr = s.options.OutputStream
+		cmd.Stdout = s.options.OutputStream
+	} else {
+		cmd.Stdout = nil
+		cmd.Stderr = nil
+	}
+	err = cmd.Run()
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (s *k3sPrepare) Verify() error {
+	return nil
+}

mdz/pkg/server/registries.yaml

@@ -0,0 +1,4 @@
+mirrors:
+  {{ .Name }}:
+    endpoint:
+      {{ range $endpoint := .Endpoints }}- "{{ $endpoint }}"{{ end }}