HugeGraph-ML

HugeGraph-ML reads graph data from HugeGraph and converts it to DGL graphs for tasks such as node embedding, node classification, graph classification, link prediction and fraud detection. Model implementations are under hugegraph-ml/src/hugegraph_ml/models/.

Requirements

• Python 3.10 or later
• HugeGraph Server 1.0 or later; 1.5 or later is recommended
• uv 0.7 or later

All server access goes through hugegraph-python-client (the pyhugegraph package) from the same repository. HugeGraph2DGL pulls vertices and edges over the Gremlin endpoint with g.V().hasLabel(...) and g.E().hasLabel(...), and the dataset importers write through the schema and batch vertex/edge APIs in batches of 500.

The ML stack is version pinned at the repository root under [tool.uv] constraint-dependencies:

Those pins install CPU builds. Every task accepts a gpu argument that defaults to -1, meaning CPU; pass a device index only after installing CUDA builds of torch and dgl yourself.

Installation

git clone https://github.com/apache/hugegraph-ai.git
cd hugegraph-ai
uv sync --extra ml
source .venv/bin/activate
cd hugegraph-ml/src

HugeGraph-ML is a path dependency of the root project but is not a uv workspace member. Select the ml extra at the repository root instead of creating another lock file in the subdirectory.

Implemented Models

Every module below lives in hugegraph-ml/src/hugegraph_ml/models/. models/__init__.py re-exports nothing, so import from the module file directly.

GIN accepts pooling values sum (default), mean, max, global_attention and set2set.

Reading Graph Data

HugeGraph2DGL in hugegraph-ml/src/hugegraph_ml/data/hugegraph2dgl.py opens a PyHugeClient and converts query results into DGL objects:

from hugegraph_ml.data.hugegraph2dgl import HugeGraph2DGL

hg2d = HugeGraph2DGL(
    url="http://127.0.0.1:8080",
    graph="hugegraph",
    user="",
    pwd="",
    graphspace=None,
)

Node features land in ndata["feat"], labels in ndata["label"] and each mask in ndata[<mask key>]. NodeEmbed requires feat only; NodeClassify, NodeClassifyWithEdge and NodeClassifyWithSample require feat, label, train_mask, val_mask and test_mask and raise ValueError when one is missing.

Importing Sample Datasets

hugegraph_ml.utils.dgl2hugegraph_utils writes DGL, OGB and NetworkX datasets into HugeGraph so the conversion layer has something to read. Every function takes the same url, graph, user, pwd and graphspace arguments as HugeGraph2DGL, and most upper-case the dataset name before matching it.

Any other name raises ValueError("dataset not supported"). import_split_edge_from_ogb additionally requires the idx_to_vertex_id mapping and a max_nodes cap returned by the vertex import.

clear_all_data() drops every vertex and edge in the target graph. The test fixture calls it, loads CORA, MUTAG and ACM, and calls it again on teardown.

AMAZONGATNE and AVAZU are not fetched automatically. Their archive URLs are recorded in comments above import_hetero_graph_from_dgl_no_feat and import_hetero_graph_from_dgl_bgnn.

Tasks

Task classes live in hugegraph-ml/src/hugegraph_ml/tasks/. Each one takes the converted graph and a model instance.

patience defaults to float("inf"). EarlyStopping in utils/early_stopping.py monitors either loss or accuracy, keeps a copy of the best weights and restores them when training stops.

Runnable Examples

Scripts sit in hugegraph-ml/src/hugegraph_ml/examples/. From hugegraph-ml/src, run one with:

python ./hugegraph_ml/examples/dgi_example.py

Each script also exposes a function of the same name, so it can be imported and called with a smaller epoch count.

DGI Node Embedding Example

First import DGL’s Cora dataset into HugeGraph. The name is upper-cased before use, so cora and CORA both produce the CORA_vertex and CORA_edge labels:

from hugegraph_ml.utils.dgl2hugegraph_utils import import_graph_from_dgl

import_graph_from_dgl("cora")

Read the graph and train DGI:

from hugegraph_ml.data.hugegraph2dgl import HugeGraph2DGL
from hugegraph_ml.models.dgi import DGI
from hugegraph_ml.models.mlp import MLPClassifier
from hugegraph_ml.tasks.node_classify import NodeClassify
from hugegraph_ml.tasks.node_embed import NodeEmbed

hg2d = HugeGraph2DGL()
graph = hg2d.convert_graph(vertex_label="CORA_vertex", edge_label="CORA_edge")

embed_model = DGI(n_in_feats=graph.ndata["feat"].shape[1])
embed_task = NodeEmbed(graph=graph, model=embed_model)
embedded_graph = embed_task.train_and_embed(
    add_self_loop=True, n_epochs=300, patience=30
)

classifier = MLPClassifier(
    n_in_feat=embedded_graph.ndata["feat"].shape[1],
    n_out_feat=embedded_graph.ndata["label"].unique().shape[0],
)
classify_task = NodeClassify(graph=embedded_graph, model=classifier)
classify_task.train(lr=1e-3, n_epochs=400, patience=40)
print(classify_task.evaluate())

evaluate() returns a dictionary such as {'accuracy': 0.82, 'loss': 0.5714246034622192}. The complete script is hugegraph-ml/src/hugegraph_ml/examples/dgi_example.py.

GRAND Node Classification Example

from hugegraph_ml.data.hugegraph2dgl import HugeGraph2DGL
from hugegraph_ml.models.grand import GRAND
from hugegraph_ml.tasks.node_classify import NodeClassify

hg2d = HugeGraph2DGL()
graph = hg2d.convert_graph(vertex_label="CORA_vertex", edge_label="CORA_edge")
model = GRAND(
    n_in_feats=graph.ndata["feat"].shape[1],
    n_out_feats=graph.ndata["label"].unique().shape[0],
)
task = NodeClassify(graph, model)
task.train(lr=1e-2, weight_decay=5e-4, n_epochs=2000, patience=100)
print(task.evaluate())

GRAND returns a list of logits per augmentation sample, and NodeClassify masks each element of that list before computing the loss. The complete script is hugegraph-ml/src/hugegraph_ml/examples/grand_example.py.

Troubleshooting

• Connection failures: check the HugeGraph Server address, port, and credentials.
• Schema mismatches: the examples use CORA_vertex and CORA_edge; pass the actual labels for your own data.
• ValueError: Graph is missing required node attribute ...: the node classification tasks need feat, label, train_mask, val_mask and test_mask in ndata. Import a dataset that carries masks, or pass your own mask_keys to convert_graph.
• ValueError: dataset not supported: the importer only accepts the names in the table above, and import_graph_from_ogb matches ogbl-collab without upper-casing.
• DGL or PyTorch import failures: rerun uv sync --extra ml from the repository root and confirm that Python comes from the root .venv.
• bgrl_example.py currently fails on import: it asks for MLP_Predictor from hugegraph_ml.models.bgrl, but that module defines the class as MLPPredictor.
• care_gnn_example.py reads AMAZON_user_v and the three AMAZON_net_*_e edge labels. No bundled importer creates them, so load that dataset yourself before running the script.