Skip to content

Commit a897b43

Browse files
authored
MLFLow update for dependabot (#5187)
* MLFLow update for dependabot * Update lower bound * Unit test fixes
1 parent 6fb3b81 commit a897b43

File tree

2 files changed

+10
-3
lines changed

2 files changed

+10
-3
lines changed

requirements/extras/test_requirements.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ nbformat>=5.9,<6
4444
accelerate>=0.24.1,<=0.27.0
4545
schema==0.7.5
4646
tensorflow>=2.16.2,<=2.18.0
47-
mlflow>=2.12.2,<2.13
47+
mlflow>=2.14.2,<3
4848
huggingface_hub==0.26.2
4949
uvicorn>=0.30.1
5050
fastapi==0.115.4

tests/unit/sagemaker/mlflow/test_forward_sagemaker_metrics.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ def mock_mlflow_client():
4848
def test_encode():
4949
existing_names = set()
5050
assert encode("test-name", existing_names) == "test-name"
51-
assert encode("test:name", existing_names) == "test_3a_name"
51+
assert encode("test:name", existing_names) == "test:name"
5252
assert encode("test-name", existing_names) == "test-name_1"
5353

5454

@@ -183,6 +183,7 @@ def getenv_side_effect(arg, default=None):
183183
spec=requests.Response
184184
),
185185
"https://test.sagemaker.aws/api/2.0/mlflow/runs/create": Mock(spec=requests.Response),
186+
"https://test.sagemaker.aws/api/2.0/mlflow/runs/update": Mock(spec=requests.Response),
186187
"https://test.sagemaker.aws/api/2.0/mlflow/runs/log-batch": [
187188
Mock(spec=requests.Response),
188189
Mock(spec=requests.Response),
@@ -211,6 +212,11 @@ def getenv_side_effect(arg, default=None):
211212
{"run_id": "test_run_id"}
212213
)
213214

215+
mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/update"].status_code = 200
216+
mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/update"].text = json.dumps(
217+
{"run_id": "test_run_id"}
218+
)
219+
214220
for mock_response in mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/log-batch"]:
215221
mock_response.status_code = 200
216222
mock_response.text = json.dumps({})
@@ -221,6 +227,7 @@ def getenv_side_effect(arg, default=None):
221227
mock_request.side_effect = [
222228
mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/experiments/get-by-name"],
223229
mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/create"],
230+
mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/update"],
224231
*mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/log-batch"],
225232
mock_responses["https://test.sagemaker.aws/api/2.0/mlflow/runs/terminate"],
226233
]
@@ -231,7 +238,7 @@ def getenv_side_effect(arg, default=None):
231238

232239
log_to_mlflow(metrics, params, tags)
233240

234-
assert mock_request.call_count == 6 # Total number of API calls
241+
assert mock_request.call_count == 7 # Total number of API calls
235242

236243

237244
@patch("sagemaker.mlflow.forward_sagemaker_metrics.get_training_job_details")

0 commit comments

Comments
 (0)