More bug fixes for PR checks

Author: aviruthen

sagemaker-serve/tox.ini

@@ -87,8 +87,8 @@ 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.8.0' 'torchvision==0.21.0'
-    pip install 'onnx>=1.17.0' 'onnxruntime>=1.20.0'
+    pip install 'torch==2.8.0' 'torchvision==0.23.0'
+    pip install 'onnx>=1.16.0,<1.17.0' 'onnxruntime>=1.19.0,<1.20.0'
     pip install 'dill>=0.3.9'
     pip install 'tensorflow==2.16.2'
 

sagemaker-train/pyproject.toml

@@ -72,6 +72,9 @@ version = { file = "VERSION"}
 [tool.pytest.ini_options]
 addopts = ["-vv"]
 testpaths = ["tests"]
+markers = [
+    "serial: marks tests that must run serially (not in parallel)",
+]
 
 [tool.black]
 line-length = 100

sagemaker-train/src/sagemaker/ai_registry/dataset.py

@@ -179,7 +179,6 @@ def _validate_dataset_file(cls, file_path: str) -> None:
                 max_size_mb = DATASET_MAX_FILE_SIZE_BYTES / (1024 * 1024)
                 raise ValueError(f"File size {file_size_mb:.2f} MB exceeds maximum allowed size of {max_size_mb:.0f} MB")
 
-    @classmethod
     @classmethod
     @_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="DataSet.get")
     def get(cls, name: str, sagemaker_session=None) -> "DataSet":

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

@@ -19,6 +19,7 @@
 from sagemaker.ai_registry.air_constants import HubContentStatus, REWARD_FUNCTION, REWARD_PROMPT
 
 
+@pytest.mark.serial
 class TestEvaluatorIntegration:
     """Integration tests for Evaluator operations."""
 

sagemaker-train/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 =