Simulator code

.gitignore

@@ -175,4 +175,6 @@ Thumbs.db
 src/cpp/third_party
 test_datasets
 simulator/datasets
-simulator/images
+simulator/images
+simulator/metis-5.1.0.tar
+simulator/metis-5.1.0

simulator/configs/arvix_linear.json

@@ -5,5 +5,6 @@
         "page_size" : "16.384 KB",
         "feature_dimension" : 128,
         "feature_size" : "float32"
-    }
+    }, 
+    "sampling_depth" : 2
 }

simulator/configs/arvix_linear_in_mem.json

@@ -0,0 +1,10 @@
+{
+    "dataset_name" : "ogbn_arxiv",
+    "features_stats" : {
+        "featurizer_type" : "default",
+        "page_size" : "16.384 KB",
+        "feature_dimension" : 128,
+        "feature_size" : "float32"
+    }, 
+    "top_percent_in_mem" : 1
+}

simulator/configs/arvix.json → simulator/configs/arvix_metis.json

@@ -1,8 +1,8 @@
 {
     "dataset_name" : "ogbn_arxiv",
     "features_stats" : {
-        "feature_layout" : "random",
-        "page_size" : "16 KB",
+        "featurizer_type" : "metis_partition",
+        "page_size" : "16.384 KB",
         "feature_dimension" : 128,
         "feature_size" : "float32"
     }

simulator/configs/products_linear.json

@@ -0,0 +1,10 @@
+{
+    "dataset_name" : "ogbn_products",
+    "features_stats" : {
+        "featurizer_type" : "default",
+        "page_size" : "16.384 KB",
+        "feature_dimension" : 100,
+        "feature_size" : "float32"
+    },
+    "sampling_depth" : 1
+}

simulator/main.py

@@ -2,13 +2,13 @@
 import json
 import argparse
 import random
+import time
 
 from src.dataset_loader import *
 from src.features_loader import *
 from src.sampler import *
 from src.visualizer import *
 
-
 def read_config_file(config_file):
     with open(config_file, "r") as reader:
         return json.load(reader)
@@ -19,40 +19,61 @@ def read_arguments():
     parser.add_argument("--config_file", type=str, help="The config file containing the details for the simulation")
     parser.add_argument("--save_path", required=True, type=str, help="The path to save the resulting image to")
     parser.add_argument("--graph_title", required=True, type=str, help="The title of the saved graph")
+    parser.add_argument("--num_nodes", default = -1, type = int, help = "The number of nodes we want in our sample")
+    parser.add_argument("--log_rate", type=int, default=20, help="Log rate of the nodes processed")
     return parser.parse_args()
 
-
 def main():
+    start_time = time.time()
     arguments = read_arguments()
     config = read_config_file(arguments.config_file)
 
     # Create the loaders
-    data_loader = DatasetLoader(config["dataset_name"])
+    data_loader = DatasetLoader(config)
+    print(data_loader.get_num_nodes(), data_loader.get_num_edges())
+    '''
     features_loader = get_featurizer(data_loader, config["features_stats"])
-    sampler = SubgraphSampler(data_loader, features_loader)
+    sampler = SubgraphSampler(data_loader, features_loader, config)
+    print("Finished loading all objects")
 
     # Perform sampling
     nodes_to_sample = [i for i in range(data_loader.get_num_nodes())]
     random.shuffle(nodes_to_sample)
+    if arguments.num_nodes > 0:
+        nodes_to_sample = nodes_to_sample[ : arguments.num_nodes]
+    log_rate = int(len(nodes_to_sample) / arguments.log_rate)
 
     pages_loaded = []
     for curr_node in nodes_to_sample:
         num_pages_read = sampler.perform_sampling_for_node(curr_node)
         if num_pages_read > 0:
             pages_loaded.append(num_pages_read)
 
+        if len(pages_loaded) > 0 and len(pages_loaded) % log_rate == 0:
+            percentage_finished = (100.0 * len(pages_loaded)) / len(nodes_to_sample)
+            print("Finished processing", round(percentage_finished), "percent of nodes")
+
+    # Get the arguments to log
+    vals_to_log = dict()
+    for curr_obj in [data_loader, features_loader, sampler]:
+        vals_to_log.update(curr_obj.get_values_to_log())
+    
+    # Log the time taken
+    total_time = time.time() - start_time
+    print("Processed all", len(nodes_to_sample), "nodes in", total_time, "seconds")
+
     # Save the histogram
     os.makedirs(os.path.dirname(arguments.save_path), exist_ok=True)
     visualize_arguments = {
         "pages_loaded": pages_loaded,
         "save_path": arguments.save_path,
         "graph_title": arguments.graph_title,
-        "total_space": features_loader.get_total_file_size(),
+        "depth" : config["sampling_depth"],
         "dataset_name": config["dataset_name"],
-        "nodes_per_page": features_loader.nodes_per_page,
+        "values_to_log": vals_to_log,
     }
     visualize_results(visualize_arguments)
-
+    '''
 
 if __name__ == "__main__":
-    main()
+    main()

simulator/src/dataset_loader.py

@@ -1,15 +1,17 @@
 import subprocess
 import os
 import numpy as np
+import torch
 from collections import defaultdict
-
+import marius.storage
 
 class DatasetLoader:
     SAVE_DIR = "datasets"
     EDGES_PATH = "edges/train_edges.bin"
 
-    def __init__(self, name):
-        self.name = name
+    def __init__(self, config):
+        self.name = config["dataset_name"]
+        self.sampling_depth = config["sampling_depth"]
         os.makedirs(DatasetLoader.SAVE_DIR, exist_ok=True)
         self.save_dir = os.path.join(DatasetLoader.SAVE_DIR, self.name)
         if not os.path.exists(self.save_dir):
@@ -29,24 +31,52 @@ def load_dataset(self):
 
         # Create the adjacency map
         edges_flaten_arr = np.frombuffer(edges_bytes, dtype=np.int32)
-        self.nodes = set(edges_flaten_arr)
         edges_arr = edges_flaten_arr.reshape((-1, 2))
-        self.num_edges = len(edges_arr)
 
-        self.adjacency_map = {}
-        for source, target in edges_arr:
-            if source not in self.adjacency_map:
-                self.adjacency_map[source] = set()
-            self.adjacency_map[source].add(target)
+        '''
+        # Create the graph
+        self.edge_list = torch.from_numpy(edges_arr)
+        self.nodes = torch.unique(self.edge_list)
+        self.current_graph = MariusGraph(self.edge_list, self.edge_list[torch.argsort(edge_list[:, -1])], self.get_num_nodes())
+        self.sampler = LayeredNeighborSampler(full_graph, [-1 for _ in range(self.sampling_depth)])
+
+        # Neighbors cache
+        '''
 
     def get_num_nodes(self):
-        return len(self.nodes)
+        return self.nodes.shape(0)
 
-    def get_neigbhors_for_node(self, node_id):
+    def get_neigbhors_for_node(self, node_id, all_depths = False):
         if node_id not in self.adjacency_map:
             return []
 
         return list(self.adjacency_map[node_id])
 
+    def get_incoming_neighbors(self, node_id):
+        if node_id not in self.num_incoming_edges:
+            return 0
+
+        return self.num_incoming_edges[node_id]
+
     def get_num_edges(self):
-        return self.num_edges
+        return self.edge_list.shape(0)
+
+    def get_average_neighbors(self):
+        neighbors_count = []
+        for node_neighbors in self.adjacency_map.values():
+            neighbors_count.append(len(node_neighbors))
+
+        return np.mean(np.array(neighbors_count))
+
+    def get_average_incoming(self):
+        incoming_counts = []
+        for num_incoming in self.num_incoming_edges.values():
+            incoming_counts.append(num_incoming)
+
+        return np.mean(np.array(incoming_counts))
+
+    def get_values_to_log(self):
+        return {
+            "Average Node Out Degree": str(round(self.get_average_neighbors(), 2)),
+            "Average Node In Degree": str(round(self.get_average_incoming(), 2)),
+        }

simulator/src/features_loader.py

@@ -3,6 +3,8 @@
 import math
 import random
 import numpy as np
+import metis
+import networkx as nx
 
 
 class FeaturesLoader:
@@ -17,23 +19,34 @@ def __init__(self, data_loader, features_stat):
 
     def initialize(self):
         total_nodes = self.data_loader.get_num_nodes()
+        print("Total nodes of", total_nodes)
         self.total_pages = int(math.ceil(total_nodes / (1.0 * self.nodes_per_page)))
-        self.node_location_map = [i for i in range(total_nodes)]
+        self.node_location_map = [i for i in range(self.total_pages)]
         if "feature_layout" in self.features_stat and self.features_stat["feature_layout"] == "random":
             random.shuffle(self.node_location_map)
 
     def get_node_page(self, src_node, neighbor_node):
+        print("Querying for node", neighbor_node, "with total of", len(self.node_location_map), "nodes")
         start_node = int(self.node_location_map[neighbor_node] / self.nodes_per_page)
-        curr_page_nodes = set()
-        for node_id in range(start_node, start_node + self.nodes_per_page):
-            curr_page_nodes.add(node_id)
-
+        curr_page_nodes = set(range(start_node, start_node + self.nodes_per_page))
         return curr_page_nodes
 
+    def get_single_node_feature_size(self):
+        return self.node_feature_size
+
+    def get_nodes_per_page(self):
+        return self.nodes_per_page
+
+    def get_page_size(self):
+        return self.page_size
+
     def get_total_file_size(self):
         total_bytes = self.page_size * self.total_pages
         return humanfriendly.format_size(total_bytes)
 
+    def get_values_to_log(self):
+        return {"Nodes per Page": self.nodes_per_page, "Total File Size": self.get_total_file_size()}
+
 
 class NeighborFeaturesLoader(FeaturesLoader):
     def __init__(self, data_loader, features_stat):
@@ -63,8 +76,55 @@ def get_node_page(self, src_node, neighbor_node):
         return self.get_page_for_node(neighbor_node)
 
 
-features_class_map = {"default": FeaturesLoader, "neighbors_nearby": NeighborFeaturesLoader}
+class MetisPartitionLoader(FeaturesLoader):
+    def __init__(self, data_loader, features_stat):
+        super().__init__(data_loader, features_stat)
+
+    def initialize(self):
+        total_nodes = self.data_loader.get_num_nodes()
+        num_partitions = int(math.ceil(total_nodes / (1.0 * self.nodes_per_page)))
+
+        # Create the adjancency list
+        edge_list = []
+        for node_id in range(total_nodes):
+            for neighbor in self.data_loader.get_neigbhors_for_node(node_id):
+                edge_list.append((node_id, neighbor))
+        graph = nx.Graph(edge_list)
+
+        # Partition the graph and store the result
+        (_, parts) = metis.part_graph(graph, num_partitions)
+        self.node_to_page = parts
+        self.page_nodes = {}
+        self.additional_page = max(self.node_to_page) + 1
+
+        # Determine the nodes in each page
+        for node_idx, node_page in enumerate(self.node_to_page):
+            # Determine if we need to spill over
+            if node_page in self.page_nodes and len(self.page_nodes[node_page]) >= self.nodes_per_page:
+                node_page = self.additional_page
+                self.node_to_page[node_idx] = node_page
+
+            # Record the node for this page
+            if node_page not in self.page_nodes:
+                self.page_nodes[node_page] = set()
+            self.page_nodes[node_page].add(node_idx)
+            assert len(self.page_nodes[node_page]) <= self.nodes_per_page
+
+            # Move to the next additonal page
+            if node_page == self.additional_page and len(self.page_nodes[node_page]) >= self.nodes_per_page:
+                self.additional_page += 1
+
+        self.total_pages = len(self.page_nodes)
+        print("Finished metis partitioning")
+
+    def get_node_page(self, src_node, neighbor_node):
+        return self.page_nodes[self.node_to_page[neighbor_node]]
 
+features_class_map = {
+    "default": FeaturesLoader,
+    "neighbors_nearby": NeighborFeaturesLoader,
+    "metis_partition": MetisPartitionLoader,
+}
 
 def get_featurizer(data_loader, features_stat):
     featurizer_type = features_stat["featurizer_type"]

simulator/src/in_mem_storage.py

@@ -0,0 +1,26 @@
+import heapq
+
+
+class InMemoryStorage:
+    def __init__(self, data_loader, percent_in_memory):
+        self.percent_in_memory = percent_in_memory
+        total_nodes = data_loader.get_num_nodes()
+        nodes_in_mem = int((total_nodes * self.percent_in_memory) / 100.0)
+
+        # Get the top nodes based on incoming neighbors
+        heap = []
+        for node_id in range(total_nodes):
+            num_incoming = data_loader.get_incoming_neighbors(node_id)
+            heapq.heappush(heap, (num_incoming, node_id))
+
+        top_pairs = heapq.nlargest(nodes_in_mem, heap)
+        self.in_memory_nodes = set([pair[1] for pair in top_pairs])
+
+    def node_in_mem_storage(self, node_id):
+        return node_id in self.in_memory_nodes
+
+    def get_percentage_in_mem(self):
+        return self.percent_in_memory
+
+    def in_mem_nodes_count(self):
+        return len(self.in_memory_nodes)