diff --git a/pkgs/development/python-modules/bytewax/default.nix b/pkgs/development/python-modules/bytewax/default.nix
index 05c00906f173..a1963deb433f 100644
--- a/pkgs/development/python-modules/bytewax/default.nix
+++ b/pkgs/development/python-modules/bytewax/default.nix
@@ -17,6 +17,7 @@
 
   # dependencies
   jsonpickle,
+  prometheus-client,
 
   # optional dependencies
   confluent-kafka,
@@ -24,11 +25,12 @@
   # test
   myst-docutils,
   pytestCheckHook,
+  pytest-benchmark,
 }:
 
 buildPythonPackage rec {
   pname = "bytewax";
-  version = "0.17.2";
+  version = "0.21.1";
   format = "pyproject";
 
   disabled = pythonOlder "3.7";
@@ -37,20 +39,16 @@ buildPythonPackage rec {
     owner = "bytewax";
     repo = pname;
     tag = "v${version}";
-    hash = "sha256-BecZvBJsaTHIhJhWM9GZldSL6Irrc7fiedulTN9e76I=";
+    hash = "sha256-O5q1Jd3AMUaQwfQM249CUnkjqEkXybxtM9SOISoULZk=";
   };
 
   env = {
     OPENSSL_NO_VENDOR = true;
   };
 
-  # Remove docs tests, myst-docutils in nixpkgs is not compatible with package requirements.
-  # Package uses old version.
-  patches = [ ./remove-docs-test.patch ];
-
   cargoDeps = rustPlatform.fetchCargoVendor {
     inherit pname version src;
-    hash = "sha256-osLr4u5vQa3/2whytC9PsXKURAwMCNObykY5Us7P3kQ=";
+    hash = "sha256-TTB1//Xza47rnfvlIs9qMvwHPj/U3w2cGTmWrEokriQ=";
   };
 
   nativeBuildInputs = [
@@ -68,7 +66,10 @@ buildPythonPackage rec {
     protobuf
   ];
 
-  propagatedBuildInputs = [ jsonpickle ];
+  dependencies = [
+    jsonpickle
+    prometheus-client
+  ];
 
   optional-dependencies = {
     kafka = [ confluent-kafka ];
@@ -78,10 +79,16 @@ buildPythonPackage rec {
     export PY_IGNORE_IMPORTMISMATCH=1
   '';
 
-  checkInputs = [
+  nativeCheckInputs = [
     myst-docutils
     pytestCheckHook
-  ] ++ lib.flatten (builtins.attrValues optional-dependencies);
+    pytest-benchmark
+  ] ++ lib.flatten (lib.attrValues optional-dependencies);
+
+  pytestFlagsArray = [
+    "--benchmark-disable"
+    "pytests"
+  ];
 
   disabledTestPaths = [
     # dependens on an old myst-docutils version