add: port gen, port check, show on run

internal/misc.go

@@ -90,46 +90,47 @@ func (p *Path) Join(subpath string) *Path {
 	return &Path{path: filepath.Join(p.path, subpath)}
 }
 
-func makeDefaultDirectories(projectName, experimentName, runName string) (string, string, error) {
+func makeDefaultDirectories(projectName, experimentName, runName string) (string, string, string, error) {
 	home, err := os.UserHomeDir()
 	if err != nil {
-		return "", "", errors.WithMessage(err, "failed to get user home directory")
+		return "", "", "", errors.WithMessage(err, "failed to get user home directory")
 	}
 
 	cacheDir := Path{path: filepath.Join(home, ".cache")}
 	if err = cacheDir.mkdirIfNotExists(); err != nil {
-		return "", "", errors.WithMessage(err, "failed to create cache directory")
+		return "", "", "", errors.WithMessage(err, "failed to create cache directory")
 	}
 
 	projectDir := cacheDir.Join(projectName)
 	if err = projectDir.mkdirIfNotExists(); err != nil {
-		return "", "", errors.WithMessage(err, "failed to create project directory")
+		return "", "", "", errors.WithMessage(err, "failed to create project directory")
 	}
 
 	experimentsDir := projectDir.Join("experiments")
 	if err = experimentsDir.mkdirIfNotExists(); err != nil {
-		return "", "", errors.WithMessage(err, "failed to create experiments directory")
+		return "", "", "", errors.WithMessage(err, "failed to create experiments directory")
 	}
 
 	experimentDir := experimentsDir.Join(experimentName)
 	if err = experimentDir.mkdirIfNotExists(); err != nil {
-		return "", "", errors.WithMessage(err, "failed to create experiment directory")
+		return "", "", "", errors.WithMessage(err, "failed to create experiment directory")
 	}
 
 	med := []string{"checkpoints", "sharded-checkpoints", "lr-schedules", "logs", "plots", "results"}
 	for _, dir := range med {
 		if err = experimentDir.Join(dir).Join(runName).mkdirIfNotExists(); err != nil {
-			return "", "", errors.WithMessagef(err, "failed to create %s directory for experiment %s and run name %s", dir, experimentName, runName)
+			return "", "", "", errors.WithMessagef(err, "failed to create %s directory for experiment %s and run name %s", dir, experimentName, runName)
 		}
 	}
 
-  guestPath := Path{path: filepath.Join("/home/nonroot/", ".cache", projectName, "experiments")}
-  return cacheDir.path, guestPath.Join("logs").Join(runName).path, nil
+	checkpointDir := cacheDir.Join("higgsfield").Join(projectName).Join("experiments").Join(experimentName).Join(runName)
+	if err = checkpointDir.mkdirIfNotExists(); err != nil {
+		return "", "", "", errors.WithMessagef(err, "failed to create checkpoint directory for experiment %s and run name %s", experimentName, runName)
+	}
+	guestPath := Path{path: filepath.Join("/home/nonroot/", ".cache", projectName, "experiments")}
+	return cacheDir.path, guestPath.Join("logs").Join(runName).path, checkpointDir.path, nil
 }
 
-
-
-
 func exitIfError(flag string, err error) {
 	if err != nil {
 		fmt.Printf("cannot parse %s: %v\n", flag, err)
@@ -138,30 +139,30 @@ func exitIfError(flag string, err error) {
 
 }
 
-func ParseOrExit[T ~string| ~int | ~[]string](cmd *cobra.Command, flag string) T {
-  got := parseOrExitInternal[T](cmd, flag)
-  return got.(T)
+func ParseOrExit[T ~string | ~int | ~[]string](cmd *cobra.Command, flag string) T {
+	got := parseOrExitInternal[T](cmd, flag)
+	return got.(T)
 }
 
-func parseOrExitInternal[T ~string| ~int | ~[]string](cmd *cobra.Command, flag string) interface{} {
+func parseOrExitInternal[T ~string | ~int | ~[]string](cmd *cobra.Command, flag string) interface{} {
 	var value T
 	switch v := any(value).(type) {
 	case string:
 		v, err := cmd.Flags().GetString(flag)
-    exitIfError(flag, err)
-    return v
+		exitIfError(flag, err)
+		return v
 	case int:
 		v, err := cmd.Flags().GetInt(flag)
-    exitIfError(flag, err)
-     return v
+		exitIfError(flag, err)
+		return v
 	case []string:
 		v, err := cmd.Flags().GetStringSlice(flag)
-    exitIfError(flag, err)
-    return v
+		exitIfError(flag, err)
+		return v
 	default:
 		fmt.Printf("cannot parse %s: unknown type %T\n", flag, v)
-    os.Exit(1)
+		os.Exit(1)
 	}
 
-  return nil
+	return nil
 }

internal/network.go

@@ -0,0 +1,26 @@
+package internal
+
+import (
+	"fmt"
+	"math/rand"
+	"net"
+)
+
+func isPortAvailable(port int) bool {
+	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
+	if err != nil {
+		return false
+	}
+	defer listener.Close()
+
+	return true
+}
+
+func GeneratePort() int {
+	port := rand.Intn(65535-1024) + 1024
+	for !isPortAvailable(port) {
+		port = rand.Intn(65535-1024) + 1024
+	}
+
+	return port
+}

internal/run.go

@@ -26,12 +26,30 @@ func Run(args RunArgs) {
 	portIsAvailable(args.Port)
 	nodeNum := len(args.Hosts)
 
-	hostCachePath, guestLogPath, err := makeDefaultDirectories(args.ProjectName, args.ExperimentName, args.RunName)
+	if !isPortAvailable(args.Port) {
+		fmt.Printf("port %d is not available\n", args.Port)
+		os.Exit(1)
+	}
+
+	hostCachePath, guestLogPath, checkpointDir, err := makeDefaultDirectories(args.ProjectName, args.ExperimentName, args.RunName)
 	if err != nil {
 		fmt.Printf("failed to create directories: %v\n", err)
 		os.Exit(1)
 	}
 
+	fmt.Printf(`
+|=============================================================================================================
+|=
+|=  Training info:
+|=  🛠🛠🛠
+|=
+|=  EXPERIMENT NAME =       %s
+|=  RUN NAME =              %s
+|=  MODEL CHECKPOINT PATH = %s
+|=
+|============================================================================================================="""
+`, args.ExperimentName, args.RunName, checkpointDir)
+
 	cmd, cmdArgs := buildArgs(
 		nodeNum,
 		rank,

main.go

@@ -75,23 +75,36 @@ func decodeSecrets() *cobra.Command {
 }
 
 func randomName() *cobra.Command {
-  cmd := &cobra.Command{
-    Use: "random-name",
-    Short: "Generate a random name",
-    Run: func(cmd *cobra.Command, args []string) {
-      fmt.Print(namesgenerator.GetRandomName(0))
-    },
-  }
-
-  return cmd
+	cmd := &cobra.Command{
+		Use:   "random-name",
+		Short: "Generate a random name",
+		Run: func(cmd *cobra.Command, args []string) {
+			fmt.Print(namesgenerator.GetRandomName(0))
+		},
+	}
+
+	return cmd
+}
+
+func randomPort() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "random-port",
+		Short: "Generate a random port",
+		Run: func(cmd *cobra.Command, args []string) {
+			fmt.Print(internal.GeneratePort())
+		},
+	}
+
+	return cmd
 }
 
 func main() {
 	experimentCmd.AddCommand(runCmdFunc())
 	experimentCmd.AddCommand(killCmdFunc())
 
 	rootCmd.AddCommand(decodeSecrets())
-  rootCmd.AddCommand(randomName())
+	rootCmd.AddCommand(randomName())
+	rootCmd.AddCommand(randomPort())
 	rootCmd.AddCommand(experimentCmd)
 
 	if err := rootCmd.Execute(); err != nil {