Add MetricOps struct and reflection-based flag parsing

Author: annapoornanarayan

internal/e2e/flags.go

@@ -4,21 +4,33 @@ package e2e
 
 import (
 	"flag"
+	"fmt"
 	"github.com/octago/sflags/gen/gpflag"
 	"github.com/spf13/pflag"
+	"reflect"
 )
 
-// Bridging for test execution => InitFlags parses struct tags into command-line flags and registers them with Go's flag package
-func InitFlags(config interface{}, metricDimensions *map[string]string) {
+// For CloudWatch metric dimension flag
+type MetricOps struct {
+	MetricDimensions map[string]string `flag:"metricDimensions" desc:"CloudWatch metric dimensions as comma-separated key=value pairs"`
+}
+
+func ParseFlags(config interface{}) (*pflag.FlagSet, error) {
 	flags, err := gpflag.Parse(config)
 	if err != nil {
-		panic(err)
+		return nil, fmt.Errorf("failed to parse flags: %w", err)
 	}
 
-	// Manually add flags for data types that gpflag cannot parse
-	flags.StringToStringVar(metricDimensions, "metricDimensions", nil, "CloudWatch metric dimensions as comma-separated key=value pairs")
+	// Handle MetricDimensions map that gpflag doesn't support
+	if _, hasField := reflect.TypeOf(config).Elem().FieldByName("MetricDimensions"); hasField {
+		field := reflect.ValueOf(config).Elem().FieldByName("MetricDimensions")
+		metricDims := field.Addr().Interface().(*map[string]string)
+		flags.StringToStringVar(metricDims, "metricDimensions", nil, "CloudWatch metric dimensions as comma-separated key=value pairs")
+	}
 
 	flags.VisitAll(func(pf *pflag.Flag) {
 		flag.CommandLine.Var(pf.Value, pf.Name, pf.Usage)
 	})
+
+	return flags, nil
 }

test/cases/nvidia-training/main_test.go

@@ -27,7 +27,10 @@ import (
 )
 
 func TestMain(m *testing.M) {
-	fwext.InitFlags(&testConfig, &testConfig.MetricDimensions)
+	_, err := fwext.ParseFlags(&testConfig)
+	if err != nil {
+		log.Fatalf("failed to parse flags: %v", err)
+	}
 	cfg, err := envconf.NewFromFlags()
 	if err != nil {
 		log.Fatalf("failed to initialize test environment: %v", err)
@@ -38,7 +41,7 @@ func TestMain(m *testing.M) {
 	testenv = env.NewWithConfig(cfg).WithContext(ctx)
 
 	// Render CloudWatch Agent manifest with dynamic dimensions
-	renderedCloudWatchAgentManifest, err := common.RenderCloudWatchAgentManifest(manifests.CloudWatchAgentManifest, testConfig.MetricDimensions)
+	renderedCloudWatchAgentManifest, err := manifests.RenderCloudWatchAgentManifest(testConfig.MetricDimensions)
 	if err != nil {
 		log.Printf("Warning: failed to render CloudWatch Agent manifest: %v", err)
 	}
@@ -84,10 +87,10 @@ func TestMain(m *testing.M) {
 		},
 
 		// Wait for DaemonSets using helper
-		common.DeployDaemonSet("nvidia-device-plugin-daemonset", "kube-system"),
-		common.DeployDaemonSet("aws-efa-k8s-device-plugin-daemonset", "kube-system"),
-		common.DeployDaemonSet("dcgm-exporter", "kube-system"),
-		common.DeployDaemonSet("cwagent", "amazon-cloudwatch"),
+		common.DeployDaemonSet("nvidia-device-plugin-daemonset", "kube-system", false),
+		common.DeployDaemonSet("aws-efa-k8s-device-plugin-daemonset", "kube-system", false),
+		common.DeployDaemonSet("dcgm-exporter", "kube-system", true),
+		common.DeployDaemonSet("cwagent", "amazon-cloudwatch", true),
 		checkNodeTypes, // Dynamically check node types and capacities after device plugins are ready
 	)
 

test/cases/nvidia-training/vars.go

@@ -3,14 +3,15 @@
 package training
 
 import (
+	fwext "github.com/aws/aws-k8s-tester/internal/e2e"
 	"sigs.k8s.io/e2e-framework/pkg/env"
 )
 
 type Config struct {
-	MetricDimensions  map[string]string `flag:"metricDimensions" desc:"CloudWatch metric dimensions as comma-separated key=value pairs"`
-	BertTrainingImage string            `flag:"bertTrainingImage" desc:"Docker image used for BERT training workload"`
-	EfaEnabled        bool              `flag:"efaEnabled" desc:"Enable Elastic Fabric Adapter (EFA)"`
-	NodeType          string            `flag:"nodeType" desc:"Instance type for cluster nodes"`
+	fwext.MetricOps
+	BertTrainingImage string `flag:"bertTrainingImage" desc:"Docker image used for BERT training workload"`
+	EfaEnabled        bool   `flag:"efaEnabled" desc:"Enable Elastic Fabric Adapter (EFA)"`
+	NodeType          string `flag:"nodeType" desc:"Instance type for cluster nodes"`
 }
 
 // Shared global variables

test/common/helper.go renamed to test/common/resources.go

@@ -5,9 +5,7 @@ package common
 import (
 	"context"
 	"fmt"
-	"html/template"
 	"log"
-	"strings"
 	"time"
 
 	fwext "github.com/aws/aws-k8s-tester/internal/e2e"
@@ -18,21 +16,9 @@ import (
 	"sigs.k8s.io/e2e-framework/pkg/envconf"
 )
 
-// RenderCloudWatchAgentManifest renders the CloudWatch Agent manifest with dynamic dimensions
-func RenderCloudWatchAgentManifest(cloudWatchAgentManifest []byte, metricDimensions map[string]string) ([]byte, error) {
-	var keys []string
-	for key := range metricDimensions {
-		keys = append(keys, `"`+key+`"`)
-	}
-	dimensionsStr := strings.Join(keys, ", ")
-	return fwext.RenderManifests(cloudWatchAgentManifest, map[string]interface{}{
-		"MetricDimensions": metricDimensions,
-		"DimensionKeys":    template.HTML(dimensionsStr),
-	})
-}
-
 // DeployDaemonSet returns a function to deploy and wait for a DaemonSet to be ready
-func DeployDaemonSet(name, namespace string) env.Func {
+// If isOptional is true, failures will log warnings instead of returning errors
+func DeployDaemonSet(name, namespace string, isOptional bool) env.Func {
 	return func(ctx context.Context, config *envconf.Config) (context.Context, error) {
 		log.Printf("Waiting for %s daemonset to be ready.", name)
 		daemonset := appsv1.DaemonSet{
@@ -44,8 +30,7 @@ func DeployDaemonSet(name, namespace string) env.Func {
 			wait.WithContext(ctx),
 		)
 		if err != nil {
-			// Do not fail test for optional daemonset failures
-			if name == "cwagent" || name == "dcgm-exporter" {
+			if isOptional {
 				log.Printf("Warning: %s daemonset is not ready: %v", name, err)
 				return ctx, nil
 			}

test/manifests/raw.go

@@ -22,5 +22,5 @@ var (
 	DCGMExporterManifest []byte
 
 	//go:embed assets/cloudwatch-agent.yaml
-	CloudWatchAgentManifest []byte
+	cloudWatchAgentManifestTemplate []byte
 )

test/manifests/rendered.go

@@ -0,0 +1,21 @@
+package manifests
+
+import (
+	"html/template"
+	"strings"
+
+	fwext "github.com/aws/aws-k8s-tester/internal/e2e"
+)
+
+// RenderCloudWatchAgentManifest renders the CloudWatch Agent manifest with dynamic dimensions
+func RenderCloudWatchAgentManifest(metricDimensions map[string]string) ([]byte, error) {
+	var keys []string
+	for key := range metricDimensions {
+		keys = append(keys, `"`+key+`"`)
+	}
+	dimensionsStr := strings.Join(keys, ", ")
+	return fwext.RenderManifests(cloudWatchAgentManifestTemplate, map[string]interface{}{
+		"MetricDimensions": metricDimensions,
+		"DimensionKeys":    template.HTML(dimensionsStr),
+	})
+}