From 686484946d52a08b2c9419bc508cbda321a018c8 Mon Sep 17 00:00:00 2001
From: Nick Cao <nickcao@nichi.co>
Date: Sun, 30 Jul 2023 10:59:01 +0800
Subject: [PATCH] python3Packages.sagemaker: 2.135.0 -> 2.173.0

---
 .../python-modules/sagemaker/default.nix      | 63 ++++++++++++-------
 1 file changed, 41 insertions(+), 22 deletions(-)

diff --git a/pkgs/development/python-modules/sagemaker/default.nix b/pkgs/development/python-modules/sagemaker/default.nix
index 6bee251ff3d..c45a7291646 100644
--- a/pkgs/development/python-modules/sagemaker/default.nix
+++ b/pkgs/development/python-modules/sagemaker/default.nix
@@ -1,65 +1,84 @@
 { lib
 , buildPythonPackage
-, fetchPypi
+, pythonOlder
+, fetchFromGitHub
 , pythonRelaxDepsHook
 , attrs
 , boto3
+, cloudpickle
 , google-pasta
-, importlib-metadata
 , numpy
 , protobuf
-, protobuf3-to-dict
 , smdebug-rulesconfig
+, importlib-metadata
+, packaging
 , pandas
 , pathos
-, packaging
-, pythonOlder
+, schema
+, pyyaml
+, jsonschema
+, platformdirs
+, tblib
+, urllib3
+, docker
+, scipy
 }:
 
 buildPythonPackage rec {
   pname = "sagemaker";
-  version = "2.135.0";
+  version = "2.173.0";
   format = "setuptools";
 
-  disabled = pythonOlder "3.7";
+  disabled = pythonOlder "3.6";
 
-  src = fetchPypi {
-    inherit pname version;
-    hash = "sha256-ypdcqEYLxHbfnq1ycq3hVLThhIIs3pq29Fv33Ly2hbE=";
+  src = fetchFromGitHub {
+    owner = "aws";
+    repo = "sagemaker-python-sdk";
+    rev = "refs/tags/v${version}";
+    hash = "sha256-Fzkw37bRBbCD7VBIsN4Qkk6dI0Qh4Gvt5TJbnoUuPCs=";
   };
 
-  nativeBuildInputs = [ pythonRelaxDepsHook ];
+  nativeBuildInputs = [
+    pythonRelaxDepsHook
+  ];
+
   pythonRelaxDeps = [
-    # FIXME: Remove when >= 2.111.0
     "attrs"
-    "protobuf"
+    "boto3"
   ];
 
   propagatedBuildInputs = [
     attrs
     boto3
+    cloudpickle
     google-pasta
-    importlib-metadata
     numpy
-    packaging
-    pathos
     protobuf
-    protobuf3-to-dict
     smdebug-rulesconfig
+    importlib-metadata
+    packaging
     pandas
+    pathos
+    schema
+    pyyaml
+    jsonschema
+    platformdirs
+    tblib
   ];
 
-  postFixup = ''
-    [ "$($out/bin/sagemaker-upgrade-v2 --help 2>&1 | grep -cim1 'pandas failed to import')" -eq "0" ]
-  '';
-
-  doCheck = false;
+  doCheck = false; # many test dependencies are not available in nixpkgs
 
   pythonImportsCheck = [
     "sagemaker"
     "sagemaker.lineage.visualizer"
   ];
 
+  passthru.optional-dependencies = {
+    local = [ urllib3 docker pyyaml ];
+    scipy = [ scipy ];
+    # feature-processor = [ pyspark sagemaker-feature-store-pyspark ]; # not available in nixpkgs
+  };
+
   meta = with lib; {
     description = "Library for training and deploying machine learning models on Amazon SageMaker";
     homepage = "https://github.com/aws/sagemaker-python-sdk/";