feat: add krun & bug fixes

Author: xtlsoft

PROJECT

@@ -37,4 +37,12 @@ resources:
   webhooks:
     defaulting: true
     webhookVersion: v1
+- core: true
+  group: batch
+  kind: Job
+  path: k8s.io/api/batch/v1
+  version: v1
+  webhooks:
+    defaulting: true
+    webhookVersion: v1
 version: "3"

cmd/krun/init.go

@@ -0,0 +1,291 @@
+package main
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"log"
+	"net"
+	"os"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/urfave/cli/v2"
+)
+
+func determineServiceName() (string, error) {
+	// Firstly get from environment variable KRUN_SERVICE
+	if serviceName := os.Getenv("KRUN_SERVICE"); serviceName != "" {
+		return serviceName, nil
+	}
+	// Secondly get from current hostname
+	hostname, err := os.Hostname()
+	if err != nil {
+		return "", wrapError("determine_service_name", err)
+	}
+	firstElem := strings.Split(hostname, ".")[0]
+	lastStashPos := strings.LastIndex(firstElem, "-")
+	if lastStashPos != -1 {
+		suffix := firstElem[lastStashPos+1:]
+		if _, err := strconv.ParseInt(suffix, 10, 64); err != nil {
+			return "", wrapError("determine_service_name", errors.New("hostname not indexed"))
+		}
+		firstElem = firstElem[:lastStashPos]
+		return firstElem, nil
+	}
+	return "", wrapError("determine_service_name", errors.New("hostname not indexed"))
+}
+
+func getIPsOfService(serviceName string) ([]string, error) {
+	ips, err := net.LookupIP(serviceName)
+	if err != nil {
+		return nil, wrapError("get_ips_of_service", err)
+	}
+	ipStrings := make([]string, 0, len(ips))
+	for _, ip := range ips {
+		ipStrings = append(ipStrings, ip.String())
+	}
+	return ipStrings, nil
+}
+
+func getRealHostnameAndIPOfPod(origIP string) (hostname, ip string, err error) {
+	krunPath := os.Args[0]
+	rawStr, err := waitGetOutput(sshRunner.Command(origIP, []string{krunPath, "_init"}))
+	if err != nil {
+		return "", "", wrapError("get_real_hostname_and_ip_of_pod", err)
+	}
+	var peerInfo InitPeerInfo
+	err = json.Unmarshal([]byte(rawStr), &peerInfo)
+	if err != nil {
+		return "", "", wrapError("get_real_hostname_and_ip_of_pod", err)
+	}
+	return peerInfo.Hostname, peerInfo.IP, nil
+}
+
+type InitPeerInfo struct {
+	Hostname string `json:"h"`
+	IP       string `json:"i"`
+}
+
+func initPeerInfoHandler(c *cli.Context) error {
+	h, err := getShortHostname()
+	if err != nil {
+		return err
+	}
+	i, err := getFastestIP()
+	if err != nil {
+		return err
+	}
+	peerInfo := &InitPeerInfo{
+		Hostname: h,
+		IP:       i,
+	}
+	peerInfoBytes, err := json.Marshal(peerInfo)
+	if err != nil {
+		return err
+	}
+	fmt.Println(string(peerInfoBytes))
+	return nil
+}
+
+// 200ms, adjust to your need
+const ReadinessWaitInterval = 200 * time.Millisecond
+const ReadinessDelay = 100 * time.Millisecond
+
+func initCommandHandler(c *cli.Context) error {
+	// waitMinStr := os.Getenv("KRUN_WAIT_MIN")
+	// if waitMinStr == "" {
+	// 	waitMinStr = "1"
+	// }
+	// waitMin, err := strconv.Atoi(waitMinStr)
+	// if err != nil {
+	// 	return wrapError("init_command_handler", err)
+	// }
+	waitMin := c.Int("wait-min")
+
+	serviceName, err := determineServiceName()
+	if err != nil {
+		return err
+	}
+	log.Println("Service name:", serviceName)
+	log.Println("Wait min:", waitMin)
+
+	var ipsOfService []string
+	for len(ipsOfService) < waitMin {
+		ipsOfService, err = getIPsOfService(serviceName)
+		if err != nil {
+			return err
+		}
+		log.Println("Discovered IPs of service:", ipsOfService)
+		if len(ipsOfService) < waitMin {
+			log.Printf("%v out of %v are ready. Still waiting...\n", len(ipsOfService), waitMin)
+			time.Sleep(ReadinessWaitInterval)
+		}
+	}
+
+	// Such delay is used to make sure underlying services are started up correctly
+	time.Sleep(ReadinessDelay)
+
+	results, err := PoolExecute(func(ip string) (*InitPeerInfo, error) {
+		hostname, realIP, err := getRealHostnameAndIPOfPod(ip)
+		return &InitPeerInfo{
+			Hostname: hostname,
+			IP:       realIP,
+		}, err
+	}, ipsOfService, c.Int("threads"))
+	if err != nil {
+		return err
+	}
+
+	for _, result := range results {
+		log.Printf("Peer (Hostname: %v, IP: %v)\n", result.Hostname, result.IP)
+	}
+
+	noHosts := c.Bool("no-hosts")
+
+	if !noHosts {
+		err = writeHosts(results)
+		if err != nil {
+			return err
+		}
+	}
+
+	log.Println("Saving state file to", c.String("state-file")+"...")
+	r, err := json.Marshal(results)
+	if err != nil {
+		return wrapError("init_command_handler", err)
+	}
+	err = os.WriteFile(c.String("state-file"), r, 0644)
+	if err != nil {
+		return wrapError("init_command_handler", err)
+	}
+
+	return nil
+}
+
+func writeHosts(results []*InitPeerInfo) error {
+	log.Println("Generating hosts file...")
+
+	myHostname, err := getShortHostname()
+	if err != nil {
+		return err
+	}
+
+	orig, err := os.ReadFile("/etc/hosts")
+	if err != nil {
+		return wrapError("write_hosts", err)
+	}
+
+	origStr := string(orig)
+
+	startBanner := "\n# --- KRUN GENERATED HOSTS START ---"
+	startBanner += "\n# Generated by krun, do not modify manually\n\n"
+	endedBanner := "\n# --- KRUN GENERATED HOSTS ENDED ---\n\n"
+
+	content := ""
+	for _, result := range results {
+		// FIX: skip the current pod, otherwise
+		//      intel MPI hydra fails to launch
+		//      processes across nodes.
+		if result.Hostname == myHostname {
+			continue
+		}
+		content += fmt.Sprintf("%v\t%v\n", result.IP, result.Hostname)
+	}
+
+	startPos := strings.Index(origStr, startBanner)
+	endPos := strings.Index(origStr, endedBanner)
+	if startPos == -1 || endPos == -1 {
+		origStr += startBanner + content + endedBanner
+	} else {
+		origStr = origStr[:startPos] + startBanner + content + origStr[endPos:]
+	}
+
+	err = os.WriteFile("/etc/hosts", []byte(origStr), 0644)
+	if err != nil {
+		return wrapError("write_hosts", err)
+	}
+
+	return nil
+}
+
+func loadState(c *cli.Context) ([]*InitPeerInfo, error) {
+	r, err := os.ReadFile(c.String("state-file"))
+	if err != nil {
+		if os.IsNotExist(err) {
+			return nil, errors.New("failed to load state file, please run `krun init` first")
+		}
+		return nil, wrapError("load_state", err)
+	}
+	var results []*InitPeerInfo
+	err = json.Unmarshal(r, &results)
+	if err != nil {
+		return nil, wrapError("load_state", err)
+	}
+	return results, nil
+}
+
+func hostsCommandHandler(c *cli.Context) error {
+	results, err := loadState(c)
+	if err != nil {
+		return err
+	}
+
+	short := c.Bool("short")
+	noHeaders := c.Bool("no-headers")
+	hostname := c.Bool("hostname")
+	ip := c.Bool("ip")
+	suffix := c.String("suffix")
+
+	if !hostname && !ip {
+		hostname = true
+		if !short {
+			ip = true
+		}
+	}
+
+	fieldDelim := "\t"
+	itemDelim := "\n"
+
+	if short {
+		fieldDelim = " "
+		itemDelim = ","
+	}
+
+	if !short && !noHeaders {
+		if hostname {
+			fmt.Print("HOSTNAME")
+		}
+		if hostname && ip {
+			fmt.Print(fieldDelim)
+		}
+		if ip {
+			fmt.Print("IP")
+		}
+	}
+
+	first := true
+
+	for _, result := range results {
+		if first && (short || noHeaders) {
+			first = false
+		} else {
+			fmt.Print(itemDelim)
+		}
+		if hostname {
+			fmt.Print(result.Hostname)
+		}
+		if hostname && ip {
+			fmt.Print(fieldDelim)
+		}
+		if ip {
+			fmt.Print(result.IP)
+		}
+		fmt.Print(suffix)
+	}
+
+	fmt.Println()
+
+	return nil
+}

cmd/krun/log.go

@@ -0,0 +1,10 @@
+package main
+
+import "fmt"
+
+func wrapError(wrapper string, err error) error {
+	if err != nil {
+		return fmt.Errorf("%s: %w", wrapper, err)
+	}
+	return nil
+}