diff --git a/pkgs/development/python-modules/sagemaker-mlflow/default.nix b/pkgs/development/python-modules/sagemaker-mlflow/default.nix
new file mode 100644
index 000000000000..152ef3b78eb3
--- /dev/null
+++ b/pkgs/development/python-modules/sagemaker-mlflow/default.nix
@@ -0,0 +1,74 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+
+  # build-system
+  setuptools,
+
+  # dependencies
+  boto3,
+  mlflow,
+
+  # tests
+  pytestCheckHook,
+  scikit-learn,
+}:
+
+buildPythonPackage rec {
+  pname = "sagemaker-mlflow";
+  version = "0.1.0";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "aws";
+    repo = "sagemaker-mlflow";
+    rev = "refs/tags/v${version}";
+    hash = "sha256-1bonIqZ+cFxCOxoFWn1MLBOIiB1wUX69/lUTPPupJaw=";
+  };
+
+  postPatch = ''
+    substituteInPlace VERSION \
+      --replace-fail "${version}.dev0" "${version}"
+  '';
+
+  build-system = [
+    setuptools
+  ];
+
+  dependencies = [
+    boto3
+    mlflow
+  ];
+
+  pythonImportsCheck = [
+    "sagemaker_mlflow"
+  ];
+
+  nativeCheckInputs = [
+    pytestCheckHook
+    scikit-learn
+  ];
+
+  disabledTests = [
+    # AssertionError: assert 's3' in '/build/source/not implemented/0/d3c16d2bad4245bf9fc68f86d2e7599d/artifacts'
+    "test_log_metric"
+
+    # AssertionError: assert 'not implemented' == 'mw'
+    "test_request_header"
+
+    # Require internet access
+    "test_auth_provider_returns_correct_sigv4"
+    "test_log_artifact"
+    "test_presigned_url"
+    "test_presigned_url_with_fields"
+  ];
+
+  meta = {
+    description = "MLFlow plugin for SageMaker";
+    homepage = "https://github.com/aws/sagemaker-mlflow";
+    changelog = "https://github.com/aws/sagemaker-mlflow/releases/tag/v${version}";
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ GaetanLepage ];
+  };
+}
diff --git a/pkgs/development/python-modules/sagemaker/default.nix b/pkgs/development/python-modules/sagemaker/default.nix
index 7c3b26cbb25c..582ddd93f0f2 100644
--- a/pkgs/development/python-modules/sagemaker/default.nix
+++ b/pkgs/development/python-modules/sagemaker/default.nix
@@ -2,7 +2,6 @@
   lib,
   buildPythonPackage,
   fetchFromGitHub,
-  fetchpatch,
 
   # build-system
   hatchling,
@@ -25,6 +24,7 @@
   pyyaml,
   requests,
   sagemaker-core,
+  sagemaker-mlflow,
   schema,
   smdebug-rulesconfig,
   tblib,
@@ -38,29 +38,16 @@
 
 buildPythonPackage rec {
   pname = "sagemaker";
-  version = "2.232.1";
+  version = "2.232.2";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "aws";
     repo = "sagemaker-python-sdk";
     rev = "refs/tags/v${version}";
-    hash = "sha256-I+iZKx1CnZIGYgYuYhhs8BnY84KPyKOGw8M0He26DGU=";
+    hash = "sha256-q0JxQi1kUnp5L/hexxpGhR4t0v53l8iVuc9H6N0K6Y4=";
   };
 
-  patches = [
-    # Distutils removal, fix build with python 3.12
-    # https://github.com/aws/sagemaker-python-sdk/pull/4544
-    (fetchpatch {
-      url = "https://github.com/aws/sagemaker-python-sdk/commit/84447ba59e544c810aeb842fd058e20d89e3fc74.patch";
-      hash = "sha256-B8Q18ViB7xYy1F5LoL1NvXj2lnFPgt+C9wssSODyAXM=";
-    })
-    (fetchpatch {
-      url = "https://github.com/aws/sagemaker-python-sdk/commit/e9e08a30cb42d4b2d7299c1c4b42d680a8c78110.patch";
-      hash = "sha256-uGPtXSXfeaIvt9kkZZKQDuiZfoRgw3teffuxai1kKlY=";
-    })
-  ];
-
   build-system = [
     hatchling
   ];
@@ -69,6 +56,7 @@ buildPythonPackage rec {
     "boto3"
     "cloudpickle"
     "importlib-metadata"
+    "protobuf"
   ];
 
   dependencies = [
@@ -89,6 +77,7 @@ buildPythonPackage rec {
     pyyaml
     requests
     sagemaker-core
+    sagemaker-mlflow
     schema
     smdebug-rulesconfig
     tblib
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 17af78dd3618..b0d9ca65a2fe 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -13964,6 +13964,8 @@ self: super: with self; {
 
   sagemaker-core = callPackage ../development/python-modules/sagemaker-core { };
 
+  sagemaker-mlflow = callPackage ../development/python-modules/sagemaker-mlflow { };
+
   salib = callPackage ../development/python-modules/salib { };
 
   salmon-mail = callPackage ../development/python-modules/salmon-mail { };