Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / hyper_offload / api / session.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"""Activation offload session built on TorchDispatchMode.""" 

16 

17from __future__ import annotations 

18 

19import contextvars 

20import logging 

21 

22from torch.utils._python_dispatch import TorchDispatchMode 

23 

24from hyper_parallel.auto_parallel.hyper_offload.api.config import OffloadConfig 

25from hyper_parallel.auto_parallel.hyper_offload.execution.replay import ReplayExecutor 

26from hyper_parallel.auto_parallel.hyper_offload.execution.warmup import WarmupExecutor 

27from hyper_parallel.auto_parallel.hyper_offload.planning.greedy import GreedyResidencyPlanner 

28from hyper_parallel.auto_parallel.hyper_offload.runtime.residency import ResidencyManager 

29 

30logger = logging.getLogger(__name__) 

31 

32# Context variable tracking the active OffloadSession (set on __enter__). 

33_active_session: contextvars.ContextVar[OffloadSession | None] = contextvars.ContextVar("_active_session", default=None) 

34 

35 

36class ActivationDispatchMode(TorchDispatchMode): 

37 """Dispatch mode that delegates to a :class:`BaseExecutor`. 

38 

39 The executor (warmup or replay) is selected by 

40 :class:`OffloadSession` and never swapped while this mode 

41 is active. 

42 """ 

43 

44 def __init__(self, session: OffloadSession): 

45 super().__init__() 

46 self.session = session 

47 

48 def __torch_dispatch__(self, func, types, args=(), kwargs=None): # pylint: disable=unused-argument 

49 """Dispatch a torch operation.""" 

50 return self.session.executor.dispatch(func, args, kwargs) 

51 

52 

53class OffloadSession: 

54 """Manage warmup and replay for activation residency control. 

55 

56 Holds lifecycle state and a single active executor (:class:`WarmupExecutor` 

57 or :class:`ReplayExecutor`). The executor is switched during the 

58 warmup → replay transition in :meth:`__exit__`. 

59 """ 

60 

61 def __init__(self, config: OffloadConfig): 

62 self.mode: str = "warmup" 

63 

64 self.planner = config.planner or GreedyResidencyPlanner() 

65 self.residency_manager = ResidencyManager( 

66 max_host_bytes=config.max_offload_activation_mb * 1024**2, 

67 ) 

68 self.executor: WarmupExecutor | ReplayExecutor = WarmupExecutor( 

69 residency_manager=self.residency_manager, 

70 memory_limit_bytes=config.max_resident_activation_mb * 1024**2, 

71 ) 

72 

73 self._dispatch_context = None 

74 self._session_token = None 

75 

76 # ------------------------------------------------------------------ 

77 # Active session lookup (used by skip_offload decorator) 

78 # ------------------------------------------------------------------ 

79 

80 @staticmethod 

81 def get_active() -> OffloadSession | None: 

82 """Return the currently active OffloadSession, or ``None``.""" 

83 return _active_session.get() 

84 

85 # ------------------------------------------------------------------ 

86 # Context manager 

87 # ------------------------------------------------------------------ 

88 

89 def __enter__(self): 

90 """Enter the runtime context.""" 

91 self.executor.reset() 

92 

93 self._session_token = _active_session.set(self) 

94 self._dispatch_context = ActivationDispatchMode(self) 

95 self._dispatch_context.__enter__() 

96 return self 

97 

98 def __exit__(self, exc_type, exc_val, exc_tb): 

99 """Exit the runtime context.""" 

100 if self._dispatch_context is not None: 

101 self._dispatch_context.__exit__(exc_type, exc_val, exc_tb) 

102 self._dispatch_context = None 

103 if self._session_token is not None: 

104 _active_session.reset(self._session_token) 

105 self._session_token = None 

106 

107 if exc_type is not None: 

108 self.residency_manager.sync_all_transfers() 

109 self.executor.reset() 

110 return False 

111 

112 self.residency_manager.wait_for_transfers() 

113 

114 if self.mode == "warmup": 

115 trace, guide = self.executor.finish() 

116 schedule = self.planner.build(trace) 

117 self.executor = ReplayExecutor(self.residency_manager, schedule, guide) 

118 self.mode = "replay" 

119 

120 return False