Bug fixes for unit and integ tests

Author: aviruthen

sagemaker-core/pyproject.toml

@@ -34,6 +34,9 @@ dependencies = [
     "omegaconf>=2.1.0",
     "torch>=1.9.0",
     "scipy>=1.5.0",
+    # Remote function dependencies
+    "cloudpickle>=2.0.0",
+    "paramiko>=2.11.0",
 ]
 requires-python = ">=3.9"
 classifiers = [

sagemaker-core/src/sagemaker/core/training/configs.py

@@ -257,15 +257,16 @@ class InputData(BaseConfig):
     Parameters:
         channel_name (StrPipeVar):
             The name of the input data source channel.
-        data_source (Union[str, S3DataSource, FileSystemDataSource, DatasetSource]):
+        data_source (Union[StrPipeVar, S3DataSource, FileSystemDataSource, DatasetSource]):
             The data source for the channel. Can be an S3 URI string, local file path string,
-            S3DataSource object, or FileSystemDataSource object.
+            S3DataSource object, FileSystemDataSource object, DatasetSource object, or a
+            pipeline variable (Properties) from a previous step.
         content_type (StrPipeVar):
             The MIME type of the data.
     """
 
     channel_name: StrPipeVar = None
-    data_source: Union[str, FileSystemDataSource, S3DataSource, DatasetSource] = None
+    data_source: Union[StrPipeVar, FileSystemDataSource, S3DataSource, DatasetSource] = None
     content_type: StrPipeVar = None
 
 

sagemaker-core/tests/integ/jumpstart/test_search_integ.py

@@ -19,6 +19,7 @@
 from sagemaker.core.resources import HubContent
 
 
+@pytest.mark.serial
 @pytest.mark.integ
 def test_search_public_hub_models_default_args():
     # Only query, uses default hub name and session
@@ -30,6 +31,7 @@ def test_search_public_hub_models_default_args():
     assert len(results) > 0, "Expected at least one matching model from the public hub"
 
 
+@pytest.mark.serial
 @pytest.mark.integ
 def test_search_public_hub_models_custom_session():
     # Provide a custom SageMaker session
@@ -41,6 +43,7 @@ def test_search_public_hub_models_custom_session():
     assert all(isinstance(m, HubContent) for m in results)
 
 
+@pytest.mark.serial
 @pytest.mark.integ
 def test_search_public_hub_models_custom_hub_name():
     # Using the default public hub but provided explicitly
@@ -51,6 +54,7 @@ def test_search_public_hub_models_custom_hub_name():
     assert all(isinstance(m, HubContent) for m in results)
 
 
+@pytest.mark.serial
 @pytest.mark.integ
 def test_search_public_hub_models_all_args():
     # Provide both hub_name and session explicitly

sagemaker-core/tests/unit/telemetry/test_telemetry_logging.py

@@ -30,7 +30,18 @@
     PYTHON_VERSION,
 )
 from sagemaker.core.user_agent import SDK_VERSION, process_studio_metadata_file
-from sagemaker.serve.utils.exceptions import ModelBuilderException, LocalModelOutOfMemoryException
+
+# Try to import sagemaker-serve exceptions, skip tests if not available
+try:
+    from sagemaker.serve.utils.exceptions import ModelBuilderException, LocalModelOutOfMemoryException
+    SAGEMAKER_SERVE_AVAILABLE = True
+except ImportError:
+    SAGEMAKER_SERVE_AVAILABLE = False
+    # Create mock exceptions for type hints
+    class ModelBuilderException(Exception):
+        pass
+    class LocalModelOutOfMemoryException(Exception):
+        pass
 
 MOCK_SESSION = Mock()
 MOCK_EXCEPTION = LocalModelOutOfMemoryException("mock raise ex")
@@ -147,6 +158,10 @@ def test_telemetry_emitter_decorator_success(
             1, [1, 2], MOCK_SESSION, None, None, expected_extra_str
         )
 
+    @pytest.mark.skipif(
+        not SAGEMAKER_SERVE_AVAILABLE,
+        reason="Requires sagemaker-serve package"
+    )
     @patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request")
     @patch("sagemaker.core.telemetry.telemetry_logging.resolve_value_from_config")
     def test_telemetry_emitter_decorator_handle_exception_success(

sagemaker-core/tox.ini

@@ -63,6 +63,7 @@ markers =
     release
     image_uris_unit_test
     timeout: mark a test as a timeout.
+    serial: marks tests that must run serially (not in parallel)
 
 [testenv]
 setenv =

sagemaker-mlops/tests/integ/test_pipeline_train_registry.py

@@ -215,7 +215,19 @@ def test_pipeline_with_train_and_registry(sagemaker_session, pipeline_session, r
                 assert execution_status == "Succeeded"
                 break
             elif execution_status in ["Failed", "Stopped"]:
-                pytest.fail(f"Pipeline execution {execution_status}")
+                # Get detailed failure information
+                steps = sagemaker_session.sagemaker_client.list_pipeline_execution_steps(
+                    PipelineExecutionArn=execution_desc["PipelineExecutionArn"]
+                )["PipelineExecutionSteps"]
+                
+                failed_steps = []
+                for step in steps:
+                    if step.get("StepStatus") == "Failed":
+                        failure_reason = step.get("FailureReason", "Unknown reason")
+                        failed_steps.append(f"{step['StepName']}: {failure_reason}")
+                
+                failure_details = "\n".join(failed_steps) if failed_steps else "No detailed failure information available"
+                pytest.fail(f"Pipeline execution {execution_status}. Failed steps:\n{failure_details}")
 
             time.sleep(60)
         else:

sagemaker-mlops/tox.ini

@@ -87,8 +87,7 @@ allowlist_externals =
 commands =
     python -c "import os; os.system('install-custom-pkgs --install-boto-wheels')"
     pip install 'apache-airflow==2.10.4' --constraint "https://raw.githubusercontent.com/apache/airflow/constraints-2.10.4/constraints-3.9.txt"
-    pip install 'torch==2.3.1+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
-    pip install 'torchvision==0.18.1+cpu' -f 'https://download.pytorch.org/whl/torch_stable.html'
+    pip install 'torch==2.8.0' 'torchvision==0.23.0'
     pip install 'dill>=0.3.9'
 
     pytest {posargs}

sagemaker-train/src/sagemaker/train/evaluate/benchmark_evaluator.py

@@ -323,12 +323,10 @@ class BenchMarkEvaluator(BaseEvaluator):
     """
 
     benchmark: _Benchmark
+    dataset: Union[str, Any]  # Required field, must come before optional fields
     subtasks: Optional[Union[str, List[str]]] = None
-    _hyperparameters: Optional[Any] = None
-    
-    # Template-required fields
-    dataset: Union[str, Any]
     evaluate_base_model: bool = True
+    _hyperparameters: Optional[Any] = None
 
     @validator('dataset', pre=True)
     def _resolve_dataset(cls, v):

sagemaker-train/tests/integ/ai_registry/test_dataset.py

@@ -21,6 +21,7 @@
 from sagemaker.ai_registry.air_constants import HubContentStatus
 
 
+@pytest.mark.serial
 class TestDataSetIntegration:
     """Integration tests for DataSet operations."""