Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / encoder_parallel / contracts.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"""Model-neutral contracts for visual encoder parallel execution.""" 

16 

17from __future__ import annotations 

18 

19from dataclasses import dataclass 

20from typing import Any, Callable, Mapping 

21 

22import torch # pylint: disable=forbidden-backend-import 

23from torch import nn # pylint: disable=forbidden-backend-import 

24 

25 

26@dataclass(frozen=True) 

27class EncoderWorkItem: 

28 """One image and the decoder token span that consumes its features.""" 

29 

30 item_id: int 

31 batch_index: int 

32 sequence_length: int 

33 output_rows: int 

34 decoder_token_start: int 

35 decoder_token_length: int 

36 payload_start: int 

37 payload_stop: int 

38 metadata: Mapping[str, Any] 

39 

40 

41@dataclass(frozen=True) 

42class EncoderSegment: 

43 """One padded image shard in a packed ViT call.""" 

44 

45 item: EncoderWorkItem | None 

46 local_start: int 

47 local_valid_length: int 

48 local_padded_length: int 

49 packed_offset: int 

50 

51 @property 

52 def is_dummy(self) -> bool: 

53 """A missing image identifies the synchronization-only segment.""" 

54 return self.item is None 

55 

56 

57EncoderFeature = tuple[EncoderWorkItem, torch.Tensor] 

58"""One image job paired with its aligned feature rows.""" 

59 

60 

61@dataclass(frozen=True) 

62class EncoderParallelAdapterSpec: 

63 """Model-owned visual boundaries and a runtime bridge factory. 

64 

65 ``bind(model)`` returns an object with ``input_fields``, ``head_counts``, 

66 ``validate_training``, ``parse_images``, ``encode_local`` and ``project``. 

67 The runtime uses these operations without inspecting model-specific state. 

68 Module replacement is selected separately by recipe ``replace_module``. 

69 """ 

70 

71 isolated_subtrees: tuple[str, ...] 

72 attention_pattern: str 

73 replacement_attention_type: type[nn.Module] 

74 bind: Callable[[nn.Module], Any]