Add Numpy 2.0 support (#5311)

rsareddy0329 · Roja Reddy Sareddy · sage-maker · web-flow · commit 99210b2cc042 · 2025-10-28T11:39:04.000-07:00
* Add numpy 2.0 support

* Add numpy 2.0 support

* Add numpy 2.0 support

* Add numpy 2.0 support

* Add numpy 2.0 support

* Fix incompatible_dependecies test

* Fix incompatible_dependecies test

* Fix incompatible_dependecies test

* Fix incompatible_dependecies test

* Fix incompatible_dependecies test

* update tensorflow artifacts

* update tensorflow artifacts

* update tensorflow artifacts

* testfile codestyle fixes

* testfile codestyle fixes

* update SKLearn image URI config

* update SKLearn image URI config

* docstyle fixes

* docstyle fixes

* numpy fixes

* numpy fixes

* numpy fixes

* numpy fixes

* numpy fixes

* numpy fixes

* numpy fixes

* numpy fixes

* numpy fix for slow test

* numpy fix for slow test

* numpy fix for slow test

* numpy fix for slow test

* Revert 'Add numpy 2.0 support'

* Revert 'Add numpy 2.0 support'

* Revert 'Add numpy 2.0 support'

* Add numpy 2.0 support

* Add numpy 2.0 support

* Add numpy 2.0 support

* Add numpy 2.0 support

* Add numpy 2.0 support

* Add numpy 2.0 support

* Add numpy 2.0 support

* Add numpy 2.0 support

* Add numpy 2.0 support

---------

Co-authored-by: Roja Reddy Sareddy &lt;rsareddy@amazon.com&gt;
Co-authored-by: parknate@ &lt;parknate@amazon.com&gt;
Co-authored-by: Gokul Anantha Narayanan &lt;166456257+nargokul@users.noreply.github.com&gt;
diff --git a/pyproject.toml b/pyproject.toml
@@ -39,10 +39,10 @@ dependencies = [
   "google-pasta",
   "importlib-metadata>=1.4.0,<7.0",
   "jsonschema",
-  "numpy==1.26.4",
+  "numpy>=1.26.4,<3.0",
   "omegaconf>=2.2,<3",
   "packaging>=23.0,<25",
-  "pandas",
+  "pandas>=2.3.0",
   "pathos",
   "platformdirs",
   "protobuf>=3.12,<6.32",
diff --git a/requirements/extras/scipy_requirements.txt b/requirements/extras/scipy_requirements.txt
@@ -1 +1 @@
-scipy==1.11.3
+scipy==1.13.0
diff --git a/requirements/extras/test_requirements.txt b/requirements/extras/test_requirements.txt
@@ -1,5 +1,5 @@
 tox==3.24.5
-numpy==1.26.4
+numpy>=2.0.0, <3.0
 build[virtualenv]==1.2.1
 flake8==7.1.2
 pytest==6.2.5
@@ -23,8 +23,8 @@ requests==2.32.2
 sagemaker-experiments==0.1.35
 Jinja2==3.1.6
 pyvis==0.2.1
-pandas==1.4.4
-scikit-learn==1.3.0
+pandas>=2.3.0
+scikit-learn==1.6.1
 cloudpickle==2.2.1
 jsonpickle<4.0.0
 PyYAML>=6.0.1
@@ -44,7 +44,7 @@ onnx==1.17.0
 nbformat>=5.9,<6
 accelerate>=0.24.1,<=0.27.0
 schema==0.7.5
-tensorflow>=2.16.2,<=2.18.0
+tensorflow>=2.16.2,<=2.19.0
 mlflow>=2.14.2,<3
 huggingface_hub==0.26.2
 uvicorn>=0.30.1
diff --git a/src/sagemaker/image_uri_config/sklearn.json b/src/sagemaker/image_uri_config/sklearn.json
@@ -388,6 +388,54 @@
                     "us-west-2": "246618743249"
                 },
                 "repository": "sagemaker-scikit-learn"
+            },
+            "1.4-2": {
+                "processors": [
+                    "cpu"
+                ],
+                "py_versions": [
+                    "py3"
+                ],
+                "registries": {
+                    "af-south-1": "510948584623",
+                    "ap-east-1": "651117190479",
+                    "ap-northeast-1": "354813040037",
+                    "ap-northeast-2": "366743142698",
+                    "ap-northeast-3": "867004704886",
+                    "ap-south-1": "720646828776",
+                    "ap-south-2": "628508329040",
+                    "ap-southeast-1": "121021644041",
+                    "ap-southeast-2": "783357654285",
+                    "ap-southeast-3": "951798379941",
+                    "ap-southeast-4": "106583098589",
+                    "ca-central-1": "341280168497",
+                    "ca-west-1": "190319476487",
+                    "cn-north-1": "450853457545",
+                    "cn-northwest-1": "451049120500",
+                    "eu-central-1": "492215442770",
+                    "eu-central-2": "680994064768",
+                    "eu-north-1": "662702820516",
+                    "eu-south-1": "978288397137",
+                    "eu-south-2": "104374241257",
+                    "eu-west-1": "141502667606",
+                    "eu-west-2": "764974769150",
+                    "eu-west-3": "659782779980",
+                    "il-central-1": "898809789911",
+                    "me-central-1": "272398656194",
+                    "me-south-1": "801668240914",
+                    "sa-east-1": "737474898029",
+                    "us-east-1": "683313688378",
+                    "us-east-2": "257758044811",
+                    "us-gov-east-1": "237065988967",
+                    "us-gov-west-1": "414596584902",
+                    "us-iso-east-1": "833128469047",
+                    "us-isob-east-1": "281123927165",
+                    "us-isof-east-1": "108575199400",
+                    "us-isof-south-1": "124985052026",
+                    "us-west-1": "746614075791",
+                    "us-west-2": "246618743249"
+                },
+                "repository": "sagemaker-scikit-learn"
             }
         }
     },
