Repository: volcengine/verl
File: verl/utils/profiler/mstx_profile.py (Python)

Code before:
        analysis (bool):
            Whether to enables automatic data parsing.
        role (str, optional):
            The role of the current data collection. Defaults to None.
        profile_step(str, optional):
            The current training step. Defaults to None.
    """
    if profile_level == "level_none":
        level = torch_npu.profiler.ProfilerLevel.Level_none
    elif profile_level == "level0":
        level = torch_npu.profiler.ProfilerLevel.Level0
    elif profile_level == "level1":
        level = torch_npu.profiler.ProfilerLevel.Level1
    elif profile_level == "level2":
        level = torch_npu.profiler.ProfilerLevel.Level2
    else:
        raise ValueError(f"level only supports level0, 1, 2, and level_none, but gets {profile_level}")

    if profile_step:
        profile_save_path = os.path.join(profile_save_path, profile_step)
    if role:
        profile_save_path = os.path.join(profile_save_path, role)

    # The ability to filter communication via mstx_domain_exclude requires torch_npu 2.1 or higher. Compatibility issues may arise with older versions.
    if version.parse(torch_npu.__version__) < version.parse('2.1'):
        raise RuntimeError("Torch_npu 2.1 or higher is required to use mstx_domain_exclude")

    experimental_config = torch_npu.profiler._ExperimentalConfig(
        profiler_level=level,
        export_type=torch_npu.profiler.ExportType.Db,
        data_simplification=True,
        msprof_tx=True,
        mstx_domain_exclude=["communication"]
    )

    activites = []
    if contents is None or "npu" in contents:
        activites.append(torch_npu.profiler.ProfilerActivity.NPU)
    if contents is None or "cpu" in contents:
        activites.append(torch_npu.profiler.ProfilerActivity.CPU)

    prof = torch_npu.profiler.profile(
        with_modules=contents is None or "module" in contents,
        with_stack=contents is None or "stack" in contents,
        record_shapes=contents is None or "shapes" in contents,
        profile_memory=contents is None or "memory" in contents,
        activities=activites,
        on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profile_save_path, analyse_flag=analysis),
        experimental_config=experimental_config,
    )
    return prof

Review comment (suggestion):
package name should be `torch_npu`, not `Torch_npu`

```suggestion
        raise RuntimeError("torch_npu==2.1 or higher is required to use mstx_domain_exclude")
```

Code after:
            Whether to enables automatic data parsing.
        role (str, optional):
            The role of the current data collection. Defaults to None.
        profile_step(str, optional):
            The current training step. Defaults to None.
    """
    if profile_level == "level_none":
        level = torch_npu.profiler.ProfilerLevel.Level_none
    elif profile_level == "level0":
        level = torch_npu.profiler.ProfilerLevel.Level0
    elif profile_level == "level1":
        level = torch_npu.profiler.ProfilerLevel.Level1
    elif profile_level == "level2":
        level = torch_npu.profiler.ProfilerLevel.Level2
    else:
        raise ValueError(f"level only supports level0, 1, 2, and level_none, but gets {profile_level}")

    if profile_step:
        profile_save_path = os.path.join(profile_save_path, profile_step)
    if role:
        profile_save_path = os.path.join(profile_save_path, role)

    # The ability to filter communication via mstx_domain_exclude requires torch_npu==2.1 or higher.
    if version.parse(torch_npu.__version__) < version.parse("2.1"):
        raise RuntimeError("torch_npu==2.1 or higher is required to use mstx_domain_exclude")

    experimental_config = torch_npu.profiler._ExperimentalConfig(
        profiler_level=level,
        export_type=torch_npu.profiler.ExportType.Db,
        data_simplification=True,
        msprof_tx=True,
        mstx_domain_exclude=["communication"],
    )

    activites = []
    if contents is None or "npu" in contents:
        activites.append(torch_npu.profiler.ProfilerActivity.NPU)
    if contents is None or "cpu" in contents:
        activites.append(torch_npu.profiler.ProfilerActivity.CPU)

    prof = torch_npu.profiler.profile(
        with_modules=contents is None or "module" in contents,
        with_stack=contents is None or "stack" in contents,
        record_shapes=contents is None or "shapes" in contents,
        profile_memory=contents is None or "memory" in contents,
        activities=activites,
        on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profile_save_path, analyse_flag=analysis),
        experimental_config=experimental_config,
    )
    return prof