Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / encoder_parallel / batch_input.py: 96%

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"""Compose Encoder Parallel planning with ordinary forward-input providers.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Mapping 

20from typing import Any 

21 

22from hyper_parallel.data.batching.runtime_input import RuntimeInputAdapter, RuntimeInputContext 

23from hyper_parallel.distributed.encoder_parallel.runtime import EncoderParallelRuntime 

24 

25 

26class EncoderParallelInputAdapter(RuntimeInputAdapter): 

27 """Add one microbatch plan after local validation and other runtime inputs.""" 

28 

29 def __init__(self, *, runtime: EncoderParallelRuntime, 

30 inner: RuntimeInputAdapter | None) -> None: 

31 """Combine visual planning with the batch's existing input provider.""" 

32 if inner is not None and "encoder_parallel_plan" in inner.runtime_input_fields(): 

33 raise ValueError("encoder_parallel_plan has two owners") 

34 self.runtime = runtime 

35 self.inner = inner 

36 

37 def runtime_input_fields(self) -> tuple[str, ...]: 

38 """Declare the existing fields and the plan field as one input contract.""" 

39 fields = () if self.inner is None else self.inner.runtime_input_fields() 

40 return (*fields, "encoder_parallel_plan") 

41 

42 def build_runtime_inputs( 

43 self, *, batch: Mapping[str, Any], context: RuntimeInputContext, 

44 ) -> Mapping[str, Any]: 

45 """Validate locally, build other inputs, then agree on the visual plan.""" 

46 prepared = self.runtime.preflight_batch(batch) 

47 generated = {} if self.inner is None else self.inner.build_runtime_inputs(batch=batch, context=context) 

48 if not isinstance(generated, Mapping): 

49 raise TypeError("RuntimeInputAdapter.build_runtime_inputs must return a mapping") 

50 inputs = dict(generated) 

51 if "encoder_parallel_plan" in batch or "encoder_parallel_plan" in inputs or set(batch).intersection(inputs): 

52 raise ValueError("runtime inputs cannot replace Omni model inputs") 

53 inputs["encoder_parallel_plan"] = self.runtime.finalize_plan(*prepared) 

54 return inputs