diff --git a/src/sagemaker/serve/utils/conda_in_process.yml b/src/sagemaker/serve/utils/conda_in_process.yml
@@ -12,7 +12,7 @@ dependencies:
     - boto3>=1.34.142,<2.0
     - cloudpickle==2.2.1
     - google-pasta
-    - numpy==1.26.4
+    - numpy>=2.0.0,<3.0
     - protobuf>=3.12,<5.0
     - smdebug_rulesconfig==1.0.1
     - importlib-metadata>=1.4.0,<7.0
@@ -64,7 +64,7 @@ dependencies:
             - multiprocess>=0.70.14
             - networkx>=3.1
             - packaging>=23.1
-            - pandas>=1.5.3
+            - pandas>=2.3.0
             - pathos>=0.3.0
             - pillow>=9.5.0
             - platformdirs>=3.2.0
diff --git a/tests/data/remote_function/requirements.txt b/tests/data/remote_function/requirements.txt
@@ -1 +1 @@
-scipy==1.11.3
+scipy==1.13.0
diff --git a/tests/data/serve_resources/mlflow/pytorch/conda.yaml b/tests/data/serve_resources/mlflow/pytorch/conda.yaml
@@ -2,23 +2,23 @@ channels:
 - conda-forge
 dependencies:
 - python=3.10.13
-- pip<=23.3.1
+- pip<=24.3
 - pip:
-  - mlflow==2.10.2
+  - mlflow>=2.16.1
   - astunparse==1.6.3
   - cffi==1.16.0
   - cloudpickle==2.2.1
   - defusedxml==0.7.1
   - dill==0.3.9
   - gmpy2==2.1.2
-  - numpy==1.26.4
+  - numpy>=2.0.0,<3.0
   - opt-einsum==3.3.0
   - packaging==24.0
-  - pandas==2.2.1
+  - pandas>=2.3.0
   - pyyaml==6.0.1
   - requests==2.31.0
   - torch>=2.6.0
   - torchvision>=0.17.0
   - tqdm==4.66.2
-  - scikit-learn==1.3.2
+  - scikit-learn==1.6.1
 name: mlflow-env
diff --git a/tests/data/serve_resources/mlflow/pytorch/requirements.txt b/tests/data/serve_resources/mlflow/pytorch/requirements.txt
@@ -5,10 +5,10 @@ cloudpickle==2.2.1
 defusedxml==0.7.1
 dill==0.3.9
 gmpy2==2.1.2
-numpy==1.26.4
+numpy>=2.0.0,<3.0
 opt-einsum==3.3.0
 packaging>=23.0,<25
-pandas==2.2.1
+pandas>=2.3.0
 pyyaml==6.0.1
 requests==2.32.4
 torch>=2.6.0
diff --git a/tests/data/serve_resources/mlflow/tensorflow/MLmodel b/tests/data/serve_resources/mlflow/tensorflow/MLmodel
@@ -10,7 +10,7 @@ flavors:
     code: null
     model_type: tf2-module
     saved_model_dir: tf2model
