Merge pull request #136398 from dims/use-pytorch-wide-deep-in-tests

Switch node perf tests from TensorFlow to PyTorch Wide-Deep

test/e2e_node/image_list.go

@@ -21,7 +21,6 @@ import (
 	"fmt"
 	"os"
 	"os/user"
-	"runtime"
 	"sync"
 	"time"
 
@@ -61,6 +60,7 @@ var NodePrePullImageList = sets.NewString(
 	imageutils.GetPauseImageName(),
 	imageutils.GetE2EImage(imageutils.NodePerfNpbEp),
 	imageutils.GetE2EImage(imageutils.NodePerfNpbIs),
+	imageutils.GetE2EImage(imageutils.NodePerfPytorchWideDeep),
 	imageutils.GetE2EImage(imageutils.Etcd),
 )
 
@@ -69,11 +69,6 @@ var NodePrePullImageList = sets.NewString(
 // 2. the ones passed in from framework.TestContext.ExtraEnvs
 // So this function needs to be called after the extra envs are applied.
 func updateImageAllowList(ctx context.Context) {
-	// Architecture-specific image
-	if !isRunningOnArm64() {
-		// NodePerfTfWideDeep is only supported on x86_64, pulling in arm64 will fail
-		NodePrePullImageList = NodePrePullImageList.Insert(imageutils.GetE2EImage(imageutils.NodePerfTfWideDeep))
-	}
 	// Union NodePrePullImageList and PrePulledImages into the framework image pre-pull list.
 	e2epod.ImagePrePullList = NodePrePullImageList.Union(commontest.PrePulledImages)
 	// Images from extra envs
@@ -95,10 +90,6 @@ func updateImageAllowList(ctx context.Context) {
 	}
 }
 
-func isRunningOnArm64() bool {
-	return runtime.GOARCH == "arm64"
-}
-
 func getNodeProblemDetectorImage() string {
 	const defaultImage string = "registry.k8s.io/node-problem-detector/node-problem-detector:v1.34.0"
 	image := os.Getenv("NODE_PROBLEM_DETECTOR_IMAGE")

test/e2e_node/perf/workloads/pytorch_wide_deep.go

@@ -1,5 +1,5 @@
 /*
-Copyright 2018 The Kubernetes Authors.
+Copyright 2025 The Kubernetes Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -29,22 +29,22 @@ import (
 	imageutils "k8s.io/kubernetes/test/utils/image"
 )
 
-// tfWideDeepWorkload defines a workload to run
-// https://github.com/tensorflow/models/tree/master/official/r1/wide_deep.
-type tfWideDeepWorkload struct{}
+// pytorchWideDeepWorkload defines a workload to run a PyTorch Wide-Deep
+// model training benchmark for CPU Manager validation.
+type pytorchWideDeepWorkload struct{}
 
-// Ensure tfWideDeepWorkload implements NodePerfWorkload interface.
-var _ NodePerfWorkload = &tfWideDeepWorkload{}
+// Ensure pytorchWideDeepWorkload implements NodePerfWorkload interface.
+var _ NodePerfWorkload = &pytorchWideDeepWorkload{}
 
-func (w tfWideDeepWorkload) Name() string {
-	return "tensorflow-wide-deep"
+func (w pytorchWideDeepWorkload) Name() string {
+	return "pytorch-wide-deep"
 }
 
-func (w tfWideDeepWorkload) PodSpec() v1.PodSpec {
+func (w pytorchWideDeepWorkload) PodSpec() v1.PodSpec {
 	var containers []v1.Container
 	ctn := v1.Container{
 		Name:  fmt.Sprintf("%s-ctn", w.Name()),
-		Image: imageutils.GetE2EImage(imageutils.NodePerfTfWideDeep),
+		Image: imageutils.GetE2EImage(imageutils.NodePerfPytorchWideDeep),
 		Resources: v1.ResourceRequirements{
 			Requests: v1.ResourceList{
 				v1.ResourceName(v1.ResourceCPU):    resource.MustParse("15000m"),
@@ -55,8 +55,7 @@ func (w tfWideDeepWorkload) PodSpec() v1.PodSpec {
 				v1.ResourceName(v1.ResourceMemory): resource.MustParse("16Gi"),
 			},
 		},
-		Command: []string{"/bin/sh"},
-		Args:    []string{"-c", "time -p python ./wide_deep.py --model_type=wide_deep --train_epochs=300 --epochs_between_evals=300 --batch_size=32561"},
+		// The container entrypoint already runs: time -p python /train_wide_deep.py
 	}
 	containers = append(containers, ctn)
 
@@ -66,11 +65,11 @@ func (w tfWideDeepWorkload) PodSpec() v1.PodSpec {
 	}
 }
 
-func (w tfWideDeepWorkload) Timeout() time.Duration {
+func (w pytorchWideDeepWorkload) Timeout() time.Duration {
 	return 15 * time.Minute
 }
 
-func (w tfWideDeepWorkload) KubeletConfig(oldCfg *kubeletconfig.KubeletConfiguration) (newCfg *kubeletconfig.KubeletConfiguration, err error) {
+func (w pytorchWideDeepWorkload) KubeletConfig(oldCfg *kubeletconfig.KubeletConfiguration) (newCfg *kubeletconfig.KubeletConfiguration, err error) {
 	// Enable CPU Manager in Kubelet with static policy.
 	newCfg = oldCfg.DeepCopy()
 	// Set the CPU Manager policy to static.
@@ -92,7 +91,7 @@ func (w tfWideDeepWorkload) KubeletConfig(oldCfg *kubeletconfig.KubeletConfigura
 	return newCfg, nil
 }
 
-func (w tfWideDeepWorkload) PreTestExec() error {
+func (w pytorchWideDeepWorkload) PreTestExec() error {
 	cmd := "/bin/sh"
 	args := []string{"-c", "rm -f /var/lib/kubelet/cpu_manager_state"}
 	err := runCmd(cmd, args)
@@ -100,7 +99,7 @@ func (w tfWideDeepWorkload) PreTestExec() error {
 	return err
 }
 
-func (w tfWideDeepWorkload) PostTestExec() error {
+func (w pytorchWideDeepWorkload) PostTestExec() error {
 	cmd := "/bin/sh"
 	args := []string{"-c", "rm -f /var/lib/kubelet/cpu_manager_state"}
 	err := runCmd(cmd, args)
@@ -108,7 +107,7 @@ func (w tfWideDeepWorkload) PostTestExec() error {
 	return err
 }
 
-func (w tfWideDeepWorkload) ExtractPerformanceFromLogs(logs string) (perf time.Duration, err error) {
+func (w pytorchWideDeepWorkload) ExtractPerformanceFromLogs(logs string) (perf time.Duration, err error) {
 	perfLine, err := getMatchingLineFromLog(logs, "real")
 	if err != nil {
 		return perf, err

test/e2e_node/perf/workloads/workloads.go

@@ -50,4 +50,4 @@ type NodePerfWorkload interface {
 }
 
 // NodePerfWorkloads is the collection of all node performance testing workloads.
-var NodePerfWorkloads = []NodePerfWorkload{npbISWorkload{}, npbEPWorkload{}, tfWideDeepWorkload{}}
+var NodePerfWorkloads = []NodePerfWorkload{npbISWorkload{}, npbEPWorkload{}, pytorchWideDeepWorkload{}}

test/images/.permitted-images

@@ -21,7 +21,7 @@ registry.k8s.io/e2e-test-images/nautilus
 registry.k8s.io/e2e-test-images/nginx
 registry.k8s.io/e2e-test-images/node-perf/npb-ep
 registry.k8s.io/e2e-test-images/node-perf/npb-is
-registry.k8s.io/e2e-test-images/node-perf/tf-wide-deep
+registry.k8s.io/e2e-test-images/node-perf/pytorch-wide-deep
 registry.k8s.io/e2e-test-images/nonewprivs
 registry.k8s.io/e2e-test-images/nonroot
 registry.k8s.io/e2e-test-images/perl

test/images/node-perf/tf-wide-deep/BASEIMAGE

@@ -1 +0,0 @@
-linux/amd64=python:3.11-slim-bookworm

test/images/node-perf/tf-wide-deep/Dockerfile

@@ -1,36 +0,0 @@
-# Copyright 2018 The Kubernetes Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-ARG BASEIMAGE
-FROM $BASEIMAGE
-
-CROSS_BUILD_COPY qemu-QEMUARCH-static /usr/bin/
-
-RUN apt-get update && apt-get install -y wget time
-
-RUN pip install tensorflow==2.20.0
-
-# Use models v1.9.0 which contains the wide_deep implementation.
-# The wide_deep directory was removed in later versions of tensorflow/models.
-# The v1.9.0 code uses tf.estimator APIs which are still available in TF 2.x.
-RUN wget https://github.com/tensorflow/models/archive/v1.9.0.tar.gz \
-&& tar xzf v1.9.0.tar.gz \
-&& rm -f v1.9.0.tar.gz
-
-RUN python /models-1.9.0/official/wide_deep/data_download.py
-
-WORKDIR $HOME/models-1.9.0/official/wide_deep
-ENV PYTHONPATH=$PYTHONPATH:$HOME/models-1.9.0
-
-ENTRYPOINT python ./wide_deep.py

test/images/node-perf/tf-wide-deep/README.md

@@ -1,17 +0,0 @@
-## Tensorflow Official Wide Deep Model
-
-The container image described here predicts the income using the census income dataset in Tensorflow. For more
-information, see
-[https://github.com/tensorflow/models/tree/v2.0/official/r1/wide_deep](https://github.com/tensorflow/models/tree/v2.0/official/r1/wide_deep).
-This image is used as a workload in node performance testing.
-
-## How to release:
-```
-# Build
-$ cd $K8S_ROOT/test/images
-$ make all WHAT=node-perf/tf-wide-deep
-
-# Push
-$ cd $K8S_ROOT/test/images
-$ make all-push WHAT=node-perf/tf-wide-deep
-```

test/images/node-perf/tf-wide-deep/VERSION

@@ -1 +0,0 @@
-1.5

test/utils/image/manifest.go

@@ -185,8 +185,8 @@ const (
 	NodePerfNpbEp
 	// NodePerfNpbIs image
 	NodePerfNpbIs
-	// NodePerfTfWideDeep image
-	NodePerfTfWideDeep
+	// NodePerfPytorchWideDeep image
+	NodePerfPytorchWideDeep
 	// Nonewprivs image
 	Nonewprivs
 	// NonRoot runs with a default user of 1234
@@ -226,7 +226,7 @@ func initImageConfigs(list RegistryList) (map[ImageID]Config, map[ImageID]Config
 	configs[NginxNew] = Config{list.PromoterE2eRegistry, "nginx", "1.15-4"}
 	configs[NodePerfNpbEp] = Config{list.PromoterE2eRegistry, "node-perf/npb-ep", "1.2"}
 	configs[NodePerfNpbIs] = Config{list.PromoterE2eRegistry, "node-perf/npb-is", "1.2"}
-	configs[NodePerfTfWideDeep] = Config{list.PromoterE2eRegistry, "node-perf/tf-wide-deep", "1.3"}
+	configs[NodePerfPytorchWideDeep] = Config{list.PromoterE2eRegistry, "node-perf/pytorch-wide-deep", "1.0.0"}
 	configs[Nonewprivs] = Config{list.PromoterE2eRegistry, "nonewprivs", "1.3"}
 	configs[NonRoot] = Config{list.PromoterE2eRegistry, "nonroot", "1.4"}
 	// Pause - when these values are updated, also update cmd/kubelet/app/options/container_runtime.go