fix: address review comments (iteration #1)

Author: aviruthen

sagemaker-core/pyproject.toml

@@ -14,7 +14,6 @@ dependencies = [
   # Add your dependencies here (Include lower and upper bounds as applicable)
     "boto3>=1.42.2,<2.0.0",
     "pydantic>=2.10.0,<3.0.0",
-    "pydantic-core>=2.27.0,<3.0.0",
     "PyYAML>=6.0, <7.0",
     "jsonschema<5.0.0",
     "platformdirs>=4.0.0, <5.0.0",

sagemaker-core/src/sagemaker/core/__init__.py

@@ -3,10 +3,10 @@
     from sagemaker.core._pydantic_compat import check_pydantic_compatibility
     check_pydantic_compatibility()
 except ImportError as e:
-    if "pydantic" in str(e).lower() and ("incompatible" in str(e).lower() or "mismatch" in str(e).lower()):
+    if "pydantic" in str(e).lower():
         raise
-    # If it's a different ImportError (e.g., pydantic not installed yet), let it pass
-    # and fail later with a more standard error
+    # If it's a different ImportError (e.g., module not found for _pydantic_compat itself),
+    # let it pass and fail later with a more standard error
     pass
 
 from sagemaker.core.utils.utils import enable_textual_rich_console_and_traceback

sagemaker-core/src/sagemaker/core/_pydantic_compat.py

@@ -15,12 +15,24 @@
 This module provides an early check for pydantic/pydantic-core version
 compatibility to give users a clear error message with fix instructions
 instead of a cryptic SystemError.
+
+When pydantic-core is upgraded independently of pydantic (e.g., via
+``pip install --force-reinstall --no-deps``), Python raises a ``SystemError``
+at import time.  This module catches that error and converts it into a
+helpful ``ImportError`` with remediation steps.
 """
 
 
 def check_pydantic_compatibility():
     """Check that pydantic and pydantic-core versions are compatible.
 
+    Pydantic internally requires an exact matching pydantic-core version
+    (e.g., pydantic 2.11.5 requires pydantic-core==2.41.5).  If the two
+    packages are out of sync, ``import pydantic`` raises a ``SystemError``.
+
+    This function catches that ``SystemError`` and re-raises it as an
+    ``ImportError`` with clear instructions on how to fix the issue.
+
     Raises:
         ImportError: If pydantic and pydantic-core versions are incompatible,
             with instructions on how to fix the issue.
@@ -37,44 +49,3 @@ def check_pydantic_compatibility():
             "    pip install pydantic pydantic-core --force-reinstall\n\n"
             "This will ensure both packages are installed at compatible versions."
         ) from e
-
-    try:
-        import pydantic_core  # noqa: F401
-    except ImportError:
-        # pydantic_core not installed separately is fine;
-        # pydantic manages it as a dependency
-        return
-
-    # Additional version check: pydantic declares the exact pydantic-core
-    # version it requires. Verify they match.
-    try:
-        pydantic_version = pydantic.VERSION
-        pydantic_core_version = pydantic_core.VERSION
-
-        # pydantic >= 2.x stores the required core version
-        expected_core_version = getattr(pydantic, '__pydantic_core_version__', None)
-        if expected_core_version is None:
-            # Try alternative attribute name used in some pydantic versions
-            expected_core_version = getattr(
-                pydantic, '_internal', None
-            ) and getattr(
-                getattr(pydantic, '_internal', None),
-                '_generate_schema',
-                None,
-            )
-            # If we can't determine the expected version, skip the check
-            return
-
-        if pydantic_core_version != expected_core_version:
-            raise ImportError(
-                f"Pydantic/pydantic-core version mismatch detected: "
-                f"pydantic {pydantic_version} requires pydantic-core=={expected_core_version}, "
-                f"but pydantic-core {pydantic_core_version} is installed.\n\n"
-                "To fix this, run:\n"
-                "    pip install pydantic pydantic-core --force-reinstall\n\n"
-                "This will ensure both packages are installed at compatible versions."
-            )
-    except (AttributeError, TypeError):
-        # If we can't determine versions, skip the check
-        # The SystemError catch above will handle the most common case
-        pass

sagemaker-core/tests/unit/test_pydantic_compat.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 """Tests for pydantic compatibility check."""
 
+import builtins
 import sys
 from unittest import mock
 
@@ -26,6 +27,18 @@ def test_check_pydantic_compatibility_passes_with_matching_versions():
     check_pydantic_compatibility()
 
 
+def _make_pydantic_system_error_import(error_msg):
+    """Create a mock import function that raises SystemError for pydantic."""
+    _real_import = builtins.__import__
+
+    def _mock_import(name, *args, **kwargs):
+        if name == "pydantic":
+            raise SystemError(error_msg)
+        return _real_import(name, *args, **kwargs)
+
+    return _mock_import
+
+
 def test_check_pydantic_compatibility_raises_on_system_error():
     """Mock pydantic import to raise SystemError and verify a clear ImportError is raised."""
     from sagemaker.core._pydantic_compat import check_pydantic_compatibility
@@ -35,20 +48,15 @@ def test_check_pydantic_compatibility_raises_on_system_error():
         "with the current pydantic version, which requires 2.41.5."
     )
 
