Adding agent-index (#189)

* Adding agent-index

* review changes

* review comments addressed

* minor fix

* fix CI failure

* review comments

* Fix agent index test and address review comments

* Build Fixes

---------

Co-authored-by: Benjamin Welton <bewelton@amd.com>

tests/conversion-script/validate.py

@@ -38,7 +38,8 @@ counters_list = ["SQ_WAVES", "GRBM_GUI_ACTIVE"]
 def test_validate_counter_collection_pmc1(input_data: pd.DataFrame):
     df = input_data
     assert not df.empty
-    assert (df["Agent_Id"].astype(int).values > 0).all()
+    df_agent_id = df["Agent_Id"].str.split(" ").str[-1]
+    assert (df_agent_id.astype(int).values >= 0).all()
     assert (df["Queue_Id"].astype(int).values > 0).all()
     assert len(df["Kernel_Name"]) > 0
 

tests/rocprofv3/CMakeLists.txt

@@ -41,3 +41,4 @@ add_subdirectory(rocjpeg-trace)
 if(TARGET att_decoder_testing)
     add_subdirectory(advanced-thread-trace)
 endif()
+add_subdirectory(agent-index)

tests/rocprofv3/agent-index/CMakeLists.txt

@@ -0,0 +1,124 @@
+#
+# rocprofv3 tool test
+#
+cmake_minimum_required(VERSION 3.21.0 FATAL_ERROR)
+
+project(
+    rocprofiler-tests-rocprofv3-agent-index
+    LANGUAGES CXX
+    VERSION 0.0.0)
+
+find_package(rocprofiler-sdk REQUIRED)
+
+string(REPLACE "LD_PRELOAD=" "ROCPROF_PRELOAD=" PRELOAD_ENV
+               "${ROCPROFILER_MEMCHECK_PRELOAD_ENV}")
+
+set(tracing-env "${PRELOAD_ENV}")
+
+rocprofiler_configure_pytest_files(
+    CONFIG pytest.ini
+    COPY validate.py conftest.py input_absolute_agent_index.yaml
+         input_type_relative_agent_index.yaml input_relative_agent_index.yaml)
+if(ROCPROFILER_MEMCHECK STREQUAL "LeakSanitizer")
+    set(LOG_LEVEL "warning") # info produces memory leak
+else()
+    set(LOG_LEVEL "info")
+endif()
+
+set(execute_tests
+    rocprofv3-test-agent-index-json-absolute-execute
+    rocprofv3-test-agent-index-json-type-relative-execute
+    rocprofv3-test-agent-index-json-relative-execute)
+set(input_files
+    ${CMAKE_CURRENT_BINARY_DIR}/input_absolute_agent_index.yaml
+    ${CMAKE_CURRENT_BINARY_DIR}/input_type_relative_agent_index.yaml
+    ${CMAKE_CURRENT_BINARY_DIR}/input_relative_agent_index.yaml)
+set(output_dir
+    ${CMAKE_CURRENT_BINARY_DIR}/%argt%-agent-index-absolute
+    ${CMAKE_CURRENT_BINARY_DIR}/%argt%-agent-index-type-relative
+    ${CMAKE_CURRENT_BINARY_DIR}/%argt%-agent-index-relative)
+
+list(LENGTH execute_tests execute_test_length)
+foreach(i RANGE 0 ${execute_test_length})
+    if(${i} EQUAL ${execute_test_length})
+        break()
+    endif()
+    list(GET execute_tests ${i} test_name)
+    list(GET input_files ${i} input_file)
+    list(GET output_dir ${i} output)
+    add_test(NAME ${test_name}
+             COMMAND $<TARGET_FILE:rocprofiler-sdk::rocprofv3> -i ${input_file} -d
+                     ${output} --log-level ${LOG_LEVEL} -- $<TARGET_FILE:transpose>)
+    set_tests_properties(
+        ${test_name}
+        PROPERTIES TIMEOUT 45 LABELS "integration-tests" ENVIRONMENT "${tracing-env}"
+                   FAIL_REGULAR_EXPRESSION "${ROCPROFILER_DEFAULT_FAIL_REGEX}")
+endforeach()
+
+set(agent_index_types absolute relative type-relative)
+set(validate_tests
+    rocprofv3-test-agent-index-json-absolute-validate
+    rocprofv3-test-agent-index-json-relative-validate
+    rocprofv3-test-agent-index-json-type-relative-validate)
+set(output_dir_name transpose-agent-index-absolute transpose-agent-index-relative
+                    transpose-agent-index-type-relative)
+
+list(LENGTH validate_tests validate_test_length)
+foreach(i RANGE 0 ${validate_test_length})
+    if(${i} EQUAL ${validate_test_length})
+        break()
+    endif()
+    list(GET validate_tests ${i} test_name)
+    list(GET agent_index_types ${i} agent_index)
+    list(GET output_dir_name ${i} output)
+    add_test(
+        NAME ${test_name}
+        COMMAND
+            ${Python3_EXECUTABLE} ${CMAKE_CURRENT_BINARY_DIR}/validate.py --agent-index
+            ${agent_index} --csv-kernel-input
+            ${CMAKE_CURRENT_BINARY_DIR}/${output}/out_kernel_trace.csv
+            --csv-memory-allocation-input
+            ${CMAKE_CURRENT_BINARY_DIR}/${output}/out_memory_allocation_trace.csv
+            --csv-memory-copy-input
+            ${CMAKE_CURRENT_BINARY_DIR}/${output}/out_memory_copy_trace.csv --json-input
+            ${CMAKE_CURRENT_BINARY_DIR}/transpose-agent-index-absolute/out_results.json)
+endforeach()
+
+set(VALIDATION_FILES
+    ${CMAKE_CURRENT_BINARY_DIR}/transpose-agent-index-absolute/out_results.json
+    ${CMAKE_CURRENT_BINARY_DIR}/transpose-agent-index-absolute/out_kernel_trace.csv
+    ${CMAKE_CURRENT_BINARY_DIR}/transpose-agent-index-absolute/out_memory_allocation_trace.csv
+    ${CMAKE_CURRENT_BINARY_DIR}/transpose-agent-index-absolute/out_memory_copy_trace.csv
+    ${CMAKE_CURRENT_BINARY_DIR}/transpose-agent-index-relative/out_kernel_trace.csv
+    ${CMAKE_CURRENT_BINARY_DIR}/transpose-agent-index-relative/out_memory_allocation_trace.csv
+    ${CMAKE_CURRENT_BINARY_DIR}/transpose-agent-index-relative/out_memory_copy_trace.csv
+    ${CMAKE_CURRENT_BINARY_DIR}/transpose-agent-index-type-relative/out_kernel_trace.csv
+    ${CMAKE_CURRENT_BINARY_DIR}/transpose-agent-index-type-relative/out_memory_allocation_trace.csv
+    ${CMAKE_CURRENT_BINARY_DIR}/transpose-agent-index-type-relative/out_memory_copy_trace.csv
+    )
+
+set(dependency_tests
+    rocprofv3-test-agent-index-json-absolute-execute
+    rocprofv3-test-agent-index-json-relative-execute;rocprofv3-test-agent-index-json-absolute-execute
+    rocprofv3-test-agent-index-json-type-relative-execute;rocprofv3-test-agent-index-json-absolute-execute
+    )
+
+foreach(i RANGE 0 ${validate_test_length})
+    if(${i} EQUAL ${validate_test_length})
+        break()
+    endif()
+    list(GET validate_tests ${i} test_name)
+    list(GET dependency_tests ${i} dep)
+    set_tests_properties(
+        ${test_name}
+        PROPERTIES TIMEOUT
+                   45
+                   LABELS
+                   "integration-tests"
+                   DEPENDS
+                   ${dep}
+                   FAIL_REGULAR_EXPRESSION
+                   "AssertionError"
+                   ATTACHED_FILES_ON_FAIL
+                   "${VALIDATION_FILES}")
+endforeach()

tests/rocprofv3/agent-index/conftest.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+
+# MIT License
+#
+# Copyright (c) 2024-2025 Advanced Micro Devices, Inc. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+import csv
+import pytest
+import json
+
+from rocprofiler_sdk.pytest_utils.dotdict import dotdict
+from rocprofiler_sdk.pytest_utils import collapse_dict_list
+from rocprofiler_sdk.pytest_utils.perfetto_reader import PerfettoReader
+
+
+def pytest_addoption(parser):
+
+    parser.addoption(
+        "--agent-index",
+        choices=("absolute", "relative", "type-relative"),
+        help="...",
+    )
+    parser.addoption(
+        "--csv-kernel-input",
+        action="store",
+        help="Path to kernel tracing CSV file",
+    )
+    parser.addoption(
+        "--csv-memory-allocation-input",
+        action="store",
+        help="Path to memory allocation tracing CSV file",
+    )
+    parser.addoption(
+        "--csv-memory-copy-input",
+        action="store",
+        help="Path to memory allocation tracing CSV file",
+    )
+    parser.addoption(
+        "--json-input",
+        action="store",
+        help="Path to JSON file.",
+    )
+
+
+def read_csv(filename):
+    data = []
+    with open(filename, "r") as inp:
+        reader = csv.DictReader(inp)
+        for row in reader:
+            data.append(row)
+    return data
+
+
+@pytest.fixture
+def csv_kernel_input(request):
+    filename = request.config.getoption("--csv-kernel-input")
+    return read_csv(filename)
+
+
+@pytest.fixture
+def csv_memory_copy_input(request):
+    filename = request.config.getoption("--csv-memory-copy-input")
+    return read_csv(filename)
+
+
+@pytest.fixture
+def csv_memory_allocation_input(request):
+    filename = request.config.getoption("--csv-memory-allocation-input")
+    return read_csv(filename)
+
+
+@pytest.fixture
+def json_data(request):
+    filename = request.config.getoption("--json-input")
+    with open(filename, "r") as inp:
+        return dotdict(collapse_dict_list(json.load(inp)))
+
+
+@pytest.fixture
+def agent_index(request):
+    return request.config.getoption("--agent-index")

tests/rocprofv3/agent-index/input_absolute_agent_index.yaml

@@ -0,0 +1,7 @@
+jobs:
+  - output_file: "out"
+    output_format: [csv, json]
+    sys_trace: True
+    agent_index: absolute
+
+    

tests/rocprofv3/agent-index/input_relative_agent_index.yaml

@@ -0,0 +1,6 @@
+jobs:
+  - output_file: "out"
+    output_format: [csv]
+    sys_trace: True
+    agent_index: relative
+

tests/rocprofv3/agent-index/input_type_relative_agent_index.yaml

@@ -0,0 +1,6 @@
+jobs:
+  - output_file: "out"
+    output_format: [csv]
+    sys_trace: True
+    agent_index: type-relative
+

tests/rocprofv3/agent-index/pytest.ini

@@ -0,0 +1,5 @@
+
+[pytest]
+addopts = --durations=20 -rA -s -vv
+testpaths = validate.py
+pythonpath = @ROCPROFILER_SDK_TESTS_BINARY_DIR@/pytest-packages

tests/rocprofv3/agent-index/validate.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+
+# MIT License
+#
+# Copyright (c) 2024-2025 Advanced Micro Devices, Inc. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+import sys
+import pytest
+import re
+
+
+def get_type(id_type):
+    if id_type == "absolute":
+        return "node_id"
+    elif id_type == "type-relative":
+        return "logical_node_type_id"
+    else:
+        return "logical_node_id"
+
+
+def check_kernel_traces(json_data, kernels_data, id_type):
+    data = json_data["rocprofiler-sdk-tool"]
+    agents = data["agents"]
+    kernels = data["buffer_records"]["kernel_dispatch"]
+    _type = get_type(id_type)
+
+    for row in kernels_data:
+        found = any(
+            int(agent[_type]) == int(row["Agent_Id"].split(" ")[-1]) for agent in agents
+        )
+
+        assert found, f"Agent ID {row['Agent_Id']} of Type {id_type} not found"
+
+
+def check_memory_allocation_traces(json_data, mem_alloc, id_type):
+    data = json_data["rocprofiler-sdk-tool"]
+    agents = data["agents"]
+    mem_alloc_record = data["buffer_records"]["memory_allocation"]
+    _type = get_type(id_type)
+
+    for row in mem_alloc:
+
+        found = row["Agent_Id"] != "" and any(
+            int(agent[_type]) == int(row["Agent_Id"].split(" ")[-1]) for agent in agents
+        )
+
+        if row["Agent_Id"] != "":
+            assert found, f"Agent ID {row['Agent_Id']} of Type {id_type} not found"
+
+
+def check_memory_copy_traces(json_data, mem_copy, id_type):
+    data = json_data["rocprofiler-sdk-tool"]
+    agents = data["agents"]
+    mem_copy_record = data["buffer_records"]["memory_copy"]
+    _type = get_type(id_type)
+
+    for row in mem_copy:
+
+        found = any(
+            int(agent[_type]) == int(row["Source_Agent_Id"].split(" ")[-1])
+            for agent in agents
+        )
+
+        assert (
+            found
+        ), f"Src Agent ID {row['Destination_Agent_Id']} of Type {id_type} not found"
+
+        found = any(
+            int(agent[_type]) == int(row["Destination_Agent_Id"].split(" ")[-1])
+            for agent in agents
+        )
+
+        assert (
+            found
+        ), f"Destination Agent ID {row['Destiantion_Agent_Id']} of Type {id_type} not found"
+
+        if (
+            _type == "logical_node_type_id"
+            and row["Direction"] == "MEMORY_COPY_HOST_TO_DEVICE"
+        ):
+            assert row["Source_Agent_Id"].split(" ")[0] == "CPU"
+            assert row["Destination_Agent_Id"].split(" ")[0] == "GPU"
+
+        elif (
+            _type == "logical_node_type_id"
+            and row["Direction"] == "MEMORY_COPY_DEVICE_TO_HOST"
+        ):
+            assert row["Source_Agent_Id"].split(" ")[0] == "GPU"
+            assert row["Destination_Agent_Id"].split(" ")[0] == "CPU"
+
+
+def test_validate(
+    csv_kernel_input,
+    json_data,
+    agent_index,
+    csv_memory_allocation_input,
+    csv_memory_copy_input,
+):
+    check_kernel_traces(json_data, csv_kernel_input, agent_index)
+    check_memory_copy_traces(json_data, csv_memory_copy_input, agent_index)
+    check_memory_allocation_traces(json_data, csv_memory_allocation_input, agent_index)
+
+
+if __name__ == "__main__":
+    exit_code = pytest.main(["-x", __file__] + sys.argv[1:])
+    sys.exit(exit_code)

tests/rocprofv3/counter-collection/extra_counters/validate.py

@@ -41,7 +41,8 @@ def test_validate_counter_collection_pmc1_extra_counters(input_data: pd.DataFram
     df = input_data
 
     assert not df.empty
-    assert (df["Agent_Id"].astype(int).values > 0).all()
+    df_agent_id = df["Agent_Id"].str.split(" ").str[-1]
+    assert (df_agent_id.astype(int).values >= 0).all()
     assert (df["Queue_Id"].astype(int).values > 0).all()
     assert (df["Process_Id"].astype(int).values > 0).all()
     assert len(df["Kernel_Name"]) > 0

tests/rocprofv3/counter-collection/input1/validate.py

@@ -41,7 +41,8 @@ def test_validate_counter_collection_pmc1(input_data: pd.DataFrame):
     df = input_data
 
     assert not df.empty
-    assert (df["Agent_Id"].astype(int).values > 0).all()
+    df_agent_id = df["Agent_Id"].str.split(" ").str[-1]
+    assert (df_agent_id.astype(int).values >= 0).all()
     assert (df["Queue_Id"].astype(int).values > 0).all()
     assert (df["Process_Id"].astype(int).values > 0).all()
     assert len(df["Kernel_Name"]) > 0

tests/rocprofv3/counter-collection/input2/validate.py

@@ -49,7 +49,7 @@ def test_validate_counter_collection_pmc2(counter_input_data):
     di_list = []
 
     for row in counter_input_data:
-        assert int(row["Agent_Id"]) > 0
+        assert int(row["Agent_Id"].split(" ")[-1]) >= 0
         assert int(row["Queue_Id"]) > 0
         assert int(row["Process_Id"]) > 0
         assert len(row["Kernel_Name"]) > 0

tests/rocprofv3/counter-collection/input3/validate.py

@@ -51,7 +51,7 @@ def test_validate_counter_collection_yml_pmc(counter_input_data):
     di_list = []
 
     for row in counter_input_data:
-        assert int(row["Agent_Id"]) > 0
+        assert int(row["Agent_Id"].split(" ")[-1]) >= 0
         assert int(row["Queue_Id"]) > 0
         assert int(row["Process_Id"]) > 0
         assert len(row["Kernel_Name"]) > 0

tests/rocprofv3/counter-collection/kernel_filtering/validate.py

@@ -36,7 +36,8 @@ def unique(lst):
 def validate_csv(df, kernel_list, counter_name):
 
     assert not df.empty
-    assert (df["Agent_Id"].astype(int).values > 0).all()
+    df_Agent_id = df["Agent_Id"].str.split(" ").str[-1]
+    assert (df_Agent_id.astype(int).values >= 0).all()
     assert (df["Queue_Id"].astype(int).values > 0).all()
     assert (df["Process_Id"].astype(int).values > 0).all()
     assert len(df["Kernel_Name"]) > 0

tests/rocprofv3/hsa-queue-dependency/validate.py

@@ -73,7 +73,7 @@ def test_kernel_trace(kernel_trace_input_data):
     assert len(kernel_trace_input_data) == 3
     for row in kernel_trace_input_data:
         assert row["Kind"] == "KERNEL_DISPATCH"
-        assert int(row["Agent_Id"]) > 0
+        assert int(row["Agent_Id"].split(" ")[-1]) >= 0
         assert int(row["Queue_Id"]) > 0
         assert int(row["Kernel_Id"]) > 0
         assert row["Kernel_Name"] in valid_kernel_names

tests/rocprofv3/pc-sampling/host-trap/transpose-multiple-agents/CMakeLists.txt

@@ -23,7 +23,7 @@ add_test(
     NAME rocprofv3-test-pc-sampling-host-trap-transpose-multiple-agents-input-cmd-execute
     COMMAND
         $<TARGET_FILE:rocprofiler-sdk::rocprofv3> --kernel-trace --pc-sampling-unit time
-        --pc-sampling-method host_trap --pc-sampling-interval 1 -d
+        --pc-sampling-method host_trap --pc-sampling-interval 1 --agent-index absolute -d
         ${CMAKE_CURRENT_BINARY_DIR}/pc_sampling_cmd_input -o out --output-format csv --
         $<TARGET_FILE:transpose> ${TRANSPOSE_NUM_THREADS} ${TRANSPOSE_NUM_ITERATIONS})
 

tests/rocprofv3/pc-sampling/host-trap/transpose-multiple-agents/input.json

@@ -8,7 +8,8 @@
             "output_file": "out",
             "output_format": [
                 "csv"
-            ]
+            ],
+            "agent_index": "type-relative" 
         }
     ]
 }

