path overlay, rename list cmd to manage (#92)

* path overlay, rename list cmd to manage * make workload name RFC-compliant * Give Ray head pod 1Gi storage --------- Co-authored-by: Antti-Ville Suni <[email protected]>
silogen · Feb 4, 2025 · 5459734 · 5459734
1 parent 8ed7d83
commit 5459734
Show file tree

Hide file tree

Showing 15 changed files with 142 additions and 982 deletions.
diff --git a/pkg/cli/apply/run.go b/pkg/cli/apply/run.go
@@ -50,16 +50,21 @@ func RunApply(workload workloads.Workload, workloadMeta any) error {
 		return nil
 	}
 
+	// Discover workload files
+	if err := loadFileList(&execFlags); err != nil {
+		return fmt.Errorf("error loading file list: %w", err)
+	}
+
 	// Generate workload configuration
 	workloadConfig, err := workload.GenerateTemplateContext(execFlags)
 	if err != nil {
 		return fmt.Errorf("error generating workload config: %w", err)
 	}
 
-	// Load custom configuration, if provided
+	// Load custom configuration, if available
 	var customConfig any
-	if execFlags.CustomConfigPath != "" {
-		customConfig, err = loadCustomConfig(execFlags.CustomConfigPath)
+	if customConfigFile, ok := execFlags.WorkloadFiles[workloads.CustomTemplateValuesFilename]; ok {
+		customConfig, err = loadCustomConfig(customConfigFile)
 		if err != nil {
 			return fmt.Errorf("error loading custom config: %w", err)
 		}
@@ -77,14 +82,8 @@ func RunApply(workload workloads.Workload, workloadMeta any) error {
 		logrus.Debugf("No explicit name provided, using name: %s", metaFlags.Name)
 	}
 
-	// Parse environment variables
-	if execFlags.EnvFilePath == "" {
-		envFilePath = filepath.Join(execFlags.Path, workloads.EnvFilename)
-	} else {
-		envFilePath = execFlags.EnvFilePath
-	}
-
-	if err := parseEnvFile(envFilePath, &metaFlags); err != nil {
+	// Parse environment variables, if any
+	if err := parseEnvFile(execFlags.WorkloadFiles[workloads.EnvFilename], &metaFlags); err != nil {
 		return fmt.Errorf("error parsing environment: %w", err)
 	}
 
@@ -158,6 +157,48 @@ func RunApply(workload workloads.Workload, workloadMeta any) error {
 	return nil
 }
 
+func loadFileList(execFlags *workloads.ExecFlags) error {
+	if execFlags.Path == "" && execFlags.OverlayPath != "" {
+		return fmt.Errorf("cannot load workload with an overlay path without base path")
+	}
+	if execFlags.Path == "" {
+		return nil
+	}
+
+	files := map[string]string{}
+
+	entries, err := os.ReadDir(execFlags.Path)
+	if err != nil {
+		return fmt.Errorf("error reading directory: %w", err)
+	}
+
+	for _, entry := range entries {
+		if !entry.IsDir() {
+			files[entry.Name()] = filepath.Join(execFlags.Path, entry.Name())
+		}
+	}
+
+	if execFlags.OverlayPath != "" {
+		overlayFiles, err := os.ReadDir(execFlags.OverlayPath)
+		if err != nil {
+			return fmt.Errorf("error reading directory: %w", err)
+		}
+		for _, overlayFile := range overlayFiles {
+			if !overlayFile.IsDir() {
+				files[overlayFile.Name()] = filepath.Join(execFlags.OverlayPath, overlayFile.Name())
+			}
+		}
+	}
+
+	for k, v := range files {
+		logrus.Debugf("Discovered file %s from %s", k, v)
+	}
+
+	execFlags.WorkloadFiles = files
+
+	return nil
+}
+
 func findDefaultStorageFlags(ctx context.Context, clientset kubernetes.Clientset, namespace string) (*workloads.StorageSchedulingFlags, error) {
 	namespaceObject, err := clientset.CoreV1().Namespaces().Get(ctx, namespace, metav1.GetOptions{})
 	if err != nil {
@@ -244,7 +285,7 @@ func makeWorkloadName(path string, image string, version string, currentUser str
 		components = append(components, version)
 	}
 
-	return strings.Join(components, "-")
+	return baseutils.MakeRFC1123Compliant(strings.Join(components, "-"))
 }
 
 func sanitizeStringForKubernetes(path string) string {

diff --git a/pkg/cli/apply/utils.go b/pkg/cli/apply/utils.go
@@ -33,25 +33,23 @@ const defaultGpuNodeLabelKey = "beta.amd.com/gpu.family.AI"
 
 // Exec flags
 var (
-	dryRun           bool
-	createNamespace  bool
-	template         string
-	path             string
-	gpuNodeLabelKey  string
-	customConfigPath string
-	envFilePath      string
+	dryRun          bool
+	createNamespace bool
+	template        string
+	path            string
+	overlayPath     string
+	gpuNodeLabelKey string
 )
 
 // AddExecFlags adds flags that are needed for the execution of apply functions
 func AddExecFlags(cmd *cobra.Command) {
 	cmd.Flags().BoolVarP(&createNamespace, "create-namespace", "", false, "Create namespace if it does not exist")
 	cmd.Flags().BoolVarP(&dryRun, "dry-run", "d", false, "Print the generated workload manifest without submitting it")
 	cmd.Flags().StringVarP(&path, "path", "p", "", "Path to directory for workload code, entrypoint/serveconfig, env-file, etc. Either image or path is mandatory")
+	cmd.Flags().StringVarP(&overlayPath, "overlay-path", "o", "", "Additional overlay path. Files from both path and overlay-path are combined, if the file exists in both, the one from overlay-path is used")
 	// TODO: remove gpuNodeLabelKey and have this logic be handled by the operator
 	cmd.Flags().StringVarP(&gpuNodeLabelKey, "gpu-node-label-key", "", defaultGpuNodeLabelKey, fmt.Sprintf("Optional node label key used to specify the resource flavor GPU count. Defaults to %s", defaultGpuNodeLabelKey))
 	cmd.Flags().StringVarP(&template, "template", "t", "", "Optional path to a custom template to use for the workload. If not provided, a default template will be used unless template file found in workload directory")
-	cmd.Flags().StringVarP(&customConfigPath, "custom-config", "c", "", "Optional path to a custom YAML configuration file whose contents are made available in the template")
-	cmd.Flags().StringVarP(&envFilePath, "env-file", "", "", "Optional path to env file. Defaults to 'env' in workload code directory")
 }
 
 func GetExecFlags() workloads.ExecFlags {
@@ -60,9 +58,8 @@ func GetExecFlags() workloads.ExecFlags {
 		DryRun:                        dryRun,
 		Template:                      template,
 		Path:                          path,
+		OverlayPath:                   overlayPath,
 		ResourceFlavorGpuNodeLabelKey: gpuNodeLabelKey,
-		CustomConfigPath:              customConfigPath,
-		EnvFilePath:                   envFilePath,
 	}
 }
 
@@ -188,10 +185,9 @@ type Config struct {
 	DryRun                  bool   `yaml:"dryRun"`
 	CreateNamespace         bool   `yaml:"createNamespace"`
 	Path                    string `yaml:"path"`
+	OverlayPath             string `yaml:"overlayPath"`
 	GpuNodeLabelKey         string `yaml:"gpuNodeLabelKey"`
 	Template                string `yaml:"template"`
-	CustomConfig            string `yaml:"customConfig"`
-	EnvFile                 string `yaml:"envFile"`
 	Name                    string `yaml:"name"`
 	Namespace               string `yaml:"namespace"`
 	Image                   string `yaml:"image"`
@@ -237,10 +233,9 @@ func ApplyConfigToFlags(cmd *cobra.Command, config *Config) {
 	setFlag("dry-run", fmt.Sprintf("%v", config.DryRun))
 	setFlag("create-namespace", fmt.Sprintf("%v", config.CreateNamespace))
 	setFlag("path", config.Path)
+	setFlag("overlay-path", config.OverlayPath)
 	setFlag("gpu-node-label-key", config.GpuNodeLabelKey)
 	setFlag("template", config.Template)
-	setFlag("custom-config", config.CustomConfig)
-	setFlag("env-file", config.EnvFile)
 
 	// MetaFlags
 	setFlag("name", config.Name)
@@ -256,15 +251,25 @@ func ApplyConfigToFlags(cmd *cobra.Command, config *Config) {
 }
 
 func PreRunLoadConfig(cmd *cobra.Command, _ []string) error {
-	if path == "" {
+	if path == "" && overlayPath == "" {
 		return nil
 	}
 
-	config, err := LoadConfigFromPath(path)
+	// First try to load config from the overlay path
+	config, err := LoadConfigFromPath(overlayPath)
 	if err != nil {
 		return fmt.Errorf("failed to load config: %w", err)
 	}
 
+	if config == nil {
+		// If no overlay, try the main path
+		config, err = LoadConfigFromPath(path)
+		if err != nil {
+			return fmt.Errorf("failed to load config: %w", err)
+		}
+		logrus.Infof("Loaded config from %s", path)
+	}
+
 	if config != nil {
 		ApplyConfigToFlags(cmd, config)
 		logrus.Infof("Configuration loaded from %s", filepath.Join(path, workloads.KaiwoconfigFilename))

diff --git a/pkg/cli/list.go b/pkg/cli/list.go
@@ -32,9 +32,9 @@ var (
 
 func BuildListCmd() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "list [workloadType] [workloadName]",
+		Use:   "manage [workloadType] [workloadName]",
 		Args:  cobra.MaximumNArgs(2),
-		Short: "List workloads",
+		Short: "Manage workloads",
 		RunE: func(cmd *cobra.Command, args []string) error {
 			workloadType := ""
 			workloadName := ""

diff --git a/pkg/k8s/utils.go b/pkg/k8s/utils.go
@@ -19,66 +19,13 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
-	"slices"
 
 	"github.com/sirupsen/logrus"
 	"gopkg.in/yaml.v3"
 	corev1 "k8s.io/api/core/v1"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 )
 
-func isBinaryFile(content []byte) bool {
-	return bytes.Contains(content, []byte{0})
-}
-
-// GenerateConfigMapFromDir generates a ConfigMap from a directory
-func GenerateConfigMapFromDir(dir string, name string, namespace string, skipFiles []string) (*corev1.ConfigMap, error) {
-	files, err := os.ReadDir(dir)
-	if err != nil {
-		return nil, fmt.Errorf("failed to read directory: %w", err)
-	}
-
-	data := make(map[string]string)
-
-	for _, file := range files {
-		if file.IsDir() {
-			continue
-		}
-
-		if slices.Contains(skipFiles, file.Name()) {
-			continue
-		}
-
-		filePath := filepath.Join(dir, file.Name())
-		content, err := os.ReadFile(filePath)
-		if err != nil {
-			return nil, fmt.Errorf("failed to read file %s: %w", filePath, err)
-		}
-
-		// Skip binary files
-		if isBinaryFile(content) {
-			logrus.Warnf("Skipping binary file: %s", file.Name())
-			continue
-		}
-		data[file.Name()] = string(content)
-	}
-
-	if len(data) == 0 {
-		return nil, nil
-	}
-
-	configMap := &corev1.ConfigMap{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      name,
-			Namespace: namespace,
-		},
-		Data: data,
-	}
-
-	return configMap, nil
-}
-
 type SecretVolume struct {
 	Name       string
 	SecretName string

diff --git a/pkg/utils/user.go b/pkg/utils/user.go
@@ -40,7 +40,7 @@ func GetCurrentUser() (string, error) {
 		parts := strings.Split(userEmail, "@")
 		username := strings.Split(parts[0], "-")[0]
 		domain := strings.ReplaceAll(parts[1], ".", "-")
-		return makeRFC1123Compliant(fmt.Sprintf("%s-%s", username, domain)), nil
+		return MakeRFC1123Compliant(fmt.Sprintf("%s-%s", username, domain)), nil
 	}
 
 	logrus.Warn("USER_EMAIL not set. Falling back to UNIX username and hostname")
@@ -55,10 +55,10 @@ func GetCurrentUser() (string, error) {
 
 	k8sCompatibleHostname := strings.ReplaceAll(hostname, ".", "-")
 
-	return makeRFC1123Compliant(fmt.Sprintf("%s-%s", currentUser.Username, k8sCompatibleHostname)), nil
+	return MakeRFC1123Compliant(fmt.Sprintf("%s-%s", currentUser.Username, k8sCompatibleHostname)), nil
 }
 
-func makeRFC1123Compliant(input string) string {
+func MakeRFC1123Compliant(input string) string {
 	input = strings.ToLower(input)
 
 	rfc1123Regex := regexp.MustCompile(`[^a-z0-9.-]`)

diff --git a/pkg/workloads/apply.go b/pkg/workloads/apply.go
@@ -15,11 +15,13 @@
 package workloads
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
 	"os"
 	"path"
+	"slices"
 	"strings"
 	"text/template"
 
@@ -64,14 +66,16 @@ func ApplyWorkload(
 	}
 
 	if execFlags.Path != "" {
-		configMapResource, err = generateConfigMapManifest(execFlags.Path, workload, templateContext.Meta)
+		configMapResource, err = generateConfigMapManifest(execFlags.WorkloadFiles, workload, templateContext.Meta)
 		if err != nil {
 			return fmt.Errorf("failed to generate configmap resource: %w", err)
 		}
 		if configMapResource != nil {
 			resources = append(resources, configMapResource)
 			templateContext.Meta.HasConfigMap = true
 		}
+	} else if execFlags.OverlayPath != "" {
+		return fmt.Errorf("overlay path set without setting the path")
 	}
 
 	workloadTemplate, err := getWorkloadTemplate(execFlags, workload)
@@ -199,15 +203,54 @@ func generateNamespaceManifestIfNotExists(
 	}, nil
 }
 
+func isBinaryFile(content []byte) bool {
+	return bytes.Contains(content, []byte{0})
+}
+
 // generateConfigMapManifest adds a config map resource
-func generateConfigMapManifest(path string, workload Workload, metaConfig MetaFlags) (*corev1.ConfigMap, error) {
-	configMap, err := k8s.GenerateConfigMapFromDir(path, metaConfig.Name, metaConfig.Namespace, workload.IgnoreFiles())
-	if err != nil {
-		return nil, fmt.Errorf("failed to generate ConfigMap: %w", err)
+func generateConfigMapManifest(files map[string]string, workload Workload, metaConfig MetaFlags) (*corev1.ConfigMap, error) {
+	configMap := &corev1.ConfigMap{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      metaConfig.Name,
+			Namespace: metaConfig.Namespace,
+		},
+		Data: map[string]string{},
+	}
+
+	skipFiles := workload.IgnoreFiles()
+
+	for fileName, filePath := range files {
+		if slices.Contains(skipFiles, fileName) {
+			continue
+		}
+
+		info, err := os.Stat(filePath)
+		if err != nil {
+			return nil, fmt.Errorf("failed to stat file %s: %w", filePath, err)
+		}
+
+		if info.Size() > 950e3 {
+			logrus.Warnf("Skipping file %s in %s as it is too large", fileName, filePath)
+			continue
+		}
+
+		content, err := os.ReadFile(filePath)
+		if err != nil {
+			return nil, fmt.Errorf("failed to read file %s: %w", filePath, err)
+		}
+
+		// Skip binary files
+		if isBinaryFile(content) {
+			logrus.Warnf("Skipping binary file %s in %s", fileName, filePath)
+			continue
+		}
+		configMap.Data[fileName] = string(content)
 	}
-	if configMap != nil {
+
+	if len(configMap.Data) > 0 {
 		return configMap, nil
 	}
+
 	return nil, nil
 }