Refactor chakra_trace_link and Add Tests

[TaekyungHeo]

Summary

Refactor chakra_trace_link and add tests

Test Plan

1. Pytest

$ python -m pytest -vv tests
========================================================= test session starts =========================================================
platform darwin -- Python 3.10.2, pytest-8.1.1, pluggy-1.5.0 -- /Users/theo/venv/bin/python
cachedir: .pytest_cache
rootdir: /Users/theo/chakra
configfile: pyproject.toml
plugins: typeguard-2.13.3, mock-3.14.0, anyio-4.2.0
collected 15 items                                                                                                                    

tests/converter/test_pytorch_node.py::test_pytorch_node_parsing PASSED                                                          [  6%]
tests/converter/test_pytorch_tensor.py::test_pytorch_tensor_initialization PASSED                                               [ 13%]
tests/converter/test_pytorch_tensor.py::test_pytorch_tensor_is_valid PASSED                                                     [ 20%]
tests/converter/test_pytorch_tensor.py::test_pytorch_tensor_properties PASSED                                                   [ 26%]
tests/converter/test_pytorch_tensor.py::test_pytorch_tensor_has_valid_storage_id PASSED                                         [ 33%]
tests/converter/test_pytorch_tensor.py::test_list_to_pytorch_tensor PASSED                                                      [ 40%]
tests/trace_link/test_kineto_operator.py::test_init_kineto_operator PASSED                                                      [ 46%]
tests/trace_link/test_kineto_operator.py::test_repr_method PASSED                                                               [ 53%]
tests/trace_link/test_kineto_operator.py::test_is_valid_method PASSED                                                           [ 60%]
tests/trace_link/test_unique_id_assigner.py::test_assign_or_retrieve_id_new PASSED                                              [ 66%]
tests/trace_link/test_unique_id_assigner.py::test_assign_or_retrieve_id_existing PASSED                                         [ 73%]
tests/trace_link/test_unique_id_assigner.py::test_assign_or_retrieve_id_distinct PASSED                                         [ 80%]
tests/trace_link/test_unique_id_assigner.py::test_generate_new_id_sequence PASSED                                               [ 86%]
tests/trace_link/test_unique_id_assigner.py::test_lookup_new_id_assigned PASSED                                                 [ 93%]
tests/trace_link/test_unique_id_assigner.py::test_lookup_new_id_unassigned PASSED                                               [100%]

========================================================= 15 passed in 0.05s ==========================================================

2. Correlation
2.1 Run trace_link

chakra_trace_link --pytorch-et-file /Users/theo/Downloads/llama_pytorch24.05/megatron_et_0.json --kineto-file /Users/theo/Downloads/llama_pytorch24.05/megatron_kineto_0.json --output-file ~/megatron_0.json &
chakra_trace_link --pytorch-et-file /Users/theo/Downloads/llama_pytorch24.05/megatron_et_1.json --kineto-file /Users/theo/Downloads/llama_pytorch24.05/megatron_kineto_1.json --output-file ~/megatron_1.json &
chakra_trace_link --pytorch-et-file /Users/theo/Downloads/llama_pytorch24.05/megatron_et_2.json --kineto-file /Users/theo/Downloads/llama_pytorch24.05/megatron_kineto_2.json --output-file ~/megatron_2.json &
chakra_trace_link --pytorch-et-file /Users/theo/Downloads/llama_pytorch24.05/megatron_et_3.json --kineto-file /Users/theo/Downloads/llama_pytorch24.05/megatron_kineto_3.json --output-file ~/megatron_3.json &
chakra_trace_link --pytorch-et-file /Users/theo/Downloads/llama_pytorch24.05/megatron_et_4.json --kineto-file /Users/theo/Downloads/llama_pytorch24.05/megatron_kineto_4.json --output-file ~/megatron_4.json &
chakra_trace_link --pytorch-et-file /Users/theo/Downloads/llama_pytorch24.05/megatron_et_5.json --kineto-file /Users/theo/Downloads/llama_pytorch24.05/megatron_kineto_5.json --output-file ~/megatron_5.json &
chakra_trace_link --pytorch-et-file /Users/theo/Downloads/llama_pytorch24.05/megatron_et_6.json --kineto-file /Users/theo/Downloads/llama_pytorch24.05/megatron_kineto_6.json --output-file ~/megatron_6.json &
chakra_trace_link --pytorch-et-file /Users/theo/Downloads/llama_pytorch24.05/megatron_et_7.json --kineto-file /Users/theo/Downloads/llama_pytorch24.05/megatron_kineto_7.json --output-file ~/megatron_7.json &

2.2 Run et_converter

chakra_converter --input_filename ~/megatron_0.json --output_filename megatron_0.chakra --input_type PyTorch > /tmp/rank_0 &
chakra_converter --input_filename ~/megatron_1.json --output_filename megatron_1.chakra --input_type PyTorch > /tmp/rank_1 &
chakra_converter --input_filename ~/megatron_2.json --output_filename megatron_2.chakra --input_type PyTorch > /tmp/rank_2 &
chakra_converter --input_filename ~/megatron_3.json --output_filename megatron_3.chakra --input_type PyTorch > /tmp/rank_3 &
chakra_converter --input_filename ~/megatron_4.json --output_filename megatron_4.chakra --input_type PyTorch > /tmp/rank_4 &
chakra_converter --input_filename ~/megatron_5.json --output_filename megatron_5.chakra --input_type PyTorch > /tmp/rank_5 &
chakra_converter --input_filename ~/megatron_6.json --output_filename megatron_6.chakra --input_type PyTorch > /tmp/rank_6 &
chakra_converter --input_filename ~/megatron_7.json --output_filename megatron_7.chakra --input_type PyTorch > /tmp/rank_7 &

2.3 Results

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