Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / debug / _dispatch_logger.py: 28%

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"""Debug logging helpers for OpDispatcher dispatch enter/exit tracing.""" 

16# pylint: disable=R0401 

17import logging 

18from typing import Any 

19 

20import torch 

21 

22logger = logging.getLogger(__name__) 

23_Tensor = torch.Tensor 

24 

25 

26def log_dispatch_enter(op_name: str, args: tuple, kwargs: dict) -> None: 

27 """Log debug information before dispatching an op. Caller must guard with isEnabledFor.""" 

28 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=C0415 

29 dtensor_args_summary = [] 

30 for i, a in enumerate(args): 

31 if isinstance(a, DTensor): 

32 dtensor_args_summary.append( 

33 f"args[{i}]: DTensor(shape={tuple(a.shape)}, " 

34 f"placements={tuple(a.placements)}, " 

35 f"mesh_shape={tuple(a.device_mesh.shape)})" 

36 ) 

37 elif isinstance(a, _Tensor): 

38 dtensor_args_summary.append( 

39 f"args[{i}]: Tensor(shape={tuple(a.shape)}, dtype={a.dtype})" 

40 ) 

41 logger.debug( 

42 "dispatch enter: op=%s, num_args=%d, num_kwargs=%d%s", 

43 op_name, len(args), len(kwargs), 

44 (", " + ", ".join(dtensor_args_summary)) if dtensor_args_summary else "", 

45 ) 

46 

47 

48def log_dispatch_exit(op_name: str, result: Any) -> None: 

49 """Log debug information after dispatching an op. Caller must guard with isEnabledFor.""" 

50 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=C0415 

51 result_summary = "" 

52 if isinstance(result, DTensor): 

53 result_summary = ( 

54 f", result: DTensor(shape={tuple(result.shape)}, " 

55 f"placements={tuple(result.placements)}, " 

56 f"mesh_shape={tuple(result.device_mesh.shape)})" 

57 ) 

58 elif isinstance(result, _Tensor): 

59 result_summary = ( 

60 f", result: Tensor(shape={tuple(result.shape)}, " 

61 f"dtype={result.dtype})" 

62 ) 

63 elif isinstance(result, (tuple, list)): 

64 type_name = type(result).__name__ 

65 result_summary = f", result: {type_name}(len={len(result)})" 

66 logger.debug("dispatch exit: op=%s%s", op_name, result_summary)