已合并
test(fx):Add validation cases for torch._C._distributed_rpc._is_current_rpc_agent_set on NPU #43170
olpk创建于 7月29日
test(fx):Add validation cases for torch._C._distributed_rpc._is_current_rpc_agent_set on NPU #43170
已合并
共 1 个文件变更+59-0
| @@ -0,0 +1,59 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd | ||
| 2 | +# All rights reserved. | ||
| 3 | +# | ||
| 4 | +# Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +# you may not use this file except in compliance with the License. | ||
| 6 | +# You may obtain a copy of the License at | ||
| 7 | +# | ||
| 8 | +# https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +# | ||
| 10 | +# Unless required by applicable law or agreed to in writing, software | ||
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +# See the License for the specific language governing permissions and | ||
| 14 | +# limitations under the License. | ||
| 15 | + | ||
| 16 | +""" | ||
| 17 | +Add validation cases for torch._C._distributed_rpc APIs on NPU: | ||
| 18 | +1. PyTorch community lacks direct and sufficient validation for some APIs, | ||
| 19 | + so this file is added. | ||
| 20 | +2. This file validates torch._C._distributed_rpc._is_current_rpc_agent_set (extendable). | ||
| 21 | +""" | ||
| 22 | + | ||
| 23 | +import torch # noqa: F401 | ||
| 24 | +from torch_npu.testing.testcase import run_tests, TestCase | ||
| 25 | + | ||
| 26 | + | ||
| 27 | +class TestIsCurrentRpcAgentSet(TestCase): | ||
| 28 | + """Test cases for torch._C._distributed_rpc._is_current_rpc_agent_set.""" | ||
| 29 | + | ||
| 30 | + def test_is_current_rpc_agent_set_import(self): | ||
| 31 | + """Verify that _is_current_rpc_agent_set is importable and callable.""" | ||
| 32 | + from torch._C._distributed_rpc import _is_current_rpc_agent_set | ||
| 33 | + self.assertTrue(callable(_is_current_rpc_agent_set)) | ||
| 34 | + | ||
| 35 | + def test_is_current_rpc_agent_set_default(self): | ||
| 36 | + """Verify that _is_current_rpc_agent_set returns False when RPC is not initialized.""" | ||
| 37 | + from torch._C._distributed_rpc import _is_current_rpc_agent_set | ||
| 38 | + self.assertFalse(_is_current_rpc_agent_set()) | ||
| 39 | + | ||
| 40 | + def test_is_current_rpc_agent_set_after_init(self): | ||
| 41 | + """Verify that _is_current_rpc_agent_set returns True after init_rpc, and False after shutdown.""" | ||
| 42 | + import os | ||
| 43 | + import torch.distributed.rpc as rpc | ||
| 44 | + from torch._C._distributed_rpc import _is_current_rpc_agent_set | ||
| 45 | + | ||
| 46 | + os.environ.setdefault("MASTER_ADDR", "localhost") | ||
| 47 | + os.environ.setdefault("MASTER_PORT", "29500") | ||
| 48 | + | ||
| 49 | + self.assertFalse(_is_current_rpc_agent_set()) | ||
| 50 | + rpc.init_rpc("worker0", rank=0, world_size=1) | ||
| 51 | + try: | ||
| 52 | + self.assertTrue(_is_current_rpc_agent_set()) | ||
| 53 | + finally: | ||
| 54 | + rpc.shutdown() | ||
| 55 | + self.assertFalse(_is_current_rpc_agent_set()) | ||
 | |||
| 56 | + | ||
| 57 | + | ||
| 58 | +if __name__ == "__main__": | ||
| 59 | + run_tests() | ||
🟡 Medium Priority
在
test_is_current_rpc_agent_set_after_init方法中（第 46-48 行），rpc.init_rpc(...)与rpc.shutdown()之间夹了一个self.assertTrue(...)断言。如果该断言失败（抛出AssertionError），则rpc.shutdown()永远不会被执行，导致：test_is_current_rpc_agent_set_default）会读取到残留的 RPC agent 状态，导致self.assertFalse(_is_current_rpc_agent_set())意外失败，产生难以定位的级联失败。触发条件：
_is_current_rpc_agent_set()在init_rpc之后返回了False（例如框架 Bug 或环境问题），断言触发后即泄漏。建议使用
try/finally确保shutdown()始终在断言之后被调用。建议：使用 try/finally 包裹 init_rpc 之后的断言，确保 shutdown 总是被调用。同时将 shutdown 后的断言放在 finally 之外，避免 shutdown 异常时跳过该断言。