Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / local_compute_context.py: 100%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Region state for trainer-managed local compute execution. 

16 

17Keeps only the ``_LOCAL_COMPUTE_ACTIVE`` ContextVar and the 

18``local_compute_context`` manager; the descendant forward adapters that 

19consult this state live in ``distributed/_builder/forward_rewriter.py`` 

20(the single forward-assignment site, AST-gated per 05 §15.2.3). 

21""" 

22 

23from contextlib import contextmanager 

24from contextvars import ContextVar 

25from typing import Iterator 

26 

27 

28_LOCAL_COMPUTE_ACTIVE = ContextVar("trainer_local_compute_active", default=False) 

29 

30 

31@contextmanager 

32def local_compute_context() -> Iterator[None]: 

33 """Mark execution of a trainer-managed local computation.""" 

34 token = _LOCAL_COMPUTE_ACTIVE.set(True) 

35 try: 

36 yield 

37 finally: 

38 _LOCAL_COMPUTE_ACTIVE.reset(token)