Implement graph_net/torch/op_names_extractor.py as a subclass of SamplePass.

graph_net/test/typical_sequence_decomposer_test.sh

@@ -14,7 +14,7 @@ python3 -m graph_net.model_path_handler \
     --model-path-list $model_list \
     --handler-config=$(base64 -w 0 <<EOF
 {
-    "handler_path": "$GRAPH_NET_ROOT/graph_net/torch/op_names_extractor.py",
+    "handler_path": "$GRAPH_NET_ROOT/graph_net/torch/sample_pass/op_names_extractor.py",
     "handler_class_name": "OpNamesExtractor",
     "handler_config": {
         "resume": true,

graph_net/tools/typical_sequence_decompose.sh

@@ -10,7 +10,7 @@ model_list="$GRAPH_NET_ROOT/graph_net/config/small10_torch_samples_list.txt"
 
 python3 -m graph_net.apply_sample_pass \
     --model-path-list $model_list \
-    --sample-pass-file-path $GRAPH_NET_ROOT/graph_net/torch/op_names_extractor.py \
+    --sample-pass-file-path $GRAPH_NET_ROOT/graph_net/torch/sample_pass/op_names_extractor.py \
     --sample-pass-class-name OpNamesExtractor \
     --sample-pass-config=$(base64 -w 0 <<EOF
 {
@@ -21,6 +21,7 @@ python3 -m graph_net.apply_sample_pass \
 EOF
 )
 
+
 TYPICAL_SEQUENCE_RANGES_WORKSPACE=$DECOMPOSE_WORKSPACE/workspace_typical_sequence_ranges
 
 python3 -m graph_net.torch.typical_sequence_split_points \

graph_net/torch/op_names_extractor.py → graph_net/torch/sample_pass/op_names_extractor.py

@@ -1,19 +1,72 @@
-import argparse
-import json
 import os
 from pathlib import Path
 from typing import Any, Dict, List
 import torch
 import torch.nn as nn
 
-from graph_net.torch.rp_expr.rp_expr_parser import RpExprParser
-from graph_net.torch.rp_expr.rp_expr_util import (
-    MakeNestedIndexRangeFromLetsListTokenRpExpr,
-)
+from graph_net.sample_pass.resumable_sample_pass_mixin import ResumableSamplePassMixin
 from graph_net.torch.fx_graph_module_util import get_torch_module_and_inputs
 from graph_net.torch.fx_graph_parse_util import parse_sole_graph_module_without_varify
+from graph_net.sample_pass.sample_pass import SamplePass
 
 
+class OpNamesExtractor(SamplePass, ResumableSamplePassMixin):
+    def __init__(self, config=None):
+        super().__init__(config)
+        if config is None:
+            config = {}
+
+        self.config = self._make_config(**config)
+
+    def declare_config(
+        self,
+        model_path_prefix: str,
+        output_dir: str,
+        resume: bool = False,
+        limits_handled_models: int = None,
+    ):
+        pass
+
+    def sample_handled(self, rel_model_path: str) -> bool:
+        return self.naive_sample_handled(
+            rel_model_path, search_file_name="op_names.txt"
+        )
+
+    def resume(self, rel_model_path: str):
+        torch.cuda.empty_cache()
+        model_path = os.path.join(self.config["model_path_prefix"], rel_model_path)
+        op_names = self._extract_ops(model_path)
+        output_path = self._get_output_path(rel_model_path)
+        output_path.write_text("\n".join(op_names))
+        print(f"Save op-names to {str(output_path)}")
+
+    def __call__(self, rel_model_path: str):
+        self.resumable_handle_sample(rel_model_path)
+
+    def _make_config(
+        self, model_path_prefix: str, output_dir: str, resume: bool = False
+    ):
+        return {
+            "model_path_prefix": model_path_prefix,
+            "resume": resume,
+            "output_dir": output_dir,
+        }
+
+    def _get_output_path(self, rel_model_path: str):
+        output_path_dir = Path(self.config["output_dir"]) / rel_model_path
+        output_path_dir.mkdir(parents=True, exist_ok=True)
+        output_path = output_path_dir / "op_names.txt"
+        return output_path
+
+    def _extract_ops(self, model_path: str) -> List[str]:
+        extractor = TypicalSequenceExtractor()
+        model, inputs = get_torch_module_and_inputs(model_path)
+        compiled_model, _ = parse_sole_graph_module_without_varify(model, inputs)
+        extractor.extract_compiler(compiled_model, inputs)
+        ops_info = extractor.extract_node
+
+        return [op["target_name"] for op in ops_info]
+
 
 class TypicalSequenceExtractor:
     def __init__(self):
@@ -52,46 +105,3 @@ def extract_compiler(self, gm: torch.fx.GraphModule, inputs: List[torch.Tensor])
         operator = self._extract_operators_from_graph(gm, inputs)
         self.extract_node = operator
         return gm.forward
-
-
-
-class OpNamesExtractor:
-    def __init__(self, config=None):
-        if config is None:  
-            config = {}
-
-        self.config = self._make_config(**config)  
-
-    def _make_config(
-        self, model_path_prefix: str, output_dir: str, resume: bool = False
-    ):
-        return {
-            "model_path_prefix": model_path_prefix,
-            "resume": resume,
-            "output_dir": output_dir,
-        }
-
-    def __call__(self, rel_model_path: str):
-        torch.cuda.empty_cache()
-        model_path = os.path.join(self.config["model_path_prefix"], rel_model_path)
-        output_path = self._get_output_path(rel_model_path)
-        if self.config["resume"] and output_path.exists():
-            return
-        op_names = self._extract_ops(model_path)
-        output_path.write_text("\n".join(op_names))
-        print(f"Save op-names to {str(output_path)}")
-
-    def _get_output_path(self, rel_model_path: str):
-        output_path_dir = Path(self.config["output_dir"]) / rel_model_path
-        output_path_dir.mkdir(parents=True, exist_ok=True)
-        output_path = output_path_dir / "op_names.txt"
-        return output_path
-
-    def _extract_ops(self, model_path: str) -> List[str]:
-        extractor = TypicalSequenceExtractor()
-        model, inputs = get_torch_module_and_inputs(model_path)
-        compiled_model, _ = parse_sole_graph_module_without_varify(model, inputs)
-        extractor.extract_compiler(compiled_model, inputs)
-        ops_info = extractor.extract_node
-
-        return [op["target_name"] for op in ops_info]