Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / hyper_offload / execution / replay / executor.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"""Replay-phase executor. 

16 

17The executor executes residency schedule actions (D2H/H2D, device/host 

18release) and validates replayed outputs against the warmup trace. 

19Physical state transitions are delegated to 

20:class:`~offload.runtime.residency.ResidencyManager`. 

21""" 

22 

23from __future__ import annotations 

24 

25import logging 

26from typing import Any 

27 

28from torch.utils._pytree import tree_flatten 

29 

30from hyper_parallel.auto_parallel.hyper_offload.execution.base import BaseExecutor 

31from hyper_parallel.auto_parallel.hyper_offload.ir.replay import OpGuide 

32from hyper_parallel.auto_parallel.hyper_offload.ir.schedule import ResidencyActionType, ResidencySchedule 

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

34 

35logger = logging.getLogger(__name__) 

36 

37 

38class ReplayExecutor(BaseExecutor): 

39 """Executor for the replay phase. 

40 

41 Executes residency schedule actions (D2H/H2D, device/host release) 

42 and validates replayed outputs against the warmup trace. 

43 """ 

44 

45 def __init__( 

46 self, 

47 residency_manager: ResidencyManager, 

48 schedule: ResidencySchedule, 

49 guide: list[OpGuide], 

50 ) -> None: 

51 super().__init__(residency_manager) 

52 self._schedule = schedule 

53 self._guide = guide 

54 

55 # ------------------------------------------------------------------ 

56 # Lifecycle hooks 

57 # ------------------------------------------------------------------ 

58 

59 def on_op_begin(self, func, args, kwargs) -> None: 

60 """Before op: execute pre-actions.""" 

61 super().on_op_begin(func, args, kwargs) 

62 if self.op_idx >= len(self._guide): 

63 raise RuntimeError( 

64 "replay op count exceeds warmup trace: " 

65 f"op_idx={self.op_idx}, " 

66 f"trace_ops={len(self._guide)}, " 

67 f"op={func.__name__}" 

68 ) 

69 

70 for action in self._schedule.pre_actions(self.op_idx): 

71 if action.kind == ResidencyActionType.COPY_H2D: 

72 self.residency_manager.copy_h2d(action.storage_id) 

73 else: 

74 raise RuntimeError(f"unsupported pre action {action} at op={self.op_idx}") 

75 

76 def on_op_end(self, result) -> Any: 

77 """After op: validate output structure, shadow outputs, then execute post-actions. 

78 

79 The order is critical: :meth:`apply_shadows` must run **before** 

80 post-actions so that output tensors are registered in the residency 

81 table (via :meth:`bind`) before the scheduler tries to copy or 

82 release them. The previous order (post-actions → apply_shadows) 

83 caused ``COPY_D2H`` / ``RELEASE_DEVICE`` to silently skip output 

84 sids that had not yet been bound, leaving stale device buffers 

85 and missing host copies. 

86 """ 

87 func = self._last_func 

88 op_guide = self._guide[self.op_idx] 

89 leaves, _ = tree_flatten(result) 

90 if len(leaves) != op_guide.output_leaf_count: 

91 raise RuntimeError( 

92 "replay output structure differs from warmup trace: " 

93 f"op_idx={self.op_idx}, " 

94 f"name={op_guide.name}, " 

95 f"expected_leaves={op_guide.output_leaf_count}, " 

96 f"actual_leaves={len(leaves)}, " 

97 f"op={func.__name__}" 

98 ) 

99 

100 # 1. Shadow outputs first — binds PhysicalBuffers so that 

101 # post-actions can find and operate on them. 

102 result = self.apply_shadows(result, op_guide.output_bindings) 

103 

104 # 2. Execute post-actions (COPY_D2H, RELEASE_DEVICE, RELEASE_HOST). 

105 for action in self._schedule.post_actions(self.op_idx): 

106 sid = action.storage_id 

107 if action.kind == ResidencyActionType.COPY_D2H: 

108 self.residency_manager.copy_d2h(sid) 

109 elif action.kind == ResidencyActionType.RELEASE_DEVICE: 

110 self.residency_manager.release_device(sid) 

111 elif action.kind == ResidencyActionType.RELEASE_HOST: 

112 self.residency_manager.release_host(sid) 

113 else: 

114 raise RuntimeError(f"unsupported post action {action} at op={self.op_idx}") 

115 

116 return result