-mlflow_version: 2.20.3
+mlflow_version: 2.11.1
 model_size_bytes: 23823
 model_uuid: 40d2323944294fce898d8693455f60e8
 run_id: 592132312fb84935b201de2c027c54c6
diff --git a/tests/data/serve_resources/mlflow/tensorflow/conda.yaml b/tests/data/serve_resources/mlflow/tensorflow/conda.yaml
@@ -2,7 +2,7 @@ channels:
 - conda-forge
 dependencies:
 - python=3.10.13
-- pip<=23.3.1
+- pip<=24.3
 - pip:
   - mlflow==2.11.1
   - cloudpickle==2.2.1
diff --git a/tests/data/serve_resources/mlflow/tensorflow_numpy2/MLmodel b/tests/data/serve_resources/mlflow/tensorflow_numpy2/MLmodel
@@ -0,0 +1,13 @@
+artifact_path: model
+flavors:
+  python_function:
+    env:
+      conda: conda.yaml
+      virtualenv: python_env.yaml
+    loader_module: mlflow.tensorflow
+    python_version: 3.10.0
+  tensorflow:
+    saved_model_dir: tf2model
+    model_type: tf2-module
+mlflow_version: 2.20.3
+model_uuid: test-uuid-numpy2
diff --git a/tests/data/serve_resources/mlflow/tensorflow_numpy2/conda.yaml b/tests/data/serve_resources/mlflow/tensorflow_numpy2/conda.yaml
@@ -0,0 +1,11 @@
+channels:
+- conda-forge
+dependencies:
+- python=3.10
+- pip
+- pip:
+  - numpy>=2.0.0
+  - tensorflow==2.19.0
+  - scikit-learn
+  - mlflow
+name: mlflow-env
diff --git a/tests/data/serve_resources/mlflow/tensorflow_numpy2/data/keras_module.txt b/tests/data/serve_resources/mlflow/tensorflow_numpy2/data/keras_module.txt
@@ -0,0 +1 @@
+tensorflow.keras
diff --git a/tests/data/serve_resources/mlflow/tensorflow_numpy2/data/model.keras b/tests/data/serve_resources/mlflow/tensorflow_numpy2/data/model.keras
diff --git a/tests/data/serve_resources/mlflow/tensorflow_numpy2/data/save_format.txt b/tests/data/serve_resources/mlflow/tensorflow_numpy2/data/save_format.txt
@@ -0,0 +1 @@
+tf
diff --git a/tests/data/serve_resources/mlflow/tensorflow_numpy2/python_env.yaml b/tests/data/serve_resources/mlflow/tensorflow_numpy2/python_env.yaml
@@ -0,0 +1,5 @@
+python: 3.10.0
+build_dependencies:
+- pip
+dependencies:
+- -r requirements.txt
diff --git a/tests/data/serve_resources/mlflow/tensorflow_numpy2/requirements.txt b/tests/data/serve_resources/mlflow/tensorflow_numpy2/requirements.txt
@@ -0,0 +1,3 @@
+numpy>=2.0.0
+tensorflow==2.19.0
+scikit-learn
diff --git a/tests/data/serve_resources/mlflow/xgboost/conda.yaml b/tests/data/serve_resources/mlflow/xgboost/conda.yaml
@@ -2,14 +2,14 @@ channels:
 - conda-forge
 dependencies:
 - python=3.10.13
-- pip<=23.3.1
+- pip<=24.3
 - pip:
-  - mlflow==2.11.1
+  - mlflow>=2.16.1
   - lz4==4.3.2
-  - numpy==1.26.4
-  - pandas==2.2.1
+  - numpy>=1.26.4,<3.0
+  - pandas>=2.3.0
   - psutil==5.9.8
-  - scikit-learn==1.3.2
-  - scipy==1.11.3
+  - scikit-learn==1.6.1
+  - scipy==1.13.0
   - xgboost==1.7.1
 name: mlflow-env
diff --git a/tests/data/serve_resources/mlflow/xgboost/requirements.txt b/tests/data/serve_resources/mlflow/xgboost/requirements.txt
@@ -1,8 +1,8 @@
 mlflow==3.1.0
 lz4==4.3.2
-numpy==1.26.4
-pandas==2.0.3
+numpy>=1.26.4,<3.0
+pandas>=2.3.0
 psutil==5.9.8
-scikit-learn==1.5.1
-scipy==1.11.3
+scikit-learn==1.6.1
+scipy==1.13.0
 xgboost==1.7.1
