Revert "fix: Move sagemaker-mlflow to extras (aws#4903)"

This reverts commit 292a00d.

Author: Ashish Gupta

hatch_build.py

@@ -20,7 +20,7 @@ def read_feature_deps(feature):
 
     optional_dependencies = {"all": []}
 
-    for feature in ("feature-processor", "huggingface", "local", "scipy", "sagemaker-mlflow"):
+    for feature in ("feature-processor", "huggingface", "local", "scipy"):
         dependencies = read_feature_deps(feature)
         optional_dependencies[feature] = dependencies
         optional_dependencies["all"].extend(dependencies)

pyproject.toml

@@ -48,6 +48,7 @@ dependencies = [
   "PyYAML~=6.0",
   "requests",
   "sagemaker-core>=1.0.0,<2.0.0",
+  "sagemaker-mlflow",
   "schema",
   "smdebug_rulesconfig==1.0.1",
   "tblib>=1.7.0,<4",

requirements/extras/sagemaker-mlflow_requirements.txt

@@ -44,4 +44,3 @@ huggingface_hub>=0.23.4
 uvicorn>=0.30.1
 fastapi>=0.111.0
 nest-asyncio
-sagemaker-mlflow>=0.1.0

requirements/extras/test_requirements.txt

@@ -107,6 +107,7 @@
 from sagemaker.workflow.parameters import ParameterString
 from sagemaker.workflow.pipeline_context import PipelineSession, runnable_by_pipeline
 
+from sagemaker.mlflow.forward_sagemaker_metrics import log_sagemaker_job_to_mlflow
 
 logger = logging.getLogger(__name__)
 
@@ -1373,14 +1374,8 @@ def fit(
             forward_to_mlflow_tracking_server = True
         if wait:
             self.latest_training_job.wait(logs=logs)
-        try:
-            if forward_to_mlflow_tracking_server:
-                from sagemaker.mlflow.forward_sagemaker_metrics import log_sagemaker_job_to_mlflow
-
-                log_sagemaker_job_to_mlflow(self.latest_training_job.name)
-        except ImportError:
-            if forward_to_mlflow_tracking_server:
-                raise ValueError("Unable to import mlflow, check if sagemaker-mlflow is installed")
+        if forward_to_mlflow_tracking_server:
+            log_sagemaker_job_to_mlflow(self.latest_training_job.name)
 
     def _compilation_job_name(self):
         """Placeholder docstring"""

src/sagemaker/estimator.py

@@ -20,11 +20,7 @@
 import re
 from typing import Set, Tuple, List, Dict, Generator
 import boto3
-
-try:
-    import mlflow
-except ImportError:
-    raise ValueError("Unable to import mlflow, check if sagemaker-mlflow is installed.")
+import mlflow
 from mlflow import MlflowClient
 from mlflow.entities import Metric, Param, RunTag
 

src/sagemaker/mlflow/forward_sagemaker_metrics.py

@@ -5924,7 +5924,7 @@ def test_estimator_get_app_url_fail(sagemaker_session):
     assert "does not support URL retrieval." in str(error)
 
 
-@patch("sagemaker.mlflow.forward_sagemaker_metrics.log_sagemaker_job_to_mlflow")
+@patch("sagemaker.estimator.log_sagemaker_job_to_mlflow")
 def test_forward_sagemaker_metrics(mock_log_to_mlflow, sagemaker_session):
     f = DummyFramework(
         entry_point=SCRIPT_PATH,
@@ -5943,7 +5943,7 @@ def test_forward_sagemaker_metrics(mock_log_to_mlflow, sagemaker_session):
     mock_log_to_mlflow.assert_called_once()
 
 
-@patch("sagemaker.mlflow.forward_sagemaker_metrics.log_sagemaker_job_to_mlflow")
+@patch("sagemaker.estimator.log_sagemaker_job_to_mlflow")
 def test_no_forward_sagemaker_metrics(mock_log_to_mlflow, sagemaker_session):
     f = DummyFramework(
         entry_point=SCRIPT_PATH,