Sql example

README.md

@@ -111,6 +111,19 @@ The outputs will be stored next to the inputs in the `data/FB15k` directory.
 
 This simple utility is only suitable for small graphs that fit entirely in memory. To handle larger data one will have to implement their own custom preprocessor.
 
+### Checking the data
+
+It is advised that you check your using PBG's checking script. It catches common errors that our developers have run into before,
+and can be hard to debug. This command is run by invoking the following:
+
+```bash
+torchbiggraph_check \
+  --lhs-col=0 --rel-col=1 --rhs-col=2 \
+  torchbiggraph/examples/configs/fb15k_config_cpu.py
+```
+
+This command will throw run-time errors with informative messages if it encounters problems. You will see ''Found no errors in the input'' logged if there are no errors.
+
 ### Training
 
 The `torchbiggraph_train` command is used to launch training. The training parameters are tucked away in a configuration file, whose path is given to the command. They can however be overridden from the command line with the `--param` flag. The sample config is used for both training and evaluation, so we will have to use the override to specify the edge set to use.

setup.cfg

@@ -53,6 +53,7 @@ parquet = parquet
 
 [options.entry_points]
 console_scripts =
+    torchbiggraph_check = torchbiggraph.check:main
     torchbiggraph_config = torchbiggraph.config:main
     torchbiggraph_eval = torchbiggraph.eval:main
     torchbiggraph_example_fb15k = torchbiggraph.examples.fb15k:main

torchbiggraph/check.py

@@ -0,0 +1,106 @@
+import argparse
+import logging
+import sys
+import torch
+
+from torchbiggraph.train_cpu import (
+    IterationManager,
+    get_num_edge_chunks,
+)
+from torchbiggraph.graph_storages import EDGE_STORAGES, ENTITY_STORAGES
+from torchbiggraph.config import ConfigFileLoader, ConfigSchema
+from torchbiggraph.types import Bucket
+from torchbiggraph.util import EmbeddingHolder
+from torchbiggraph.util import (
+    set_logging_verbosity,
+    setup_logging,
+)
+
+
+logger = logging.getLogger("torchbiggraph")
+
+class Checker:
+    def __init__(self, config):
+        entity_storage = ENTITY_STORAGES.make_instance(config.entity_path)
+        entity_counts = {}
+        for entity, econf in config.entities.items():
+            entity_counts[entity] = []
+            for part in range(econf.num_partitions):
+                entity_counts[entity].append(entity_storage.load_count(entity, part))
+        self.entity_counts = entity_counts
+        self.config = config
+        holder = self.holder = EmbeddingHolder(config)
+
+
+    def check_all_edges(self):
+        num_edge_chunks = get_num_edge_chunks(self.config)
+
+        iteration_manager = IterationManager(
+            1,
+            self.config.edge_paths,
+            num_edge_chunks,
+            iteration_idx=0,
+        )
+        edge_storage = EDGE_STORAGES.make_instance(iteration_manager.edge_path)
+
+        for _, _, edge_chunk_idx in iteration_manager:
+            for lhs in range(self.holder.nparts_lhs):
+                for rhs in range(self.holder.nparts_rhs):
+                    cur_b = Bucket(lhs, rhs)
+                    logging.info(f"Checking edge chunk: {edge_chunk_idx} for edges_{cur_b.lhs}_{cur_b.rhs}.h5")
+                    edges = edge_storage.load_chunk_of_edges(
+                        cur_b.lhs,
+                        cur_b.rhs,
+                        edge_chunk_idx,
+                        iteration_manager.num_edge_chunks,
+                        shared=True,
+                    )
+                    self.check_edge_chunk(cur_b, edges)
+
+    def check_edge_chunk(self, cur_b, edges):
+        rhs = edges.rhs.to_tensor()
+        lhs = edges.lhs.to_tensor()
+        rel_lhs_entity_counts = torch.tensor(
+            [self.entity_counts[r.lhs][cur_b.lhs] for r in self.config.relations]
+        )
+        #Check LHS         
+        edge_lhs_entity_count = rel_lhs_entity_counts[edges.rel]
+
+        if any(lhs >= edge_lhs_entity_count):
+            _, worst_edge_idx = (lhs - edge_lhs_entity_count).max(0)
+            raise RuntimeError(f"edge {worst_edge_idx} has LHS entity of "
+                                f"{lhs[worst_edge_idx]} but rel "
+                                f"{edges.rel[worst_edge_idx]} only has "
+                                f"{edge_lhs_entity_count[worst_edge_idx]} "
+                                "entities "
+                                f" with r.name: {self.config.relations[edges.rel[worst_edge_idx]].name}. "
+                                "Preprocessing bug?")
+        #Check RHS
+        rel_rhs_entity_counts = torch.tensor(
+            [self.entity_counts[r.rhs][cur_b.rhs] for r in self.config.relations]
+        )
+        edge_rhs_entity_count = rel_rhs_entity_counts[edges.rel]
+        if any(rhs >= edge_rhs_entity_count):
+            _, worst_edge_idx = (rhs - edge_rhs_entity_count).max(0)
+            raise RuntimeError(f"edge {worst_edge_idx} has RHS entity of "
+                                f"{rhs[worst_edge_idx]} but rel "
+                                f"{edges.rel[worst_edge_idx]} only has "
+                                f"{edge_rhs_entity_count[worst_edge_idx]} "
+                                "entities "
+                                f" with r.name: {self.config.relations[edges.rel[worst_edge_idx]].name}. "
+                                "Preprocessing bug?")       
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("config", help="Path to config file")
+    parser.add_argument("-p", "--param", action="append", nargs="*")
+    opt = parser.parse_args()
+
+    loader = ConfigFileLoader()
+    config = loader.load_config(opt.config, opt.param)
+
+    set_logging_verbosity(config.verbose)
+    setup_logging(config.verbose)
+
+    Checker(config).check_all_edges()
+    logging.info("Found no errors in the input")

