From a54545f724ebaaa3464dfa8edd00d9dc188ec8bd Mon Sep 17 00:00:00 2001 From: Taekyung Heo <7621438+TaekyungHeo@users.noreply.github.com> Date: Mon, 27 Nov 2023 09:52:51 -0500 Subject: [PATCH 1/3] Refactor directories --- .github/workflows/main.yml | 8 ++++---- INSTALL.md | 14 +++++++------- pyproject.toml | 9 ++------- setup.cfg | 4 ++-- {et_converter => src}/et_converter.py | 0 {et_def => src}/et_def.proto | 0 {et_feeder => src/et_feeder}/et_feeder.cpp | 0 {et_feeder => src/et_feeder}/et_feeder.h | 0 {et_feeder => src/et_feeder}/et_feeder_node.cpp | 0 {et_feeder => src/et_feeder}/et_feeder_node.h | 0 {utils/et_generator => src}/et_generator.py | 2 +- {utils/et_jsonizer => src}/et_jsonizer.py | 2 +- {et_visualizer => src}/et_visualizer.py | 2 +- {et_converter => src}/flexflow2chakra_converter.py | 2 +- {et_converter => src}/pytorch2chakra_converter.py | 2 +- {et_converter => src}/text2chakra_converter.py | 2 +- {third_party => src/third_party}/utils/protoio.cc | 0 {third_party => src/third_party}/utils/protoio.hh | 0 {third_party => src/third_party}/utils/protolib.py | 0 .../timeline_visualizer.py | 0 20 files changed, 21 insertions(+), 26 deletions(-) rename {et_converter => src}/et_converter.py (100%) rename {et_def => src}/et_def.proto (100%) rename {et_feeder => src/et_feeder}/et_feeder.cpp (100%) rename {et_feeder => src/et_feeder}/et_feeder.h (100%) rename {et_feeder => src/et_feeder}/et_feeder_node.cpp (100%) rename {et_feeder => src/et_feeder}/et_feeder_node.h (100%) rename {utils/et_generator => src}/et_generator.py (99%) rename {utils/et_jsonizer => src}/et_jsonizer.py (96%) rename {et_visualizer => src}/et_visualizer.py (98%) rename {et_converter => src}/flexflow2chakra_converter.py (99%) rename {et_converter => src}/pytorch2chakra_converter.py (99%) rename {et_converter => src}/text2chakra_converter.py (99%) rename {third_party => src/third_party}/utils/protoio.cc (100%) rename {third_party => src/third_party}/utils/protoio.hh (100%) rename {third_party => src/third_party}/utils/protolib.py (100%) rename {timeline_visualizer => src}/timeline_visualizer.py (100%) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 7fd79eb..2af3094 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -12,7 +12,7 @@ jobs: version: 1.0 - run: sudo apt install -y graphviz - uses: lukka/get-cmake@latest - - run: protoc et_def.proto --proto_path et_def --cpp_out et_def + - run: protoc et_def.proto --proto_path src --cpp_out src - uses: actions/setup-python@v2 with: python-version: 3.7 @@ -23,10 +23,10 @@ jobs: - run: conda install pytorch - run: pip install fbgemm-gpu-cpu - run: pip install . - - run: python3 -m chakra.et_generator.et_generator --num_npus 1 + - run: python3 -m chakra.et_generator --num_npus 1 - run: | - python3 -m chakra.et_visualizer.et_visualizer --help + python3 -m chakra.et_visualizer --help for f in *.et do - python3 -m chakra.et_visualizer.et_visualizer --input_filename ${f} --output_filename "${f%.*}".dot + python3 -m chakra.et_visualizer --input_filename ${f} --output_filename "${f%.*}".dot done diff --git a/INSTALL.md b/INSTALL.md index 4c74ca0..4f31053 100644 --- a/INSTALL.md +++ b/INSTALL.md @@ -23,7 +23,7 @@ You can use the following commands for each input type. ### ASTRA-sim Text Files ```shell -$ python -m chakra.et_converter.et_converter\ +$ python -m chakra.et_converter\ --input_type Text\ --input_filename \ --output_filename \ @@ -34,7 +34,7 @@ $ python -m chakra.et_converter.et_converter\ ### FlexFlow Execution Graphs ```shell -$ python -m chakra.et_converter.et_converter\ +$ python -m chakra.et_converter\ --input_type FlexFlow\ --input_filename \ --output_filename \ @@ -44,7 +44,7 @@ $ python -m chakra.et_converter.et_converter\ ### PyTorch Execution Graphs ```shell -$ python -m chakra.et_converter.et_converter\ +$ python -m chakra.et_converter\ --input_type PyTorch\ --input_filename \ --output_filename \ @@ -56,7 +56,7 @@ This is an execution trace generator that generates synthetic execution traces. A user can define a new function in the generator to generate new synthetic execution traces. You can follow the commands below to run the generator. ```shell -$ python -m chakra.et_generator.et_generator\ +$ python -m chakra.et_generator\ --num_npus \ --num_dims ``` @@ -72,7 +72,7 @@ For visualizing GraphML files, you can use Gephi (https://gephi.org/). Run the tool with the following command: ```shell -$ python -m chakra.et_visualizer.et_visualizer\ +$ python -m chakra.et_visualizer\ --input_filename \ --output_filename ``` @@ -124,7 +124,7 @@ $ cd - $ ./build/astra_analytical/build.sh -c $ cd extern/graph_frontend/chakra/ -$ python -m chakra.et_generator.et_generator\ +$ python -m chakra.et_generator\ --num_npus \ --num_dims @@ -136,7 +136,7 @@ $ ./run.sh This tool prints the nodes within execution traces for better comprehension. The printed information includes the node's id, name, type, and any associated metadata, which are all outputted in a user-friendly text format. ``` -$ python -m chakra.et_jsonizer.et_jsonizer\ +$ python -m chakra.et_jsonizer\ --input_filename \ --output_filename ``` diff --git a/pyproject.toml b/pyproject.toml index 932dbd7..ca8f451 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,13 +19,8 @@ Documentation = "https://github.com/mlcommons/chakra/README.md" Repository = "https://github.com/mlcommons/chakra.git" [tool.setuptools.package-dir] -"chakra.et_def" = "et_def" -"chakra.et_converter" = "et_converter" -"chakra.et_visualizer" = "et_visualizer" -"chakra.timeline_visualizer" = "timeline_visualizer" -"chakra.et_generator" = "utils/et_generator" -"chakra.et_jsonizer" = "utils/et_jsonizer" -"chakra.third_party.utils" = "third_party/utils" +"chakra" = "src" +"chakra.third_party.utils" = "src/third_party/utils" [tool.setuptools.package-data] "chakra.et_def" = ["et_def_pb2.pyi", "et_def.proto"] diff --git a/setup.cfg b/setup.cfg index 545ad57..9afb748 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,5 +1,5 @@ [build_grpc] proto_files = et_def.proto grpc_files = et_def.proto -proto_path = et_def -output_path = et_def +proto_path = src +output_path = src diff --git a/et_converter/et_converter.py b/src/et_converter.py similarity index 100% rename from et_converter/et_converter.py rename to src/et_converter.py diff --git a/et_def/et_def.proto b/src/et_def.proto similarity index 100% rename from et_def/et_def.proto rename to src/et_def.proto diff --git a/et_feeder/et_feeder.cpp b/src/et_feeder/et_feeder.cpp similarity index 100% rename from et_feeder/et_feeder.cpp rename to src/et_feeder/et_feeder.cpp diff --git a/et_feeder/et_feeder.h b/src/et_feeder/et_feeder.h similarity index 100% rename from et_feeder/et_feeder.h rename to src/et_feeder/et_feeder.h diff --git a/et_feeder/et_feeder_node.cpp b/src/et_feeder/et_feeder_node.cpp similarity index 100% rename from et_feeder/et_feeder_node.cpp rename to src/et_feeder/et_feeder_node.cpp diff --git a/et_feeder/et_feeder_node.h b/src/et_feeder/et_feeder_node.h similarity index 100% rename from et_feeder/et_feeder_node.h rename to src/et_feeder/et_feeder_node.h diff --git a/utils/et_generator/et_generator.py b/src/et_generator.py similarity index 99% rename from utils/et_generator/et_generator.py rename to src/et_generator.py index cc2a94e..601294e 100644 --- a/utils/et_generator/et_generator.py +++ b/src/et_generator.py @@ -3,7 +3,7 @@ import argparse from chakra.third_party.utils.protolib import encodeMessage as encode_message -from chakra.et_def.et_def_pb2 import ( +from chakra.et_def_pb2 import ( Node as ChakraNode, DoubleList, FloatList, diff --git a/utils/et_jsonizer/et_jsonizer.py b/src/et_jsonizer.py similarity index 96% rename from utils/et_jsonizer/et_jsonizer.py rename to src/et_jsonizer.py index 5dc5c08..78f91a3 100644 --- a/utils/et_jsonizer/et_jsonizer.py +++ b/src/et_jsonizer.py @@ -9,7 +9,7 @@ decodeMessage as decode_message ) -from chakra.et_def.et_def_pb2 import ( +from chakra.et_def_pb2 import ( Node as ChakraNode, ) diff --git a/et_visualizer/et_visualizer.py b/src/et_visualizer.py similarity index 98% rename from et_visualizer/et_visualizer.py rename to src/et_visualizer.py index 589db44..4ab0202 100644 --- a/et_visualizer/et_visualizer.py +++ b/src/et_visualizer.py @@ -8,7 +8,7 @@ openFileRd as open_file_rd, decodeMessage as decode_message ) -from chakra.et_def.et_def_pb2 import Node +from chakra.et_def_pb2 import Node def main() -> None: diff --git a/et_converter/flexflow2chakra_converter.py b/src/flexflow2chakra_converter.py similarity index 99% rename from et_converter/flexflow2chakra_converter.py rename to src/flexflow2chakra_converter.py index abb7884..f3fdf8e 100644 --- a/et_converter/flexflow2chakra_converter.py +++ b/src/flexflow2chakra_converter.py @@ -6,7 +6,7 @@ from typing import Any from chakra.third_party.utils.protolib import encodeMessage as encode_message -from chakra.et_def.et_def_pb2 import ( +from chakra.et_def_pb2 import ( Node as ChakraNode, AttributeProto as ChakraAttr, COMP_NODE, diff --git a/et_converter/pytorch2chakra_converter.py b/src/pytorch2chakra_converter.py similarity index 99% rename from et_converter/pytorch2chakra_converter.py rename to src/pytorch2chakra_converter.py index 2a1ece7..c2a3301 100644 --- a/et_converter/pytorch2chakra_converter.py +++ b/src/pytorch2chakra_converter.py @@ -9,7 +9,7 @@ from typing import Any, Dict, List from chakra.third_party.utils.protolib import encodeMessage as encode_message -from chakra.et_def.et_def_pb2 import ( +from chakra.et_def_pb2 import ( GlobalMetadata, Node as ChakraNode, AttributeProto as ChakraAttr, diff --git a/et_converter/text2chakra_converter.py b/src/text2chakra_converter.py similarity index 99% rename from et_converter/text2chakra_converter.py rename to src/text2chakra_converter.py index c5fb826..639ec83 100644 --- a/et_converter/text2chakra_converter.py +++ b/src/text2chakra_converter.py @@ -5,7 +5,7 @@ from io import TextIOWrapper from typing import Any, List from chakra.third_party.utils.protolib import encodeMessage as encode_message -from chakra.et_def.et_def_pb2 import ( +from chakra.et_def_pb2 import ( Node, AttributeProto as ChakraAttr, COMP_NODE, diff --git a/third_party/utils/protoio.cc b/src/third_party/utils/protoio.cc similarity index 100% rename from third_party/utils/protoio.cc rename to src/third_party/utils/protoio.cc diff --git a/third_party/utils/protoio.hh b/src/third_party/utils/protoio.hh similarity index 100% rename from third_party/utils/protoio.hh rename to src/third_party/utils/protoio.hh diff --git a/third_party/utils/protolib.py b/src/third_party/utils/protolib.py similarity index 100% rename from third_party/utils/protolib.py rename to src/third_party/utils/protolib.py diff --git a/timeline_visualizer/timeline_visualizer.py b/src/timeline_visualizer.py similarity index 100% rename from timeline_visualizer/timeline_visualizer.py rename to src/timeline_visualizer.py From b20737487f9904ccde2c4231538e582bdfb9e633 Mon Sep 17 00:00:00 2001 From: Taekyung Heo <7621438+TaekyungHeo@users.noreply.github.com> Date: Mon, 27 Nov 2023 10:18:05 -0500 Subject: [PATCH 2/3] Add test_et_visualizer.py --- src/et_visualizer.py | 2 +- src/third_party/utils/protolib.py | 3 +-- tests/test_et_visualizer.py | 33 +++++++++++++++++++++++++++++++ 3 files changed, 35 insertions(+), 3 deletions(-) create mode 100644 tests/test_et_visualizer.py diff --git a/src/et_visualizer.py b/src/et_visualizer.py index 4ab0202..11ba26e 100644 --- a/src/et_visualizer.py +++ b/src/et_visualizer.py @@ -72,7 +72,7 @@ def main() -> None: nx.write_graphml(G, args.output_filename) else: - print("Unknown output file extension. Must be one of pdf, dot, graphml.") + raise ValueError("Unsupported file extension. Must be one of: pdf, dot, graphml.") et.close() diff --git a/src/third_party/utils/protolib.py b/src/third_party/utils/protolib.py index ea0ff09..6a34bcf 100644 --- a/src/third_party/utils/protolib.py +++ b/src/third_party/utils/protolib.py @@ -91,8 +91,7 @@ def openFileRd(in_file): except IOError: proto_in = open(in_file, 'rb') except IOError: - print("Failed to open ", in_file, " for reading") - exit(-1) + raise FileNotFoundError("Failed to open ", in_file, " for reading") return proto_in def _DecodeVarint32(in_file): diff --git a/tests/test_et_visualizer.py b/tests/test_et_visualizer.py new file mode 100644 index 0000000..f133e64 --- /dev/null +++ b/tests/test_et_visualizer.py @@ -0,0 +1,33 @@ +import pytest +from chakra.et_visualizer import main +from unittest.mock import patch + +# Assuming 'example_input.et' is the example input file you provided +example_input_file = 'example_input.et' + +def test_run_with_pdf_output(): + test_args = ["et_visualizer.py", "--input_filename", example_input_file, "--output_filename", "output.pdf"] + with patch('sys.argv', test_args): + main() # No assertion needed; we're checking if it runs without error + +def test_run_with_dot_output(): + test_args = ["et_visualizer.py", "--input_filename", example_input_file, "--output_filename", "output.dot"] + with patch('sys.argv', test_args): + main() # No assertion needed; we're checking if it runs without error + +def test_run_with_graphml_output(): + test_args = ["et_visualizer.py", "--input_filename", example_input_file, "--output_filename", "output.graphml"] + with patch('sys.argv', test_args): + main() # No assertion needed; we're checking if it runs without error + +def test_incorrect_file_path(): + test_args = ["et_visualizer.py", "--input_filename", "nonexistent_input.et", "--output_filename", "output.pdf"] + with patch('sys.argv', test_args): + with pytest.raises(FileNotFoundError): + main() + +def test_unsupported_file_extension(): + test_args = ["et_visualizer.py", "--input_filename", example_input_file, "--output_filename", "output.xyz"] + with patch('sys.argv', test_args): + with pytest.raises(ValueError): + main() From 6238674d13b90c17faab5146f48ab4bc286609cf Mon Sep 17 00:00:00 2001 From: Taekyung Heo <7621438+TaekyungHeo@users.noreply.github.com> Date: Mon, 27 Nov 2023 10:25:14 -0500 Subject: [PATCH 3/3] Add test_et_jsonizer.py --- tests/test_et_jsonizer.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 tests/test_et_jsonizer.py diff --git a/tests/test_et_jsonizer.py b/tests/test_et_jsonizer.py new file mode 100644 index 0000000..9b24046 --- /dev/null +++ b/tests/test_et_jsonizer.py @@ -0,0 +1,21 @@ +import pytest +from chakra.et_jsonizer import main +from unittest.mock import patch, mock_open + +def test_run_with_valid_input(): + test_args = ["et_jsonizer.py", "--input_filename", "valid_input.et", "--output_filename", "output.json"] + with patch('sys.argv', test_args), \ + patch('chakra.et_jsonizer.open_file_rd') as mock_open_file_rd, \ + patch('builtins.open', mock_open()): + main() # No assertion needed; we're checking if it runs without error + +def test_missing_arguments(): + with pytest.raises(SystemExit): + with patch('sys.argv', ['et_jsonizer.py']): + main() + +def test_missing_file_path(): + test_args = ["et_jsonizer.py", "--input_filename", "nonexistent_input.et", "--output_filename", "output.json"] + with patch('sys.argv', test_args): + with pytest.raises(FileNotFoundError): + main()