From 7fdb79397c1c186d686200f67a73437584e49df6 Mon Sep 17 00:00:00 2001 From: Honglei-Qiu <1044497581@qq.com> Date: Thu, 26 Feb 2026 06:27:56 +0000 Subject: [PATCH 1/3] add table SampleInputTensorMeta --- sqlite/graphsample_delete.py | 4 + sqlite/graphsample_insert.py | 94 +++++++++++++++++++ sqlite/merge_db.py | 22 +++++ .../create_main_tables_2026-02-25-065451.sql | 16 ++++ sqlite/orm_models.py | 27 ++++++ 5 files changed, 163 insertions(+) mode change 100644 => 100755 sqlite/merge_db.py create mode 100755 sqlite/migrates/create_main_tables_2026-02-25-065451.sql diff --git a/sqlite/graphsample_delete.py b/sqlite/graphsample_delete.py index 3da9571e1d..5cc3e676dd 100755 --- a/sqlite/graphsample_delete.py +++ b/sqlite/graphsample_delete.py @@ -8,6 +8,7 @@ DataTypeGeneralizationSource, SampleOpNameList, SampleOpName, + SampleInputTensorMeta, ) @@ -71,6 +72,9 @@ def delete_graph_sample(db_path: str, relative_model_path: str, repo_uid: str = session.query(SampleOpName).filter( SampleOpName.sample_uuid == graph_sample.uuid ).update({"deleted": True, "delete_at": delete_at}) + session.query(SampleInputTensorMeta).filter( + SampleInputTensorMeta.sample_uuid == graph_sample.uuid + ).update({"deleted": True, "delete_at": delete_at}) session.commit() print(f"Successfully deleted: {relative_model_path}") diff --git a/sqlite/graphsample_insert.py b/sqlite/graphsample_insert.py index 4048196014..a31e346060 100755 --- a/sqlite/graphsample_insert.py +++ b/sqlite/graphsample_insert.py @@ -13,6 +13,7 @@ DataTypeGeneralizationSource, SampleOpName, SampleOpNameList, + SampleInputTensorMeta, ) from sqlalchemy import delete as sql_delete from sqlalchemy.exc import IntegrityError @@ -370,6 +371,93 @@ def insert_sample_op_name_list( session.close() +# SampleInputTensorMeta insert func +def insert_sample_input_tensor_meta( + sample_uuid: str, + model_path_prefix: str, + relative_model_path: str, + db_path: str, +): + model_path = Path(model_path_prefix) / relative_model_path + weight_meta_file = model_path / "weight_meta.py" + + try: + with open(weight_meta_file) as f: + content = f.read() + except Exception as e: + print(f"Warning: Failed to read {weight_meta_file}: {e}") + return + + input_tensor_metas = [] + input_idx = 0 + original_name_pattern = re.compile(r'original_name\s*=\s*"(.+?)"') + shape_pattern = re.compile(r"shape\s*=\s*\[(.+?)\]") + dtype_pattern = re.compile(r'dtype\s*=\s*"(.+?)"') + class_blocks = re.split(r"\nclass Program_weight_tensor_meta_", content) + + for block in class_blocks: + original_name_match = original_name_pattern.search(block) + if not original_name_match: + continue + + original_name = original_name_match.group(1) + + shape_match = shape_pattern.search(block) + if not shape_match: + print(f"Warning: No shape found for {original_name} in {weight_meta_file}") + continue + + shape_str = f"[{shape_match.group(1)}]" + + dtype_match = dtype_pattern.search(block) + dtype = dtype_match.group(1) if dtype_match else "" + + input_tensor_metas.append( + { + "input_name": original_name, + "input_idx": input_idx, + "shape": shape_str, + "dtype": dtype, + } + ) + input_idx += 1 + + if not input_tensor_metas: + print(f"No input tensor meta found in {weight_meta_file}") + return + + session = get_session(db_path) + try: + session.execute( + sql_delete(SampleInputTensorMeta).where( + SampleInputTensorMeta.sample_uuid == sample_uuid + ) + ) + for meta in input_tensor_metas: + sample_input_tensor_meta = SampleInputTensorMeta( + sample_uuid=sample_uuid, + input_name=meta["input_name"], + input_idx=meta["input_idx"], + shape=meta["shape"], + dtype=meta["dtype"], + create_at=datetime.now(), + deleted=False, + delete_at=None, + ) + session.add(sample_input_tensor_meta) + + session.commit() + print( + f"Inserted {len(input_tensor_metas)} input tensor meta(s) for sample_uuid={sample_uuid}" + ) + except IntegrityError as e: + session.rollback() + print(f"Error inserting input tensor meta: {e}") + raise e + finally: + session.close() + + # main func def main(args): data = get_graph_sample_data( @@ -397,6 +485,12 @@ def main(args): relative_model_path=args.relative_model_path, db_path=args.db_path, ) + insert_sample_input_tensor_meta( + sample_uuid=data["uuid"], + model_path_prefix=args.model_path_prefix, + relative_model_path=args.relative_model_path, + db_path=args.db_path, + ) if args.sample_type in ["fusible_graph"]: insert_dimension_generalization_source( subgraph_source_data["subgraph_uuid"], diff --git a/sqlite/merge_db.py b/sqlite/merge_db.py old mode 100644 new mode 100755 index 7c0e6fd0bf..8ac4d1c495 --- a/sqlite/merge_db.py +++ b/sqlite/merge_db.py @@ -10,6 +10,7 @@ BackwardGraphSource, SampleOpName, SampleOpNameList, + SampleInputTensorMeta, ) from sqlalchemy.exc import IntegrityError @@ -27,6 +28,7 @@ def merge_databases(main_db_path: str, new_db_path: str): "backward_graph_source": 0, "sample_op_name": 0, "sample_op_name_list": 0, + "sample_input_tensor_meta": 0, } try: @@ -70,6 +72,11 @@ def merge_databases(main_db_path: str, new_db_path: str): sample_op_name_list_map = { s.sample_uuid: s for s in new_session.query(SampleOpNameList).all() } + sample_input_tensor_meta_list_map = {} + for s in new_session.query(SampleInputTensorMeta).all(): + if s.sample_uuid not in sample_input_tensor_meta_list_map: + sample_input_tensor_meta_list_map[s.sample_uuid] = [] + sample_input_tensor_meta_list_map[s.sample_uuid].append(s) existing_graph_uuids = { g.uuid for g in main_session.query(GraphSample.uuid).all() @@ -187,6 +194,21 @@ def merge_databases(main_db_path: str, new_db_path: str): main_session.add(new_op_name) stats["sample_op_name"] += 1 + if sample.uuid in sample_input_tensor_meta_list_map: + for src in sample_input_tensor_meta_list_map[sample.uuid]: + new_input_tensor_meta = SampleInputTensorMeta( + sample_uuid=src.sample_uuid, + input_name=src.input_name, + input_idx=src.input_idx, + shape=src.shape, + dtype=src.dtype, + create_at=src.create_at, + deleted=src.deleted, + delete_at=src.delete_at, + ) + main_session.add(new_input_tensor_meta) + stats["sample_input_tensor_meta"] += 1 + main_session.commit() total_merged = sum(stats.values()) diff --git a/sqlite/migrates/create_main_tables_2026-02-25-065451.sql b/sqlite/migrates/create_main_tables_2026-02-25-065451.sql new file mode 100755 index 0000000000..b1f3c1a122 --- /dev/null +++ b/sqlite/migrates/create_main_tables_2026-02-25-065451.sql @@ -0,0 +1,16 @@ +CREATE TABLE IF NOT EXISTS sample_input_tensor_meta ( + sample_uuid VARCHAR(255) NOT NULL, + input_name VARCHAR(255) NOT NULL, + input_idx INTEGER NOT NULL, + shape TEXT NOT NULL, + dtype VARCHAR(50) NOT NULL, + create_at DATETIME DEFAULT CURRENT_TIMESTAMP, + delete_at DATETIME, + deleted BOOLEAN DEFAULT FALSE, + FOREIGN KEY (sample_uuid) REFERENCES graph_sample(uuid), + PRIMARY KEY (sample_uuid, input_name, shape, dtype) +); + +CREATE INDEX IF NOT EXISTS idx_sample_input_tensor_meta_sample_uid ON sample_input_tensor_meta(sample_uuid); +CREATE INDEX IF NOT EXISTS idx_sample_input_tensor_meta_shape ON sample_input_tensor_meta(shape); +CREATE INDEX IF NOT EXISTS idx_sample_input_tensor_meta_dtype ON sample_input_tensor_meta(dtype); diff --git a/sqlite/orm_models.py b/sqlite/orm_models.py index c48278f954..84e368c009 100755 --- a/sqlite/orm_models.py +++ b/sqlite/orm_models.py @@ -107,6 +107,10 @@ class GraphSample(Base): back_populates="sample", uselist=False, ) + sample_input_tensor_metas = relationship( + "SampleInputTensorMeta", + back_populates="sample", + ) class SubgraphSource(Base): @@ -278,6 +282,29 @@ class SampleOpNameList(Base): sample = relationship("GraphSample", back_populates="sample_op_name_list") +class SampleInputTensorMeta(Base): + __tablename__ = "sample_input_tensor_meta" + + sample_uuid = Column( + String(255), ForeignKey("graph_sample.uuid"), nullable=False, primary_key=True + ) + input_name = Column(String(255), nullable=False, primary_key=True) + input_idx = Column(Integer, nullable=False) + shape = Column(Text, nullable=False, primary_key=True) + dtype = Column(String(50), nullable=False, primary_key=True) + create_at = Column(DateTime, default=datetime.now) + delete_at = Column(DateTime) + deleted = Column(Boolean, default=False) + + __table_args__ = ( + Index("idx_sample_input_tensor_meta_sample_uid", "sample_uuid"), + Index("idx_sample_input_tensor_meta_shape", "shape"), + Index("idx_sample_input_tensor_meta_dtype", "dtype"), + ) + + sample = relationship("GraphSample", back_populates="sample_input_tensor_metas") + + def get_session(db_path: str, echo: bool = False): engine = create_engine(f"sqlite:///{db_path}", echo=echo) Session = sessionmaker(bind=engine) From ee4bcbf05cc026ffe400a489729043e21cf1d589 Mon Sep 17 00:00:00 2001 From: Honglei-Qiu <1044497581@qq.com> Date: Thu, 26 Feb 2026 06:50:23 +0000 Subject: [PATCH 2/3] add table SampleInputTensorMeta --- sqlite/graphsample_insert.py | 40 ++++++++---------------------------- 1 file changed, 9 insertions(+), 31 deletions(-) diff --git a/sqlite/graphsample_insert.py b/sqlite/graphsample_insert.py index a31e346060..e430db05e5 100755 --- a/sqlite/graphsample_insert.py +++ b/sqlite/graphsample_insert.py @@ -378,52 +378,30 @@ def insert_sample_input_tensor_meta( relative_model_path: str, db_path: str, ): + from graph_net.tensor_meta import TensorMeta + model_path = Path(model_path_prefix) / relative_model_path weight_meta_file = model_path / "weight_meta.py" try: - with open(weight_meta_file) as f: - content = f.read() + tensor_metas = TensorMeta.unserialize_from_py_file(str(weight_meta_file)) except Exception as e: - print(f"Warning: Failed to read {weight_meta_file}: {e}") + print(f"Warning: Failed to parse {weight_meta_file}: {e}") return input_tensor_metas = [] - input_idx = 0 - original_name_pattern = re.compile(r'original_name\s*=\s*"(.+?)"') - shape_pattern = re.compile(r"shape\s*=\s*\[(.+?)\]") - dtype_pattern = re.compile(r'dtype\s*=\s*"(.+?)"') - class_blocks = re.split(r"\nclass Program_weight_tensor_meta_", content) - - for block in class_blocks: - original_name_match = original_name_pattern.search(block) - if not original_name_match: - continue - - original_name = original_name_match.group(1) - - shape_match = shape_pattern.search(block) - if not shape_match: - print(f"Warning: No shape found for {original_name} in {weight_meta_file}") - continue - - shape_str = f"[{shape_match.group(1)}]" - - dtype_match = dtype_pattern.search(block) - dtype = dtype_match.group(1) if dtype_match else "" - + for input_idx, tensor_meta in enumerate(tensor_metas): input_tensor_metas.append( { - "input_name": original_name, + "input_name": tensor_meta.original_name or tensor_meta.name, "input_idx": input_idx, - "shape": shape_str, - "dtype": dtype, + "shape": str(tensor_meta.shape), + "dtype": tensor_meta.dtype, } ) - input_idx += 1 if not input_tensor_metas: - print(f"No input tensor meta found in {weight_meta_file}") + print(f"No tensor meta found in {weight_meta_file}") return session = get_session(db_path) From 9b2caa5db998ceec313e2a1b21ecf87e2ca83bfd Mon Sep 17 00:00:00 2001 From: Honglei-Qiu <1044497581@qq.com> Date: Sun, 1 Mar 2026 03:40:12 +0000 Subject: [PATCH 3/3] add graph_net_sample_groups_insert.py --- sqlite/download.py | 22 + sqlite/graph_net_sample_groups_insert.py | 388 ++++++++++++++++++ .../create_main_tables_2026-02-28-085350.sql | 34 ++ sqlite/orm_models.py | 57 +++ 4 files changed, 501 insertions(+) create mode 100755 sqlite/download.py create mode 100644 sqlite/graph_net_sample_groups_insert.py create mode 100755 sqlite/migrates/create_main_tables_2026-02-28-085350.sql diff --git a/sqlite/download.py b/sqlite/download.py new file mode 100755 index 0000000000..d9cec2c722 --- /dev/null +++ b/sqlite/download.py @@ -0,0 +1,22 @@ +import os +from datasets import load_dataset +from huggingface_hub import hf_hub_download + +REPO_ID = "PaddlePaddle/GraphNet" +REVISION = "20260224" +SAVE_DIR = "./workspace" + +ds = load_dataset(REPO_ID, split="GraphNet", revision=REVISION) +for item in ds: + full_path = os.path.join(SAVE_DIR, item["path"]) + os.makedirs(os.path.dirname(full_path), exist_ok=True) + with open(full_path, "w", encoding="utf-8") as f: + f.write(item["content"]) + +hf_hub_download( + repo_id=REPO_ID, + filename="GraphNet.db", + repo_type="dataset", + revision=REVISION, + local_dir=SAVE_DIR, +) diff --git a/sqlite/graph_net_sample_groups_insert.py b/sqlite/graph_net_sample_groups_insert.py new file mode 100644 index 0000000000..4b71e49d82 --- /dev/null +++ b/sqlite/graph_net_sample_groups_insert.py @@ -0,0 +1,388 @@ +import argparse +import json +import uuid as uuid_module +from dataclasses import dataclass, field +from typing import List, Dict, Callable, Optional +from datetime import datetime + +from orm_models import ( + get_session, + GraphNetSampleBucket, + GraphNetSampleGroup, +) + + +GraphNetSampleUid = str +GraphNetSampleType = str +BucketId = str + + +@dataclass +class SampleBucketInfo: + sample_uid: str = "" + op_seq_bucket_id: List[str] = field(default_factory=list) + input_shapes_bucket_id: List[List[int]] = field(default_factory=list) + input_dtypes_bucket_id: List[str] = field(default_factory=list) + + def get_bucket_key(self) -> str: + return ".".join(self.op_seq_bucket_id) if self.op_seq_bucket_id else "unknown" + + +@dataclass +class Ai4cSample: + graph_net_sample_uids: List[GraphNetSampleUid] = field(default_factory=list) + shape_diversity_min_bucket_size: int = 0 + shape_diversity_max_bucket_size: int = 0 + dtype_diversity_min_bucket_size: int = 0 + dtype_diversity_max_bucket_size: int = 0 + op_seq_bucket_id: List[str] = field(default_factory=list) + group_type: str = "" + + def __post_init__(self): + self._validate() + + def _validate(self): + if self.shape_diversity_min_bucket_size > self.shape_diversity_max_bucket_size: + raise ValueError( + f"shape_diversity_min_bucket_size ({self.shape_diversity_min_bucket_size}) " + f"> shape_diversity_max_bucket_size ({self.shape_diversity_max_bucket_size})" + ) + if self.dtype_diversity_min_bucket_size > self.dtype_diversity_max_bucket_size: + raise ValueError( + f"dtype_diversity_min_bucket_size ({self.dtype_diversity_min_bucket_size}) " + f"> dtype_diversity_max_bucket_size ({self.dtype_diversity_max_bucket_size})" + ) + + +class Ai4cSampleGenerator: + def __init__( + self, + shape_diversity_min: int = 1, + shape_diversity_max: int = 2, + dtype_diversity_min: int = 1, + dtype_diversity_max: int = 2, + ): + self.shape_diversity_min = shape_diversity_min + self.shape_diversity_max = shape_diversity_max + self.dtype_diversity_min = dtype_diversity_min + self.dtype_diversity_max = dtype_diversity_max + + # Inline function hooks (can be overridden) + self.get_max_num_graph_net_samples: Optional[ + Callable[[Ai4cSample, BucketId], int] + ] = None + self.get_min_num_graph_net_samples: Optional[ + Callable[[Ai4cSample, BucketId], int] + ] = None + self._shape_diversity_sample_fn: Optional[ + Callable[[List[SampleBucketInfo], int, int], Ai4cSample] + ] = None + self._dtype_diversity_sample_fn: Optional[ + Callable[[List[SampleBucketInfo], int, int], Ai4cSample] + ] = None + + def __call__( + self, + bucket_infos: List[SampleBucketInfo], + ) -> List[Ai4cSample]: + results: List[Ai4cSample] = [] + + same_op_seq_buckets: Dict[str, List[SampleBucketInfo]] = {} + + for bucket_info in bucket_infos: + bucket_key = bucket_info.get_bucket_key() + if bucket_key not in same_op_seq_buckets: + same_op_seq_buckets[bucket_key] = [] + same_op_seq_buckets[bucket_key].append(bucket_info) + + for bucket_key, bucket_info_list in same_op_seq_buckets.items(): + op_seq = bucket_info_list[0].op_seq_bucket_id if bucket_info_list else [] + temp_sample = Ai4cSample(op_seq_bucket_id=op_seq) + max_samples = self._get_max_num_graph_net_samples(temp_sample, bucket_key) + min_samples = self._get_min_num_graph_net_samples(temp_sample, bucket_key) + + shape_sample = self._shape_diversity_sample( + bucket_info_list, max_samples, min_samples + ) + self._validate_shape_diversity_sample(shape_sample) + if shape_sample.graph_net_sample_uids: + results.append(shape_sample) + + dtype_sample = self._dtype_diversity_sample( + bucket_info_list, max_samples, min_samples + ) + self._validate_dtype_diversity_sample(dtype_sample) + if dtype_sample.graph_net_sample_uids: + results.append(dtype_sample) + + return results + + def _get_max_num_graph_net_samples( + self, sample: Ai4cSample, bucket_key: str + ) -> int: + if self.get_max_num_graph_net_samples: + return self.get_max_num_graph_net_samples(sample, bucket_key) + return 10 + + def _get_min_num_graph_net_samples( + self, sample: Ai4cSample, bucket_key: str + ) -> int: + if self.get_min_num_graph_net_samples: + return self.get_min_num_graph_net_samples(sample, bucket_key) + return 1 + + def _shape_diversity_sample( + self, + bucket_infos: List[SampleBucketInfo], + max_samples: int, + min_samples: int, + ) -> Ai4cSample: + if self._shape_diversity_sample_fn: + return self._shape_diversity_sample_fn( + bucket_infos, max_samples, min_samples + ) + + shape_groups: Dict[str, List[GraphNetSampleUid]] = {} + + for bucket_info in bucket_infos[:max_samples]: + shape_sig = str(bucket_info.input_shapes_bucket_id) + if shape_sig not in shape_groups: + shape_groups[shape_sig] = [] + shape_groups[shape_sig].append(bucket_info.sample_uid) + + filtered_uids: List[GraphNetSampleUid] = [] + bucket_sizes: List[int] = [] + + for shape_sig, uids in shape_groups.items(): + if len(uids) < self.shape_diversity_min: + continue + + truncated = uids[: self.shape_diversity_max] + filtered_uids.extend(truncated) + bucket_sizes.append(len(truncated)) + + op_seq = bucket_infos[0].op_seq_bucket_id if bucket_infos else [] + + return Ai4cSample( + graph_net_sample_uids=filtered_uids, + shape_diversity_min_bucket_size=min(bucket_sizes) if bucket_sizes else 0, + shape_diversity_max_bucket_size=max(bucket_sizes) if bucket_sizes else 0, + dtype_diversity_min_bucket_size=0, + dtype_diversity_max_bucket_size=0, + op_seq_bucket_id=op_seq, + group_type="shape_diversity", + ) + + def _dtype_diversity_sample( + self, + bucket_infos: List[SampleBucketInfo], + max_samples: int, + min_samples: int, + ) -> Ai4cSample: + if self._dtype_diversity_sample_fn: + return self._dtype_diversity_sample_fn( + bucket_infos, max_samples, min_samples + ) + + dtype_groups: Dict[str, List[GraphNetSampleUid]] = {} + + for bucket_info in bucket_infos[:max_samples]: + dtype_sig = str(bucket_info.input_dtypes_bucket_id) + if dtype_sig not in dtype_groups: + dtype_groups[dtype_sig] = [] + dtype_groups[dtype_sig].append(bucket_info.sample_uid) + + filtered_uids: List[GraphNetSampleUid] = [] + bucket_sizes: List[int] = [] + + for dtype_sig, uids in dtype_groups.items(): + if len(uids) < self.dtype_diversity_min: + continue + + truncated = uids[: self.dtype_diversity_max] + filtered_uids.extend(truncated) + bucket_sizes.append(len(truncated)) + + op_seq = bucket_infos[0].op_seq_bucket_id if bucket_infos else [] + + return Ai4cSample( + graph_net_sample_uids=filtered_uids, + shape_diversity_min_bucket_size=0, + shape_diversity_max_bucket_size=0, + dtype_diversity_min_bucket_size=min(bucket_sizes) if bucket_sizes else 0, + dtype_diversity_max_bucket_size=max(bucket_sizes) if bucket_sizes else 0, + op_seq_bucket_id=op_seq, + group_type="dtype_diversity", + ) + + def _validate_shape_diversity_sample(self, sample: Ai4cSample) -> None: + pass + + def _validate_dtype_diversity_sample(self, sample: Ai4cSample) -> None: + pass + + +def load_bucket_infos_from_db(session) -> List[SampleBucketInfo]: + bucket_infos = [] + + db_buckets = ( + session.query(GraphNetSampleBucket) + .filter(GraphNetSampleBucket.deleted.is_(False)) + .all() + ) + + for db_bucket in db_buckets: + op_seq = ( + json.loads(db_bucket.op_seq_bucket_id) if db_bucket.op_seq_bucket_id else [] + ) + input_shapes = ( + json.loads(db_bucket.input_shapes_bucket_id) + if db_bucket.input_shapes_bucket_id + else [] + ) + input_dtypes = ( + json.loads(db_bucket.input_dtypes_bucket_id) + if db_bucket.input_dtypes_bucket_id + else [] + ) + + bucket_info = SampleBucketInfo( + sample_uid=db_bucket.sample_uid, + op_seq_bucket_id=op_seq, + input_shapes_bucket_id=input_shapes, + input_dtypes_bucket_id=input_dtypes, + ) + bucket_infos.append(bucket_info) + + return bucket_infos + + +def save_groups_to_db( + session, + ai4c_samples: List[Ai4cSample], + policy_version: str = "v0.1", +) -> int: + count = 0 + + for sample in ai4c_samples: + group_uid = str(uuid_module.uuid4()) + + for sample_uid in sample.graph_net_sample_uids: + group_record = GraphNetSampleGroup( + sample_uid=sample_uid, + group_uid=group_uid, + group_type=sample.group_type, + group_policy="by_bucket", + policy_version=policy_version, + create_at=datetime.now(), + deleted=False, + ) + session.add(group_record) + count += 1 + + session.commit() + return count + + +def main(): + parser = argparse.ArgumentParser( + description="Generate graph_net_sample_groups from graph_net_sample_buckets" + ) + parser.add_argument( + "--db_path", + type=str, + required=True, + help="Path to the SQLite database file", + ) + parser.add_argument( + "--shape_diversity_min", + type=int, + default=2, + help="Minimum shape diversity (default: 2)", + ) + parser.add_argument( + "--shape_diversity_max", + type=int, + default=3, + help="Maximum shape diversity (default: 3)", + ) + parser.add_argument( + "--dtype_diversity_min", + type=int, + default=2, + help="Minimum dtype diversity (default: 2)", + ) + parser.add_argument( + "--dtype_diversity_max", + type=int, + default=3, + help="Maximum dtype diversity (default: 3)", + ) + parser.add_argument( + "--policy_version", + type=str, + default="v0.1", + help="Policy version for the generated groups (default: v0.1)", + ) + 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: Loading bucket info from database...") + bucket_infos = load_bucket_infos_from_db(session) + print(f" Loaded {len(bucket_infos)} bucket records") + + print("=" * 70) + print("Step 2: Generating Ai4cSample groups...") + generator = Ai4cSampleGenerator( + shape_diversity_min=args.shape_diversity_min, + shape_diversity_max=args.shape_diversity_max, + dtype_diversity_min=args.dtype_diversity_min, + dtype_diversity_max=args.dtype_diversity_max, + ) + ai4c_samples = generator(bucket_infos) + print(f" Generated {len(ai4c_samples)} Ai4cSample groups") + + print("=" * 70) + print("Generated Ai4cSample groups:") + for i, sample in enumerate(ai4c_samples): + print(f"\n[{i}] op_seq: {sample.op_seq_bucket_id}") + print(f" group_type: {sample.group_type}") + print(f" sample_count: {len(sample.graph_net_sample_uids)}") + print( + f" sample_uids: {sample.graph_net_sample_uids[:5]}{'...' if len(sample.graph_net_sample_uids) > 5 else ''}" + ) + if sample.group_type == "shape_diversity": + print( + f" shape_diversity: min={sample.shape_diversity_min_bucket_size}, max={sample.shape_diversity_max_bucket_size}" + ) + else: + print( + f" dtype_diversity: min={sample.dtype_diversity_min_bucket_size}, max={sample.dtype_diversity_max_bucket_size}" + ) + + print("=" * 70) + if args.dry_run: + print("Dry run mode - skipping database insert") + total_records = sum(len(s.graph_net_sample_uids) for s in ai4c_samples) + print(f" Would insert {total_records} records into graph_net_sample_groups") + else: + print("Step 3: Saving to database...") + count = save_groups_to_db(session, ai4c_samples, args.policy_version) + print(f" Inserted {count} records into graph_net_sample_groups") + + print("=" * 70) + print("Done!") + + session.close() + + +if __name__ == "__main__": + main() diff --git a/sqlite/migrates/create_main_tables_2026-02-28-085350.sql b/sqlite/migrates/create_main_tables_2026-02-28-085350.sql new file mode 100755 index 0000000000..20c1d5bf21 --- /dev/null +++ b/sqlite/migrates/create_main_tables_2026-02-28-085350.sql @@ -0,0 +1,34 @@ +-- create graph_net_sample_buckets table +CREATE TABLE IF NOT EXISTS graph_net_sample_buckets ( + sample_uid VARCHAR(255) NOT NULL PRIMARY KEY, + op_seq_bucket_id TEXT NOT NULL, + input_shapes_bucket_id TEXT NOT NULL, + input_dtypes_bucket_id TEXT NOT NULL, + create_at DATETIME DEFAULT CURRENT_TIMESTAMP, + delete_at DATETIME, + deleted BOOLEAN DEFAULT FALSE, + FOREIGN KEY (sample_uid) REFERENCES graph_sample(uuid) +); +CREATE INDEX IF NOT EXISTS idx_buckets_sample_uid ON graph_net_sample_buckets (sample_uid); +CREATE INDEX IF NOT EXISTS idx_buckets_op_seq_bucket_id ON graph_net_sample_buckets (op_seq_bucket_id); +CREATE INDEX IF NOT EXISTS idx_buckets_input_shapes_bucket_id ON graph_net_sample_buckets (input_shapes_bucket_id); +CREATE INDEX IF NOT EXISTS idx_buckets_input_dtypes_bucket_id ON graph_net_sample_buckets (input_dtypes_bucket_id); + +-- create graph_net_sample_groups table +CREATE TABLE IF NOT EXISTS graph_net_sample_groups ( + sample_uid VARCHAR(255) NOT NULL, + group_uid VARCHAR(255) NOT NULL, + group_type VARCHAR(50) NOT NULL, + group_policy VARCHAR(50) NOT NULL, + policy_version VARCHAR(20) NOT NULL, + create_at DATETIME DEFAULT CURRENT_TIMESTAMP, + delete_at DATETIME, + deleted BOOLEAN DEFAULT FALSE, + PRIMARY KEY (sample_uid, group_uid), + FOREIGN KEY (sample_uid) REFERENCES graph_sample(uuid) +); +CREATE INDEX IF NOT EXISTS idx_groups_sample_uid ON graph_net_sample_groups (sample_uid); +CREATE INDEX IF NOT EXISTS idx_groups_group_type ON graph_net_sample_groups (group_type); +CREATE INDEX IF NOT EXISTS idx_groups_group_policy ON graph_net_sample_groups (group_policy); +CREATE INDEX IF NOT EXISTS idx_groups_group_uid ON graph_net_sample_groups (group_uid); +CREATE INDEX IF NOT EXISTS idx_groups_policy_version ON graph_net_sample_groups (policy_version); diff --git a/sqlite/orm_models.py b/sqlite/orm_models.py index 84e368c009..7445dda664 100755 --- a/sqlite/orm_models.py +++ b/sqlite/orm_models.py @@ -111,6 +111,15 @@ class GraphSample(Base): "SampleInputTensorMeta", back_populates="sample", ) + sample_buckets = relationship( + "GraphNetSampleBucket", + back_populates="sample", + uselist=False, + ) + sample_groups = relationship( + "GraphNetSampleGroup", + back_populates="sample", + ) class SubgraphSource(Base): @@ -305,6 +314,54 @@ class SampleInputTensorMeta(Base): sample = relationship("GraphSample", back_populates="sample_input_tensor_metas") +class GraphNetSampleBucket(Base): + __tablename__ = "graph_net_sample_buckets" + + sample_uid = Column( + String(255), ForeignKey("graph_sample.uuid"), nullable=False, primary_key=True + ) + op_seq_bucket_id = Column(Text, nullable=False) + input_shapes_bucket_id = Column(Text, nullable=False) + input_dtypes_bucket_id = Column(Text, nullable=False) + create_at = Column(DateTime, default=datetime.now) + delete_at = Column(DateTime) + deleted = Column(Boolean, default=False) + + __table_args__ = ( + Index("idx_buckets_sample_uid", "sample_uid"), + Index("idx_buckets_op_seq_bucket_id", "op_seq_bucket_id"), + Index("idx_buckets_input_shapes_bucket_id", "input_shapes_bucket_id"), + Index("idx_buckets_input_dtypes_bucket_id", "input_dtypes_bucket_id"), + ) + + sample = relationship("GraphSample", back_populates="sample_buckets") + + +class GraphNetSampleGroup(Base): + __tablename__ = "graph_net_sample_groups" + + sample_uid = Column( + String(255), ForeignKey("graph_sample.uuid"), nullable=False, primary_key=True + ) + group_uid = Column(String(255), nullable=False, primary_key=True) + group_type = Column(String(50), nullable=False) + group_policy = Column(String(50), nullable=False) + policy_version = Column(String(20), nullable=False) + create_at = Column(DateTime, default=datetime.now) + delete_at = Column(DateTime) + deleted = Column(Boolean, default=False) + + __table_args__ = ( + Index("idx_groups_sample_uid", "sample_uid"), + Index("idx_groups_group_type", "group_type"), + Index("idx_groups_group_policy", "group_policy"), + Index("idx_groups_group_uid", "group_uid"), + Index("idx_groups_policy_version", "policy_version"), + ) + + sample = relationship("GraphSample", back_populates="sample_groups") + + def get_session(db_path: str, echo: bool = False): engine = create_engine(f"sqlite:///{db_path}", echo=echo) Session = sessionmaker(bind=engine)