Add pytest unit tests

.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

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 <input_filename>\
     --output_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 <input_filename>\
     --output_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 <input_filename>\
     --output_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_npus>\
     --num_dims <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 <input_filename>\
     --output_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_npus>\
     --num_dims <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 <input_filename>\
     --output_filename <output_filename>
 ```

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"]

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

utils/et_generator/et_generator.py → 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,

utils/et_jsonizer/et_jsonizer.py → 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,
 )
 

et_visualizer/et_visualizer.py → 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:
@@ -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()
 

et_converter/flexflow2chakra_converter.py → 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,

et_converter/pytorch2chakra_converter.py → 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,

et_converter/text2chakra_converter.py → 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,

third_party/utils/protolib.py → 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):

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()

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()