Fix ruff check errors in src/python/converter/pytorch2chakra_converte…

…r.py
mlcommons · May 9, 2024 · 7ab1796 · 7ab1796
1 parent 1fb4e8c
commit 7ab1796
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 22 deletions.
diff --git a/src/python/converter/pytorch2chakra_converter.py b/src/python/converter/pytorch2chakra_converter.py
@@ -101,7 +101,7 @@ def convert(self) -> None:
 
         self.open_chakra_execution_trace()
 
-        for pytorch_nid, pytorch_node in self.pytorch_nodes.items():
+        for _, pytorch_node in self.pytorch_nodes.items():
             if (pytorch_node.get_op_type() == PyTorchNodeType.CPU_OP) or (
                 pytorch_node.get_op_type() == PyTorchNodeType.LABEL
             ):
@@ -155,7 +155,7 @@ def load_pytorch_execution_traces(self) -> None:
             self._parse_and_instantiate_nodes(pytorch_et_data)
         except IOError as e:
             self.logger.error(f"Error opening file {self.input_filename}: {e}")
-            raise Exception(f"Could not open file {self.input_filename}")
+            raise Exception(f"Could not open file {self.input_filename}") from e
 
     def _parse_and_instantiate_nodes(self, pytorch_et_data: Dict) -> None:
         """
@@ -180,7 +180,7 @@ def _parse_and_instantiate_nodes(self, pytorch_et_data: Dict) -> None:
         }
         self._establish_parent_child_relationships(pytorch_node_objects)
 
-    def _establish_parent_child_relationships(self, pytorch_node_objects: Dict[int, PyTorchNode]) -> None:
+    def _establish_parent_child_relationships(self, pytorch_node_objects: Dict[int, PyTorchNode]) -> None:  # noqa: C901
         """
         Establishes parent-child relationships among PyTorch nodes and counts
         the node types.
@@ -248,11 +248,11 @@ def open_chakra_execution_trace(self) -> None:
         """
         self.logger.info(f"Opening Chakra execution trace file: {self.output_filename}")
         try:
-            self.chakra_et = open(self.output_filename, "wb")
+            self.chakra_et = open(self.output_filename, "wb")  # noqa: SIM115
         except IOError as e:
             err_msg = f"Error opening file {self.output_filename}: {e}"
             self.logger.error(err_msg)
-            raise Exception(err_msg)
+            raise Exception(err_msg) from e
 
     def convert_to_chakra_node(self, pytorch_node: PyTorchNode) -> ChakraNode:
         """
@@ -363,7 +363,7 @@ def is_root_node(self, node):
         if node.name in ["[pytorch|profiler|execution_graph|thread]", "[pytorch|profiler|execution_trace|thread]"]:
             return True
 
-    def convert_ctrl_dep_to_data_dep(self, chakra_node: ChakraNode) -> None:
+    def convert_ctrl_dep_to_data_dep(self, chakra_node: ChakraNode) -> None:  # noqa: C901
         """
         Traverses nodes based on control dependencies (parent nodes) and encodes
         data dependencies appropriately. This method is crucial for converting the
@@ -431,13 +431,11 @@ def convert_ctrl_dep_to_data_dep(self, chakra_node: ChakraNode) -> None:
             node_op_type = pytorch_node.get_op_type()
 
             if node_op_type == PyTorchNodeType.GPU_OP:
-                if last_visited_any:
-                    if last_visited_any.id not in current_node.data_deps:
-                        current_node.data_deps.append(last_visited_any.id)
-                        self.logger.debug(
-                            f"GPU Node ID {current_node.id} now has a data "
-                            f"dependency on Node ID {last_visited_any.id}"
-                        )
+                if (last_visited_any) and (last_visited_any.id not in current_node.data_deps):
+                    current_node.data_deps.append(last_visited_any.id)
+                    self.logger.debug(
+                        f"GPU Node ID {current_node.id} now has a data " f"dependency on Node ID {last_visited_any.id}"
+                    )
 
                 last_visited_any = last_visited_non_gpu
             else:
@@ -449,13 +447,12 @@ def convert_ctrl_dep_to_data_dep(self, chakra_node: ChakraNode) -> None:
                             f"CPU Node ID {current_node.id} now has an inter-thread data dependency on Node ID {id}"
                         )
 
-                if last_visited_non_gpu:
-                    if last_visited_non_gpu.id not in current_node.data_deps:
-                        current_node.data_deps.append(last_visited_non_gpu.id)
-                        self.logger.debug(
-                            f"CPU Node ID {current_node.id} now has a data "
-                            f"dependency on non-GPU Node ID {last_visited_non_gpu.id}"
-                        )
+                if (last_visited_non_gpu) and (last_visited_non_gpu.id not in current_node.data_deps):
+                    current_node.data_deps.append(last_visited_non_gpu.id)
+                    self.logger.debug(
+                        f"CPU Node ID {current_node.id} now has a data "
+                        f"dependency on non-GPU Node ID {last_visited_non_gpu.id}"
+                    )
                 last_visited_non_gpu = current_node
                 last_visited_any = current_node
 

diff --git a/src/python/converter/pytorch_node.py b/src/python/converter/pytorch_node.py
@@ -233,5 +233,5 @@ def get_data_type_size(data_type: str) -> int:
         }
         try:
             return data_type_size_map[data_type]
-        except KeyError:
-            raise ValueError(f"Unsupported data type: {data_type}")
+        except KeyError as e:
+            raise ValueError(f"Unsupported data type: {data_type}") from e