diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
index 7c9483343626..383691285228 100644
--- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
+++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
@@ -3004,7 +3004,7 @@ class BeamModulePlugin implements Plugin<Project> {
 
       // Transform service delivers transforms that refer to SDK harness containers with following sufixes.
       def transformServiceJavaContainerSuffix = 'java11'
-      def transformServicePythonContainerSuffix = '310'
+      def transformServicePythonContainerSuffix = pythonContainerSuffix
 
       def setupTask = project.tasks.register(config.name+"Setup", Exec) {
         // Containers for main SDKs when running tests.
diff --git a/sdks/python/expansion-service-container/Dockerfile b/sdks/python/expansion-service-container/Dockerfile
index e825d852479d..1f325bb81030 100644
--- a/sdks/python/expansion-service-container/Dockerfile
+++ b/sdks/python/expansion-service-container/Dockerfile
@@ -18,7 +18,8 @@
 
 # We just need to support one Python version supported by Beam.
 # Picking the current default Beam Python version which is Python 3.10.
-FROM python:3.10-bookworm as expansion-service
+ARG py_version=3.10
+FROM python:${py_version}-bookworm as expansion-service
 LABEL Author "Apache Beam <dev@beam.apache.org>"
 ARG TARGETOS
 ARG TARGETARCH
diff --git a/sdks/python/expansion-service-container/build.gradle b/sdks/python/expansion-service-container/build.gradle
index 4e46f060e59f..2ec6f7a44122 100644
--- a/sdks/python/expansion-service-container/build.gradle
+++ b/sdks/python/expansion-service-container/build.gradle
@@ -40,7 +40,8 @@ task copyDockerfileDependencies(type: Copy) {
 }
 
 task copyRequirementsFile(type: Copy) {
-    from project(':sdks:python:container:py39').fileTree("./")
+    def pythonVersionSuffix = project.ext.pythonVersion.replace('.', '')
+    from project(":sdks:python:container:py${pythonVersionSuffix}").fileTree("./")
     include 'base_image_requirements.txt'
     rename 'base_image_requirements.txt', 'requirements.txt'
     setDuplicatesStrategy(DuplicatesStrategy.INCLUDE)
@@ -70,6 +71,7 @@ docker {
   // tags used by dockerTag task
   tags containerImageTags()
   files "./build"
+  buildArgs(['py_version': "${project.ext.pythonVersion}"])
   buildx project.useBuildx()
   platform(*project.containerPlatforms())
   load project.useBuildx() && !pushContainers