Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / hyper_offload / runtime / timer.py: 59%

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"""Device timer using asynchronous events. 

16 

17:class:`DeviceTimer` encapsulates the ``torch.Event`` lifecycle so 

18that callers do not need to deal with API details directly. 

19""" 

20 

21from __future__ import annotations 

22 

23import torch 

24 

25 

26class DeviceTimer: 

27 """Asynchronous op timer backed by events. 

28 

29 Usage:: 

30 

31 timer = DeviceTimer() 

32 timer.start() 

33 ... # Device work 

34 elapsed_ms = timer.stop() 

35 

36 Only the **end** event is synchronised (light-weight); the start 

37 and end events are recorded on ``torch.accelerator.current_stream()``. 

38 """ 

39 

40 def __init__(self) -> None: 

41 self._start_event: torch.Event | None = None 

42 

43 def start(self) -> None: 

44 """Record the start event on the current accelerator stream. 

45 

46 Raises: 

47 RuntimeError: If no accelerator is available. 

48 """ 

49 if not torch.accelerator.is_available(): 

50 raise RuntimeError( 

51 "Cannot start DeviceTimer: no accelerator available" 

52 ) 

53 ev = torch.Event(enable_timing=True) 

54 ev.record(torch.accelerator.current_stream()) 

55 self._start_event = ev 

56 

57 def stop(self) -> float: 

58 """Record the end event, synchronise, and return elapsed milliseconds. 

59 

60 Raises: 

61 RuntimeError: If ``start()`` was not called or no accelerator is 

62 available. 

63 """ 

64 if self._start_event is None: 

65 raise RuntimeError( 

66 "Cannot stop DeviceTimer: start() was not called" 

67 ) 

68 if not torch.accelerator.is_available(): 

69 raise RuntimeError( 

70 "Cannot stop DeviceTimer: no accelerator available" 

71 ) 

72 end_event = torch.Event(enable_timing=True) 

73 end_event.record(torch.accelerator.current_stream()) 

74 end_event.synchronize() 

75 duration_ms = self._start_event.elapsed_time(end_event) 

76 self._start_event = None 

77 return duration_ms