Don't run torch.compile on runtime images in docker validations

.github/workflows/validate_docker_images.yml → .github/workflows/validate-docker-images.yml

@@ -1,4 +1,4 @@
-name: Validate Docker Images (with Matrix Generation)
+name: Validate Nightly Docker Images
 on:
   workflow_call:
     inputs:
@@ -12,6 +12,11 @@ on:
         default: false
         required: false
         type: boolean
+      ref:
+        description: 'Reference to checkout, defaults to empty'
+        default: ""
+        required: false
+        type: string
   workflow_dispatch:
     inputs:
       channel:
@@ -28,6 +33,11 @@ on:
         default: false
         required: false
         type: boolean
+      ref:
+        description: 'Reference to checkout, defaults to empty'
+        default: ""
+        required: false
+        type: string
 
 jobs:
   generate-matrix:
@@ -54,5 +64,10 @@ jobs:
         set -ex
         export MATRIX_GPU_ARCH_TYPE="cuda"
         export MATRIX_GPU_ARCH_VERSION="${{ matrix.cuda }}"
+        export MATRIX_IMAGE_TYPE="${{ matrix.image_type }}"
         export TARGET_OS="linux"
-        python test/smoke_test/smoke_test.py --package torchonly --runtime-error-check enabled
+        TORCH_COMPILE_CHECK="--torch-compile-check enabled"
+        if[[ ${MATRIX_IMAGE_TYPE} == "runtime" ]]; then
+          TORCH_COMPILE_CHECK="--torch-compile-check disabled"
+        fi
+        python test/smoke_test/smoke_test.py --package torchonly --runtime-error-check enabled ${TORCH_COMPILE_CHECK}

test/smoke_test/smoke_test.py

@@ -146,7 +146,7 @@ def test_cuda_runtime_errors_captured() -> None:
         raise RuntimeError("Expected CUDA RuntimeError but have not received!")
 
 
-def smoke_test_cuda(package: str, runtime_error_check: str) -> None:
+def smoke_test_cuda(package: str, runtime_error_check: str, torch_compile_check: str) -> None:
     if not torch.cuda.is_available() and is_cuda_system:
         raise RuntimeError(f"Expected CUDA {gpu_arch_ver}. However CUDA is not loaded.")
 
@@ -163,7 +163,7 @@ def smoke_test_cuda(package: str, runtime_error_check: str) -> None:
             print(f"{module['name']} CUDA: {version}")
 
      # torch.compile is available on macos-arm64 and Linux for python 3.8-3.11
-    if sys.version_info < (3, 12, 0) and (
+    if torch_compile_check == "enabled" and sys.version_info < (3, 12, 0) and (
         (target_os == "linux" and torch.cuda.is_available()) or
         target_os == "macos-arm64"):
         smoke_test_compile()
@@ -310,6 +310,13 @@ def main() -> None:
         choices=["enabled", "disabled"],
         default="enabled",
     )
+    parser.add_argument(
+        "--torch-compile-check",
+        help="Check torch compile",
+        type=str,
+        choices=["enabled", "disabled"],
+        default="enabled",
+    )
     options = parser.parse_args()
     print(f"torch: {torch.__version__}")
 
@@ -323,7 +330,7 @@ def main() -> None:
     if options.package == "all":
         smoke_test_modules()
 
-    smoke_test_cuda(options.package, options.runtime_error_check)
+    smoke_test_cuda(options.package, options.runtime_error_check, options.torch_compile_check)
 
 
 if __name__ == "__main__":