diff --git a/tests/data/workflow/requirements.txt b/tests/data/workflow/requirements.txt
@@ -1 +1 @@
-scipy==1.11.3
+scipy==1.13.0
diff --git a/tests/integ/sagemaker/experiments/test_run.py b/tests/integ/sagemaker/experiments/test_run.py
@@ -693,7 +693,7 @@ def _generate_estimator(
         sagemaker_client_config=sagemaker_client_config,
     )
     return SKLearn(
-        framework_version="1.2-1",
+        framework_version="1.4-2",
         entry_point=_ENTRY_POINT_PATH,
         dependencies=[sdk_tar],
         role=execution_role,
diff --git a/tests/integ/sagemaker/feature_store/feature_processor/test_feature_processor_integ.py b/tests/integ/sagemaker/feature_store/feature_processor/test_feature_processor_integ.py
@@ -1108,15 +1108,15 @@ def get_expected_dataframe():
     expected_dataframe = pd.read_csv(os.path.join(_FEATURE_PROCESSOR_DIR, "car-data.csv"))
     expected_dataframe["Model"].replace("^\d\d\d\d\s", "", regex=True, inplace=True)  # noqa: W605
     expected_dataframe["Mileage"].replace("(,)|(mi\.)", "", regex=True, inplace=True)  # noqa: W605
-    expected_dataframe["Mileage"].replace("Not available", np.NaN, inplace=True)
+    expected_dataframe["Mileage"].replace("Not available", np.nan, inplace=True)
     expected_dataframe["Price"].replace("\$", "", regex=True, inplace=True)  # noqa: W605
     expected_dataframe["Price"].replace(",", "", regex=True, inplace=True)
     expected_dataframe["MSRP"].replace(
         "(^MSRP\s\\$)|(,)", "", regex=True, inplace=True  # noqa: W605
     )
-    expected_dataframe["MSRP"].replace("Not specified", np.NaN, inplace=True)
+    expected_dataframe["MSRP"].replace("Not specified", np.nan, inplace=True)
     expected_dataframe["MSRP"].replace(
-        "\\$\d+[a-zA-Z\s]+", np.NaN, regex=True, inplace=True  # noqa: W605
+        "\\$\d+[a-zA-Z\s]+", np.nan, regex=True, inplace=True  # noqa: W605
     )
     expected_dataframe["Mileage"] = expected_dataframe["Mileage"].astype(float)
     expected_dataframe["Price"] = expected_dataframe["Price"].astype(float)
diff --git a/tests/integ/sagemaker/remote_function/test_decorator.py b/tests/integ/sagemaker/remote_function/test_decorator.py
@@ -315,6 +315,9 @@ def divide(x, y):
         divide(10, 2)
 
 
+@pytest.mark.skip(
+    reason="Test only valid for numpy < 2.0 due to serialization compatibility changes",
+)
 def test_with_incompatible_dependencies(
     sagemaker_session, dummy_container_without_error, cpu_instance_type
 ):
diff --git a/tests/integ/sagemaker/serve/test_serve_mlflow_tensorflow_flavor_happy.py b/tests/integ/sagemaker/serve/test_serve_mlflow_tensorflow_flavor_happy.py
@@ -105,7 +105,9 @@ def tensorflow_schema_builder(custom_request_translator, custom_response_transla
 
 @pytest.mark.skipif(
     PYTHON_VERSION_IS_NOT_310,
-    reason="The goal of these test are to test the serving components of our feature",
+    np.__version__ >= "2.0.0",
+    reason="The goal of these test are to test the serving components of our feature and \
+       the input model artifacts used in this specific test are generated with py310 and numpy<2.",
 )
 def test_happy_tensorflow_sagemaker_endpoint_with_tensorflow_serving(
     sagemaker_session,
diff --git a/tests/integ/sagemaker/serve/test_tensorflow_serving_numpy2.py b/tests/integ/sagemaker/serve/test_tensorflow_serving_numpy2.py
diff --git a/tests/integ/test_dependency_compatibility.py b/tests/integ/test_dependency_compatibility.py
diff --git a/tests/integ/test_dependency_resolution.py b/tests/integ/test_dependency_resolution.py
diff --git a/tests/unit/sagemaker/jumpstart/constants.py b/tests/unit/sagemaker/jumpstart/constants.py
diff --git a/tests/unit/sagemaker/serve/detector/test_dependency_manager.py b/tests/unit/sagemaker/serve/detector/test_dependency_manager.py