torchbiggraph/examples/sql_end2end/README.md

@@ -0,0 +1,56 @@
+# SQL End to End Example
+
+This is intended as a simple end-to-end example of how to get your data into
+the format that PyTorch BigGraph expects using SQL. It's implemented in SQLite
+for portability, but similar techniques scale to billions of edges using cloud
+databases such as BigQuery or SnowFlake. This pipeline can be split into three
+different components:
+
+1. Data preparation
+2. Data verification/checking
+3. Training
+
+To run the pipeline, you'll first need to download the edges.csv file,
+available HERE (TODO: INSERT LINK). This graph was constructed by
+taking the [ogbl-citation2](https://github.com/snap-stanford/ogb) graph, and
+adding edges for both paper-citations and years-published. While this graph
+might not make a huge amount of sense, it's intended to largely fulfill a
+pedagogical purpose.
+
+In the data preparation stage, we first load the graph
+into a SQLite database and then we transform and partition it. The transformation
+can be understood as first partitioning the entities, then generating a mapping
+between the graph-ids and ordinal ids per-type that PBG will expect, and finally
+writing out all the files required to train, including the config file. By
+keeping track of the vertex types, we're able to specifically verify our mappings
+in a fully self consistent fashion.
+
+Once the data has been prepared and generated, we're ready to embed the graph. We
+do this by passing the generated config to `torchbiggraph_train` in the following
+way:
+
+```
+torchbiggraph_train \
+  path/to/generated/config.py
+```
+
+The `data_prep.py` script will also compute the approximate amount of shared memory
+that will be needed for training. If the training demands are more than the
+available shared memory, you'll need to regenerate your data with more partitions
+than what you currently have. If you're seeing either a bus error or a OOM kill
+message in the kernel ring buffer but your machine has enough ram, you'll want to
+verify that `/dev/shm` is large enough to accomodate your embedding table.
+
+# Extensions
+
+A few changes will need to be made to use this at scale in production environment.
+First, this pipeline is brittle and simplistic. For production workloads it's
+probably better to use a tool like DBT or dataflow to create independent tables
+in parallel. It's also important to be careful with our indices to make our joins
+performant.
+
+When it comes time to map your buckets to hdf5 files it's almost certainly more
+performant to dump them to chunked parquet/avro files and merge those files together
+in parallel. Every company's compute infrastructure is going to be different
+enough that this piece of code will have to be custom written. Fortunately, this
+code can be written once and reused.

torchbiggraph/examples/sql_end2end/config_template.py

@@ -0,0 +1,42 @@
+CONFIG_TEMPLATE = """
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE.txt file in the root directory of this source tree.
+import torch
+
+def get_torchbiggraph_config():
+
+    config = dict(  # noqa
+        # I/O data
+        entity_path="{TRAINING_DIR}",
+        edge_paths=[
+            "{TRAINING_DIR}",
+        ],
+        checkpoint_path="{MODEL_PATH}",
+        # Graph structure
+        entities={ENTITIES_DICT},
+        relations=[
+            {RELN_DICT}
+        ],
+        # Scoring model
+        dimension=200,
+        comparator="dot",
+        # Training
+        num_epochs=50,
+        num_uniform_negs=1000,
+        num_batch_negs=1000,
+        batch_size=150_000,
+        loss_fn="softmax",
+        lr=0.05,
+        regularization_coef=1e-3,
+        num_gpus=torch.cuda.device_count(),
+        # Evaluation during training
+        eval_fraction=0,  # to reproduce results, we need to use all training data
+    )
+
+    return config
+"""