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..e430db05e5 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,71 @@ 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, +): + 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: + tensor_metas = TensorMeta.unserialize_from_py_file(str(weight_meta_file)) + except Exception as e: + print(f"Warning: Failed to parse {weight_meta_file}: {e}") + return + + input_tensor_metas = [] + for input_idx, tensor_meta in enumerate(tensor_metas): + input_tensor_metas.append( + { + "input_name": tensor_meta.original_name or tensor_meta.name, + "input_idx": input_idx, + "shape": str(tensor_meta.shape), + "dtype": tensor_meta.dtype, + } + ) + + if not input_tensor_metas: + print(f"No 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 +463,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)