From 61c506349134db0a0a2fd6fb2eff8e29a2f84e79 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Wed, 2 Jun 2021 12:06:37 -0700 Subject: [PATCH] [deepspeed] add nvme test skip rule (#11997) * add nvme skip rule * fix --- tests/deepspeed/test_deepspeed.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py index 98dc185888..149cbd4d20 100644 --- a/tests/deepspeed/test_deepspeed.py +++ b/tests/deepspeed/test_deepspeed.py @@ -69,6 +69,22 @@ def require_deepspeed(test_case): return test_case +def require_deepspeed_aio(test_case): + """ + Decorator marking a test that requires deepspeed aio (nvme) + """ + if not is_deepspeed_available(): + return unittest.skip("test requires deepspeed")(test_case) + + import deepspeed + from deepspeed.ops.aio import AsyncIOBuilder + + if not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]: + return unittest.skip("test requires deepspeed async-io")(test_case) + else: + return test_case + + if is_deepspeed_available(): from deepspeed.utils import logger as deepspeed_logger # noqa from transformers.deepspeed import deepspeed_config, is_deepspeed_zero3_enabled # noqa @@ -235,6 +251,7 @@ class TrainerIntegrationDeepSpeed(TestCasePlus, TrainerIntegrationCommon): f"got exception: {context.exception}", ) + @require_deepspeed_aio def test_stage3_nvme_offload(self): with mockenv_context(**self.dist_env_1_gpu): # this actually doesn't have to be on NVMe, any storage will do since this test only