Diff Coverage

                        continue
                    with open(spec.origin, "r", encoding="utf-8") as mf:
                        source = mf.read()
                        tree = ast.parse(source)
                        mod_cls = None
                        for node in ast.walk(tree):
                            if isinstance(node, ast.ClassDef) and node.name == hook_cls:
                                mod_cls = node
                                break
                        if mod_cls:
                            target_mod_path = mod_path
                            break
            if target_mod_path:

            result = self._fetch_node_and_switch_env(
                stages, record_lay_types, stage_id, 0, lay_id
            )
            if result is None:
                raise RuntimeError("_fetch_node_and_switch_env returned None.")
            fwd_first += sum(self._inner_dynamic_mem(default_micro_factor=1))
        for lay_id, lay in enumerate(stages[stage_id][1]):
            self._ctx.current_node = LayerType.NOT_REC_LAYER
            result = self._fetch_node_and_switch_env(

            result = self._fetch_node_and_switch_env(
                stages, record_lay_types, stage_id, 1, lay_id
            )
            if result is None:
                raise RuntimeError("_fetch_node_and_switch_env returned None.")
            if lay == LayerType.FULL_REC_LAYER:
                bwd_last = max(
                    bwd_last,
                    sum(self._inner_dynamic_mem(default_micro_factor=1)),

            result = self._fetch_node_and_switch_env(
                stages, record_lay_types, stage_id, 1, lay_id
            )
            if result is None:
                raise RuntimeError("_fetch_node_and_switch_env returned None.")
            fwd_last = sum(self._inner_dynamic_mem(default_micro_factor=1))
        for lay_id, lay in enumerate(stages[stage_id][0]):
            self._ctx.current_node = LayerType.NOT_REC_LAYER
            result = self._fetch_node_and_switch_env(

            result = self._fetch_node_and_switch_env(
                stages, record_lay_types, stage_id, 0, lay_id
            )
            if result is None:
                raise RuntimeError("_fetch_node_and_switch_env returned None.")
            if lay == LayerType.FULL_REC_LAYER:
                bwd_first = max(
                    bwd_first,
                    sum(self._inner_dynamic_mem(default_micro_factor=1)),

                        node = self.fetch_node_from_lineno(frame.f_lineno, co)
                        if not node:
                            return self.line_tracer
                        line = ast.unparse(node)
                    sign = next(op for op in op_equals if op in line)
                    left, right = line.split(sign, 1)
                    left, right = left.strip(), right.strip()
                    if sign != "=":
                        right = f"{left} {sign[0]} {right}"

# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Walkthrough of the EvaluatorV2 memory estimation API, see README.md."""
from typing import Any

from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.estimate_v2 import EvaluatorV2
from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.hooks.template import Template
from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger


def my_attn_num_param(ccfg: Any, ctx: Any) -> float:
    """Attention parameter count that overrides the default formula."""
    del ctx
    return 10 * ccfg.h * ccfg.h


def custom(ccfg: Any) -> None:
    """Cost model variables that override the parsed configuration."""
    ccfg.bytes_compute = 1
    ccfg.s = 1024
    ccfg.n_attMM = 5

    ccfg.s = 1024
    ccfg.n_attMM = 5


def main() -> None:
    """Estimate, inspect and customize the memory of the bundled test cases."""
    # Instantiate evaluator with a model configuration,
    #  log_level=0 removes warning messages
    e = EvaluatorV2("./test_cases/mixtral/default.yaml", log_level=0)

    # Check all defined node type
    logger.output("%s", list(LayerType))

    # Estimate peak memory (in Megabytes)
    peak_mem = e.estimate_peak(verbose=True)
    # Check whether estimation fits in device's max memory
    e.mem_fit(peak_mem)

    # Estimate static memory of a specific pipeline stage (in Megabytes)
    logger.output("%s", e.static_mem_stage(1))
    # Estimate dynamic memory of a specific pipeline stage (in Megabytes)
    logger.output("%s", e.dynamic_mem_stage(1))
    # Estimate static memory of a specific layer and stage (in Megabytes)
    logger.output("%s", e.static_mem_layer(LayerType.FULL_REC_LAYER, 1))
    # Estimate dynamic memory of a specific layer and stage (in Megabytes)
    logger.output("%s", e.dynamic_mem_layer(LayerType.FULL_REC_LAYER, 1))
    # Retrieve the memory estimation logs of a specific stage (in Megabytes)
    logger.output("%s", e.logs_mem_stage(1))
    # Fetch memory insights from each pipeline stage
    logger.output("%s", e.estimate_peak_insight())
    # PPB Input
    logger.output("%s", e.estimate_layer_memory())

    # Inspect a specific stage (here is the first one)
    e.estimate_peak(spec_stage_id=0, verbose=True)

    # Plot
    e.estimate_peak(plot=True)

    e = EvaluatorV2("./test_cases/deepseek3/default.yaml", log_level=0)

    # Overwriting context function
    e.set_attn_eval_fun(num_p=my_attn_num_param)

    # Overwriting a training feature
    e.set_passes(swap_os=True)

    # Overwriting cost model variables
    e.set_ccfg(custom)

    # Overwriting strategy
    logger.output("%s", e.get_strategy())
    e.set_strategy(dp=8, tp=8, m=128)
    logger.output("%s", e.get_strategy())

    e.estimate_peak(verbose=True)
    # Inspect ccfg object (cost model variables)
    e.print_ccfg()
    # Inspect ctx object (evaluation variables and functions)
    e.print_ctx()

    # Load a hook class
    # ... when declaring an Evaluator
    e = EvaluatorV2(
        "./test_cases/deepseek3/default.yaml", log_level=0, hook_cls=Template()
    )
    # ... by using load_hook_cls()
    e.load_hook_cls(Template())


if __name__ == "__main__":
    main()

            chunk_lay_occ = []
            if chunk:
                layer, count = chunk[0], 1
                for lay_id in range(1, len(chunk)):
                    if chunk[lay_id] == layer:
                        count += 1
                    else:
                        chunk_lay_occ += [f"{count}{layer.name[0]}"]
                        layer, count = chunk[lay_id], 1
                chunk_lay_occ += [f"{count}{layer.name[0]}"]
            stage_layers += [chunk_lay_occ]
        logger.info("stage _%s : %s", stage_id, stage_layers)

            self.ccfg.specs = specs
            self.__parse_toml()
        else:
            now = time.time()
            home_path = os.path.expanduser("~")
            if home_path not in sys.path:
                sys.path.append(home_path)
            spec_torch = importlib.util.find_spec("torch")
            spec_torchtitan = importlib.util.find_spec("torchtitan")
            if spec_torch is not None and spec_torchtitan is not None:
                # existing torchtitan package

            spec_torchtitan = importlib.util.find_spec("torchtitan")
            if spec_torch is not None and spec_torchtitan is not None:
                # existing torchtitan package
                spec_path = spec_torchtitan.submodule_search_locations[0]
                if spec_path not in sys.path:
                    sys.path.append(spec_path)
                try:
                    logger.info(
                        "found torchtitan package from homedir: %s",
                        spec_path,

        self.name = name
        self.level_bound_number = bounds
        self.level_bandwidth = bandwidths
        if len(bounds) != len(bandwidths):
            raise ValueError("bounds and bandwidths must have the same length")
        self.levels = len(bounds)

    def __str__(self):
        return self.name

            str(self.config.dimensions),
        )
        if self.enable_debug:
            file_path = os.path.dirname(os.path.realpath(__file__))
            output_path = os.path.join(file_path, "output")
            if scored_space:
                Debug.plot_nd(
                    scored_space,
                    output_path,

    non_steady_perf = 0
    steady_perf = 0
    if cfg.p == 1:
        if len(stage_perfs) != 1:
            raise ValueError("Expected exactly one stage performance")
        steady_perf = sum_time * cfg.m
    elif cfg.vp == 1:
        non_steady_perf = sum_time
        if GENERALIZE_PIPELINE_CALCULATION:

        return None

    def _current_layer_sum(self, variables, layer, interleave, stage_range):
        """Sum current interleave variables over a stage range."""
        considered = Recompute.get_used_list(self.recompute_considered_)
        return lpSolver.lpSum(
            variables[layer][rec][interleave][stage]
            for rec in considered for stage in stage_range
        )

        )

    def _previous_layer_sum(self, variables, layer, interleave):
        """Sum variables from previous interleaves."""
        considered = Recompute.get_used_list(self.recompute_considered_)
        previous_variables = []
        for rec in considered:
            for prev_interleave in range(interleave):
                previous_variables.extend(
                    variables[layer][rec][prev_interleave][stage] for stage in range(self.num_of_stage_)
                )
        return lpSolver.lpSum(previous_variables)

    def _add_frontier_upper_bounds(self, prob, variables, sorted_layers):
        """Prevent previous body layer types after each multimodal frontier."""
        for frontier in range(1, len(sorted_layers[Layer.type_enum.BODY])):

        block_index = self._get_block_indices(blocks, mode=mode, equal_wide=equal_wide)
        width = max(np.max(block_index[p]) for p in range(pp)) if blocks[0][-1].end is None \
            else max(blocks[p][-1].end for p in range(pp))
        if self.draw_block(block_index, blocks, ax_index, equal_wide, width, phase=phase) is not self:
            raise ValueError("Drawing blocks must return the current plot manager")
        if comm:
            if self.draw_comm(block_index, blocks, ax_index, equal_wide, mode) is not self:
                raise ValueError("Drawing communications must return the current plot manager")
        if connect:
            if self.draw_connect(block_index, blocks, ax_index, equal_wide, mode) is not self:
                raise ValueError("Drawing connections must return the current plot manager")
        self._set_block_ax(self.ax[ax_index], pp)
        self.ax[ax_index].set_xlim(0, width)
        self.ax[ax_index].set_xticks(np.linspace(0, width, 8))
        return self

                  ax_index: int = 0, comm: bool = False, connect: bool = False,
                  equal_wide: bool = False) -> "PlotMgr":
        """Highlight a dependency loop (non-comm) with red arrows and a textual trace."""
        if self.draw(blocks, ax_index, comm, connect, equal_wide, phase=True) is not self:
            raise ValueError("Drawing a dependency loop must return the current plot manager")
        block_index = self._get_block_indices(blocks, equal_wide=equal_wide)
        msg = 'dependency loop: '
        for b in range(len(loop) - 1):
            p = loop[b].stage

    def draw_comm_loop(self, lines: List[List[BlockSim]], loop: List[BlockSim],
                       ax_index: int = 0) -> "PlotMgr":
        """Highlight a dependency loop in the send-receive graph."""
        if self.draw(lines, ax_index, True, True, True, 'joint', phase=True) is not self:
            raise ValueError("Drawing a communication loop must return the current plot manager")
        block_index = self._get_block_indices(lines, mode='joint', equal_wide=True)
        msg = 'dependency loop: '
        for b in range(len(loop) - 1):
            p = loop[b].stage

            if p % 2 == 0 and lines[p][i_b + 1].type == 'r' and lines[p][i_b + 2].type == 's':
                lines[p][i_b + 1], lines[p][i_b + 2] = lines[p][i_b + 2], lines[p][i_b + 1]
            if p % 2 == 1 and lines[p][i_b + 1].type == 's' and lines[p][i_b + 2].type == 'r':
                if block.phase == 'warmup' and self.blocks[p][b + 1].phase == 'cooldown':
                    return
                lines[p][i_b + 1], lines[p][i_b + 2] = lines[p][i_b + 2], lines[p][i_b + 1]
            if lines[p][i_b + 1].dual.stage == lines[p][i_b + 2].dual.stage:
                pd = lines[p][i_b + 1].dual.stage
                j_b1 = lines[pd].index(lines[p][i_b + 1].dual)