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.""" 

16import logging 

17from typing import Any 

18 

19from hyper_parallel.platform import get_platform 

20 

21logger = logging.getLogger(__name__) 

22_Tensor = get_platform().Tensor 

23 

24 

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

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

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

28 dtensor_args_summary = [] 

29 for i, a in enumerate(args): 

30 if isinstance(a, DTensor): 

31 dtensor_args_summary.append( 

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

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

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

35 ) 

36 elif isinstance(a, _Tensor): 

37 dtensor_args_summary.append( 

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

39 ) 

40 logger.debug( 

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

42 op_name, len(args), len(kwargs), 

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

44 ) 

45 

46 

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

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

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

50 result_summary = "" 

51 if isinstance(result, DTensor): 

52 result_summary = ( 

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

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

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

56 ) 

57 elif isinstance(result, _Tensor): 

58 result_summary = ( 

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

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

61 ) 

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

63 type_name = type(result).__name__ 

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

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