sklearn: add sagemaker tests + version-drift check

.github/workflows/sklearn.dispatch-release.yml

@@ -4,12 +4,13 @@ on:
   workflow_dispatch:
     inputs:
       config-file:
-        description: "Path to image config YAML (e.g., .github/config/image/sklearn/sagemaker-py312.yml)"
+        description: "Path to image config YAML"
         required: true
         type: string
+        default: ".github/config/image/sklearn/sagemaker-py312.yml"
 
 concurrency:
-  group: ${{ github.workflow }}
+  group: ${{ github.workflow }}-${{ inputs.config-file }}
   cancel-in-progress: false
 
 permissions:

.github/workflows/sklearn.pipeline.yml

@@ -41,6 +41,11 @@ on:
         required: false
         type: boolean
         default: true
+      run-sagemaker-test:
+        description: "Run SageMaker tests (real training jobs + endpoints)"
+        required: false
+        type: boolean
+        default: true
       release:
         description: "Run release job after tests pass"
         required: false
@@ -131,15 +136,27 @@ jobs:
       config-file: ${{ inputs.config-file }}
       image-uri: ${{ needs.build.outputs.image-uri || '' }}
 
+  # SageMaker tests — real training jobs + endpoints
+  sagemaker-test:
+    if: ${{ always() && !failure() && !cancelled() && inputs.run-sagemaker-test }}
+    needs: [build]
+    concurrency:
+      group: ${{ github.workflow }}-sagemaker-test-${{ inputs.config-file }}-${{ github.ref }}
+      cancel-in-progress: true
+    uses: ./.github/workflows/sklearn.tests-sagemaker.yml
+    with:
+      config-file: ${{ inputs.config-file }}
+      image-uri: ${{ needs.build.outputs.image-uri || '' }}
+
   release-gate:
     if: >-
       github.ref == 'refs/heads/main' &&
-      (contains(github.workflow_ref, '.autorelease') || contains(github.workflow_ref, '.dispatch-release')) &&
+      contains(github.workflow_ref, '.dispatch-release') &&
       !contains(github.workflow_ref, '.pr') &&
       inputs.release &&
       !failure() && !cancelled() &&
       needs.build.result == 'success'
-    needs: [build, sanity-test, security-test, unit-test, integ-test]
+    needs: [build, sanity-test, security-test, unit-test, integ-test, sagemaker-test]
     runs-on: ubuntu-latest
     outputs:
       environment: ${{ steps.check.outputs.environment }}

.github/workflows/sklearn.pr-1.4-py312.yml

@@ -11,13 +11,15 @@ on:
       - ".github/workflows/sklearn.pr-1.4-py312.yml"
       - ".github/workflows/sklearn.tests-unit.yml"
       - ".github/workflows/sklearn.tests-integ-local.yml"
+      - ".github/workflows/sklearn.tests-sagemaker.yml"
       - "docker/sklearn/1.4-2-py312/**"
       - "docker/sklearn/resources/**"
       - "scripts/ci/build/sklearn/**"
       - "scripts/docker/common/**"
       - "scripts/docker/telemetry/**"
       - "test/sanity/**"
       - "test/security/data/ecr_scan_allowlist/sklearn/**"
+      - "test/sklearn/**"
       - "test/telemetry/**"
       - "!docs/**"
 
@@ -41,6 +43,7 @@ jobs:
     outputs:
       build-change: ${{ steps.changes.outputs.build-change }}
       integ-test-change: ${{ steps.changes.outputs.integ-test-change }}
+      sagemaker-test-change: ${{ steps.changes.outputs.sagemaker-test-change }}
       sanity-test-change: ${{ steps.changes.outputs.sanity-test-change }}
       telemetry-test-change: ${{ steps.changes.outputs.telemetry-test-change }}
       unit-test-change: ${{ steps.changes.outputs.unit-test-change }}
@@ -60,6 +63,9 @@ jobs:
               - "test/security/data/ecr_scan_allowlist/sklearn/**"
             integ-test-change:
               - "test/sklearn/**"
+            sagemaker-test-change:
+              - "test/sklearn/sagemaker/**"
+              - ".github/workflows/sklearn.tests-sagemaker.yml"
             sanity-test-change:
               - "test/sanity/**"
             telemetry-test-change:
@@ -91,6 +97,7 @@ jobs:
       tag-suffix: pr-${{ github.event.pull_request.number }}
       build: ${{ needs.check-changes.outputs.build-change == 'true' }}
       run-integ-test: ${{ needs.check-changes.outputs.build-change == 'true' || needs.check-changes.outputs.integ-test-change == 'true' }}
+      run-sagemaker-test: ${{ needs.check-changes.outputs.build-change == 'true' || needs.check-changes.outputs.sagemaker-test-change == 'true' }}
       run-sanity-test: ${{ needs.check-changes.outputs.build-change == 'true' || needs.check-changes.outputs.sanity-test-change == 'true' }}
       run-security-test: ${{ needs.check-changes.outputs.build-change == 'true' }}
       run-telemetry-test: false

.github/workflows/sklearn.tests-integ-local.yml

@@ -35,7 +35,7 @@ jobs:
         with:
           image-uri: ${{ inputs.image-uri }}
           ci-aws-account-id: ${{ vars.CI_AWS_ACCOUNT_ID }}
-          prod-aws-account-id: ${{ vars.PROD_AWS_ACCOUNT_ID }}
+          prod-aws-account-id: ${{ vars.PROD_AWS_ACCOUNT_ID_SAGEMAKER }}
           aws-region: ${{ vars.AWS_REGION }}
           config-file: ${{ inputs.config-file }}
 