tests/rocprofv3/pc-sampling/host-trap/transpose-multiple-agents/input.yml

@@ -3,3 +3,4 @@ jobs:
     pc_sampling_unit: "time"
     pc_sampling_method: "host_trap"
     pc_sampling_interval: 1
+    agent-index: relative

tests/rocprofv3/pc-sampling/host-trap/transpose-multiple-agents/validate.py

@@ -49,7 +49,11 @@ def test_multi_agent_support(
     # Determine the agent on which sample was generated
     samples_df["Agent_Id"] = (
         samples_df["Dispatch_Id"]
-        .map(input_kernel_trace_csv.set_index("Dispatch_Id")["Agent_Id"])
+        .map(
+            input_kernel_trace_csv.set_index("Dispatch_Id")["Agent_Id"]
+            .str.split(" ")
+            .str[1]
+        )
         .astype(np.uint64)
     )
     sampled_agents = samples_df["Agent_Id"].unique()

tests/rocprofv3/scratch-memory/validate.py

@@ -127,7 +127,7 @@ def test_scratch_memory(json_input_data, csv_input_data):
 
         assert row["Kind"] == "SCRATCH_MEMORY"
         assert row["Operation"] in ["SCRATCH_MEMORY_ALLOC", "SCRATCH_MEMORY_FREE"]
-        assert int(row["Agent_Id"]) > 0
+        assert int(row["Agent_Id"].split(" ")[-1]) >= 0
         assert int(row["Queue_Id"]) > 0
         assert int(row["Thread_Id"]) > 0
         assert int(row["Start_Timestamp"]) > 0

tests/rocprofv3/tracing-hip-in-libraries/validate.py

@@ -300,7 +300,7 @@ def test_kernel_trace(kernel_input_data, kernel_stats_data):
             continue
 
         assert row["Kind"] == "KERNEL_DISPATCH"
-        assert int(row["Agent_Id"]) > 0
+        assert int(row["Agent_Id"].split(" ")[-1]) >= 0
         assert int(row["Queue_Id"]) > 0
         assert int(row["Kernel_Id"]) > 0
         assert int(row["Correlation_Id"]) > 0
@@ -419,8 +419,8 @@ def test_memory_copy_trace(
             "MEMORY_COPY_DEVICE_TO_HOST",
         )
 
-        src_agent = get_agent(row["Source_Agent_Id"])
-        dst_agent = get_agent(row["Destination_Agent_Id"])
+        src_agent = get_agent(row["Source_Agent_Id"].split(" ")[-1])
+        dst_agent = get_agent(row["Destination_Agent_Id"].split(" ")[-1])
         assert src_agent is not None and dst_agent is not None, f"{agent_info_input_data}"
 
         if row["Direction"] == "MEMORY_COPY_HOST_TO_DEVICE":

tests/rocprofv3/tracing/validate.py

@@ -129,7 +129,7 @@ def test_kernel_trace(kernel_input_data):
     assert len(kernel_input_data) == 1
     for row in kernel_input_data:
         assert row["Kind"] == "KERNEL_DISPATCH"
-        assert int(row["Agent_Id"]) > 0
+        assert int(row["Agent_Id"].split(" ")[-1]) >= 0
         assert int(row["Queue_Id"]) > 0
         assert int(row["Kernel_Id"]) > 0
         assert row["Kernel_Name"] in valid_kernel_names
@@ -223,8 +223,8 @@ def test_memory_copy_trace(agent_info_input_data, memory_copy_input_data):
         assert direction in ("MEMORY_COPY_HOST_TO_DEVICE", "MEMORY_COPY_DEVICE_TO_HOST")
         row = memory_copy_input_data[idx]
         assert row["Direction"] == direction
-        src_agent = get_agent(row["Source_Agent_Id"])
-        dst_agent = get_agent(row["Destination_Agent_Id"])
+        src_agent = get_agent(row["Source_Agent_Id"].split(" ")[-1])
+        dst_agent = get_agent(row["Destination_Agent_Id"].split(" ")[-1])
         assert src_agent is not None and dst_agent is not None, f"{agent_info_input_data}"
         if direction == "MEMORY_COPY_HOST_TO_DEVICE":
             assert src_agent["Agent_Type"] == "CPU"