-    with mock.patch.dict(sys.modules, {"pydantic": None}):
-        original_import = __builtins__.__import__ if hasattr(__builtins__, '__import__') else __import__
-
-        def mock_import(name, *args, **kwargs):
-            if name == "pydantic":
-                raise SystemError(error_msg)
-            return original_import(name, *args, **kwargs)
+    mock_import = _make_pydantic_system_error_import(error_msg)
 
+    with mock.patch.dict(sys.modules, {"pydantic": None}):
         with mock.patch("builtins.__import__", side_effect=mock_import):
             with pytest.raises(ImportError) as exc_info:
                 check_pydantic_compatibility()
 
-            assert "incompatibility detected" in str(exc_info.value).lower() or \
-                   "incompatible" in str(exc_info.value).lower()
+            error_str = str(exc_info.value).lower()
+            assert "incompatibility detected" in error_str
 
 
 def test_pydantic_import_error_message_contains_instructions():
@@ -60,17 +68,32 @@ def test_pydantic_import_error_message_contains_instructions():
         "with the current pydantic version, which requires 2.41.5."
     )
 
-    with mock.patch.dict(sys.modules, {"pydantic": None}):
-        original_import = __builtins__.__import__ if hasattr(__builtins__, '__import__') else __import__
-
-        def mock_import(name, *args, **kwargs):
-            if name == "pydantic":
-                raise SystemError(error_msg)
-            return original_import(name, *args, **kwargs)
+    mock_import = _make_pydantic_system_error_import(error_msg)
 
+    with mock.patch.dict(sys.modules, {"pydantic": None}):
         with mock.patch("builtins.__import__", side_effect=mock_import):
             with pytest.raises(ImportError) as exc_info:
                 check_pydantic_compatibility()
 
             error_str = str(exc_info.value)
             assert "pip install pydantic pydantic-core --force-reinstall" in error_str
+
+
+def test_pydantic_import_error_chains_original_system_error():
+    """Verify the ImportError chains the original SystemError as __cause__."""
+    from sagemaker.core._pydantic_compat import check_pydantic_compatibility
+
+    error_msg = (
+        "The installed pydantic-core version (2.42.0) is incompatible "
+        "with the current pydantic version, which requires 2.41.5."
+    )
+
+    mock_import = _make_pydantic_system_error_import(error_msg)
+
+    with mock.patch.dict(sys.modules, {"pydantic": None}):
+        with mock.patch("builtins.__import__", side_effect=mock_import):
+            with pytest.raises(ImportError) as exc_info:
+                check_pydantic_compatibility()
+
+            assert isinstance(exc_info.value.__cause__, SystemError)
+            assert "incompatible" in str(exc_info.value.__cause__).lower()