|
47 | 47 | cpu_only = unittest.skipIf(DEVICE != "cpu", "CPU-only test") |
48 | 48 | sagemaker_only = unittest.skipIf(CUSTOMER != "sagemaker", "SageMaker-only test") |
49 | 49 | tensorflow_only = unittest.skipIf(FRAMEWORK != "tensorflow", "TF-only test") |
50 | | -# Tests gated with this decorator assume PT/TF-style training DLC |
51 | | -# (SSH cluster, MPI, entrypoint script). xgboost is a SageMaker |
52 | | -# algorithm container and does not honor this contract. |
| 50 | +# Tests gated with this decorator assume a training DLC with the SSH+MPI cluster |
| 51 | +# stack and /opt/venv. Covers PyTorch, TensorFlow, and RayTrain (all reuse the |
| 52 | +# PyTorch EFA/NCCL/SSH install scripts). xgboost is a SageMaker algorithm |
| 53 | +# container and does not honor this contract. |
53 | 54 | training_cluster_only = unittest.skipUnless( |
54 | | - FRAMEWORK in {"tensorflow", "pytorch_runtime"}, |
| 55 | + FRAMEWORK in {"tensorflow", "pytorch_runtime", "ray"}, |
55 | 56 | "training-cluster-only test (requires SSH+MPI stack; xgboost is algorithm container)", |
56 | 57 | ) |
| 58 | +# RayTrain reuses EFA's bundled OpenMPI and a passive/KubeRay entrypoint — it does |
| 59 | +# NOT rebuild OpenMPI from source and has no /usr/local/bin/entrypoint.sh, so the |
| 60 | +# double-wrap and entrypoint-script checks below do not apply to it. |
| 61 | +pt_tf_only = unittest.skipUnless( |
| 62 | + FRAMEWORK in {"tensorflow", "pytorch_runtime"}, |
| 63 | + "PyTorch/TensorFlow-only test (source-built OpenMPI + entrypoint.sh contract)", |
| 64 | +) |
57 | 65 |
|
58 | 66 |
|
59 | 67 | class TestContainerEnv(unittest.TestCase): |
@@ -204,11 +212,13 @@ class TestOpenMPI(unittest.TestCase): |
204 | 212 | def test_openmpi_binary_exists(self): |
205 | 213 | self.assertTrue(os.access("/opt/amazon/openmpi/bin/mpirun", os.X_OK)) |
206 | 214 |
|
| 215 | + @pt_tf_only |
207 | 216 | def test_openmpi_double_wrap(self): |
208 | 217 | """`mpirun` is a wrapper that exec's `mpirun.real --allow-run-as-root`. Verify the |
209 | 218 | wrapper-vs-real split is in place — single grep -c match means it's NOT |
210 | 219 | double-wrapped (which would happen if EFA's bundled OMPI wasn't wiped before |
211 | | - the from-source build).""" |
| 220 | + the from-source build). RayTrain uses EFA's bundled OpenMPI as-is, so this |
| 221 | + source-build-specific check does not apply.""" |
212 | 222 | self.assertTrue(os.access("/opt/amazon/openmpi/bin/mpirun.real", os.X_OK)) |
213 | 223 | out = subprocess.check_output( |
214 | 224 | ["grep", "-c", "mpirun.real", "/opt/amazon/openmpi/bin/mpirun"], text=True |
@@ -315,13 +325,53 @@ def test_venv_bin_exists(self): |
315 | 325 | self.assertTrue(os.path.isdir("/opt/venv/bin")) |
316 | 326 |
|
317 | 327 |
|
318 | | -@training_cluster_only |
| 328 | +@pt_tf_only |
319 | 329 | class TestEntrypoint(unittest.TestCase): |
320 | | - """Entrypoint script is executable (universal).""" |
| 330 | + """Entrypoint script is executable. PyTorch/TensorFlow ship |
| 331 | + /usr/local/bin/entrypoint.sh; RayTrain uses a passive/KubeRay entrypoint and |
| 332 | + does not, so this is gated to PT/TF only.""" |
321 | 333 |
|
322 | 334 | def test_entrypoint_executable(self): |
323 | 335 | self.assertTrue(os.access("/usr/local/bin/entrypoint.sh", os.X_OK)) |
324 | 336 |
|
325 | 337 |
|
| 338 | +ray_only = unittest.skipUnless(FRAMEWORK == "ray", "RayTrain-only test") |
| 339 | + |
| 340 | + |
| 341 | +@ray_only |
| 342 | +class TestRayTrain(unittest.TestCase): |
| 343 | + """RayTrain-specific contract: Ray training stack present, training-scoped |
| 344 | + (no Ray Serve extra). Gated on EXPECTED_FRAMEWORK == 'ray'.""" |
| 345 | + |
| 346 | + def test_ray_train_stack_imports(self): |
| 347 | + import accelerate # noqa: F401 |
| 348 | + import datasets # noqa: F401 |
| 349 | + import pytorch_lightning # noqa: F401 |
| 350 | + import ray.data # noqa: F401 |
| 351 | + import ray.train # noqa: F401 |
| 352 | + import ray.tune # noqa: F401 |
| 353 | + import transformers # noqa: F401 |
| 354 | + |
| 355 | + def test_ray_version_matches_expected(self): |
| 356 | + expected = os.environ.get("EXPECTED_FRAMEWORK_VERSION", "") |
| 357 | + if not expected: |
| 358 | + self.skipTest("EXPECTED_FRAMEWORK_VERSION not set") |
| 359 | + import ray |
| 360 | + |
| 361 | + self.assertEqual(ray.__version__, expected) |
| 362 | + |
| 363 | + def test_torch_is_cuda_build(self): |
| 364 | + import torch |
| 365 | + |
| 366 | + self.assertIn("cu", torch.__version__) |
| 367 | + |
| 368 | + def test_ray_serve_extra_absent(self): |
| 369 | + """RayTrain is training-scoped. The `ray.serve` module always ships in the |
| 370 | + ray package, but without the `serve` extra its deps (e.g. starlette) are |
| 371 | + missing and the import fails — a successful import means serve leaked in.""" |
| 372 | + with self.assertRaises(ImportError): |
| 373 | + import ray.serve # noqa: F401 |
| 374 | + |
| 375 | + |
326 | 376 | if __name__ == "__main__": |
327 | 377 | unittest.main(verbosity=2) |
0 commit comments