【合入来源】

如有社区issue，请关联issue链接
请勿携带内部流程信息（需求链接、问题单、内部issue等）

• 需求
• 问题单
• issue/工单
  https://gitcode.com/Ascend/pytorch/issues/1802
• 重构优化
• 资料更新

【修改方案】

对齐pytorch profiler的部分接口：
一个是增加一个skip_first_wait：

一个是custom_trace_id_callback，然后增加两个对应的set和get的接口set_custom_trace_id_callback和get_trace_id

【资料变更】

请确认是否涉及资料变更。如涉及，需要在PR中体现，并简要说明修改内容。如不涉及，需填写“不涉及”

涉及，需要和资料沟通

【接口变更】

请确认是否涉及跨代码仓或者客户面可见的接口变更。如涉及，需要详细说明接口以及对应的变更内容，同时需要在资料中体现。如不涉及，需填写“不涉及”

涉及，前面已经说明

【功能验证】

场景1
设置schedule = torch_npu.profiler.schedule(wait=2, warmup=1, active=1, repeat=2, skip_first=0, skip_first_wait=1)
之前：采集第3和第7个step

现在：采集第1和第5个step

场景2
测试get_trace_id接口：

……
    with torch_npu.profiler.profile(
        activities=[
            torch_npu.profiler.ProfilerActivity.NPU,
            torch_npu.profiler.ProfilerActivity.CPU,
        ],
        with_stack=True,
        schedule = torch_npu.profiler.schedule(wait=2, warmup=1, active=2, repeat=2, skip_first=1, skip_first_wait=1),
        on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./result_profiling"),
        experimental_config=experimental_config) as prof:

        for i in range(12):
            add(x0, x1)
            prof.step()
            print(f"step {i}: {prof.get_trace_id()}")
……

默认是一个uuid，是直接从pytorch里面拷贝过来的，现在会在profiler_metadata.json里面落盘，db里面也有

场景3
异常的skip_first_wait参数不生效，reset为0
正常：
异常：必须设置为整数，否则有警告信息，reset为0

场景4：
设置custom_trace_id_callback，这个trace_id，我们是想和每一份ascend_pt数据或者repeat参数绑定的，

class RepeatTraceIdGenerator:
    def __init__(self):
        self.repeat_count = 0  # 从 0 开始

    def __call__(self) -> str:
        # 每一轮 profile 启动，计数 +1
        current_id = str(self.repeat_count)
        self.repeat_count += 1
        return current_id

trace_id_gen = RepeatTraceIdGenerator()
……
    with torch_npu.profiler.profile(
        activities=[
            torch_npu.profiler.ProfilerActivity.NPU,
            torch_npu.profiler.ProfilerActivity.CPU,
        ],
        with_stack=True,
        schedule = torch_npu.profiler.schedule(wait=2, warmup=1, active=2, repeat=2, skip_first=1, skip_first_wait=1),
        on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./result_profiling"),
        experimental_config=experimental_config,
        custom_trace_id_callback=trace_id_gen) as prof:

        for i in range(12):
            add(x0, x1)
            prof.step()
            print(f"step {i}: {prof.get_trace_id()}")

如果call_back类型不对，会有警告信息，然后使用默认的uuid

【CheckList】

PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]

• 代码注释完备，正确记录错误日志
• 代码实现进行了返回值、空指针等校验
• PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
• PR持续集成流水线（CI）执行通过，代码检查无异常