gigl.distributed.dataset_factory#
DatasetFactory is responsible for building and returning a DistLinkPredictionDataset class or subclass. It does this by spawning a process which initializes rpc + worker group, loads and builds a partitioned dataset, and shuts down the rpc + worker group.
Functions
Launches a spawned process for building and returning a DistLinkPredictionDataset instance provided some SerializedGraphMetadata Args: serialized_graph_metadata (SerializedGraphMetadata): Metadata about TFRecords that are serialized to disk distributed_context (DistributedContext): Distributed context containing information for master_ip_address, rank, and world size sample_edge_direction (Union[Literal["in", "out"], str]): Whether edges in the graph are directed inward or outward. Note that this is listed as a possible string to satisfy type check, but in practice must be a Literal["in", "out"]. should_load_tensors_in_parallel (bool): Whether tensors should be loaded from serialized information in parallel or in sequence across the [node, edge, pos_label, neg_label] entity types. partitioner (Optional[DistLinkPredictionDataPartitioner]): Initialized partitioner to partition the graph inputs. If provided, this must be a DistLinkPredictionDataPartitioner or subclass of it. If not provided, will initialize a DistLinkPredictionDataPartitioner instance using provided edge assign strategy. dataset (Optional[DistLinkPredictionDataset]): Initialized dataset class to store the graph inputs. If provided, this must be a DistLinkPredictionDataset or subclass of it. If not provided, will initialize a DistLinkPredictionDataset instance using provided edge_dir. tf_dataset_options (TFDatasetOptions): Options provided to a tf.data.Dataset to tune how serialized data is read. splitter (Optional[NodeAnchorLinkSplitter]): Optional splitter to use for splitting the graph data into train, val, and test sets. If not provided (None), no splitting will be performed. _ssl_positive_label_percentage (Optional[float]): Percentage of edges to select as self-supervised labels. Must be None if supervised edge labels are provided in advance. Slotted for refactor once this functionality is available in the transductive splitter directly _dataset_building_port (int): WARNING: You don't need to configure this unless port conflict issues. Slotted for refactor. The RPC port to use to build the dataset. In future, the port will be automatically assigned based on availability. Currently defaults to: gigl.distributed.constants.DEFAULT_MASTER_DATA_BUILDING_PORT. |