pr feedback

Signed-off-by: Geoff Flarity <[email protected]>

Author: gflarity

operator/e2e_testing/setup/helm.go renamed to operator/e2e/setup/helm.go

@@ -24,7 +24,7 @@ import (
 	"path/filepath"
 	"strings"
 
-	"github.com/sirupsen/logrus"
+	"github.com/ai-dynamo/grove/operator/e2e/utils"
 	"helm.sh/helm/v3/pkg/action"
 	"helm.sh/helm/v3/pkg/chart/loader"
 	"helm.sh/helm/v3/pkg/cli"
@@ -57,7 +57,7 @@ type HelmInstallConfig struct {
 	// HelmLoggerFunc is called for Helm operation logging.
 	HelmLoggerFunc func(format string, v ...interface{})
 	// Logger is the full logger for component operations.
-	Logger *logrus.Logger
+	Logger *utils.Logger
 	// RepoURL is the base URL of the Helm repository (optional, for direct chart downloads).
 	RepoURL string
 }

operator/e2e_testing/setup/k8s_clusters.go renamed to operator/e2e/setup/k8s_clusters.go

@@ -24,7 +24,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/ai-dynamo/grove/operator/e2e_testing/utils"
+	"github.com/ai-dynamo/grove/operator/e2e/utils"
 	"github.com/docker/docker/api/types/container"
 	dockerclient "github.com/docker/docker/client"
 	"github.com/k3d-io/k3d/v5/pkg/client"
@@ -95,7 +95,7 @@ func DefaultClusterConfig() ClusterConfig {
 }
 
 // SetupCompleteK3DCluster creates a complete k3d cluster with Grove, Kai Scheduler, and NVIDIA GPU Operator
-func SetupCompleteK3DCluster(ctx context.Context, cfg ClusterConfig, skaffoldYAMLPath string, logger *logrus.Logger) (*rest.Config, func(), error) {
+func SetupCompleteK3DCluster(ctx context.Context, cfg ClusterConfig, skaffoldYAMLPath string, logger *utils.Logger) (*rest.Config, func(), error) {
 	restConfig, cleanup, err := SetupK3DCluster(ctx, cfg, logger)
 	if err != nil {
 		return nil, cleanup, err
@@ -216,9 +216,10 @@ func SetupCompleteK3DCluster(ctx context.Context, cfg ClusterConfig, skaffoldYAM
 }
 
 // SetupK3DCluster creates a k3d cluster and returns a REST config
-func SetupK3DCluster(ctx context.Context, cfg ClusterConfig, logger *logrus.Logger) (*rest.Config, func(), error) {
+func SetupK3DCluster(ctx context.Context, cfg ClusterConfig, logger *utils.Logger) (*rest.Config, func(), error) {
 	// k3d is very verbose, we don't want the INFO level logs unless the logger is set to DEBUG
-	if logger.GetLevel() == logrus.DebugLevel {
+	// k3d uses logrus internally, so we need to convert our level to logrus level
+	if logger.GetLevel() == utils.DebugLevel {
 		k3dlogger.Log().SetLevel(logrus.DebugLevel)
 	} else {
 		k3dlogger.Log().SetLevel(logrus.ErrorLevel)
@@ -336,7 +337,7 @@ func SetupK3DCluster(ctx context.Context, cfg ClusterConfig, logger *logrus.Logg
 }
 
 // InstallCoreComponents installs the core components (Grove via Skaffold, Kai Scheduler and NVIDIA GPU Operator via Helm)
-func InstallCoreComponents(ctx context.Context, restConfig *rest.Config, kaiConfig *HelmInstallConfig, nvidiaConfig *HelmInstallConfig, skaffoldYAMLPath string, registryPort string, logger *logrus.Logger) error {
+func InstallCoreComponents(ctx context.Context, restConfig *rest.Config, kaiConfig *HelmInstallConfig, nvidiaConfig *HelmInstallConfig, skaffoldYAMLPath string, registryPort string, logger *utils.Logger) error {
 	// use wait group to wait for all installations to complete
 	var wg sync.WaitGroup
 
@@ -443,7 +444,7 @@ func InstallCoreComponents(ctx context.Context, restConfig *rest.Config, kaiConf
 }
 
 // retryInstallation retries an installation function up to maxRetries times with a delay between attempts
-func retryInstallation(installFunc func() error, componentName string, maxRetries int, retryDelay time.Duration, logger *logrus.Logger) error {
+func retryInstallation(installFunc func() error, componentName string, maxRetries int, retryDelay time.Duration, logger *utils.Logger) error {
 	var lastErr error
 
 	for attempt := 1; attempt <= maxRetries; attempt++ {
@@ -481,7 +482,7 @@ func retryInstallation(installFunc func() error, componentName string, maxRetrie
 // 3. Deletes the not ready node from Kubernetes
 // 4. Finds and restarts the corresponding Docker container (node names match container names exactly)
 // 5. The restarted container will rejoin the cluster as a new node
-func StartNodeMonitoring(ctx context.Context, clientset *kubernetes.Clientset, logger *logrus.Logger) func() {
+func StartNodeMonitoring(ctx context.Context, clientset *kubernetes.Clientset, logger *utils.Logger) func() {
 	logger.Debug("🔍 Starting node monitoring for not ready nodes...")
 
 	// Create a context that can be cancelled to stop the monitoring
@@ -512,7 +513,7 @@ func StartNodeMonitoring(ctx context.Context, clientset *kubernetes.Clientset, l
 }
 
 // checkAndReplaceNotReadyNodes checks for nodes that are not ready and replaces them
-func checkAndReplaceNotReadyNodes(ctx context.Context, clientset *kubernetes.Clientset, logger *logrus.Logger) error {
+func checkAndReplaceNotReadyNodes(ctx context.Context, clientset *kubernetes.Clientset, logger *utils.Logger) error {
 	// List all nodes
 	nodes, err := clientset.CoreV1().Nodes().List(ctx, metav1.ListOptions{})
 	if err != nil {
@@ -554,7 +555,7 @@ func isNodeReady(node *v1.Node) bool {
 }
 
 // replaceNotReadyNode handles the process of replacing a not ready node
-func replaceNotReadyNode(ctx context.Context, node *v1.Node, clientset *kubernetes.Clientset, logger *logrus.Logger) error {
+func replaceNotReadyNode(ctx context.Context, node *v1.Node, clientset *kubernetes.Clientset, logger *utils.Logger) error {
 	nodeName := node.Name
 
 	// Step 1: Delete the node from Kubernetes
@@ -573,7 +574,7 @@ func replaceNotReadyNode(ctx context.Context, node *v1.Node, clientset *kubernet
 }
 
 // restartNodeContainer finds and restarts the Docker container corresponding to a k3d node
-func restartNodeContainer(ctx context.Context, nodeName string, logger *logrus.Logger) error {
+func restartNodeContainer(ctx context.Context, nodeName string, logger *utils.Logger) error {
 	// Create Docker client
 	dockerClient, err := dockerclient.NewClientWithOpts(dockerclient.FromEnv, dockerclient.WithAPIVersionNegotiation())
 	if err != nil {

operator/e2e_testing/setup/kai_scheduler.go renamed to operator/e2e/setup/kai_scheduler.go

@@ -19,12 +19,11 @@ package setup
 import (
 	"context"
 	"fmt"
-	"os"
 	"path/filepath"
 	"runtime"
 	"time"
 
-	"github.com/ai-dynamo/grove/operator/e2e_testing/utils"
+	"github.com/ai-dynamo/grove/operator/e2e/utils"
 	apiextensionsv1 "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1"
 	apiextensionsclientset "k8s.io/apiextensions-apiserver/pkg/client/clientset/clientset"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -95,14 +94,8 @@ func CreateDefaultKaiQueues(ctx context.Context, config *HelmInstallConfig) erro
 	_, currentFile, _, _ := runtime.Caller(0)
 	queuesPath := filepath.Join(filepath.Dir(currentFile), "../yaml/queues.yaml")
 
-	// Read the queues YAML file content
-	yamlContent, err := os.ReadFile(queuesPath)
-	if err != nil {
-		return fmt.Errorf("failed to read queues YAML file %s: %w", queuesPath, err)
-	}
-
-	// Apply the YAML content using the k8s client
-	appliedResources, err := utils.ApplyYAMLContent(ctx, string(yamlContent), "", config.RestConfig, config.Logger)
+	// Apply the YAML file using the k8s client
+	appliedResources, err := utils.ApplyYAMLFile(ctx, queuesPath, "", config.RestConfig, config.Logger)
 	if err != nil {
 		return fmt.Errorf("failed to apply queues YAML: %w", err)
 	}

operator/e2e_testing/setup/shared_cluster.go renamed to operator/e2e/setup/shared_cluster.go

@@ -23,10 +23,9 @@ import (
 	"sync"
 	"time"
 
-	"github.com/ai-dynamo/grove/operator/e2e_testing/utils"
+	"github.com/ai-dynamo/grove/operator/e2e/utils"
 	"github.com/docker/docker/api/types/image"
-	"github.com/docker/docker/client"
-	"github.com/sirupsen/logrus"
+	dockerclient "github.com/docker/docker/client"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime/schema"
@@ -35,17 +34,23 @@ import (
 	"k8s.io/client-go/rest"
 )
 
+const (
+	// relativeSkaffoldYAMLPath is the path to the skaffold.yaml file relative to the e2e/tests directory
+	relativeSkaffoldYAMLPath = "../../skaffold.yaml"
+
+	defaulPollInterval = 1 * time.Second
+)
+
 // SharedClusterManager manages a shared (singleton) k3d cluster for E2E tests
 type SharedClusterManager struct {
-	clientset                *kubernetes.Clientset
-	restConfig               *rest.Config
-	dynamicClient            dynamic.Interface
-	cleanup                  func()
-	logger                   *logrus.Logger
-	isSetup                  bool
-	workerNodes              []string
-	registryPort             string
-	relativeSkaffoldYAMLPath string
+	clientset     *kubernetes.Clientset
+	restConfig    *rest.Config
+	dynamicClient dynamic.Interface
+	cleanup       func()
+	logger        *utils.Logger
+	isSetup       bool
+	workerNodes   []string
+	registryPort  string
 }
 
 var (
@@ -54,11 +59,10 @@ var (
 )
 
 // SharedCluster returns the singleton shared cluster manager
-func SharedCluster(logger *logrus.Logger, skaffoldYAMLPath string) *SharedClusterManager {
+func SharedCluster(logger *utils.Logger) *SharedClusterManager {
 	once.Do(func() {
 		sharedCluster = &SharedClusterManager{
-			logger:                   logger,
-			relativeSkaffoldYAMLPath: skaffoldYAMLPath,
+			logger: logger,
 		}
 	})
 	return sharedCluster
@@ -110,7 +114,7 @@ func (scm *SharedClusterManager) Setup(ctx context.Context, testImages []string)
 
 	scm.logger.Info("🚀 Setting up shared k3d cluster for all e2e tests...")
 
-	restConfig, cleanup, err := SetupCompleteK3DCluster(ctx, customCfg, scm.relativeSkaffoldYAMLPath, scm.logger)
+	restConfig, cleanup, err := SetupCompleteK3DCluster(ctx, customCfg, relativeSkaffoldYAMLPath, scm.logger)
 	if err != nil {
 		return fmt.Errorf("failed to setup shared k3d cluster: %w", err)
 	}
@@ -198,8 +202,8 @@ func (scm *SharedClusterManager) CleanupWorkloads(ctx context.Context) error {
 		scm.logger.Warnf("failed to delete PodCliqueSets: %v", err)
 	}
 
-	// Step 2: Poll for all resources and pods to be cleaned up (max 15 seconds)
-	if err := scm.waitForAllResourcesAndPodsDeleted(ctx, 15*time.Second); err != nil {
+	// Step 2: Poll for all resources and pods to be cleaned up
+	if err := scm.waitForAllResourcesAndPodsDeleted(ctx, defaulPollInterval); err != nil {
 		scm.logger.Warnf("timeout waiting for resources and pods to be deleted: %v", err)
 		// List remaining resources and pods for debugging
 		scm.listRemainingResources(ctx)
@@ -292,9 +296,6 @@ func (scm *SharedClusterManager) resetNodeStates(ctx context.Context) error {
 
 // waitForAllResourcesAndPodsDeleted waits for all Grove resources and pods to be deleted
 func (scm *SharedClusterManager) waitForAllResourcesAndPodsDeleted(ctx context.Context, timeout time.Duration) error {
-	timeoutCtx, cancel := context.WithTimeout(ctx, timeout)
-	defer cancel()
-
 	// Define all resource types to check
 	resourceTypes := []struct {
 		group    string
@@ -308,59 +309,53 @@ func (scm *SharedClusterManager) waitForAllResourcesAndPodsDeleted(ctx context.C
 		{"scheduler.grove.io", "v1alpha1", "podgangs", "PodGangs"},
 	}
 
-	ticker := time.NewTicker(1 * time.Second)
-	defer ticker.Stop()
-
-	for {
-		select {
-		case <-timeoutCtx.Done():
-			return fmt.Errorf("timeout waiting for resources and pods to be deleted")
-		case <-ticker.C:
-			allResourcesDeleted := true
-			totalResources := 0
-
-			// Check Grove resources
-			for _, rt := range resourceTypes {
-				gvr := schema.GroupVersionResource{
-					Group:    rt.group,
-					Version:  rt.version,
-					Resource: rt.resource,
-				}
+	return utils.PollForCondition(ctx, timeout, defaulPollInterval, func() (bool, error) {
+		allResourcesDeleted := true
+		totalResources := 0
 
-				resourceList, err := scm.dynamicClient.Resource(gvr).List(ctx, metav1.ListOptions{})
-				if err != nil {
-					// If we can't list the resource type, assume it doesn't exist or is being deleted
-					continue
-				}
-
-				if len(resourceList.Items) > 0 {
-					allResourcesDeleted = false
-					totalResources += len(resourceList.Items)
-				}
+		// Check Grove resources
+		for _, rt := range resourceTypes {
+			gvr := schema.GroupVersionResource{
+				Group:    rt.group,
+				Version:  rt.version,
+				Resource: rt.resource,
 			}
 
-			// Check pods
-			allPodsDeleted := true
-			nonSystemPods := 0
-			pods, err := scm.clientset.CoreV1().Pods("default").List(ctx, metav1.ListOptions{})
-			if err == nil {
-				for _, pod := range pods.Items {
-					if !isSystemPod(&pod) {
-						allPodsDeleted = false
-						nonSystemPods++
-					}
-				}
+			resourceList, err := scm.dynamicClient.Resource(gvr).List(ctx, metav1.ListOptions{})
+			if err != nil {
+				// If we can't list the resource type, assume it doesn't exist or is being deleted
+				continue
 			}
 
-			if allResourcesDeleted && allPodsDeleted {
-				return nil
+			if len(resourceList.Items) > 0 {
+				allResourcesDeleted = false
+				totalResources += len(resourceList.Items)
 			}
+		}
 
-			if totalResources > 0 || nonSystemPods > 0 {
-				scm.logger.Debugf("⏳ Waiting for %d Grove resources and %d pods to be deleted...", totalResources, nonSystemPods)
+		// Check pods
+		allPodsDeleted := true
+		nonSystemPods := 0
+		pods, err := scm.clientset.CoreV1().Pods("default").List(ctx, metav1.ListOptions{})
+		if err == nil {
+			for _, pod := range pods.Items {
+				if !isSystemPod(&pod) {
+					allPodsDeleted = false
+					nonSystemPods++
+				}
 			}
 		}
-	}
+
+		if allResourcesDeleted && allPodsDeleted {
+			return true, nil
+		}
+
+		if totalResources > 0 || nonSystemPods > 0 {
+			scm.logger.Debugf("⏳ Waiting for %d Grove resources and %d pods to be deleted...", totalResources, nonSystemPods)
+		}
+
+		return false, nil
+	})
 }
 
 // listRemainingResources lists remaining Grove resources for debugging
@@ -436,7 +431,7 @@ func setupRegistryTestImages(registryPort string, images []string) error {
 	ctx := context.Background()
 
 	// Initialize Docker client
-	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
+	cli, err := dockerclient.NewClientWithOpts(dockerclient.FromEnv, dockerclient.WithAPIVersionNegotiation())
 	if err != nil {
 		return fmt.Errorf("failed to create Docker client: %w", err)
 	}

operator/e2e_testing/setup/skaffold.go renamed to operator/e2e/setup/skaffold.go

@@ -26,7 +26,7 @@ import (
 	"path/filepath"
 	"strings"
 
-	"github.com/sirupsen/logrus"
+	"github.com/ai-dynamo/grove/operator/e2e/utils"
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/clientcmd"
 	clientcmdapi "k8s.io/client-go/tools/clientcmd/api"
@@ -49,7 +49,7 @@ type SkaffoldInstallConfig struct {
 	// Env are environment variables required by skaffold.yaml (e.g., VERSION, LD_FLAGS).
 	Env map[string]string
 	// Logger is the logger for operations (optional, will use default if nil).
-	Logger *logrus.Logger
+	Logger *utils.Logger
 }
 
 // Validate validates the configuration.
@@ -69,7 +69,7 @@ func (c *SkaffoldInstallConfig) Validate() error {
 
 	// Set defaults
 	if c.Logger == nil {
-		c.Logger = logrus.New()
+		c.Logger = utils.NewTestLogger(utils.InfoLevel)
 	}
 	if c.Env == nil {
 		c.Env = make(map[string]string)
@@ -161,7 +161,7 @@ func runSkaffoldBuild(ctx context.Context, absSkaffoldPath, skaffoldDir, kubecon
 	cmd.Stdout = &stdout
 
 	// cmd.Stderr is logging, it'll get passed on at debug level only
-	cmd.Stderr = config.Logger.WriterLevel(logrus.DebugLevel)
+	cmd.Stderr = config.Logger.WriterLevel(utils.DebugLevel)
 
 	config.Logger.Debugf("   Running: skaffold %v", args)
 	if err := cmd.Run(); err != nil {
@@ -232,16 +232,16 @@ func runSkaffoldDeploy(ctx context.Context, absSkaffoldPath, skaffoldDir, kubeco
 	cmd.Env = append(cmd.Env, fmt.Sprintf("CONTAINER_REGISTRY=%s", config.PullRepo))
 
 	// log cmd.Stdout and cmd.Stderr at debug level only
-	cmd.Stdout = config.Logger.WriterLevel(logrus.DebugLevel)
-	cmd.Stderr = config.Logger.WriterLevel(logrus.DebugLevel)
+	cmd.Stdout = config.Logger.WriterLevel(utils.DebugLevel)
+	cmd.Stderr = config.Logger.WriterLevel(utils.DebugLevel)
 
 	config.Logger.Debugf("   Running: skaffold %v", args)
 	return cmd.Run()
 }
 
 // writeTemporaryKubeconfig converts a rest.Config to a kubeconfig file and writes it to a temporary location.
 // Returns the path to the temporary file and a cleanup function.
-func writeTemporaryKubeconfig(restConfig *rest.Config, logger *logrus.Logger) (string, func(), error) {
+func writeTemporaryKubeconfig(restConfig *rest.Config, logger *utils.Logger) (string, func(), error) {
 	// Create a temporary file for the kubeconfig
 	tmpFile, err := os.CreateTemp("", "kubeconfig-*.yaml")
 	if err != nil {

operator/e2e_testing/tests/gang_scheduling_test.go renamed to operator/e2e/tests/gang_scheduling_test.go

@@ -22,7 +22,7 @@ import (
 	"context"
 	"testing"
 
-	"github.com/ai-dynamo/grove/operator/e2e_testing/utils"
+	"github.com/ai-dynamo/grove/operator/e2e/utils"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )