diff --git a/MODULE.bazel b/MODULE.bazel
index 50f1376..4d2f3c6 100644
--- a/MODULE.bazel
+++ b/MODULE.bazel
@@ -78,11 +78,27 @@
 bazel_dep(name = "rules_testing", version = "0.6.0", dev_dependency = True)
 bazel_dep(name = "rules_shell", version = "0.2.0", dev_dependency = True)
 bazel_dep(name = "rules_multirun", version = "0.9.0", dev_dependency = True)
+bazel_dep(name = "bazel_ci_rules", version = "1.0.0", dev_dependency = True)
 
 # Extra gazelle plugin deps so that WORKSPACE.bzlmod can continue including it for e2e tests.
 # We use `WORKSPACE.bzlmod` because it is impossible to have dev-only local overrides.
 bazel_dep(name = "rules_go", version = "0.41.0", dev_dependency = True, repo_name = "io_bazel_rules_go")
 bazel_dep(name = "gazelle", version = "0.33.0", dev_dependency = True, repo_name = "bazel_gazelle")
+bazel_dep(name = "rules_python_gazelle_plugin", version = "0", dev_dependency = True)
+
+internal_dev_deps = use_extension(
+    "//python/private:internal_dev_deps.bzl",
+    "internal_dev_deps",
+    dev_dependency = True,
+)
+use_repo(internal_dev_deps, "buildkite_config", "wheel_for_testing")
+
+# Add gazelle plugin so that we can run the gazelle example as an e2e integration
+# test and include the distribution files.
+local_path_override(
+    module_name = "rules_python_gazelle_plugin",
+    path = "gazelle",
+)
 
 dev_python = use_extension(
     "//python/extensions:python.bzl",
diff --git a/WORKSPACE.bzlmod b/WORKSPACE.bzlmod
index ca89afe..e69de29 100644
--- a/WORKSPACE.bzlmod
+++ b/WORKSPACE.bzlmod
@@ -1,62 +0,0 @@
-# Copyright 2024 The Bazel Authors. All rights reserved.
-#
-# 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.
-
-# This file contains everything that is needed when using bzlmod
-workspace(name = "rules_python")
-
-load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive", "http_file")
-
-# Used for Bazel CI
-http_archive(
-    name = "bazelci_rules",
-    sha256 = "eca21884e6f66a88c358e580fd67a6b148d30ab57b1680f62a96c00f9bc6a07e",
-    strip_prefix = "bazelci_rules-1.0.0",
-    url = "https://github.com/bazelbuild/continuous-integration/releases/download/rules-1.0.0/bazelci_rules-1.0.0.tar.gz",
-)
-
-load("@bazelci_rules//:rbe_repo.bzl", "rbe_preconfig")
-
-# Creates a default toolchain config for RBE.
-# Use this as is if you are using the rbe_ubuntu16_04 container,
-# otherwise refer to RBE docs.
-rbe_preconfig(
-    name = "buildkite_config",
-    toolchain = "ubuntu1804-bazel-java11",
-)
-
-# Add gazelle plugin so that we can run the gazelle example as an e2e integration
-# test and include the distribution files.
-local_repository(
-    name = "rules_python_gazelle_plugin",
-    path = "gazelle",
-)
-
-#####################
-
-# This wheel is purely here to validate the wheel extraction code. It's not
-# intended for anything else.
-http_file(
-    name = "wheel_for_testing",
-    downloaded_file_path = "numpy-1.25.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl",
-    sha256 = "0d60fbae8e0019865fc4784745814cff1c421df5afee233db6d88ab4f14655a2",
-    urls = [
-        "https://files.pythonhosted.org/packages/50/67/3e966d99a07d60a21a21d7ec016e9e4c2642a86fea251ec68677daf71d4d/numpy-1.25.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl",
-    ],
-)
-
-# rules_proto expects //external:python_headers to point at the python headers.
-bind(
-    name = "python_headers",
-    actual = "//python/cc:current_py_cc_headers",
-)
diff --git a/python/private/internal_dev_deps.bzl b/python/private/internal_dev_deps.bzl
new file mode 100644
index 0000000..2a3b84e
--- /dev/null
+++ b/python/private/internal_dev_deps.bzl
@@ -0,0 +1,44 @@
+# Copyright 2024 The Bazel Authors. All rights reserved.
+#
+# 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.
+"""Module extension for internal dev_dependency=True setup."""
+
+load("@bazel_ci_rules//:rbe_repo.bzl", "rbe_preconfig")
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_file")
+
+def _internal_dev_deps_impl(mctx):
+    _ = mctx  # @unused
+
+    # This wheel is purely here to validate the wheel extraction code. It's not
+    # intended for anything else.
+    http_file(
+        name = "wheel_for_testing",
+        downloaded_file_path = "numpy-1.25.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl",
+        sha256 = "0d60fbae8e0019865fc4784745814cff1c421df5afee233db6d88ab4f14655a2",
+        urls = [
+            "https://files.pythonhosted.org/packages/50/67/3e966d99a07d60a21a21d7ec016e9e4c2642a86fea251ec68677daf71d4d/numpy-1.25.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl",
+        ],
+    )
+
+    # Creates a default toolchain config for RBE.
+    # Use this as is if you are using the rbe_ubuntu16_04 container,
+    # otherwise refer to RBE docs.
+    rbe_preconfig(
+        name = "buildkite_config",
+        toolchain = "ubuntu1804-bazel-java11",
+    )
+
+internal_dev_deps = module_extension(
+    implementation = _internal_dev_deps_impl,
+    doc = "This extension creates internal rules_python dev dependencies.",
+)