Add numpy 2.0 support

Roja Reddy Sareddy · Roja Reddy Sareddy · commit 40d8e9c5e2e1 · 2025-05-06T11:56:05.000-07:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -39,7 +39,7 @@ dependencies = [
   "google-pasta",
   "importlib-metadata>=1.4.0,<7.0",
   "jsonschema",
-  "numpy==1.26.4",
+  "numpy>=2.0.0,<2.3.0",
   "omegaconf>=2.2,<=2.3",
   "packaging>=23.0,<25",
   "pandas",
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, <2.3.0
 build[virtualenv]==1.2.1
 flake8==7.1.2
 pytest==6.2.5
@@ -22,8 +22,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.2.3
+scikit-learn==1.6.1
 cloudpickle==2.2.1
 jsonpickle<4.0.0
 PyYAML>=6.0.1
@@ -43,7 +43,7 @@ nbformat>=5.9,<6
 accelerate>=0.24.1,<=0.27.0
 schema==0.7.5
 tensorflow>=2.16.2,<=2.18.0
-mlflow>=2.12.2,<2.13
+mlflow>=2.16.1
 huggingface_hub==0.26.2
 uvicorn>=0.30.1
 fastapi==0.115.4
diff --git a/src/sagemaker/serve/utils/conda_in_process.yml b/src/sagemaker/serve/utils/conda_in_process.yml
@@ -12,12 +12,12 @@ dependencies:
     - boto3>=1.34.142,<2.0
     - cloudpickle==2.2.1
     - google-pasta
-    - numpy==1.26.4
+    - numpy>=2.0.0,<2.3.0
     - protobuf>=3.12,<5.0
     - smdebug_rulesconfig==1.0.1
     - importlib-metadata>=1.4.0,<7.0
     - packaging>=23.0,<25
-    - pandas
+    - pandas==2.2.3
     - pathos
     - schema
     - PyYAML>=6.0.1
@@ -64,7 +64,7 @@ dependencies:
             - multiprocess>=0.70.14
             - networkx>=3.1
             - packaging>=23.1
-            - pandas>=1.5.3
+            - pandas==2.2.3
             - 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
@@ -4,21 +4,21 @@ dependencies:
 - python=3.10.13
 - pip<=23.3.1
 - 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,<2.3.0
   - opt-einsum==3.3.0
   - packaging==24.0
-  - pandas==2.2.1
+  - pandas==2.2.3
   - 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,<2.3.0
 opt-einsum==3.3.0
 packaging>=23.0,<25
-pandas==2.2.1
+pandas==2.2.3
 pyyaml==6.0.1
 requests==2.32.2
 torch>=2.6.0
diff --git a/tests/data/serve_resources/mlflow/tensorflow/conda.yaml b/tests/data/serve_resources/mlflow/tensorflow/conda.yaml
@@ -2,10 +2,10 @@ channels:
 - conda-forge
 dependencies:
 - python=3.10.13
-- pip<=23.3.1
+- pip<=24.3
 - pip:
-  - mlflow==2.11.1
+  - mlflow>=2.16.1
   - cloudpickle==2.2.1
-  - numpy==1.26.4
-  - tensorflow==2.16.1
+  - numpy>=1.26.4
+  - tensorflow==2.18.0
 name: mlflow-env
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
+  - pandas==2.2.3
   - psutil==5.9.8
-  - scikit-learn==1.3.2
-  - scipy==1.11.3
-  - xgboost==1.7.1
+  - scikit-learn==1.6.1
+  - scipy==1.13.0
+  - xgboost>=2.1.0
 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==2.13.2
+mlflow>=2.16.1
 lz4==4.3.2
-numpy==1.26.4
-pandas==2.0.3
+numpy>=1.26.4
+pandas==2.2.3
 psutil==5.9.8
-scikit-learn==1.5.1
-scipy==1.11.3
-xgboost==1.7.1
+scikit-learn==1.6.1
+scipy==1.13.0
+xgboost>=2.1.0
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/unit/sagemaker/jumpstart/constants.py b/tests/unit/sagemaker/jumpstart/constants.py
@@ -5361,7 +5361,7 @@
             "safetensors==0.3.1",
             "sagemaker_jumpstart_huggingface_script_utilities==1.1.3",
             "sagemaker_jumpstart_script_utilities==1.1.9",
-            "scipy==1.11.1",
+            "scipy==1.13.0",
             "termcolor==2.3.0",
             "texttable==1.6.7",
             "tokenize-rt==5.1.0",
@@ -7870,7 +7870,7 @@
             "safetensors==0.3.1",
             "sagemaker_jumpstart_huggingface_script_utilities==1.1.3",
             "sagemaker_jumpstart_script_utilities==1.1.9",
-            "scipy==1.11.1",
+            "scipy==1.13.0",
             "termcolor==2.3.0",
             "texttable==1.6.7",
             "tokenize-rt==5.1.0",
@@ -8346,7 +8346,7 @@
             "safetensors==0.3.1",
             "sagemaker_jumpstart_huggingface_script_utilities==1.1.3",
             "sagemaker_jumpstart_script_utilities==1.1.9",
-            "scipy==1.11.1",
+            "scipy==1.13.0",
             "termcolor==2.3.0",
             "texttable==1.6.7",
             "tokenize-rt==5.1.0",
@@ -12095,7 +12095,7 @@
         "inference_vulnerabilities": [],
         "training_vulnerable": False,
         "training_dependencies": [
-            "numpy==1.23.1",
+            "numpy>=2.0.0",
             "opencv_python==4.7.0.68",
             "sagemaker_jumpstart_prepack_script_utilities==1.0.0",
         ],
@@ -14360,10 +14360,10 @@
             "jmespath==1.0.1",
             "jsonschema==4.17.3",
             "multiprocess==0.70.14",
-            "numpy==1.26.4",
+            "numpy>=2.0.0",
             "oscrypto==1.3.0",
             "packaging==23.1",
-            "pandas==2.0.2",
+            "pandas==2.2.3",
             "pathos==0.3.0",
             "pkgutil-resolve-name==1.3.10",
             "platformdirs==3.8.0",
@@ -14884,10 +14884,10 @@
             "jmespath==1.0.1",
             "jsonschema==4.17.3",
             "multiprocess==0.70.14",
-            "numpy==1.24.3",
+            "numpy>=2.0.0",
             "oscrypto==1.3.0",
             "packaging==23.1",
-            "pandas==2.0.2",
+            "pandas==2.2.3",
             "pathos==0.3.0",
             "pkgutil-resolve-name==1.3.10",
             "platformdirs==3.8.0",
@@ -17400,7 +17400,7 @@
             "safetensors==0.3.1",
             "sagemaker_jumpstart_huggingface_script_utilities==1.1.4",
             "sagemaker_jumpstart_script_utilities==1.1.9",
-            "scipy==1.11.1",
+            "scipy==1.13.0",
             "termcolor==2.3.0",
             "texttable==1.6.7",
             "tokenize-rt==5.1.0",
diff --git a/tests/unit/sagemaker/mlflow/test_forward_sagemaker_metrics.py b/tests/unit/sagemaker/mlflow/test_forward_sagemaker_metrics.py
@@ -48,7 +48,7 @@ def mock_mlflow_client():
 def test_encode():
     existing_names = set()
     assert encode("test-name", existing_names) == "test-name"
-    assert encode("test:name", existing_names) == "test_3a_name"
+    assert encode("test:name", existing_names) == "test:name"
     assert encode("test-name", existing_names) == "test-name_1"
 
 
@@ -188,6 +188,7 @@ def getenv_side_effect(arg, default=None):
             Mock(spec=requests.Response),
             Mock(spec=requests.Response),
         ],
