Add generation of buckets and groups into the build of db.

Author: Xreki

graph_net/tools/generate_subgraph_dataset.sh

@@ -93,36 +93,6 @@ function generate_subgraph_list() {
         | tee $sample_list
 }
 
-function insert_graph_sample(){
-    local target_dir="$1"
-    local repo_uid="$2"
-    local sample_type="$3"
-    local sample_list="$4"
-    echo ">>> [0] Inserting samples into database: ${DB_PATH}."
-    echo ">>>"
-
-    if [ ! -f "$DB_PATH" ]; then
-        echo "Fail ! No Database ! : $DB_PATH"
-        exit 1
-    fi
-
-    local order_value=0
-    while IFS= read -r model_rel_path; do
-        echo "insert : $model_rel_path"
-        python3 "${GRAPH_NET_ROOT}/sqlite/graphsample_insert.py" \
-            --model_path_prefix "${target_dir}" \
-            --relative_model_path "$model_rel_path" \
-            --repo_uid "${repo_uid}" \
-            --sample_type "${sample_type}" \
-            --order_value "$order_value" \
-            --db_path "$DB_PATH"
-
-        ((order_value++))
-
-    done < "$sample_list"
-}
-
-
 function rewrite_device() {
     echo ">>> [1] Rewrite devices for subgraph samples under ${GRAPH_NET_ROOT}."
     echo ">>>"
@@ -658,29 +628,6 @@ function generate_typical_subgraphs() {
     # generate_unittest_for_typical_subgraphs 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_unittests_typical_subgraphs_${suffix}.txt
 }
 
-function generate_database() {
-    timestamp=`date +%Y%m%d_%H%M`
-
-    # init database
-    if [ ! -f ${DB_PATH} ]; then
-        python ${GRAPH_NET_ROOT}/sqlite/init_db.py --db_path ${DB_PATH} 2>&1 | tee ${DECOMPOSE_WORKSPACE}/log_init_db_${timestamp}.txt
-    fi
-    
-    # full_graph
-    insert_graph_sample ${GRAPH_NET_ROOT} "hf_torch_samples" "full_graph" ${model_list}
-
-    # fusible_graph, typical_graph
-    for sample_type in fusible_graph typical_graph; do
-        insert_graph_sample $OUTPUT_DIR/$sample_type "hf_torch_samples" $sample_type $OUTPUT_DIR/${sample_type}/sample_list.txt
-    done
-
-    # insert buckets
-    python ${GRAPH_NET_ROOT}/sqlite/graph_net_sample_bucket_generator.py --db_path ${DB_PATH}
-
-    # insert groups
-    python ${GRAPH_NET_ROOT}/sqlite/graph_net_sample_groups_insert.py --db_path ${DB_PATH}
-}
-
 function main() {
     do_common_generalzation_and_decompose
 
@@ -693,8 +640,6 @@ function main() {
     generate_typical_subgraphs
     #cp -rf $DTYPE_GENERALIZED_TYPICAL_SUBGRAPH_DIR $OUTPUT_DIR/$sample_type
     #cp -rf $dtype_generalized_typical_subgraph_list $OUTPUT_DIR/$sample_type/sample_list.txt
-
-    #generate_database
 }
 
 function summary() {

sqlite/build_db.py

@@ -100,6 +100,12 @@ def main(args):
 
     print("all done")
 
+    from graph_sample_bucket_generator import generate_buckets
+    from graph_sample_groups_insert import generate_groups
+
+    generate_buckets(db_path)
+    generate_groups(db_path, num_dtypes=3)
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(

sqlite/graph_sample_bucket_generator.py

@@ -213,52 +213,36 @@ def save_bucket_results(
     return count
 
 
-def main():
-    parser = argparse.ArgumentParser(
-        description="Generate graph_net_sample_buckets from graph_sample"
-    )
-    parser.add_argument(
-        "--db_path",
-        type=str,
-        required=True,
-        help="Path to the SQLite database file",
-    )
-    parser.add_argument(
-        "--dry_run",
-        action="store_true",
-        help="Only print what would be done, don't actually insert into database",
-    )
-
-    args = parser.parse_args()
-
-    session = get_session(args.db_path)
-
-    print("=" * 70)
-    print("Step 1: Generating bucket info from graph_sample...")
-    sample_type_results, all_bucket_info_map = generate_sample_buckets(session)
-    print(f"  Total samples: {len(all_bucket_info_map)}")
-    print(f"  Number of sample_types: {len(sample_type_results)}")
-
-    print()
-    for result in sorted(sample_type_results, key=lambda x: -len(x)):
-        flag = " [sole-op]" if result.is_sole_op else ""
-        print(f"  {result.sample_type}{flag}: {len(result)} samples")
+def generate_buckets(db_path, dry_run=False):
+    """Generate buckets and save to DB."""
+    session = get_session(db_path)
+    try:
+        print("=" * 70)
+        print("Step 1: Generating bucket info from graph_sample...")
+        _, bucket_info_map = generate_sample_buckets(session)
+        if dry_run:
+            print("Dry run mode - skipping database insert")
+            print(
+                f"  Would insert {len(bucket_info_map)} records into graph_net_sample_buckets"
+            )
+            return 0
 
-    print("=" * 70)
-    if args.dry_run:
-        print("Dry run mode - skipping database insert")
-        print(
-            f"  Would insert {len(all_bucket_info_map)} records into graph_net_sample_buckets"
-        )
-    else:
         print("Step 2: Saving to database...")
-        count = save_bucket_results(session, all_bucket_info_map)
+        count = save_bucket_results(session, bucket_info_map)
         print(f"  Inserted {count} records into graph_net_sample_buckets")
+        return count
+    finally:
+        session.close()
 
-    print("=" * 70)
-    print("Done!")
 
-    session.close()
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate graph_net_sample_buckets from graph_sample"
+    )
+    parser.add_argument("--db_path", type=str, required=True)
+    parser.add_argument("--dry_run", action="store_true")
+    args = parser.parse_args()
+    generate_buckets(args.db_path, dry_run=args.dry_run)
 
 
 if __name__ == "__main__":

sqlite/graph_sample_groups_insert.py

@@ -7,8 +7,6 @@
 from orm_models import get_session, GraphNetSampleGroup
 
 
-# ── Types ──
-
 BucketGroup = namedtuple(
     "BucketGroup",
     ["head_uid", "op_seq", "shapes", "sample_type", "all_uids_csv"],
@@ -20,9 +18,6 @@
 )
 
 
-# ── Helpers ──
-
-
 def _new_group_id():
     return str(uuid_module.uuid4())
 
@@ -51,9 +46,6 @@ def _print_stats(stats):
     print(f"\n  Total: {total_records} records, {total_groups} groups.")
 
 
-# ── Database Queries ──
-
-
 class DB:
     def __init__(self, path):
         self.path = path
@@ -229,39 +221,24 @@ def _insert_groups(session, rows, policy):
     return stats
 
 
-# ═══════════════════════════════════════════════════════════════════
-# Main
-# ═══════════════════════════════════════════════════════════════════
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Generate graph_net_sample_groups (v1 + v2)"
-    )
-    parser.add_argument("--db_path", type=str, required=True)
-    parser.add_argument("--num_dtypes", type=int, default=3)
-    args = parser.parse_args()
-
-    db = DB(args.db_path)
+def generate_groups(db_path, num_dtypes=3):
+    """Generate sample groups and save to DB."""
+    db = DB(db_path)
     db.connect()
-    session = get_session(args.db_path)
-
+    session = get_session(db_path)
     all_stats = defaultdict(lambda: {"records": 0, "groups": set()})
-
     try:
-        # V1
         buckets = query_bucket_groups(db)
         print(f"Bucket groups: {len(buckets)}")
         v1 = _insert_groups(session, generate_v1_groups(buckets), "bucket_policy_v1")
         _merge_stats(all_stats, v1)
 
-        # V2
         candidates = query_v2_candidates(db)
         print(f"V2 candidates: {len(candidates)}")
         if candidates:
             v2 = _insert_groups(
                 session,
-                generate_v2_groups(candidates, args.num_dtypes),
+                generate_v2_groups(candidates, num_dtypes),
                 "bucket_policy_v2",
             )
             _merge_stats(all_stats, v2)
@@ -276,7 +253,17 @@ def main():
 
     print("=" * 60)
     _print_stats(all_stats)
-    print("\nDone!")
+    return all_stats
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate graph_net_sample_groups (v1 + v2)"
+    )
+    parser.add_argument("--db_path", type=str, required=True)
+    parser.add_argument("--num_dtypes", type=int, default=3)
+    args = parser.parse_args()
+    generate_groups(args.db_path, args.num_dtypes)
 
 
 if __name__ == "__main__":