I have written a simple Python script to load sync dependencies from a Kineto trace and then print out event IDs and external event IDs. When I print out event IDs, I don't find any self-dependencies. However, when I print out external event IDs of the start event and end event, I see self-dependencies. Is this expected behavior?
import argparse
import logging
import os
from typing import Dict, List
from hta.analyzers.critical_path_analysis import CPEdgeType
from hta.trace_analysis import TraceAnalysis
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
def load_sync_dependencies(
rank: int, kineto_file: str, annotation: str = "ProfilerStep", instance_id: int = 0
) -> Dict[int, List[int]]:
"""
Load synchronization dependencies using Holistic Trace Analysis (HTA).
Args:
rank (int): Rank for the input Kineto trace.
kineto_file (str): Path to the Kineto trace file.
annotation (str): Annotation to use for the analysis. Defaults to "ProfilerStep".
instance_id (int): Instance ID for the analysis. Defaults to 0.
Returns:
Dict[int, List[int]]: A dictionary mapping end event's external ID to a list of start event's external IDs
that have synchronization dependencies.
"""
sync_dependencies = {}
trace_analysis = TraceAnalysis(trace_dir=os.path.dirname(kineto_file))
cp_graph, success = trace_analysis.critical_path_analysis(rank=rank, annotation=annotation, instance_id=instance_id)
if not success:
logger.error("Failed to load Critical Path Graph")
return sync_dependencies
raw_events = trace_analysis.t.get_raw_trace_for_one_rank(rank=rank)["traceEvents"]
for edge in cp_graph.critical_path_edges_set:
if edge.type in [CPEdgeType.SYNC_DEPENDENCY]:
start_event_id, end_event_id = cp_graph.get_events_for_edge(edge)
start_event, end_event = raw_events[start_event_id], raw_events[end_event_id]
if "External id" in end_event["args"] and "External id" in start_event["args"]:
start_cat = start_event["cat"]
end_event_external_id = end_event["args"]["External id"]
start_event_external_id = start_event["args"]["External id"]
start_name = start_event["name"]
end_name = end_event["name"]
print(
f"THEO: start_event_id {start_event_id}, end_event_id {end_event_id}, "
f"start_event_external_id {start_event_external_id}, end_event_external_id {end_event_external_id}"
)
else:
logger.warning(
f"Synchronization dependency from event {start_event_id} to event {end_event_id} will "
"not be considered due to missing external IDs."
)
return sync_dependencies
def main() -> None:
"""
Main function to parse arguments and load synchronization dependencies.
"""
parser = argparse.ArgumentParser(description="Load and print critical paths from Kineto traces.")
parser.add_argument("--input", type=str, help="Path to the Kineto trace file.")
parser.add_argument("--rank", type=int, help="Rank for the input traces.")
args = parser.parse_args()
load_sync_dependencies(args.rank, args.input)
if __name__ == "__main__":
main()
What is your question?
I have written a simple Python script to load sync dependencies from a Kineto trace and then print out event IDs and external event IDs. When I print out event IDs, I don't find any self-dependencies. However, when I print out external event IDs of the start event and end event, I see self-dependencies. Is this expected behavior?
How to reproduce
sync_dep.py
kineto_0.json
/tmp/out
Environment