diff --git a/cmd/kperf/commands/runnergroup/run.go b/cmd/kperf/commands/runnergroup/run.go index 2380071..5fc3f51 100644 --- a/cmd/kperf/commands/runnergroup/run.go +++ b/cmd/kperf/commands/runnergroup/run.go @@ -5,6 +5,7 @@ import ( "fmt" "github.com/Azure/kperf/api/types" + "github.com/Azure/kperf/cmd/kperf/commands/utils" "github.com/Azure/kperf/runner" runnergroup "github.com/Azure/kperf/runner/group" @@ -31,6 +32,10 @@ var runCommand = cli.Command{ // Right now, we need to set image manually. Required: true, }, + cli.StringSliceFlag{ + Name: "affinity", + Usage: "Deploy server to the node with a specific labels (FORMAT: KEY=VALUE[,VALUE])", + }, }, Action: func(cliCtx *cli.Context) error { imgRef := cliCtx.String("runner-image") @@ -38,6 +43,11 @@ var runCommand = cli.Command{ return fmt.Errorf("required valid runner image") } + affinityLabels, err := utils.KeyValuesMap(cliCtx.StringSlice("affinity")) + if err != nil { + return fmt.Errorf("failed to parse affinity: %w", err) + } + specs, err := loadRunnerGroupSpec(cliCtx) if err != nil { return fmt.Errorf("failed to load runner group spec: %w", err) @@ -51,6 +61,7 @@ var runCommand = cli.Command{ kubeCfgPath, imgRef, specs[0], + affinityLabels, ) }, } diff --git a/cmd/kperf/commands/virtualcluster/utils.go b/cmd/kperf/commands/utils/helper.go similarity index 67% rename from cmd/kperf/commands/virtualcluster/utils.go rename to cmd/kperf/commands/utils/helper.go index f73f0d9..66b950e 100644 --- a/cmd/kperf/commands/virtualcluster/utils.go +++ b/cmd/kperf/commands/utils/helper.go @@ -1,12 +1,12 @@ -package virtualcluster +package utils import ( "fmt" "strings" ) -// stringSliceToMap converts key=value[,value] into map[string][]string. -func stringSliceToMap(strs []string) (map[string][]string, error) { +// KeyValuesMap converts key=value[,value] into map[string][]string. +func KeyValuesMap(strs []string) (map[string][]string, error) { res := make(map[string][]string, len(strs)) for _, str := range strs { key, valuesInStr, ok := strings.Cut(str, "=") diff --git a/cmd/kperf/commands/virtualcluster/nodepool.go b/cmd/kperf/commands/virtualcluster/nodepool.go index f1721b3..f42f01e 100644 --- a/cmd/kperf/commands/virtualcluster/nodepool.go +++ b/cmd/kperf/commands/virtualcluster/nodepool.go @@ -5,6 +5,7 @@ import ( "fmt" "strings" + "github.com/Azure/kperf/cmd/kperf/commands/utils" "github.com/Azure/kperf/virtualcluster" "github.com/urfave/cli" @@ -62,7 +63,7 @@ var nodepoolAddCommand = cli.Command{ kubeCfgPath := cliCtx.String("kubeconfig") - affinityLabels, err := stringSliceToMap(cliCtx.StringSlice("affinity")) + affinityLabels, err := utils.KeyValuesMap(cliCtx.StringSlice("affinity")) if err != nil { return fmt.Errorf("failed to parse affinity: %w", err) } diff --git a/manifests/runnergroup/server/templates/pod.yaml b/manifests/runnergroup/server/templates/pod.yaml index 6676283..b04af05 100644 --- a/manifests/runnergroup/server/templates/pod.yaml +++ b/manifests/runnergroup/server/templates/pod.yaml @@ -4,6 +4,21 @@ metadata: name: {{ .Values.name }} namespace: {{ .Release.Namespace }} spec: +{{- if .Values.nodeSelectors }} + affinity: + nodeAffinity: + requiredDuringSchedulingIgnoredDuringExecution: + nodeSelectorTerms: + - matchExpressions: + {{- range $key, $values := .Values.nodeSelectors }} + - key: "{{ $key }}" + operator: In + values: + {{- range $values }} + - {{ . }} + {{- end }} + {{- end }} +{{- end }} containers: - name: server command: diff --git a/manifests/runnergroup/server/values.yaml b/manifests/runnergroup/server/values.yaml index 96272c0..fd028bb 100644 --- a/manifests/runnergroup/server/values.yaml +++ b/manifests/runnergroup/server/values.yaml @@ -2,3 +2,4 @@ name: "" image: "" # TODO(weifu): need https://github.com/Azure/kperf/issues/25 to support list runnerGroupSpec: "" +nodeSelectors: {} diff --git a/runner/runnergroup_run.go b/runner/runnergroup_run.go index 637780e..12c6db4 100644 --- a/runner/runnergroup_run.go +++ b/runner/runnergroup_run.go @@ -17,12 +17,27 @@ import ( // TODO: // 1. create a new package to define ErrNotFound, ErrAlreadyExists, ... errors. // 2. support configurable timeout. -func CreateRunnerGroupServer(ctx context.Context, kubeconfigPath string, runnerImage string, rgSpec *types.RunnerGroupSpec) error { +func CreateRunnerGroupServer(ctx context.Context, + kubeconfigPath string, + runnerImage string, + rgSpec *types.RunnerGroupSpec, + nodeSelectors map[string][]string, +) error { specInStr, err := tweakAndMarshalSpec(rgSpec) if err != nil { return err } + nodeSelectorsInYAML, err := renderNodeSelectors(nodeSelectors) + if err != nil { + return err + } + + nodeSelectorsAppiler, err := helmcli.YAMLValuesApplier(nodeSelectorsInYAML) + if err != nil { + return fmt.Errorf("failed to prepare YAML value applier for nodeSelectors: %w", err) + } + getCli, err := helmcli.NewGetCli(kubeconfigPath, runnerGroupReleaseNamespace) if err != nil { return fmt.Errorf("failed to create helm get client: %w", err) @@ -49,6 +64,7 @@ func CreateRunnerGroupServer(ctx context.Context, kubeconfigPath string, runnerI "image="+runnerImage, "runnerGroupSpec="+specInStr, ), + nodeSelectorsAppiler, ) if err != nil { return fmt.Errorf("failed to create helm release client: %w", err) @@ -71,3 +87,17 @@ func tweakAndMarshalSpec(spec *types.RunnerGroupSpec) (string, error) { } return string(data), nil } + +// renderNodeSelectors renders labels into YAML string. +func renderNodeSelectors(labels map[string][]string) (string, error) { + // NOTE: It should be aligned with ../manifests/runnergroup/server/values.yaml. + target := map[string]interface{}{ + "nodeSelectors": labels, + } + + rawData, err := yaml.Marshal(target) + if err != nil { + return "", fmt.Errorf("failed to render nodeSelectors: %w", err) + } + return string(rawData), nil +}