+        "https://test.sagemaker.aws/api/2.0/mlflow/runs/update": Mock(spec=requests.Response),
         "https://test.sagemaker.aws/api/2.0/mlflow/runs/terminate": Mock(spec=requests.Response),
     }
 
@@ -215,13 +216,19 @@ def getenv_side_effect(arg, default=None):
         mock_response.status_code = 200
         mock_response.text = json.dumps({})
 
+    mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/update"].status_code = 200
+    mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/update"].text = json.dumps(
+        {"run_id": "test_run_id", "status": "FINISHED"}
+    )
+
     mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/terminate"].status_code = 200
     mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/terminate"].text = json.dumps({})
 
     mock_request.side_effect = [
         mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/experiments/get-by-name"],
         mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/create"],
         *mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/log-batch"],
+        mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/update"],
         mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/terminate"],
     ]
 
@@ -231,7 +238,7 @@ def getenv_side_effect(arg, default=None):
 
     log_to_mlflow(metrics, params, tags)
 
-    assert mock_request.call_count == 6  # Total number of API calls
+    assert mock_request.call_count == 7  # Total number of API calls
 
 
 @patch("sagemaker.mlflow.forward_sagemaker_metrics.get_training_job_details")
diff --git a/tests/unit/sagemaker/serve/detector/test_dependency_manager.py b/tests/unit/sagemaker/serve/detector/test_dependency_manager.py
@@ -21,7 +21,7 @@
 
 DEPENDENCY_LIST = [
     "requests==2.26.0",
-    "numpy==1.26.4",
+    "numpy>=2.0.0",
     "pandas<=1.3.3",
     "matplotlib<3.5.0",
     "scikit-learn>0.24.1",
@@ -34,7 +34,7 @@
 
 EXPECTED_DEPENDENCY_MAP = {
     "requests": "==2.26.0",
-    "numpy": "==1.26.4",
+    "numpy": ">=2.0.0",
     "pandas": "<=1.3.3",
     "matplotlib": "<3.5.0",
     "scikit-learn": ">0.24.1",