@@ -61,6 +61,25 @@ jobs:
           echo "framework-version=$(yq '.metadata.framework_version' "$CONFIG_FILE")" >> $GITHUB_OUTPUT
           echo "sklearn-branch=$(yq '.build.sklearn_container_branch // "master"' "$CONFIG_FILE")" >> $GITHUB_OUTPUT
           echo "python-version=$(yq '.build.python_version' "$CONFIG_FILE")" >> $GITHUB_OUTPUT
+          echo "dockerfile=$(yq '.build.dockerfile' "$CONFIG_FILE")" >> $GITHUB_OUTPUT
+
+      - name: Verify pinned package versions match declared requirements.txt
+        if: steps.check.outputs.exists == 'true'
+        env:
+          IMAGE_URI: ${{ steps.image.outputs.image-uri }}
+          DOCKERFILE: ${{ steps.config.outputs.dockerfile }}
+        run: |
+          REQS_FILE="$(dirname "$DOCKERFILE")/requirements.txt"
+          if [[ ! -f "$REQS_FILE" ]]; then
+            echo "No requirements.txt at $REQS_FILE — skipping version check"
+            exit 0
+          fi
+          docker run --rm \
+            -v "$PWD/$REQS_FILE:/tmp/requirements.txt:ro" \
+            -v "$PWD/test/sklearn/scripts/check_versions.py:/tmp/check_versions.py:ro" \
+            --entrypoint python3 \
+            "$IMAGE_URI" \
+            /tmp/check_versions.py /tmp/requirements.txt
 
       - name: Clone sagemaker-scikit-learn-container (tests + source)
         if: steps.check.outputs.exists == 'true'

.github/workflows/sklearn.tests-sagemaker.yml

@@ -0,0 +1,90 @@
+name: Reusable Scikit-learn SageMaker Tests
+
+on:
+  workflow_call:
+    inputs:
+      config-file:
+        description: "Path to image config YAML"
+        required: true
+        type: string
+      image-uri:
+        description: "Image URI to test. If empty, derives prod URI from config."
+        required: false
+        type: string
+        default: ""
+
+permissions:
+  contents: read
+
+jobs:
+  preflight:
+    runs-on:
+      - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }}
+        fleet:default-runner
+        buildspec-override:true
+    outputs:
+      exists: ${{ steps.check.outputs.exists }}
+      image-uri: ${{ steps.image.outputs.image-uri }}
+      aws-account-id: ${{ steps.image.outputs.aws-account-id }}
+      framework-version: ${{ steps.config.outputs.framework-version }}
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Resolve image URI
+        id: image
+        uses: ./.github/actions/resolve-image-uri
+        with:
+          image-uri: ${{ inputs.image-uri }}
+          ci-aws-account-id: ${{ vars.CI_AWS_ACCOUNT_ID }}
+          prod-aws-account-id: ${{ vars.PROD_AWS_ACCOUNT_ID_SAGEMAKER }}
+          aws-region: ${{ vars.AWS_REGION }}
+          config-file: ${{ inputs.config-file }}
+
+      - name: Check image exists
+        id: check
+        uses: ./.github/actions/check-image-exists
+        with:
+          image-uri: ${{ steps.image.outputs.image-uri }}
+
+      - name: Parse config
+        id: config
+        run: |
+          CONFIG_FILE="${{ inputs.config-file }}"
+          echo "framework-version=$(yq '.metadata.framework_version' "$CONFIG_FILE")" >> $GITHUB_OUTPUT
+
+  # SageMaker tests — real training jobs + endpoints
+  sagemaker-test:
+    if: ${{ needs.preflight.outputs.exists == 'true' }}
+    needs: [preflight]
+    timeout-minutes: 120
+    strategy:
+      fail-fast: false
+      matrix:
+        test-module:
+          - test_inference
+          - test_inference_mme
+          - test_network_isolation
+          - test_scoring_pipelines
+          - test_script_mode
+    runs-on:
+      - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }}
+        fleet:default-runner
+        buildspec-override:true
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Install test dependencies
+        run: |
+          uv venv --python 3.12
+          source .venv/bin/activate
+          uv pip install -r test/requirements.txt
+          uv pip install -r test/sklearn/requirements.txt
+
+      - name: Run ${{ matrix.test-module }}
+        run: |
+          source .venv/bin/activate
+          cd test/
+          python3 -m pytest -v --tb=short -rA --log-cli-level=INFO \
+            --image-uri ${{ needs.preflight.outputs.image-uri }} \
+            --sklearn-version ${{ needs.preflight.outputs.framework-version }} \
+            sklearn/sagemaker/${{ matrix.test-module }}.py

.github/workflows/sklearn.tests-unit.yml

@@ -32,7 +32,7 @@ jobs:
         with:
           image-uri: ${{ inputs.image-uri }}
           ci-aws-account-id: ${{ vars.CI_AWS_ACCOUNT_ID }}
-          prod-aws-account-id: ${{ vars.PROD_AWS_ACCOUNT_ID }}
+          prod-aws-account-id: ${{ vars.PROD_AWS_ACCOUNT_ID_SAGEMAKER }}
           aws-region: ${{ vars.AWS_REGION }}
           config-file: ${{ inputs.config-file }}
 

test/sklearn/conftest.py

@@ -0,0 +1,22 @@
+"""Shared pytest configuration for Scikit-learn tests."""
+
+import pytest
+from test_utils.constants import SAGEMAKER_ROLE
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--sklearn-version",
+        default="1.4.2",
+        help="Scikit-learn version under test (e.g. 1.4.2)",
+    )
+
+
+@pytest.fixture(scope="session")
+def sklearn_version(request):
+    return request.config.getoption("--sklearn-version")
+
+
+@pytest.fixture(scope="session")
+def role():
+    return SAGEMAKER_ROLE

test/sklearn/requirements.txt

@@ -0,0 +1 @@
+sagemaker>=2,<3